From a47493eb3b2988392d9eb425d28f09da0bd64cd4 Mon Sep 17 00:00:00 2001 From: Max Randhahn Date: Sat, 18 Jul 2026 22:19:32 +0200 Subject: [PATCH] fix scope proxy semaphore leaks --- .github/workflows/tests.yml | 4 ++ lib/scope_proxy.py | 31 +++++++---- tests/test_ci_workflow.py | 5 ++ tests/test_scope_proxy.py | 101 ++++++++++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 9 deletions(-) create mode 100644 tests/test_scope_proxy.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 49ac02e..40743ed 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -51,6 +51,10 @@ jobs: if: always() run: python3 tests/test_completion_guard.py + - name: Scope-proxy regression test + if: always() + run: python3 tests/test_scope_proxy.py + - name: CI workflow regression test if: always() run: python3 tests/test_ci_workflow.py diff --git a/lib/scope_proxy.py b/lib/scope_proxy.py index df24247..3ed8c2d 100755 --- a/lib/scope_proxy.py +++ b/lib/scope_proxy.py @@ -46,7 +46,7 @@ def within_window(value: str) -> bool: class ScopeEnforcer: def __init__(self) -> None: - self.error = "" + self.config_error = "" self.directory: Path | None = None self.config: dict = {} self.tool = os.environ.get("PENTEST_PROXY_TOOL", "").strip() @@ -67,7 +67,7 @@ def __init__(self) -> None: self.tokens = float(self.policy["burst"]) self.concurrency = asyncio.BoundedSemaphore(int(self.policy["max_concurrency"])) except ConfigError as exc: - self.error = str(exc) + self.config_error = str(exc) def allowed_support(self, host: str) -> bool: for pattern in self.config.get("egress_support") or []: @@ -120,10 +120,10 @@ def release(self, flow: http.HTTPFlow) -> None: async def request(self, flow: http.HTTPFlow) -> None: allowed = False - reason = self.error or "scope configuration unavailable" + reason = self.config_error or "scope configuration unavailable" try: - if self.error: - raise ConfigError(self.error) + if self.config_error: + raise ConfigError(self.config_error) if not within_window(str(self.config.get("time_window") or "any")): raise ConfigError("outside engagement time_window") allowed, host, reason = scope_decision(self.config, flow.request.pretty_url) @@ -144,10 +144,23 @@ async def request(self, flow: http.HTTPFlow) -> None: # so the header is never sent to out-of-scope hosts. for name, value in self.required_headers.items(): flow.request.headers[name] = value - if self.concurrency is not None: - await self.concurrency.acquire() - self.acquired.add(flow.id) - await self.throttle() + concurrency = self.concurrency + release_on_exit = False + registered = False + try: + if concurrency is not None: + await concurrency.acquire() + release_on_exit = True + self.acquired.add(flow.id) + registered = True + await self.throttle() + release_on_exit = False + finally: + if concurrency is not None and release_on_exit: + if registered: + self.release(flow) + else: + concurrency.release() def response(self, flow: http.HTTPFlow) -> None: self.release(flow) diff --git a/tests/test_ci_workflow.py b/tests/test_ci_workflow.py index af4ab76..8da6bfe 100644 --- a/tests/test_ci_workflow.py +++ b/tests/test_ci_workflow.py @@ -22,6 +22,11 @@ def test_completion_guard_regression_test_runs_in_ci(self) -> None: self.assertIn("python3 tests/test_completion_guard.py", workflow) + def test_scope_proxy_regression_test_runs_in_ci(self) -> None: + workflow = WORKFLOW.read_text(encoding="utf-8") + + self.assertIn("python3 tests/test_scope_proxy.py", workflow) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_scope_proxy.py b/tests/test_scope_proxy.py new file mode 100644 index 0000000..12c3319 --- /dev/null +++ b/tests/test_scope_proxy.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Cancellation and lifecycle regression tests for the scope proxy addon.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import os +import sys +import tempfile +import time +import types +import unittest +from pathlib import Path +from unittest import mock + + +ROOT = Path(__file__).resolve().parent.parent +SCOPE_PROXY = ROOT / "lib" / "scope_proxy.py" + + +class Request: + def __init__(self, url: str) -> None: + self.pretty_url = url + self.method = "GET" + self.headers: dict[str, str] = {} + + +class Flow: + def __init__(self, flow_id: str) -> None: + self.id = flow_id + self.request = Request("http://127.0.0.1/") + self.response = None + + +def load_scope_proxy(engagement: Path): + """Load the addon with the small mitmproxy surface needed by these tests.""" + http = types.ModuleType("mitmproxy.http") + http.HTTPFlow = object + http.Response = types.SimpleNamespace(make=lambda *args, **kwargs: (args, kwargs)) + mitmproxy = types.ModuleType("mitmproxy") + mitmproxy.http = http + + spec = importlib.util.spec_from_file_location("scope_proxy_under_test", SCOPE_PROXY) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + with ( + mock.patch.dict(sys.modules, {"mitmproxy": mitmproxy, "mitmproxy.http": http}), + mock.patch.dict(os.environ, {"PENTEST_ENGAGEMENT_DIR": str(engagement)}), + ): + spec.loader.exec_module(module) + return module + + +class ScopeProxyTests(unittest.TestCase): + def setUp(self) -> None: + self.tempdir = tempfile.TemporaryDirectory() + self.addCleanup(self.tempdir.cleanup) + self.engagement = Path(self.tempdir.name) + (self.engagement / "engagement.yaml").write_text( + "targets:\n" + " - 127.0.0.1\n" + "out_of_scope: []\n" + "rate_limit_enabled: true\n" + "rate_limit:\n" + " requests_per_second: 100\n" + " burst: 1\n" + " max_concurrency: 1\n" + "time_window: any\n", + encoding="utf-8", + ) + + def test_error_hook_remains_callable_after_initialization(self) -> None: + addon = load_scope_proxy(self.engagement).addons[0] + + self.assertTrue(callable(addon.error)) + + def test_cancellation_during_throttle_releases_concurrency_permit(self) -> None: + async def exercise() -> None: + addon = load_scope_proxy(self.engagement).addons[0] + assert addon.policy is not None and addon.concurrency is not None + addon.policy["requests_per_second"] = 0.01 + addon.tokens = 0.0 + addon.updated = time.monotonic() + flow = Flow("cancelled-flow") + + task = asyncio.create_task(addon.request(flow)) + while flow.id not in addon.acquired: + await asyncio.sleep(0) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + self.assertNotIn(flow.id, addon.acquired) + self.assertFalse(addon.concurrency.locked()) + + asyncio.run(exercise()) + + +if __name__ == "__main__": + unittest.main()