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

README.md

@@ -23,7 +23,7 @@ The goal is to keep it simple! For feature-rich solutions please check out [hedg
 **Available**:
 - live collab notepad
 - line numbers
-- custom path `{pad_id}` for more privacy
+- custom path `{pad_id}` for more privacy (1–64 chars, `[a-zA-Z0-9_-]`)
 - optional caching with valkey/redis
 - pad creation with HTTP post requests with curl (see *Usage*)
 - `{pad_id}/raw` HTTP endpoint 
@@ -117,6 +117,7 @@ The following environment variables can be configured:
 | `MAX_TEXT_SIZE` | `5` | Maximum text size in MB (5MB default) |
 | `MAX_CONNECTIONS_PER_IP` | `10` | Maximum concurrent connections per IP address |
 | `RETENTION_HOURS` | `48` | How long to retain pads in hours after last access |
+| `MAX_ROOMS` | `10000` | Maximum number of pads kept in memory; new pads are refused (WS close 1008) when full until cleanup reclaims space |
 | `DESCRIPTION` | `powered by aukpad.com` | Instance description shown on info page |
 
 ---

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,6 +697,11 @@ 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"
 
@@ -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)