From c113d28dd8e4ae266e04c549edcf0540b172fcb7 Mon Sep 17 00:00:00 2001 From: HC-ONLINE Date: Sat, 17 Oct 2026 11:23:02 -0500 Subject: [PATCH] fix(download): bound in-memory download registry and gate tokens when disabled --- .env.example | 2 + docs/API.md | 13 +- docs/CONFIGURATION.md | 39 ++- .../api/helpers/download_helper.py | 10 + src/ciberwebscan/config/models.py | 4 + src/ciberwebscan/services/download_service.py | 234 +++++++++++++-- .../api/test_api_export_download.py | 5 + .../integration/api/test_cleanup_scheduler.py | 3 + .../integration/api/test_download_endpoint.py | 2 + .../unit/api/helpers/test_download_helper.py | 30 +- tests/unit/api/routes/test_scrape_routes.py | 46 +++ tests/unit/services/test_download_service.py | 271 +++++++++++++++++- .../services/test_download_token_deletion.py | 12 +- 13 files changed, 621 insertions(+), 50 deletions(-) diff --git a/.env.example b/.env.example index 9de2fd4..f1c3489 100644 --- a/.env.example +++ b/.env.example @@ -52,6 +52,8 @@ CIBERWEBSCAN_CACHE_ENABLED=true # -> cache.enabled # API auth / downloads CIBERWEBSCAN_API_AUTH_SERVER_SECRET= # -> api.auth.server_secret (HMAC for logs only) CIBERWEBSCAN_DOWNLOAD_REQUIRE_SAME_USER=true # -> download.require_same_user (token bound to issuing key) +CIBERWEBSCAN_DOWNLOAD_MAX_TOKENS=100 # -> download.max_tokens (max live tokens in the in-memory registry) +CIBERWEBSCAN_DOWNLOAD_MAX_TOTAL_BYTES=1073741824 # -> download.max_total_bytes (byte budget for cached exports) # Reverse proxies (client identification). Empty = X-Forwarded-For always ignored. # Declare proxy endpoints only; the proxy must overwrite or append (rightmost) diff --git a/docs/API.md b/docs/API.md index 2881ced..dd233bc 100644 --- a/docs/API.md +++ b/docs/API.md @@ -648,6 +648,13 @@ that issued it. The `401` case above (`token belongs to different user`) is cont share tokens across API keys, because `false` disables that access control entirely. See [CONFIGURATION.md — Download](CONFIGURATION.md#download). +**When no token is issued:** export-enabled POST responses return `200` with +`download_token: null` (and no `download_url`) while the export itself stays available on +disk. This happens when downloads are disabled (`download.enabled=false`) or when the +registry budget is exhausted (`download.max_tokens` / `download.max_total_bytes` — see +[CONFIGURATION.md — Download](CONFIGURATION.md#download)). A `GET /api/download/{token}` +against a disabled service returns `503`. + **Example:** ```bash @@ -1041,9 +1048,11 @@ CIBERWEBSCAN_API_CORS_EXPOSE_HEADERS="X-Custom-Header" # Downloads CIBERWEBSCAN_DOWNLOAD_ENABLED=true -CIBERWEBSCAN_DOWNLOAD_TOKEN_EXPIRY_MINUTES=60 -CIBERWEBSCAN_DOWNLOAD_MAX_RETRY_ATTEMPTS=3 +CIBERWEBSCAN_DOWNLOAD_RETENTION_SECONDS=1800 +CIBERWEBSCAN_DOWNLOAD_MAX_RETRIES=3 CIBERWEBSCAN_DOWNLOAD_REQUIRE_SAME_USER=true +CIBERWEBSCAN_DOWNLOAD_MAX_TOKENS=100 +CIBERWEBSCAN_DOWNLOAD_MAX_TOTAL_BYTES=1073741824 ``` ### Rate Limiting diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index bc65a84..9317d73 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -481,7 +481,9 @@ Configure the download tokens returned by export-enabled API requests (`download "download": { "enabled": true, "retention_seconds": 1800, - "require_same_user": true + "require_same_user": true, + "max_tokens": 100, + "max_total_bytes": 1073741824 } } ``` @@ -505,17 +507,34 @@ The value is read from the configuration on every download request, so it applie The returned-attempt refund is capped: it can never push a token above `download.max_retries`, and it never recreates a deleted token. If you lower `max_retries` at runtime, refunds on tokens created under the previous (higher) value may be refused, leaving their counter unchanged — attempts are never inflated beyond the current configuration. +#### Registry budget: `download.max_tokens` / `download.max_total_bytes` + +Each download token keeps a full in-memory copy of its export file, so total memory grows with every live token. Two budget keys bound that: + +- **`download.max_tokens`** (default `100`, range 1-10000) — maximum number of live tokens in the registry. +- **`download.max_total_bytes`** (default `1073741824` = 1 GiB, range 1 MiB-1024 GiB) — maximum total bytes held by all live tokens. + +Admission is atomic and happens **before** the export file is read: a request whose token would exceed either limit is rejected without loading the file into memory. When a budget is exhausted: + +- The export itself still completes and is written to disk normally. +- The response returns `200` with `download_token: null` (no token, no `download_url`). +- Budget is freed when a token is downloaded (deleted), expires (`retention_seconds`), or is removed by the cleanup job — there is no eviction of live tokens. + +Raise these values only when the host has the RAM for them: `max_total_bytes` is a hard cap on cached-export memory across the whole process. The rejection is logged as a warning identifying the exhausted limit. + #### Default values (quick reference) -| Key | Default | Description | -| ----------------------------------- | --------: | ------------------------------------------------------- | -| `download.enabled` | `true` | Download service enabled (tokens + `GET /api/download`) | -| `download.retention_seconds` | `1800` | Seconds before a download token expires (60-86400) | -| `download.max_file_size_mb` | `500` | Max export file size accepted for download (MB) | -| `download.max_retries` | `3` | Max download attempts per token (1-10) | -| `download.cleanup_interval_seconds` | `300` | Cleanup scheduler interval (requires restart) | -| `download.require_same_user` | `true` | Bind each download token to the API key that issued it | -| `download.stream_chunk_size` | `1048576` | Streaming chunk size in bytes | +| Key | Default | Description | +| ----------------------------------- | -----------: | ------------------------------------------------------- | +| `download.enabled` | `true` | Download service enabled (tokens + `GET /api/download`) | +| `download.retention_seconds` | `1800` | Seconds before a download token expires (60-86400) | +| `download.max_file_size_mb` | `500` | Max export file size accepted for download (MB) | +| `download.max_retries` | `3` | Max download attempts per token (1-10) | +| `download.cleanup_interval_seconds` | `300` | Cleanup scheduler interval (requires restart) | +| `download.require_same_user` | `true` | Bind each download token to the API key that issued it | +| `download.stream_chunk_size` | `1048576` | Streaming chunk size in bytes | +| `download.max_tokens` | `100` | Max live download tokens held in memory (1-10000) | +| `download.max_total_bytes` | `1073741824` | Byte budget across all live tokens (1 MiB-1024 GiB) | ### Cache diff --git a/src/ciberwebscan/api/helpers/download_helper.py b/src/ciberwebscan/api/helpers/download_helper.py index cf17ad4..9bf75b3 100644 --- a/src/ciberwebscan/api/helpers/download_helper.py +++ b/src/ciberwebscan/api/helpers/download_helper.py @@ -13,6 +13,7 @@ from fastapi import Request +from ciberwebscan.config.loader import get_config from ciberwebscan.export import EXPORT_FORMATS from ciberwebscan.services.base import ServiceResult from ciberwebscan.services.download_service import DownloadService @@ -70,6 +71,15 @@ def enrich_response_with_token( if result.export_path is None: return result.data, None, None + # Downloads disabled is an expected state, not a failure: skip the + # service call so the (already gated) export does not emit a warning + # per request. The service remains the authoritative admission gate. + if not get_config().download.enabled: + download_service.logger.debug( + "Download token skipped: downloads are disabled (download.enabled=false)" + ) + return result.data, None, None + # Generate token from the exported file with the real export format # (previously hardcoded to "json", which mislabeled every download). token_result = download_service.generate_download_token( diff --git a/src/ciberwebscan/config/models.py b/src/ciberwebscan/config/models.py index ebb21eb..a22427a 100644 --- a/src/ciberwebscan/config/models.py +++ b/src/ciberwebscan/config/models.py @@ -327,6 +327,10 @@ class DownloadConfig(BaseModel): cleanup_interval_seconds: Annotated[int, Field(ge=60, le=3600)] = 300 require_same_user: bool = True stream_chunk_size: Annotated[int, Field(ge=1024, le=10 * 1024 * 1024)] = 1024 * 1024 + max_tokens: Annotated[int, Field(ge=1, le=10000)] = 100 + max_total_bytes: Annotated[int, Field(ge=1_048_576, le=1_099_511_627_776)] = ( + 1_073_741_824 + ) # ============================================================================= diff --git a/src/ciberwebscan/services/download_service.py b/src/ciberwebscan/services/download_service.py index d774792..125793a 100644 --- a/src/ciberwebscan/services/download_service.py +++ b/src/ciberwebscan/services/download_service.py @@ -25,12 +25,112 @@ class _DownloadRegistry: - """In-memory registry for download tokens and file data.""" + """In-memory registry for download tokens and file data. + + The registry keeps a full in-memory copy of every exported file, so it + enforces a budget (``download.max_tokens`` / ``download.max_total_bytes``) + to bound total process memory. Admission uses an atomic reserve-commit + cycle: the byte reservation is taken *before* the file is read, so + concurrent admissions can never oversubscribe the budget. + """ def __init__(self) -> None: self._tokens: dict[str, DownloadInfo] = {} self._file_data: dict[str, bytes] = {} self._lock = threading.Lock() + # Budget accounting (guarded by _lock): + # - _total_bytes: bytes held by committed tokens. + # - _reserved_bytes / _pending: bytes promised to in-flight + # admissions that have reserved budget but not yet read their file. + self._total_bytes = 0 + self._reserved_bytes = 0 + self._pending = 0 + + @property + def stats(self) -> dict[str, int]: + """Snapshot of registry occupancy (tokens, bytes, in-flight work).""" + with self._lock: + return { + "tokens": len(self._tokens), + "total_bytes": self._total_bytes, + "reserved_bytes": self._reserved_bytes, + "pending": self._pending, + } + + async def clear(self) -> None: + """Drop every token and reset budget counters. + + Test-fixture support: production cleanup happens via + ``cleanup_expired`` and ``delete``. Guards below clamp counters + at zero, so a late release from an in-flight reservation is safe. + """ + with self._lock: + self._tokens.clear() + self._file_data.clear() + self._total_bytes = 0 + self._reserved_bytes = 0 + self._pending = 0 + + async def reserve( + self, + size: int, + max_tokens: int, + max_total_bytes: int, + ) -> bool: + """Atomically reserve admission budget for a file of ``size`` bytes. + + Returns False when either budget dimension is exhausted. The + reservation must be paired with ``commit`` (on success) or + ``release_reservation`` (on any failure path). + """ + with self._lock: + if len(self._tokens) + self._pending >= max_tokens: + return False + if self._total_bytes + self._reserved_bytes + size > max_total_bytes: + return False + self._pending += 1 + self._reserved_bytes += size + return True + + async def release_reservation(self, size: int) -> None: + """Release an outstanding reservation after a failed admission.""" + with self._lock: + self._pending = max(0, self._pending - 1) + self._reserved_bytes = max(0, self._reserved_bytes - size) + + async def commit( + self, + token: str, + info: DownloadInfo, + file_data: bytes, + reserved_size: int, + max_total_bytes: int, + ) -> bool: + """Convert a reservation into a committed token entry. + + Re-checks the byte budget under the lock because the file may have + grown between the ``stat()`` reservation and ``read_bytes()``; other + concurrent commits may also have consumed the reserved slack. + Returns False (caller must ``release_reservation``) when the budget + no longer fits. + """ + with self._lock: + actual = len(file_data) + if self._total_bytes + actual > max_total_bytes: + return False + self._pending = max(0, self._pending - 1) + self._reserved_bytes = max(0, self._reserved_bytes - reserved_size) + self._insert_locked(token, info, file_data) + return True + + def _insert_locked(self, token: str, info: DownloadInfo, file_data: bytes) -> None: + """Insert or overwrite a token, keeping ``_total_bytes`` exact.""" + previous = self._file_data.get(token) + if previous is not None: + self._total_bytes -= len(previous) + self._tokens[token] = info + self._file_data[token] = file_data + self._total_bytes += len(file_data) async def store( self, @@ -38,10 +138,9 @@ async def store( info: DownloadInfo, file_data: bytes, ) -> None: - """Store token with metadata and file data.""" + """Store token with metadata and file data (no budget enforcement).""" with self._lock: - self._tokens[token] = info - self._file_data[token] = file_data + self._insert_locked(token, info, file_data) async def get_info(self, token: str) -> DownloadInfo | None: """Retrieve token metadata.""" @@ -100,7 +199,9 @@ async def delete(self, token: str) -> bool: if token not in self._tokens: return False del self._tokens[token] - self._file_data.pop(token, None) + file_data = self._file_data.pop(token, None) + if file_data is not None: + self._total_bytes = max(0, self._total_bytes - len(file_data)) return True async def cleanup_expired(self) -> int: @@ -112,7 +213,9 @@ async def cleanup_expired(self) -> int: ] for token in expired_tokens: del self._tokens[token] - self._file_data.pop(token, None) + file_data = self._file_data.pop(token, None) + if file_data is not None: + self._total_bytes = max(0, self._total_bytes - len(file_data)) return len(expired_tokens) async def get_expired_count(self) -> int: @@ -149,6 +252,18 @@ def generate_download_token( try: config = get_config() + # Single admission gate for every consumer (helper, scrape/batch, + # direct service use): no tokens may exist while downloads are off. + if not config.download.enabled: + self.logger.info( + "Download token not generated: downloads are disabled " + "(download.enabled=false)" + ) + return ServiceResult( + success=False, + error="Download service is disabled", + ) + # Validate path to prevent traversal attacks output_base = Path(config.export.output_dir) if not output_base.is_absolute(): @@ -174,35 +289,102 @@ def generate_download_token( file_path = validated_path - # Read file data - file_data = file_path.read_bytes() - file_size_mb = len(file_data) / (1024 * 1024) + # Size check on metadata first: never load a file that is already + # known to exceed the per-file limit (read_bytes() would copy the + # whole file into RAM before the check could reject it). + try: + size_bytes = file_path.stat().st_size + except OSError as e: + self.logger.error(f"Unable to stat export file: {e}") + return ServiceResult( + success=False, + error=f"Unable to read file size: {file_path}", + ) - # Validate size + file_size_mb = size_bytes / (1024 * 1024) if file_size_mb > config.download.max_file_size_mb: return ServiceResult( success=False, error=f"File size {file_size_mb:.2f}MB exceeds limit of {config.download.max_file_size_mb}MB", ) - # Generate token - token = str(uuid.uuid4()) - now = datetime.now(timezone.utc) - expires_at = now + timedelta(seconds=config.download.retention_seconds) - - # Create metadata - info = DownloadInfo( - token=token, - user_id=user_id, - file_size_bytes=len(file_data), - created_at=now, - expires_at=expires_at, - attempts_remaining=config.download.max_retries, - file_format=file_format, + # Atomic budget reservation BEFORE reading: concurrent admissions + # (up to the server threadpool width) can never jointly hold more + # than the configured memory budget. + reserved = run_async( + _registry.reserve( + size_bytes, + config.download.max_tokens, + config.download.max_total_bytes, + ) ) + if not reserved: + self.logger.warning( + f"Download token not created for user " + f"{mask_identifier(user_id)}: registry budget exhausted " + f"(max_tokens={config.download.max_tokens}, " + f"max_total_bytes={config.download.max_total_bytes})" + ) + return ServiceResult( + success=False, + error="Download storage quota exceeded: token not created", + ) - # Store token (async operation) - run_async(_registry.store(token, info, file_data)) + committed = False + try: + # Read file data (reservation covers this allocation) + file_data = file_path.read_bytes() + file_size_mb = len(file_data) / (1024 * 1024) + + # Defensive re-check: the file may have grown between stat() + # and read_bytes() (TOCTOU); commit re-checks the budget too. + if file_size_mb > config.download.max_file_size_mb: + return ServiceResult( + success=False, + error=f"File size {file_size_mb:.2f}MB exceeds limit of {config.download.max_file_size_mb}MB", + ) + + # Generate token + token = str(uuid.uuid4()) + now = datetime.now(timezone.utc) + expires_at = now + timedelta(seconds=config.download.retention_seconds) + + # Create metadata + info = DownloadInfo( + token=token, + user_id=user_id, + file_size_bytes=len(file_data), + created_at=now, + expires_at=expires_at, + attempts_remaining=config.download.max_retries, + file_format=file_format, + ) + + # Commit the reservation (re-checks budget under the lock) + committed = run_async( + _registry.commit( + token, + info, + file_data, + size_bytes, + config.download.max_total_bytes, + ) + ) + if not committed: + self.logger.warning( + f"Download token not created for user " + f"{mask_identifier(user_id)}: file grew beyond the " + f"reserved budget" + ) + return ServiceResult( + success=False, + error="Download storage quota exceeded: token not created", + ) + finally: + if not committed: + # Failure path (size re-check, read error, commit refusal): + # hand the reservation back so the budget never leaks. + run_async(_registry.release_reservation(size_bytes)) response = DownloadTokenResponse( token=token, diff --git a/tests/integration/api/test_api_export_download.py b/tests/integration/api/test_api_export_download.py index 7c5a9eb..2b4eee7 100644 --- a/tests/integration/api/test_api_export_download.py +++ b/tests/integration/api/test_api_export_download.py @@ -79,6 +79,8 @@ async def _override_user() -> AuthenticatedUser: mock_dl_cfg.return_value.download.max_retries = 3 mock_dl_cfg.return_value.download.stream_chunk_size = 1024 * 1024 mock_dl_cfg.return_value.download.enabled = True + mock_dl_cfg.return_value.download.max_tokens = 100 + mock_dl_cfg.return_value.download.max_total_bytes = 1_073_741_824 with TestClient(app) as client: yield client, mock_dl_cfg @@ -312,6 +314,9 @@ def test_token_generate_validate_delete(self, tmp_path: Path): mock_cfg.return_value.export.output_dir = str(tmp_path) mock_cfg.return_value.download.max_file_size_mb = 10 mock_cfg.return_value.download.retention_seconds = 3600 + mock_cfg.return_value.download.enabled = True + mock_cfg.return_value.download.max_tokens = 100 + mock_cfg.return_value.download.max_total_bytes = 1_073_741_824 service = DownloadService() generate_result = service.generate_download_token( diff --git a/tests/integration/api/test_cleanup_scheduler.py b/tests/integration/api/test_cleanup_scheduler.py index 72c7978..1cc3e57 100644 --- a/tests/integration/api/test_cleanup_scheduler.py +++ b/tests/integration/api/test_cleanup_scheduler.py @@ -69,6 +69,9 @@ def test_cleanup_service_integration(self, test_file: Path, tmp_path: Path): mock_cfg.return_value.export.output_dir = str(tmp_path) mock_cfg.return_value.download.max_file_size_mb = 10 mock_cfg.return_value.download.retention_seconds = 3600 + mock_cfg.return_value.download.enabled = True + mock_cfg.return_value.download.max_tokens = 100 + mock_cfg.return_value.download.max_total_bytes = 1_073_741_824 service = DownloadService() # Generate token diff --git a/tests/integration/api/test_download_endpoint.py b/tests/integration/api/test_download_endpoint.py index 1fb0759..44c1d20 100644 --- a/tests/integration/api/test_download_endpoint.py +++ b/tests/integration/api/test_download_endpoint.py @@ -56,6 +56,8 @@ async def _override_user() -> AuthenticatedUser: mock_cfg.return_value.download.max_retries = 3 mock_cfg.return_value.download.stream_chunk_size = 1024 * 1024 mock_cfg.return_value.download.enabled = True + mock_cfg.return_value.download.max_tokens = 100 + mock_cfg.return_value.download.max_total_bytes = 1_073_741_824 yield TestClient(app) diff --git a/tests/unit/api/helpers/test_download_helper.py b/tests/unit/api/helpers/test_download_helper.py index 2e200a4..ff96b10 100644 --- a/tests/unit/api/helpers/test_download_helper.py +++ b/tests/unit/api/helpers/test_download_helper.py @@ -11,7 +11,7 @@ from datetime import datetime, timezone from pathlib import Path -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest @@ -20,6 +20,7 @@ enrich_response_with_token, ) from ciberwebscan.api.models.responses import DownloadTokenResponse +from ciberwebscan.config.models import AppConfig from ciberwebscan.services.base import ServiceResult pytestmark = pytest.mark.unit @@ -118,3 +119,30 @@ def test_no_export_path_returns_no_token(self): assert token is None assert url is None service.generate_download_token.assert_not_called() + + def test_disabled_downloads_skip_token_generation(self): + """download.enabled=false short-circuits without warning per export. + + The service is the authoritative gate; the helper skips the call so + an expected "disabled" state does not emit a warning on every export. + """ + result: ServiceResult[dict] = ServiceResult.ok( + data={"ok": True}, export_path=Path("exports/report.csv") + ) + + service = self._make_service() + config = AppConfig() + config.download.enabled = False + + with patch( + "ciberwebscan.api.helpers.download_helper.get_config", + return_value=config, + ): + data, token, url = enrich_response_with_token( + result, "user-1", service, self._make_request() + ) + + assert data == {"ok": True} + assert token is None + assert url is None + service.generate_download_token.assert_not_called() diff --git a/tests/unit/api/routes/test_scrape_routes.py b/tests/unit/api/routes/test_scrape_routes.py index b287826..9c11041 100644 --- a/tests/unit/api/routes/test_scrape_routes.py +++ b/tests/unit/api/routes/test_scrape_routes.py @@ -4,12 +4,14 @@ from __future__ import annotations +from pathlib import Path from unittest.mock import MagicMock, Mock, patch import pytest from fastapi.testclient import TestClient from ciberwebscan.api.app import create_app +from ciberwebscan.config.models import AppConfig from ciberwebscan.export.models import ScrapeResult pytestmark = pytest.mark.unit @@ -161,6 +163,50 @@ def test_post_scrape_batch_empty_urls_returns_422(self, client: TestClient): response = client.post("/api/scrape/batch", json={"urls": []}) assert response.status_code == 422 + def test_post_scrape_batch_export_no_token_when_downloads_disabled( + self, + client: TestClient, + mock_scrape_result: ScrapeResult, + tmp_path: Path, + ): + """BUG-011: batch exports never mint tokens while downloads are off. + + The batch route bypasses enrich_response_with_token, so the gate must + live in the service: the response stays 200 with download_token=None + and the export file itself is untouched. + """ + export_file = tmp_path / "batch_export.json" + export_file.write_text('{"batch": true}') + + config = AppConfig() + config.download.enabled = False + config.export.output_dir = str(tmp_path) + + with ( + patch("ciberwebscan.api.routes.scrape.ScrapeService") as mock_service_cls, + patch( + "ciberwebscan.services.download_service.get_config", + return_value=config, + ), + ): + mock_service = MagicMock() + mock_service_cls.return_value = mock_service + scrape_result = _make_service_result(data=[mock_scrape_result]) + scrape_result.export_path = export_file + mock_service.scrape_multiple.return_value = scrape_result + + response = client.post( + "/api/scrape/batch", + json={"urls": ["https://example.com"]}, + ) + + assert response.status_code == 200 + body = response.json() + assert body["success"] is True + assert body["download_token"] is None + assert body["download_url"] is None + assert export_file.exists() + def _mock_http_response( status: int, text: str = "ok" diff --git a/tests/unit/services/test_download_service.py b/tests/unit/services/test_download_service.py index d07daf5..5f6ffe0 100644 --- a/tests/unit/services/test_download_service.py +++ b/tests/unit/services/test_download_service.py @@ -8,6 +8,7 @@ import concurrent.futures import logging import tempfile +from contextlib import contextmanager from datetime import datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch @@ -41,16 +42,18 @@ def test_file(tmp_path: Path) -> Path: @pytest.fixture(autouse=True) def clear_registry(): - """Purge expired tokens before and after each test. - - This does NOT empty the registry: fresh tokens (default TTL 1800s) - created by a test survive into later tests. Tests must therefore only - assert on tokens they created themselves (unique UUIDs), never on the - registry being globally empty. + """Purge all tokens before and after each test. + + The registry now enforces ``download.max_tokens`` (default 100), so + fresh tokens (default TTL 1800s) surviving from earlier tests would + consume the shared budget and make later tests order-dependent (a + single test can create up to 100 tokens). Tests must therefore only + assert on tokens they created themselves (unique UUIDs); an empty + registry between tests is strictly stronger isolation. """ - asyncio.run(_registry.cleanup_expired()) + asyncio.run(_registry.clear()) yield - asyncio.run(_registry.cleanup_expired()) + asyncio.run(_registry.clear()) @pytest.fixture(autouse=True) @@ -813,3 +816,255 @@ def test_ownership_denial_masks_owner_and_requester( assert "intruder" not in caplog.text assert "token owner apikey:" in caplog.text assert "requester apikey:" in caplog.text + + +# ============================================================================= +# Test: Registry admission budget (BUG-011) +# ============================================================================= + + +@contextmanager +def _download_config(**overrides): + """Temporarily override download.* config values, restoring them after.""" + config = get_config() + saved = {name: getattr(config.download, name) for name in overrides} + for name, value in overrides.items(): + setattr(config.download, name, value) + try: + yield config + finally: + for name, value in saved.items(): + setattr(config.download, name, value) + + +class TestDownloadRegistryBudget: + """The registry enforces max_tokens/max_total_bytes before reading files. + + Regression suite for BUG-011: unbounded in-memory token registry and + read_bytes() executed before any size/budget validation. + """ + + def test_disabled_downloads_generate_no_token( + self, + service: DownloadService, + test_file: Path, + caplog: pytest.LogCaptureFixture, + ): + """download.enabled=false never admits a token (single gate for all consumers).""" + with _download_config(enabled=False), caplog.at_level(logging.INFO): + result = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + + assert not result.success + assert "disabled" in result.error.lower() + assert _registry.stats["tokens"] == 0 + # Expected state: logged at info, not as a warning. + assert "downloads are disabled" in caplog.text + assert "WARNING" not in [r.levelname for r in caplog.records] + + def test_oversized_file_rejected_without_reading( + self, service: DownloadService, tmp_path: Path + ): + """stat() runs first: an over-limit file is never loaded into RAM.""" + big_file = tmp_path / "big_export.json" + big_file.write_bytes(b"x" * (2 * 1024 * 1024)) # 2 MiB + + with ( + _download_config(max_file_size_mb=1), + patch.object( + Path, "read_bytes", side_effect=AssertionError("file must not be read") + ), + ): + result = service.generate_download_token( + file_path=big_file, user_id="test_user", file_format="json" + ) + + assert not result.success + assert "exceeds limit" in result.error + assert _registry.stats["tokens"] == 0 + + def test_token_budget_rejected_without_reading( + self, service: DownloadService, test_file: Path, tmp_path: Path + ): + """When max_tokens is exhausted, admission fails before read_bytes().""" + with _download_config(max_tokens=1): + first = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + assert first.success, first.error + + second_file = tmp_path / "second_export.json" + second_file.write_text('{"second": true}') + + with patch.object( + Path, + "read_bytes", + side_effect=AssertionError("file must not be read"), + ): + second = service.generate_download_token( + file_path=second_file, user_id="test_user", file_format="json" + ) + + assert not second.success + assert "quota" in second.error.lower() + + stats = _registry.stats + assert stats["tokens"] == 1 + assert stats["pending"] == 0 + assert stats["reserved_bytes"] == 0 + + def test_token_budget_frees_after_delete( + self, service: DownloadService, test_file: Path + ): + """Deleting a token returns its slot to the budget immediately.""" + with _download_config(max_tokens=2): + first = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + second = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + third = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + assert first.success and second.success + assert not third.success + assert "quota" in third.error.lower() + + assert service.delete_token(first.data.token).success + + fourth = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + assert fourth.success, fourth.error + + assert _registry.stats["tokens"] == 2 + + def test_byte_budget_rejects_and_frees( + self, service: DownloadService, tmp_path: Path + ): + """max_total_bytes bounds committed bytes; delete releases them.""" + file_a = tmp_path / "budget_a.json" + file_b = tmp_path / "budget_b.json" + file_a.write_bytes(b"a" * (600 * 1024)) + file_b.write_bytes(b"b" * (600 * 1024)) + chunk = 600 * 1024 + + # 1 MiB budget: 600 KiB + 600 KiB must not fit. + with _download_config(max_total_bytes=1_048_576): + first = service.generate_download_token( + file_path=file_a, user_id="test_user", file_format="json" + ) + assert first.success, first.error + assert _registry.stats["total_bytes"] == chunk + + second = service.generate_download_token( + file_path=file_b, user_id="test_user", file_format="json" + ) + assert not second.success + assert "quota" in second.error.lower() + assert _registry.stats["total_bytes"] == chunk + assert _registry.stats["pending"] == 0 + + assert service.delete_token(first.data.token).success + + retry = service.generate_download_token( + file_path=file_b, user_id="test_user", file_format="json" + ) + assert retry.success, retry.error + assert _registry.stats["total_bytes"] == chunk + + def test_failed_read_releases_reservation( + self, service: DownloadService, test_file: Path + ): + """A read failure after reserving never leaks budget.""" + with _download_config(max_tokens=1): + with patch.object(Path, "read_bytes", side_effect=OSError("disk error")): + failed = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + + assert not failed.success + assert _registry.stats["pending"] == 0 + assert _registry.stats["reserved_bytes"] == 0 + + # A leaked reservation would keep max_tokens=1 unadmittable forever. + admitted = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + assert admitted.success, admitted.error + + def test_concurrent_admissions_never_exceed_token_budget( + self, service: DownloadService, test_file: Path + ): + """Atomic reserve-commit: 8 concurrent workers admit exactly max_tokens.""" + import threading as _threading + + max_tokens = 3 + successes: list[str] = [] + lock = _threading.Lock() + + def worker() -> None: + result = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + if result.success: + with lock: + successes.append(result.data.token) + + with ( + _download_config(max_tokens=max_tokens), + concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool, + ): + futures = [pool.submit(worker) for _ in range(8)] + for future in concurrent.futures.as_completed(futures): + future.result() + + stats = _registry.stats + assert len(successes) == max_tokens + assert stats["tokens"] == max_tokens + assert stats["pending"] == 0 + assert stats["reserved_bytes"] == 0 + + def test_byte_counter_tracks_store_overwrite_and_delete( + self, service: DownloadService, test_file: Path + ): + """_total_bytes stays exact across store(), delete_token() and TTL cleanup.""" + generated = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + assert generated.success, generated.error + token = generated.data.token + assert _registry.stats["total_bytes"] == test_file.stat().st_size + + # Direct store() overwrite must not double-count (test helper path). + info = asyncio.run(_registry.get_info(token)) + asyncio.run(_registry.store(token, info, b"new")) + assert _registry.stats["total_bytes"] == 3 + + assert service.delete_token(token).success + assert _registry.stats["total_bytes"] == 0 + assert _registry.stats["tokens"] == 0 + + def test_expired_cleanup_releases_byte_budget( + self, service: DownloadService, test_file: Path + ): + """TTL cleanup returns the token's bytes to the budget.""" + generated = service.generate_download_token( + file_path=test_file, user_id="test_user", file_format="json" + ) + assert generated.success, generated.error + token = generated.data.token + assert _registry.stats["total_bytes"] == test_file.stat().st_size + + async def expire_token() -> None: + info = await _registry.get_info(token) + info.expires_at = datetime.now(timezone.utc) - timedelta(seconds=1) + + asyncio.run(expire_token()) + assert service.cleanup_expired_tokens().data == 1 + + stats = _registry.stats + assert stats["tokens"] == 0 + assert stats["total_bytes"] == 0 diff --git a/tests/unit/services/test_download_token_deletion.py b/tests/unit/services/test_download_token_deletion.py index 6746bc2..9f87d99 100644 --- a/tests/unit/services/test_download_token_deletion.py +++ b/tests/unit/services/test_download_token_deletion.py @@ -30,10 +30,16 @@ def download_service(self): @pytest.fixture(autouse=True) def clear_registry(self): - """Clear the registry before and after each test.""" - asyncio.run(_registry.cleanup_expired()) + """Clear the registry before and after each test. + + Full purge (not just expired tokens): the registry enforces + ``download.max_tokens``, so tokens leaking from earlier tests + would consume the shared budget and make these tests fail for + reasons unrelated to deletion. + """ + asyncio.run(_registry.clear()) yield - asyncio.run(_registry.cleanup_expired()) + asyncio.run(_registry.clear()) @pytest.fixture(autouse=True) def mock_export_dir(self):