Update app.py

app.py

@@ -4,14 +4,16 @@ import time
 import uuid
 import base64
 import logging
-from typing import Optional, Dict, Set, List, Any
+import tempfile
+import shutil
+from typing import Optional, Dict, List, Any
 from contextlib import asynccontextmanager
 
 import aiosqlite
 
 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Header, Query, UploadFile, File
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import Response, JSONResponse, FileResponse
+from fastapi.responses import Response, JSONResponse, FileResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
@@ -35,6 +37,10 @@ if MESSAGE_KEY_B64 is None:
 MESSAGE_KEY = base64.urlsafe_b64decode(MESSAGE_KEY_B64)
 assert len(MESSAGE_KEY) == 32
 
+# Temp directory for chunk assembly - uses disk not RAM
+TEMP_DIR = os.environ.get("TEMP_DIR", "/tmp/infinitychat_uploads")
+os.makedirs(TEMP_DIR, exist_ok=True)
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("InfinityChat")
 
@@ -43,7 +49,6 @@ logger = logging.getLogger("InfinityChat")
 # ------------------------------------------------------------------------
 async def init_database():
     os.makedirs(os.path.dirname(DATABASE_URL), exist_ok=True)
-
     db_existed = download_database(DATABASE_URL)
     if db_existed:
         logger.info("✅ Restored database from bucket")
@@ -106,6 +111,9 @@ async def init_database():
     start_db_sync(DATABASE_URL)
     logger.info("✅ Database initialized successfully")
 
+    # Clean up any leftover temp upload dirs from previous runs
+    _cleanup_temp_dir()
+
 # ------------------------------------------------------------------------
 # FastAPI lifespan
 # ------------------------------------------------------------------------
@@ -115,6 +123,7 @@ async def lifespan(app: FastAPI):
     yield
     logger.info("🔄 Final database sync on shutdown...")
     upload_database(DATABASE_URL)
+    _cleanup_temp_dir()
 
 app = FastAPI(title="InfinityChat", version="1.0.0", lifespan=lifespan)
 app.add_middleware(
@@ -126,6 +135,22 @@ app.add_middleware(
 )
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
+# ------------------------------------------------------------------------
+# Temp directory cleanup
+# ------------------------------------------------------------------------
+def _cleanup_temp_dir():
+    """Remove all leftover temp chunk dirs older than 2 hours."""
+    try:
+        now = time.time()
+        for entry in os.scandir(TEMP_DIR):
+            if entry.is_dir():
+                age = now - entry.stat().st_mtime
+                if age > 7200:  # 2 hours
+                    shutil.rmtree(entry.path, ignore_errors=True)
+                    logger.debug(f"🧹 Cleaned up old temp dir: {entry.path}")
+    except Exception as e:
+        logger.warning(f"Temp cleanup error: {e}")
+
 # ------------------------------------------------------------------------
 # Database helper
 # ------------------------------------------------------------------------
@@ -138,8 +163,7 @@ async def get_db() -> aiosqlite.Connection:
 
 def schedule_db_sync():
     import threading
-    t = threading.Thread(target=upload_database, args=(DATABASE_URL,), daemon=True)
-    t.start()
+    threading.Thread(target=upload_database, args=(DATABASE_URL,), daemon=True).start()
 
 # ------------------------------------------------------------------------
 # Encryption helpers (AES-256-GCM)
@@ -336,12 +360,29 @@ async def upload_avatar(
     user = await authenticate_user(token)
     if not user:
         raise HTTPException(401)
-    data = await file.read()
-    if len(data) > 5 * 1024 * 1024:
-        raise HTTPException(400, "Image too large (max 5MB)")
-    ext = os.path.splitext(file.filename)[1] if file.filename else '.jpg'
-    remote_path = f"avatars/{user['username']}_{uuid.uuid4().hex}{ext}"
-    store_file(remote_path, data)
+
+    # Write to temp file on disk first, not RAM
+    suffix = os.path.splitext(file.filename)[1] if file.filename else '.jpg'
+    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, dir=TEMP_DIR) as tmp:
+        tmp_path = tmp.name
+        size = 0
+        chunk = await file.read(65536)
+        while chunk:
+            size += len(chunk)
+            if size > 5 * 1024 * 1024:
+                os.unlink(tmp_path)
+                raise HTTPException(400, "Image too large (max 5MB)")
+            tmp.write(chunk)
+            chunk = await file.read(65536)
+
+    try:
+        with open(tmp_path, 'rb') as f:
+            data = f.read()
+        remote_path = f"avatars/{user['username']}_{uuid.uuid4().hex}{suffix}"
+        store_file(remote_path, data)
+    finally:
+        os.unlink(tmp_path)
+
     db = await get_db()
     try:
         await db.execute(
@@ -355,8 +396,9 @@ async def upload_avatar(
     return {"avatar_path": remote_path}
 
 # ------------------------------------------------------------------------
-# Chunked Upload
+# Chunked Upload - assembles chunks on DISK not RAM
 # ------------------------------------------------------------------------
+# Tracks active uploads: upload_id -> metadata dict (no chunk data in RAM)
 upload_sessions: Dict[str, dict] = {}
 
 @app.post("/api/upload/chunk")
@@ -374,14 +416,22 @@ async def upload_chunk(
     if not user:
         raise HTTPException(401)
 
+    # Validate file size limit (500MB)
+    MAX_FILE_SIZE = 500 * 1024 * 1024
+    if file_size > MAX_FILE_SIZE:
+        raise HTTPException(400, "File too large (max 500MB)")
+
+    # Create session with temp dir on disk
     if upload_id not in upload_sessions:
+        session_dir = os.path.join(TEMP_DIR, f"upload_{upload_id}")
+        os.makedirs(session_dir, exist_ok=True)
         upload_sessions[upload_id] = {
-            "chunks": {},
+            "session_dir": session_dir,
             "filename": file_name,
             "file_type": file_type,
             "file_size": file_size,
             "total_chunks": total_chunks,
-            "received": 0,
+            "received_chunks": set(),
             "user_id": user['id'],
             "created_at": time.time()
         }
@@ -390,15 +440,50 @@ async def upload_chunk(
     if session['user_id'] != user['id']:
         raise HTTPException(403)
 
-    chunk_data = await file.read()
-    session['chunks'][chunk_index] = chunk_data
-    session['received'] += 1
+    # Write chunk directly to disk
+    chunk_path = os.path.join(session['session_dir'], f"chunk_{chunk_index:06d}")
+    with open(chunk_path, 'wb') as f:
+        # Stream chunk to disk in 64KB pieces to avoid RAM spike
+        data = await file.read(65536)
+        while data:
+            f.write(data)
+            data = await file.read(65536)
 
-    if session['received'] == total_chunks:
-        full_data = b"".join(session['chunks'][i] for i in sorted(session['chunks']))
+    session['received_chunks'].add(chunk_index)
+
+    if len(session['received_chunks']) == total_chunks:
+        # All chunks received - assemble on disk then stream to bucket
         remote_path = f"uploads/{user['username']}/{uuid.uuid4().hex}/{file_name}"
-        store_file(remote_path, full_data)
-        del upload_sessions[upload_id]
+        session_dir = session['session_dir']
+
+        try:
+            # Assemble chunks into single temp file on disk
+            assembled_path = os.path.join(session_dir, "assembled")
+            with open(assembled_path, 'wb') as out_f:
+                for i in range(total_chunks):
+                    chunk_file = os.path.join(session_dir, f"chunk_{i:06d}")
+                    with open(chunk_file, 'rb') as in_f:
+                        # Copy in 1MB pieces
+                        buf = in_f.read(1024 * 1024)
+                        while buf:
+                            out_f.write(buf)
+                            buf = in_f.read(1024 * 1024)
+                    os.unlink(chunk_file)  # Delete chunk immediately after use
+
+            # Read assembled file and store to bucket
+            # Note: this does load into RAM once for encryption
+            # For very large files this is unavoidable with AES-GCM
+            # as it needs to process the whole file
+            with open(assembled_path, 'rb') as f:
+                full_data = f.read()
+            store_file(remote_path, full_data)
+            del full_data  # Explicitly free RAM immediately
+
+        finally:
+            # Always clean up temp dir
+            shutil.rmtree(session_dir, ignore_errors=True)
+            del upload_sessions[upload_id]
+
         return {
             "status": "complete",
             "file_path": remote_path,
@@ -409,14 +494,16 @@ async def upload_chunk(
 
     return {
         "status": "in_progress",
-        "received": session['received'],
+        "received": len(session['received_chunks']),
         "total": total_chunks
     }
 
 # ------------------------------------------------------------------------
-# File Download - accepts token as header OR ?t= query param
-# so <img src> and window.open() work without fetch headers
+# File Download - streams from bucket to client in chunks
+# avoids loading entire file into RAM
 # ------------------------------------------------------------------------
+DOWNLOAD_CHUNK_SIZE = 1024 * 1024  # 1MB streaming chunks
+
 @app.get("/api/download/{file_path:path}")
 async def download_file(
     file_path: str,
@@ -429,19 +516,40 @@ async def download_file(
         raise HTTPException(401)
     if '..' in file_path:
         raise HTTPException(400, "Invalid path")
+
     try:
+        # retrieve_file decrypts and returns bytes
+        # For true streaming we'd need a streaming decrypt but AES-GCM
+        # requires the full ciphertext to verify the auth tag before
+        # decrypting - so one full read is required for security.
+        # We do however stream the response TO the client in chunks.
         data = retrieve_file(file_path)
     except FileNotFoundError:
         raise HTTPException(404, "File not found")
     except Exception as e:
         logger.error(f"Download error for {file_path}: {e}")
         raise HTTPException(500, "Failed to retrieve file")
+
     import mimetypes
     mt = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-    return Response(
-        content=data,
+    file_size = len(data)
+
+    # Stream response to client in chunks to avoid keeping
+    # large response in RAM on the server side
+    def iter_data():
+        offset = 0
+        while offset < len(data):
+            yield data[offset:offset + DOWNLOAD_CHUNK_SIZE]
+            offset += DOWNLOAD_CHUNK_SIZE
+
+    return StreamingResponse(
+        iter_data(),
         media_type=mt,
-        headers={"Cache-Control": "private, max-age=3600"}
+        headers={
+            "Content-Length": str(file_size),
+            "Cache-Control": "private, max-age=3600",
+            "Content-Disposition": f'inline; filename="{os.path.basename(file_path)}"'
+        }
     )
 
 # ------------------------------------------------------------------------
@@ -532,7 +640,6 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
             data = await ws.receive_json()
             mtype = data.get("type")
 
-            # -------------------- SEND MESSAGE --------------------
             if mtype == "send_message":
                 content = data.get("content", "").strip()
                 reply_to_id = data.get("reply_to_id")
@@ -543,18 +650,10 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 client_id = data.get("client_id")
 
                 if not content and not file_path:
-                    await ws.send_json({
-                        "type": "error",
-                        "code": "EMPTY",
-                        "message": "Message cannot be empty"
-                    })
+                    await ws.send_json({"type": "error", "code": "EMPTY", "message": "Message cannot be empty"})
                     continue
                 if len(content) > 10000:
-                    await ws.send_json({
-                        "type": "error",
-                        "code": "TOO_LONG",
-                        "message": "Message too long"
-                    })
+                    await ws.send_json({"type": "error", "code": "TOO_LONG", "message": "Message too long"})
                     continue
 
                 encrypted = encrypt_message(content)
@@ -567,8 +666,7 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                            (sender_id, encrypted_content, timestamp_ms, reply_to_id,
                             file_path, file_type, file_name, file_size)
                            VALUES (?,?,?,?,?,?,?,?)""",
-                        (uid, encrypted, ts, reply_to_id,
-                         file_path, file_type, file_name, file_size)
+                        (uid, encrypted, ts, reply_to_id, file_path, file_type, file_name, file_size)
                     )
                     mid = cursor.lastrowid
                     await db.commit()
@@ -605,7 +703,6 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     "message": msg_obj
                 }, exclude_user_id=uid)
 
-            # -------------------- LOAD MESSAGES --------------------
             elif mtype == "load_messages":
                 cursor_id = data.get("cursor")
                 limit = min(data.get("limit", 50), 100)
@@ -669,7 +766,6 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 finally:
                     await db.close()
 
-            # -------------------- EDIT MESSAGE --------------------
             elif mtype == "edit_message":
                 mid = data.get("message_id")
                 new_content = data.get("content", "").strip()
@@ -682,11 +778,7 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                         (mid, uid)
                     ) as cursor:
                         if not await cursor.fetchone():
-                            await ws.send_json({
-                                "type": "error",
-                                "code": "NOT_FOUND",
-                                "message": "Message not found"
-                            })
+                            await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Message not found"})
                             continue
                     new_enc = encrypt_message(new_content)
                     await db.execute(
@@ -704,7 +796,6 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 finally:
                     await db.close()
 
-            # -------------------- DELETE MESSAGE --------------------
             elif mtype == "delete_message":
                 mid = data.get("message_id")
                 if not mid:
@@ -717,11 +808,7 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     ) as cursor:
                         msg = await cursor.fetchone()
                     if not msg:
-                        await ws.send_json({
-                            "type": "error",
-                            "code": "NOT_FOUND",
-                            "message": "Message not found"
-                        })
+                        await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Message not found"})
                         continue
                     if msg['file_path']:
                         try:
@@ -732,25 +819,19 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (mid,))
                     await db.commit()
                     schedule_db_sync()
-                    await manager.broadcast({
-                        "type": "message_deleted",
-                        "message_id": mid
-                    })
+                    await manager.broadcast({"type": "message_deleted", "message_id": mid})
                 finally:
                     await db.close()
 
-            # -------------------- TYPING --------------------
             elif mtype == "typing":
-                is_typing = bool(data.get("is_typing", False))
                 await manager.broadcast({
                     "type": "typing_indicator",
                     "user_id": uid,
                     "username": username,
                     "display_name": display_name,
-                    "is_typing": is_typing
+                    "is_typing": bool(data.get("is_typing", False))
                 }, exclude_user_id=uid)
 
-            # -------------------- MARK READ --------------------
             elif mtype == "mark_read":
                 up_to = data.get("up_to_message_id")
                 if not up_to:
@@ -784,7 +865,6 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 finally:
                     await db.close()
 
-            # -------------------- GET ONLINE USERS --------------------
             elif mtype == "get_online_users":
                 await ws.send_json({
                     "type": "online_users",