diff --git a/cmd_chat/client/client.py b/cmd_chat/client/client.py index 42870b8..d2ce520 100644 --- a/cmd_chat/client/client.py +++ b/cmd_chat/client/client.py @@ -319,10 +319,21 @@ class Client: elif ft_type == "chunk": if transfer_id in self.received_chunks: chunk_data = base64.b64decode(ft_data.get("data", "")) - self.received_chunks[transfer_id].append(chunk_data) meta = self.transfer_meta.get(transfer_id, {}) - total = meta.get("size", 0) + total = meta.get("size", 0) or 0 + # Enforce the declared size (clamped to MAX_FILE_SIZE) on receipt: + # a sender can lie about `size` or just keep streaming, so abort + # the moment the accumulated bytes would exceed the cap. + cap = min(total, MAX_FILE_SIZE) if total else MAX_FILE_SIZE received = sum(len(c) for c in self.received_chunks[transfer_id]) + if received + len(chunk_data) > cap: + self.received_chunks.pop(transfer_id, None) + self.transfer_meta.pop(transfer_id, None) + self.console.print() + self.error("Transfer exceeds declared size — aborted.") + return True + self.received_chunks[transfer_id].append(chunk_data) + received += len(chunk_data) pct = int(received * 100 / total) if total else 0 self.console.print( f"\r[cyan]Receiving: {pct}% ({_human_size(received)}/{_human_size(total)})[/]", diff --git a/cmd_chat/server/helpers.py b/cmd_chat/server/helpers.py index bfd6b7f..46394d1 100644 --- a/cmd_chat/server/helpers.py +++ b/cmd_chat/server/helpers.py @@ -1,3 +1,4 @@ +import os import time from collections import defaultdict from datetime import datetime, timezone @@ -10,9 +11,16 @@ def utcnow() -> datetime: return datetime.now(timezone.utc) +# Only honor X-Forwarded-For when explicitly told we sit behind a trusted proxy +# (TRUST_PROXY=1). Otherwise a direct client can spoof the header to forge a +# source IP and dodge the per-IP rate limiter, so we use the real peer address. +_TRUST_PROXY = os.environ.get("TRUST_PROXY", "").lower() in ("1", "true", "yes") + + def get_client_ip(request: Request) -> str: - if forwarded := request.headers.get("x-forwarded-for"): - return forwarded.split(",")[0].strip() + if _TRUST_PROXY: + if forwarded := request.headers.get("x-forwarded-for"): + return forwarded.split(",")[0].strip() return request.ip diff --git a/cmd_chat/server/srp_auth.py b/cmd_chat/server/srp_auth.py index 98ccb03..d47915d 100644 --- a/cmd_chat/server/srp_auth.py +++ b/cmd_chat/server/srp_auth.py @@ -1,3 +1,4 @@ +import time from dataclasses import dataclass, field from typing import Optional from uuid import uuid4 @@ -7,6 +8,11 @@ import srp srp.rfc5054_enable() +# Half-finished handshakes (init without a matching verify) would otherwise pile +# up forever, letting an attacker exhaust memory. Evict any session that hasn't +# authenticated within this many seconds whenever a new handshake begins. +UNVERIFIED_TTL_SECONDS = 60 + @dataclass class SRPSession: @@ -15,6 +21,7 @@ class SRPSession: svr: Optional[srp.Verifier] = None session_key: Optional[bytes] = None authenticated: bool = False + created_at: float = field(default_factory=time.monotonic) class SRPAuthManager: @@ -25,9 +32,20 @@ class SRPAuthManager: b"chat", self.password, hash_alg=srp.SHA256 ) + def _evict_stale_unverified(self) -> None: + now = time.monotonic() + stale = [ + uid + for uid, s in self.sessions.items() + if not s.authenticated and now - s.created_at > UNVERIFIED_TTL_SECONDS + ] + for uid in stale: + del self.sessions[uid] + def init_auth( self, username: str, client_public: bytes ) -> tuple[str, bytes, bytes]: + self._evict_stale_unverified() session = SRPSession(username=username) svr = srp.Verifier( diff --git a/cmd_chat/server/views.py b/cmd_chat/server/views.py index df99c0e..c483d51 100644 --- a/cmd_chat/server/views.py +++ b/cmd_chat/server/views.py @@ -11,6 +11,13 @@ from .models import Message, UserSession from .helpers import get_client_ip, send_state, utcnow +# Hard cap on a single relayed WS frame. The largest legitimate frame is one +# Fernet-encrypted 64 KB file chunk (~120 KB after base64 + token overhead), so +# 256 KB leaves headroom while bounding per-message memory and the 1000-message +# store. Oversized frames are dropped, not stored or broadcast. +MAX_FRAME_SIZE = 256 * 1024 + + def generate_ws_token(user_id: str, secret: bytes) -> str: return hmac.new(secret, user_id.encode(), hashlib.sha256).hexdigest() @@ -152,10 +159,16 @@ async def chat_ws(request: Request, ws: Websocket, app: Sanic) -> None: if data is None: break + text = str(data) + # Drop oversized frames before they reach the store/broadcast: this + # bounds memory and stops a single client from flooding the room. + if len(text) > MAX_FRAME_SIZE: + continue + app.ctx.session_store.update_activity(user_id) message = Message( - text=str(data), + text=text, username=session.username, ) app.ctx.message_store.add(message) diff --git a/hh/src/app.rs b/hh/src/app.rs index 12b804c..5c5c02a 100644 --- a/hh/src/app.rs +++ b/hh/src/app.rs @@ -573,9 +573,27 @@ fn handle_ft( } } ft::Ft::Chunk { id, data } => { + // Enforce the declared size (clamped to MAX_SIZE) on receipt — a + // malicious sender can lie about `size` or just keep streaming, so + // never let an accepted transfer grow its buffer past the cap. + let mut overflow = false; if let Some(t) = app.transfers.get_mut(&id) { if t.accepted { - t.buf.extend_from_slice(&data); + let cap = (t.meta.size as usize).min(ft::MAX_SIZE); + if t.buf.len() + data.len() > cap { + overflow = true; + } else { + t.buf.extend_from_slice(&data); + } + } + } + if overflow { + if let Some(t) = app.transfers.remove(&id) { + app.err(format!( + "{} — transfer exceeds declared size (max {}), aborted", + t.meta.name, + ft::human((t.meta.size as usize).min(ft::MAX_SIZE)) + )); } } }