sec: FIX: doc_id precheck to prevent app DoS #26

app.py

@@ -1,7 +1,7 @@
 # aukpad.py
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, HTTPException
 from fastapi.responses import HTMLResponse, RedirectResponse, PlainTextResponse, FileResponse
-import json, secrets, string, time, os, threading, asyncio, hashlib
+import json, re, secrets, string, time, os, threading, asyncio, hashlib
 from collections import defaultdict
 from typing import Optional
 
@@ -14,8 +14,14 @@ VALKEY_URL = os.getenv("VALKEY_URL", "redis://localhost:6379/0")
 MAX_TEXT_SIZE = int(os.getenv("MAX_TEXT_SIZE", "5")) * 1024 * 1024  # 5MB default
 MAX_CONNECTIONS_PER_IP = int(os.getenv("MAX_CONNECTIONS_PER_IP", "10"))
 RETENTION_HOURS = int(os.getenv("RETENTION_HOURS", "48"))  # Default 48 hours
+MAX_ROOMS = int(os.getenv("MAX_ROOMS", "10000"))
 DESCRIPTION = os.getenv("DESCRIPTION", "powered by aukpad.com")
 
+DOC_ID_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")
+
+def is_valid_doc_id(doc_id: str) -> bool:
+    return bool(DOC_ID_RE.match(doc_id))
+
 # Valkey/Redis client (initialized later if enabled)
 redis_client = None
 
@@ -632,12 +638,16 @@ async def create_pad_with_content(request: Request):
 
 @app.get("/{doc_id}/", response_class=HTMLResponse)
 def pad(doc_id: str):
+    if not is_valid_doc_id(doc_id):
+        raise HTTPException(status_code=400, detail="Invalid pad ID")
     # Update access time when pad is accessed
     update_room_access_time(doc_id)
     return HTMLResponse(HTML)
 
 @app.get("/{doc_id}/raw", response_class=PlainTextResponse)
 def get_raw_pad_content(doc_id: str, pw: str = ""):
+    if not is_valid_doc_id(doc_id):
+        raise HTTPException(status_code=400, detail="Invalid pad ID")
     # Load room into memory if needed
     if doc_id not in rooms:
         cached_data = get_room_data_from_cache(doc_id)
@@ -687,14 +697,19 @@ async def _broadcast(doc_id: str, message: dict, exclude: WebSocket | None = Non
 
 @app.websocket("/ws/{doc_id}")
 async def ws(doc_id: str, ws: WebSocket):
+    # Validate doc_id format
+    if not is_valid_doc_id(doc_id):
+        await ws.close(code=1008, reason="Invalid pad ID")
+        return
+
     # Get client IP for connection limiting
     client_ip = ws.client.host if ws.client else "unknown"
-    
+
     # Check connection limit per IP
     if connections_per_ip[client_ip] >= MAX_CONNECTIONS_PER_IP:
         await ws.close(code=1008, reason="Too many connections from this IP")
         return
-    
+
     await ws.accept()
     connections_per_ip[client_ip] += 1
 
@@ -712,6 +727,12 @@ async def ws(doc_id: str, ws: WebSocket):
                 "pw_salt": cached_data.get("pw_salt"),
             }
 
+    # Refuse to create new rooms when at capacity
+    if doc_id not in rooms and len(rooms) >= MAX_ROOMS:
+        await ws.close(code=1008, reason="Server at capacity")
+        connections_per_ip[client_ip] = max(0, connections_per_ip[client_ip] - 1)
+        return
+
     room = rooms.setdefault(doc_id, {"text": "", "ver": 0, "peers": set(), "authed_peers": set(),
                                      "last_access": time.time(), "pw_hash": None, "pw_salt": None})
     room["peers"].add(ws)