fix: harden error handling, concurrency, and XML safety

- Escape XML in error responses to prevent injection via error messages
- Add retry limit (MAX_WATCH_RETRIES=5) to Redis WATCH/MULTI/EXEC loops
  to prevent unbounded recursion under high contention
- Refactor concurrency limiter from module-level globals to a
  ConcurrencyLimiter class with proper encapsulation; module-level
  functions delegate to a default instance for backward compatibility
- Clean up metrics import (remove noqa suppression, use real name)

Author: ServerSideHannes

s3proxy/app.py

@@ -8,14 +8,15 @@
 import uuid
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
+from xml.sax.saxutils import escape as xml_escape
 
 import structlog
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.responses import PlainTextResponse
 from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
 from structlog.stdlib import BoundLogger
 
-from . import metrics as _  # noqa: F401 - Import to register metrics
+from . import metrics  # Ensure Prometheus collectors are registered at import time
 from .config import Settings
 from .errors import S3Error, get_s3_error_code
 from .handlers import S3ProxyHandler
@@ -130,8 +131,8 @@ async def s3_exception_handler(request: Request, exc: HTTPException):
 
         error_xml = f"""<?xml version="1.0" encoding="UTF-8"?>
 <Error>
-    <Code>{error_code}</Code>
-    <Message>{message}</Message>
+    <Code>{xml_escape(error_code)}</Code>
+    <Message>{xml_escape(str(message))}</Message>
     <RequestId>{request_id}</RequestId>
 </Error>"""
         return Response(
@@ -161,7 +162,7 @@ async def metrics():
         "/{path:path}",
         methods=["GET", "PUT", "POST", "DELETE", "HEAD"],
     )
-    async def proxy(request: Request, path: str):  # noqa: ARG001
+    async def proxy(request: Request, path: str):  # noqa: ARG001 - required by FastAPI for {path:path}
         return await handle_proxy_request(
             request, request.app.state.handler, request.app.state.verifier
         )

s3proxy/concurrency.py

@@ -20,15 +20,6 @@
 MIN_RESERVATION = 64 * 1024  # 64KB minimum per request
 MAX_BUFFER_SIZE = 8 * 1024 * 1024  # 8MB streaming buffer size
 
-# Module-level state
-_limit_mb = int(os.environ.get("S3PROXY_MEMORY_LIMIT_MB", "64"))
-_limit_bytes = _limit_mb * 1024 * 1024
-_active_bytes = 0
-_lock: asyncio.Lock | None = None
-
-# Initialize memory limit metric
-MEMORY_LIMIT_BYTES.set(_limit_bytes)
-
 
 def _create_malloc_release() -> Callable[[], int] | None:
     """Create platform-specific function to release memory back to OS.
@@ -50,19 +41,89 @@ def _create_malloc_release() -> Callable[[], int] | None:
 _malloc_release = _create_malloc_release()
 
 
-async def _get_lock() -> asyncio.Lock:
-    global _lock
-    if _lock is None:
-        _lock = asyncio.Lock()
-    return _lock
-
-
-def get_memory_limit() -> int:
-    return _limit_bytes
+class ConcurrencyLimiter:
+    """Memory-based concurrency limiter.
 
+    Tracks reserved memory across concurrent requests and rejects new requests
+    when the configured limit would be exceeded.
+    """
 
-def get_active_memory() -> int:
-    return _active_bytes
+    def __init__(self, limit_mb: int = 64) -> None:
+        self._limit_mb = limit_mb
+        self._limit_bytes = limit_mb * 1024 * 1024
+        self._active_bytes = 0
+        self._lock = asyncio.Lock()
+        MEMORY_LIMIT_BYTES.set(self._limit_bytes)
+
+    @property
+    def limit_bytes(self) -> int:
+        return self._limit_bytes
+
+    @property
+    def active_bytes(self) -> int:
+        return self._active_bytes
+
+    @active_bytes.setter
+    def active_bytes(self, value: int) -> None:
+        """Set active memory (testing only)."""
+        self._active_bytes = value
+
+    def set_memory_limit(self, limit_mb: int) -> None:
+        """Update the memory limit."""
+        self._limit_mb = limit_mb
+        self._limit_bytes = limit_mb * 1024 * 1024
+        MEMORY_LIMIT_BYTES.set(self._limit_bytes)
+
+    async def try_acquire(self, bytes_needed: int) -> int:
+        """Reserve memory. Returns bytes reserved. Raises S3Error.slow_down if exhausted."""
+        if self._limit_bytes <= 0:
+            return 0
+
+        to_reserve = max(MIN_RESERVATION, min(bytes_needed, self._limit_bytes))
+
+        async with self._lock:
+            if self._active_bytes + to_reserve > self._limit_bytes:
+                active_mb = self._active_bytes / 1024 / 1024
+                request_mb = to_reserve / 1024 / 1024
+                limit_mb = self._limit_bytes / 1024 / 1024
+                logger.warning("MEMORY_REJECTED", active_mb=round(active_mb, 2),
+                             requested_mb=round(request_mb, 2), limit_mb=round(limit_mb, 2))
+                MEMORY_REJECTIONS.inc()
+                raise S3Error.slow_down(
+                    f"Memory limit: {active_mb:.0f}MB + {request_mb:.0f}MB > {limit_mb:.0f}MB"
+                )
+            self._active_bytes += to_reserve
+            MEMORY_RESERVED_BYTES.set(self._active_bytes)
+            return to_reserve
+
+    async def release(self, bytes_reserved: int) -> None:
+        """Release reserved memory and trigger OS memory release."""
+        if self._limit_bytes <= 0 or bytes_reserved <= 0:
+            return
+
+        async with self._lock:
+            self._active_bytes = max(0, self._active_bytes - bytes_reserved)
+            MEMORY_RESERVED_BYTES.set(self._active_bytes)
+
+        # Run garbage collection and release memory to OS
+        gc.collect(0)
+        gc.collect(1)
+        gc.collect(2)
+
+        if _malloc_release:
+            try:
+                _malloc_release()
+            except OSError:
+                pass
+
+        # Yield to allow OS memory reclaim
+        await asyncio.sleep(0)
+
+
+# Default instance used by module-level functions
+_default = ConcurrencyLimiter(
+    limit_mb=int(os.environ.get("S3PROXY_MEMORY_LIMIT_MB", "64"))
+)
 
 
 def estimate_memory_footprint(method: str, content_length: int) -> int:
@@ -78,75 +139,38 @@ def estimate_memory_footprint(method: str, content_length: int) -> int:
     return MAX_BUFFER_SIZE
 
 
-async def try_acquire_memory(bytes_needed: int) -> int:
-    """Reserve memory. Returns bytes reserved. Raises S3Error.slow_down if exhausted."""
-    global _active_bytes
+# Module-level convenience functions delegating to the default instance
 
-    if _limit_bytes <= 0:
-        return 0
-
-    to_reserve = max(MIN_RESERVATION, min(bytes_needed, _limit_bytes))
-
-    lock = await _get_lock()
-    async with lock:
-        if _active_bytes + to_reserve > _limit_bytes:
-            active_mb = _active_bytes / 1024 / 1024
-            request_mb = to_reserve / 1024 / 1024
-            limit_mb = _limit_bytes / 1024 / 1024
-            logger.warning("MEMORY_REJECTED", active_mb=round(active_mb, 2),
-                         requested_mb=round(request_mb, 2), limit_mb=round(limit_mb, 2))
-            MEMORY_REJECTIONS.inc()
-            raise S3Error.slow_down(
-                f"Memory limit: {active_mb:.0f}MB + {request_mb:.0f}MB > {limit_mb:.0f}MB"
-            )
-        _active_bytes += to_reserve
-        MEMORY_RESERVED_BYTES.set(_active_bytes)
-        return to_reserve
 
+def get_memory_limit() -> int:
+    return _default.limit_bytes
 
-async def release_memory(bytes_reserved: int) -> None:
-    """Release reserved memory and trigger OS memory release."""
-    global _active_bytes
 
-    if _limit_bytes <= 0 or bytes_reserved <= 0:
-        return
+def get_active_memory() -> int:
+    return _default.active_bytes
 
-    lock = await _get_lock()
-    async with lock:
-        _active_bytes = max(0, _active_bytes - bytes_reserved)
-        MEMORY_RESERVED_BYTES.set(_active_bytes)
 
-    # Run garbage collection and release memory to OS
-    gc.collect(0)
-    gc.collect(1)
-    gc.collect(2)
+async def try_acquire_memory(bytes_needed: int) -> int:
+    return await _default.try_acquire(bytes_needed)
 
-    if _malloc_release:
-        try:
-            _malloc_release()
-        except OSError:
-            pass
 
-    # Yield to allow OS memory reclaim
-    await asyncio.sleep(0)
+async def release_memory(bytes_reserved: int) -> None:
+    await _default.release(bytes_reserved)
 
 
 def reset_state() -> None:
-    """Reset state (testing only)."""
-    global _active_bytes, _lock
-    _active_bytes = 0
-    _lock = None
+    """Reset default instance state (testing only)."""
+    global _default
+    _default = ConcurrencyLimiter(limit_mb=_default._limit_mb)
+    # Reset reserved bytes metric to 0 for clean test state
+    MEMORY_RESERVED_BYTES.set(0)
 
 
 def set_memory_limit(limit_mb: int) -> None:
-    """Set memory limit (testing only)."""
-    global _limit_mb, _limit_bytes
-    _limit_mb = limit_mb
-    _limit_bytes = limit_mb * 1024 * 1024
-    MEMORY_LIMIT_BYTES.set(_limit_bytes)
+    """Set memory limit on default instance (testing only)."""
+    _default.set_memory_limit(limit_mb)
 
 
 def set_active_memory(bytes_val: int) -> None:
-    """Set active memory (testing only)."""
-    global _active_bytes
-    _active_bytes = bytes_val
+    """Set active memory on default instance (testing only)."""
+    _default.active_bytes = bytes_val

s3proxy/state/storage.py

@@ -21,6 +21,9 @@
 # Type alias for updater function: takes bytes, returns bytes
 Updater = Callable[[bytes], bytes]
 
+# Maximum retries for Redis optimistic locking (WATCH/MULTI/EXEC)
+MAX_WATCH_RETRIES = 5
+
 
 class StateStore(ABC):
     """Abstract interface for state storage backends."""
@@ -118,7 +121,7 @@ async def set(self, key: str, value: bytes, ttl_seconds: int) -> None:
     async def delete(self, key: str) -> None:
         await self._client.delete(self._key(key))
 
-    async def get_and_delete(self, key: str) -> bytes | None:
+    async def get_and_delete(self, key: str, _retries: int = 0) -> bytes | None:
         """Atomically get and delete using Redis transaction."""
         import redis.asyncio as redis
 
@@ -137,11 +140,14 @@ async def get_and_delete(self, key: str) -> bytes | None:
                 return data
 
             except redis.WatchError:
-                # Retry on conflict
-                return await self.get_and_delete(key)
+                if _retries >= MAX_WATCH_RETRIES:
+                    logger.error("REDIS_WATCH_RETRIES_EXHAUSTED", key=key, operation="get_and_delete")
+                    raise
+                logger.debug("REDIS_WATCH_RETRY", key=key, attempt=_retries + 1)
+                return await self.get_and_delete(key, _retries=_retries + 1)
 
     async def update(
-        self, key: str, updater: Updater, ttl_seconds: int
+        self, key: str, updater: Updater, ttl_seconds: int, _retries: int = 0
     ) -> bytes | None:
         """Atomically update using Redis WATCH/MULTI/EXEC."""
         import redis.asyncio as redis
@@ -163,5 +169,8 @@ async def update(
                 return new_data
 
             except redis.WatchError:
-                logger.debug("REDIS_WATCH_RETRY", key=key)
-                return await self.update(key, updater, ttl_seconds)
+                if _retries >= MAX_WATCH_RETRIES:
+                    logger.error("REDIS_WATCH_RETRIES_EXHAUSTED", key=key, operation="update")
+                    raise
+                logger.debug("REDIS_WATCH_RETRY", key=key, attempt=_retries + 1)
+                return await self.update(key, updater, ttl_seconds, _retries=_retries + 1)