fix scope proxy semaphore leaks

.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

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)

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()

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()