Improve encrypted DM file handling for Android, add server-side JPEG thumbnail generation

This commit is contained in:
2026-02-17 19:15:41 +03:00
Unverified
parent 618f55e057
commit 451355dc83
6 changed files with 270 additions and 35 deletions
+1
View File
@@ -568,6 +568,7 @@ buck-out/
# Custom rules (everything added below won't be overriden by 'Generate .gitignore File' if you use 'Update' option) # Custom rules (everything added below won't be overriden by 'Generate .gitignore File' if you use 'Update' option)
backend/data backend/data
backend/files
.vite .vite
*.db *.db
package-lock.json package-lock.json
+2 -1
View File
@@ -16,4 +16,5 @@ httpx>=0.27.2
rich>=13.9.4 rich>=13.9.4
slowapi>=0.1.9 slowapi>=0.1.9
firebase_admin>=7.1.0 firebase_admin>=7.1.0
PyNaCl>=1.5.0 PyNaCl>=1.5.0
numpy
+2 -2
View File
@@ -163,8 +163,8 @@ if not _running_in_docker():
app.mount("/internal/file_storage", file_storage_service_module.app) app.mount("/internal/file_storage", file_storage_service_module.app)
logger.info("Mounted messaging and file_storage services in development mode") logger.info("Mounted messaging and file_storage services in development mode")
except Exception as e: except Exception as e:
# If mounting fails, continue without blocking startup; log for debugging import traceback
logger.warning(f"Failed to mount internal services for development: {e}") logger.warning("Failed to mount internal services for development: %s\n%s", e, traceback.format_exc())
def _get_username_for_log(user) -> str | None: def _get_username_for_log(user) -> str | None:
+172 -15
View File
@@ -187,15 +187,22 @@ async def store_encrypted_file(
} }
""" """
mod = _get_file_storage_module() mod = _get_file_storage_module()
# Build allowed users list once for both in-process and HTTP modes
allowed_user_ids: list[int] = []
if sender_id is not None:
allowed_user_ids.append(sender_id)
if recipient_id is not None:
allowed_user_ids.append(recipient_id)
if mod: if mod:
try: try:
# In-process: call the upload-base64 endpoint directly # In-process: call internal function directly
return await mod.upload_base64_file( return await mod.upload_base64_internal(
None, # request - not needed for in-process
filename=filename, filename=filename,
data_b64=encrypted_file_data_b64, data_b64=encrypted_file_data_b64,
content_type=content_type, content_type=content_type,
) # type: ignore allowed_user_ids=allowed_user_ids,
)
except Exception as e: except Exception as e:
logger.error("In-process file_storage.store_encrypted_file failed: %s", e) logger.error("In-process file_storage.store_encrypted_file failed: %s", e)
raise raise
@@ -207,11 +214,6 @@ async def store_encrypted_file(
try: try:
try: try:
import httpx import httpx
allowed_user_ids = []
if sender_id is not None:
allowed_user_ids.append(sender_id)
if recipient_id is not None:
allowed_user_ids.append(recipient_id)
payload = { payload = {
"filename": filename, "filename": filename,
@@ -225,11 +227,6 @@ async def store_encrypted_file(
return r.json() return r.json()
except Exception: except Exception:
from urllib import request from urllib import request
allowed_user_ids = []
if sender_id is not None:
allowed_user_ids.append(sender_id)
if recipient_id is not None:
allowed_user_ids.append(recipient_id)
payload = { payload = {
"filename": filename, "filename": filename,
@@ -396,7 +393,7 @@ async def process_message_with_files_in_messaging_service(
compliance_public_key_b64=compliance_public_key_b64, compliance_public_key_b64=compliance_public_key_b64,
sender_public_key_b64=sender_public_key_b64, sender_public_key_b64=sender_public_key_b64,
recipient_public_key_b64=recipient_public_key_b64, recipient_public_key_b64=recipient_public_key_b64,
files=[f["encrypted_file_data_b64"] for f in transport_files], transport_files=transport_files,
) )
except Exception as e: except Exception as e:
logger.error("In-process messaging.process_message_with_files failed: %s", e) logger.error("In-process messaging.process_message_with_files failed: %s", e)
@@ -424,3 +421,163 @@ async def process_message_with_files_in_messaging_service(
raise raise
async def init_resumable_upload_in_storage(
filename: str,
total_size: int,
allowed_user_ids: list[int],
chunk_size: int | None = None,
timeout: float = 10.0,
) -> Dict[str, Any]:
mod = _get_file_storage_module()
if mod:
try:
return await mod.init_resumable_upload_internal(
filename=filename,
total_size=total_size,
allowed_user_ids=allowed_user_ids,
chunk_size=chunk_size,
)
except Exception as e:
logger.error("In-process file_storage.init_resumable_upload failed: %s", e)
raise
file_storage_url = os.getenv("FILE_STORAGE_URL") or os.getenv("FILE_STORAGE_SERVICE_URL") or "http://127.0.0.1:8302"
url = f"{file_storage_url.rstrip('/')}/uploads/resumable/init"
payload = {
"filename": filename,
"total_size": total_size,
"allowed_user_ids": allowed_user_ids,
}
if chunk_size is not None:
payload["chunk_size"] = chunk_size
import httpx
async with httpx.AsyncClient(timeout=timeout) as client:
r = await client.post(url, json=payload)
r.raise_for_status()
return r.json()
async def get_resumable_upload_status_in_storage(
upload_id: str,
user_id: int,
timeout: float = 10.0,
) -> Dict[str, Any]:
mod = _get_file_storage_module()
if mod:
try:
return await mod.get_resumable_upload_status_internal(upload_id, user_id)
except Exception as e:
logger.error("In-process file_storage.get_resumable_upload_status failed: %s", e)
raise
file_storage_url = os.getenv("FILE_STORAGE_URL") or os.getenv("FILE_STORAGE_SERVICE_URL") or "http://127.0.0.1:8302"
url = f"{file_storage_url.rstrip('/')}/uploads/resumable/{upload_id}"
import httpx
async with httpx.AsyncClient(timeout=timeout) as client:
r = await client.get(url, headers={"X-User-ID": str(user_id)})
r.raise_for_status()
return r.json()
async def upload_resumable_chunk_in_storage(
upload_id: str,
user_id: int,
offset: int,
data_b64: str,
timeout: float = 30.0,
) -> Dict[str, Any]:
mod = _get_file_storage_module()
if mod:
try:
return await mod.upload_resumable_chunk_internal(
upload_id, user_id, offset, data_b64
)
except Exception as e:
logger.error("In-process file_storage.upload_resumable_chunk failed: %s", e)
raise
file_storage_url = os.getenv("FILE_STORAGE_URL") or os.getenv("FILE_STORAGE_SERVICE_URL") or "http://127.0.0.1:8302"
url = f"{file_storage_url.rstrip('/')}/uploads/resumable/{upload_id}"
payload = {
"offset": offset,
"data_b64": data_b64,
}
import httpx
async with httpx.AsyncClient(timeout=timeout) as client:
r = await client.patch(url, json=payload, headers={"X-User-ID": str(user_id)})
r.raise_for_status()
return r.json()
async def complete_resumable_upload_in_storage(
upload_id: str,
user_id: int,
timeout: float = 10.0,
) -> Dict[str, Any]:
mod = _get_file_storage_module()
if mod:
try:
return await mod.complete_resumable_upload_internal(upload_id, user_id)
except Exception as e:
logger.error("In-process file_storage.complete_resumable_upload failed: %s", e)
raise
file_storage_url = os.getenv("FILE_STORAGE_URL") or os.getenv("FILE_STORAGE_SERVICE_URL") or "http://127.0.0.1:8302"
url = f"{file_storage_url.rstrip('/')}/uploads/resumable/{upload_id}/complete"
import httpx
async with httpx.AsyncClient(timeout=timeout) as client:
r = await client.post(url, json={"upload_id": upload_id}, headers={"X-User-ID": str(user_id)})
r.raise_for_status()
return r.json()
async def get_resumable_upload_data_in_storage(
upload_id: str,
user_id: int,
timeout: float = 30.0,
) -> Dict[str, Any]:
mod = _get_file_storage_module()
if mod:
try:
return await mod.get_resumable_upload_data_internal(upload_id, user_id)
except Exception as e:
logger.error("In-process file_storage.get_resumable_upload_data failed: %s", e)
raise
file_storage_url = os.getenv("FILE_STORAGE_URL") or os.getenv("FILE_STORAGE_SERVICE_URL") or "http://127.0.0.1:8302"
url = f"{file_storage_url.rstrip('/')}/uploads/resumable/{upload_id}/data-b64"
import httpx
async with httpx.AsyncClient(timeout=timeout) as client:
r = await client.get(url, headers={"X-User-ID": str(user_id)})
r.raise_for_status()
return r.json()
async def delete_resumable_upload_in_storage(
upload_id: str,
user_id: int,
timeout: float = 10.0,
) -> Dict[str, Any]:
mod = _get_file_storage_module()
if mod:
try:
return await mod.delete_resumable_upload_internal(upload_id, user_id)
except Exception as e:
logger.error("In-process file_storage.delete_resumable_upload failed: %s", e)
raise
file_storage_url = os.getenv("FILE_STORAGE_URL") or os.getenv("FILE_STORAGE_SERVICE_URL") or "http://127.0.0.1:8302"
url = f"{file_storage_url.rstrip('/')}/uploads/resumable/{upload_id}"
import httpx
async with httpx.AsyncClient(timeout=timeout) as client:
r = await client.delete(url, headers={"X-User-ID": str(user_id)})
r.raise_for_status()
return r.json()
+19 -9
View File
@@ -187,6 +187,7 @@ class ProcessMessageWithFilesFile(BaseModel):
A single transport-encrypted file blob (base64 of nonce||ciphertext). A single transport-encrypted file blob (base64 of nonce||ciphertext).
""" """
encrypted_file_data_b64: str encrypted_file_data_b64: str
filename: str = "file"
class ProcessMessageWithFilesRequest(ProcessMessageRequest): class ProcessMessageWithFilesRequest(ProcessMessageRequest):
@@ -200,13 +201,13 @@ class ProcessMessageWithFilesRequest(ProcessMessageRequest):
# Health Checks # Health Checks
# ============================================================================ # ============================================================================
@app.get("/health") @app.get("/health", response_model=None)
async def health_check(): async def health_check():
"""Health check endpoint for messaging service.""" """Health check endpoint for messaging service."""
return {"status": "healthy", "service": "messaging"} return {"status": "healthy", "service": "messaging"}
@app.get("/") @app.get("/", response_model=None)
async def root(): async def root():
"""Root endpoint for messaging service.""" """Root endpoint for messaging service."""
return {"message": "FromChat Messaging Service", "status": "operational"} return {"message": "FromChat Messaging Service", "status": "operational"}
@@ -216,7 +217,7 @@ async def root():
# Ephemeral Key Endpoints # Ephemeral Key Endpoints
# ============================================================================ # ============================================================================
@app.get("/key/transport/public") @app.get("/key/transport/public", response_model=None)
async def get_transport_public_key(): async def get_transport_public_key():
""" """
Return the current ephemeral transport public key for client-side message encryption. Return the current ephemeral transport public key for client-side message encryption.
@@ -304,7 +305,7 @@ async def process_message(
raise raise
@app.post("/process") @app.post("/process", response_model=None)
async def process_message_http(request: ProcessMessageRequest): async def process_message_http(request: ProcessMessageRequest):
""" """
HTTP endpoint for processing encrypted messages. HTTP endpoint for processing encrypted messages.
@@ -328,10 +329,11 @@ async def process_message_with_files(
compliance_public_key_b64: str, compliance_public_key_b64: str,
sender_public_key_b64: str, sender_public_key_b64: str,
recipient_public_key_b64: str, recipient_public_key_b64: str,
files: list[str], transport_files: list[dict],
): ):
""" """
In-process helper: process message + transport-encrypted files with one MEK. In-process helper: process message + transport-encrypted files with one MEK.
transport_files: list of {"encrypted_file_data_b64": str, "filename": str}
""" """
private_key = _get_ephemeral_private_key() private_key = _get_ephemeral_private_key()
@@ -343,8 +345,10 @@ async def process_message_with_files(
) )
plaintext_files: list[bytes] = [] plaintext_files: list[bytes] = []
for encrypted_file_data_b64 in files: filenames: list[str] = []
transport_blob = base64.b64decode(encrypted_file_data_b64) for tf in transport_files:
enc_b64 = tf.get("encrypted_file_data_b64", "")
transport_blob = base64.b64decode(enc_b64)
plaintext_files.append( plaintext_files.append(
decrypt_transport_blob( decrypt_transport_blob(
client_public_key_b64=sender_public_key_b64, client_public_key_b64=sender_public_key_b64,
@@ -352,17 +356,19 @@ async def process_message_with_files(
ephemeral_private_key=private_key, ephemeral_private_key=private_key,
) )
) )
filenames.append(tf.get("filename", "file"))
return process_encrypted_message_and_files( return process_encrypted_message_and_files(
plaintext_message=plaintext_message, plaintext_message=plaintext_message,
plaintext_files=plaintext_files, plaintext_files=plaintext_files,
filenames=filenames,
compliance_public_key_b64=compliance_public_key_b64, compliance_public_key_b64=compliance_public_key_b64,
sender_public_key_b64=sender_public_key_b64, sender_public_key_b64=sender_public_key_b64,
recipient_public_key_b64=recipient_public_key_b64, recipient_public_key_b64=recipient_public_key_b64,
) )
@app.post("/process-with-files") @app.post("/process-with-files", response_model=None)
async def process_message_with_files_http(request: ProcessMessageWithFilesRequest): async def process_message_with_files_http(request: ProcessMessageWithFilesRequest):
""" """
Process an encrypted message and its files using a single MEK. Process an encrypted message and its files using a single MEK.
@@ -373,6 +379,10 @@ async def process_message_with_files_http(request: ProcessMessageWithFilesReques
- MEK is wrapped for compliance, sender, and recipient (stored on DM envelope) - MEK is wrapped for compliance, sender, and recipient (stored on DM envelope)
""" """
try: try:
transport_files = [
{"encrypted_file_data_b64": f.encrypted_file_data_b64, "filename": f.filename}
for f in request.files
]
return await process_message_with_files( return await process_message_with_files(
client_public_key_b64=request.client_public_key_b64, client_public_key_b64=request.client_public_key_b64,
transport_nonce_b64=request.transport_nonce_b64, transport_nonce_b64=request.transport_nonce_b64,
@@ -380,7 +390,7 @@ async def process_message_with_files_http(request: ProcessMessageWithFilesReques
compliance_public_key_b64=request.compliance_public_key_b64, compliance_public_key_b64=request.compliance_public_key_b64,
sender_public_key_b64=request.sender_public_key_b64, sender_public_key_b64=request.sender_public_key_b64,
recipient_public_key_b64=request.recipient_public_key_b64, recipient_public_key_b64=request.recipient_public_key_b64,
files=[f.encrypted_file_data_b64 for f in request.files], transport_files=transport_files,
) )
except Exception as e: except Exception as e:
logger.exception("Failed to process message with files: %s", e) logger.exception("Failed to process message with files: %s", e)
+74 -8
View File
@@ -9,10 +9,12 @@ This module handles the core envelope encryption workflow:
5. Store encrypted message + wrapped keys 5. Store encrypted message + wrapped keys
""" """
import io
import logging import logging
import json import json
import time import time
import base64 import base64
from pathlib import Path
from typing import Dict, Any, Optional from typing import Dict, Any, Optional
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
@@ -147,9 +149,41 @@ def process_encrypted_message(
raise raise
_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp"})
_THUMB_SIZE = 80
def _generate_thumbnail(image_bytes: bytes) -> tuple[str | None, list[int]]:
"""Generate tiny JPEG thumbnail (Telegram-style). Returns (base64_jpeg, [w,h]) or (None, [1,1]) on error."""
try:
from math import gcd
from PIL import Image
img = Image.open(io.BytesIO(image_bytes))
img = img.convert("RGB")
if hasattr(img, "info") and img.info:
img.info.pop("icc_profile", None)
w, h = img.size
g = gcd(w, h) if h else 1
aspect_wh = [w // g, h // g] if g else [1, 1]
if w > _THUMB_SIZE or h > _THUMB_SIZE:
scale = min(_THUMB_SIZE / w, _THUMB_SIZE / h)
new_w = max(1, int(w * scale))
new_h = max(1, int(h * scale))
img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
buf = io.BytesIO()
img.save(buf, format="JPEG", quality=85, optimize=True)
jpeg_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
logger.info("THUMB: Image %dx%d -> thumb %dx%d, b64len=%d", w, h, img.width, img.height, len(jpeg_b64))
return (jpeg_b64, aspect_wh)
except Exception as e:
logger.warning("THUMB: Generation failed: %s", e)
return (None, [1, 1])
def process_encrypted_message_and_files( def process_encrypted_message_and_files(
plaintext_message: bytes, plaintext_message: bytes,
plaintext_files: list[bytes], plaintext_files: list[bytes],
filenames: list[str],
compliance_public_key_b64: str, compliance_public_key_b64: str,
sender_public_key_b64: str, sender_public_key_b64: str,
recipient_public_key_b64: str, recipient_public_key_b64: str,
@@ -160,29 +194,61 @@ def process_encrypted_message_and_files(
- Generates one random MEK - Generates one random MEK
- Encrypts message and each file with AES-GCM using that MEK (unique nonce per item) - Encrypts message and each file with AES-GCM using that MEK (unique nonce per item)
- Wraps the MEK for compliance, sender, and recipient - Wraps the MEK for compliance, sender, and recipient
Returns: Returns:
{ {
"message": {"nonce": str, "ciphertext": str}, "message": {"nonce": str, "ciphertext": str},
"files": [{"nonce": str, "ciphertext": str}, ...], "files": [{"nonce": str, "ciphertext": str}, ...],
"compliance_wrapped_mek": str, ...
"sender_wrapped_mek": str,
"recipient_wrapped_mek": str,
} }
""" """
start_time = time.time() start_time = time.time()
if len(filenames) != len(plaintext_files):
filenames = [f"file_{i}" for i in range(len(plaintext_files))]
# One MEK for everything in this envelope # One MEK for everything in this envelope
mek = generate_mek() mek = generate_mek()
# Build message plaintext: when we have files, use JSON with text + fileThumbnails + fileAspectRatios + fileSizes
file_thumbnails: list[str] = []
file_aspect_ratios: list[list[int]] = []
file_sizes: list[int] = []
for i, f_bytes in enumerate(plaintext_files):
name = filenames[i] if i < len(filenames) else ""
file_sizes.append(len(f_bytes))
if Path(name).suffix.lower() in _IMAGE_EXTENSIONS:
thumb_b64, wh = _generate_thumbnail(f_bytes)
file_thumbnails.append(thumb_b64 or "")
file_aspect_ratios.append(wh)
else:
file_thumbnails.append("")
file_aspect_ratios.append([1, 1])
if plaintext_files:
msg_obj = {
"text": plaintext_message.decode("utf-8", errors="replace"),
"fileThumbnails": file_thumbnails,
"fileAspectRatios": file_aspect_ratios,
"fileSizes": file_sizes,
}
logger.info(
"THUMB: Message with %d files, thumbnails=%s, aspectRatios=%s",
len(file_thumbnails),
[f"len={len(t)}" if t else "empty" for t in file_thumbnails],
file_aspect_ratios,
)
plaintext_to_encrypt = json.dumps(msg_obj, ensure_ascii=False).encode("utf-8")
else:
plaintext_to_encrypt = plaintext_message
# Encrypt message # Encrypt message
msg_nonce, msg_ciphertext = encrypt_message(plaintext_message, mek) msg_nonce, msg_ciphertext = encrypt_message(plaintext_to_encrypt, mek)
# Encrypt files (same MEK, per-file nonce) # Encrypt files (same MEK, per-file nonce)
files_out: list[Dict[str, str]] = [] files_out: list[Dict[str, Any]] = []
for f_bytes in plaintext_files: for i, f_bytes in enumerate(plaintext_files):
f_nonce, f_ciphertext = encrypt_message(f_bytes, mek) f_nonce, f_ciphertext = encrypt_message(f_bytes, mek)
files_out.append({"nonce": f_nonce, "ciphertext": f_ciphertext}) entry: Dict[str, Any] = {"nonce": f_nonce, "ciphertext": f_ciphertext}
files_out.append(entry)
# Derive wrap keys deterministically (same as existing flow) # Derive wrap keys deterministically (same as existing flow)
compliance_key_bytes = base64.b64decode(compliance_public_key_b64) compliance_key_bytes = base64.b64decode(compliance_public_key_b64)