diff --git a/bbot/defaults.yml b/bbot/defaults.yml index 3856a1f644..84545b2202 100644 --- a/bbot/defaults.yml +++ b/bbot/defaults.yml @@ -15,6 +15,9 @@ home: ~/.bbot keep_scans: 20 # Interval for displaying status messages status_frequency: 15 +# Pause event processing when system memory exceeds this percentage. +# Workers finish their current event then stop pulling new ones until memory drops. +max_mem_percent: 90 # Include the raw data of files (i.e. PDFs, web screenshots) as base64 in the event file_blobs: false # Include the raw data of directories (i.e. git repos) as tar.gz base64 in the event diff --git a/bbot/scanner/manager.py b/bbot/scanner/manager.py index 0ad7f59c85..1d5c8fa556 100644 --- a/bbot/scanner/manager.py +++ b/bbot/scanner/manager.py @@ -145,6 +145,14 @@ def module_priority_weights(self): return self._module_priority_weights async def get_incoming_event(self): + # memory-pressure backpressure: pause ingress while RAM is high. + # modules keep draining their queues (freeing memory via _minimize()), + # we just stop feeding new events into the pipeline. + if not self.scan._memory_ok.is_set(): + await self.scan._memory_ok.wait() + elif self.scan._memory_throttled: + await asyncio.sleep(self.scan._memory_throttle_delay) + for q in self.helpers.weighted_shuffle(self.incoming_queues, self.module_priority_weights): try: return q.get_nowait() diff --git a/bbot/scanner/scanner.py b/bbot/scanner/scanner.py index edfe5ca8db..73845e0923 100644 --- a/bbot/scanner/scanner.py +++ b/bbot/scanner/scanner.py @@ -1,4 +1,5 @@ import sys +import time import asyncio import logging import traceback @@ -264,6 +265,16 @@ def __init__( # how often to print scan status self.status_frequency = self.config.get("status_frequency", 15) + # memory-pressure backpressure: cleared when RAM is high, set when OK. + # ingress awaits _memory_ok before pulling new events. + # _memory_throttled = slow drip mode after safety valve fires. + # _memory_throttle_delay = seconds between events, scales with overshoot. + self._memory_ok = asyncio.Event() + self._memory_ok.set() + self._memory_paused_since = None + self._memory_throttled = False + self._memory_throttle_delay = 0.0 + from .stats import ScanStats self.stats = ScanStats(self) @@ -728,14 +739,52 @@ def modules_status(self, _log=False, detailed=False): modules_errored = [m for m, s in status["modules"].items() if s["errored"]] - max_mem_percent = 90 mem_status = self.helpers.memory_status() - # abort if we don't have the memory mem_percent = mem_status.percent + max_mem_percent = self.config.get("max_mem_percent", 90) + resume_threshold = max(0, max_mem_percent - 2) + # safety valve: force-resume after this many seconds to prevent stuck scans + max_pause_seconds = 120 + if mem_percent > max_mem_percent: - free_memory = mem_status.available - free_memory_human = self.helpers.bytes_to_human(free_memory) - self.warning(f"System memory is at {mem_percent:.1f}% ({free_memory_human} remaining)") + free_memory_human = self.helpers.bytes_to_human(mem_status.available) + if self._memory_ok.is_set() and not self._memory_throttled: + self.warning( + f"System memory is at {mem_percent:.1f}% ({free_memory_human} remaining); " + f"pausing event processing until it drops below {resume_threshold}%" + ) + self._memory_ok.clear() + self._memory_throttled = False + self._memory_paused_since = time.time() + elif not self._memory_ok.is_set(): + paused_for = time.time() - self._memory_paused_since + if paused_for >= max_pause_seconds: + self._memory_ok.set() + self._memory_throttled = True + self._memory_paused_since = None + self._update_throttle_delay(mem_percent, max_mem_percent) + self.warning( + f"System memory still at {mem_percent:.1f}% after {paused_for:.0f}s paused; " + f"switching to throttled ingress ({self._memory_throttle_delay:.1f}s/event)" + ) + else: + self.warning( + f"System memory still at {mem_percent:.1f}% ({free_memory_human} remaining); " + f"event processing paused ({paused_for:.0f}s / {max_pause_seconds}s)" + ) + else: + # throttled mode, memory still high — update delay and keep dripping + self._update_throttle_delay(mem_percent, max_mem_percent) + self.warning( + f"System memory still at {mem_percent:.1f}% ({free_memory_human} remaining); " + f"ingress throttled ({self._memory_throttle_delay:.1f}s/event)" + ) + elif (self._memory_throttled or not self._memory_ok.is_set()) and mem_percent <= resume_threshold: + self.hugesuccess(f"System memory dropped to {mem_percent:.1f}%; resuming full-speed event processing") + self._memory_ok.set() + self._memory_throttled = False + self._memory_throttle_delay = 0.0 + self._memory_paused_since = None if _log: modules_status = [] @@ -808,6 +857,13 @@ def modules_status(self, _log=False, detailed=False): return status + def _update_throttle_delay(self, mem_percent, max_mem_percent): + """Scale throttle delay linearly: 0s at threshold, 5s at cap (threshold+5 or 95%, whichever is lower).""" + cap = min(max_mem_percent + 5, 95) + overshoot_range = max(cap - max_mem_percent, 1) + overshoot = min(mem_percent - max_mem_percent, overshoot_range) + self._memory_throttle_delay = max(0.0, (overshoot / overshoot_range) * 5.0) + async def async_stop(self): """Stops the in-progress scan and performs necessary cleanup. diff --git a/bbot/test/test_step_1/test_scan.py b/bbot/test/test_step_1/test_scan.py index a23b3c4240..9229119935 100644 --- a/bbot/test/test_step_1/test_scan.py +++ b/bbot/test/test_step_1/test_scan.py @@ -431,3 +431,70 @@ async def test_scan_name(bbot_scanner): await scan._prep() assert scan.name == "test_scan_name" assert scan.preset.scan_name == "test_scan_name" + + +@pytest.mark.asyncio +async def test_memory_backpressure(bbot_scanner, monkeypatch): + """Test that memory pressure pauses ingress and the scan still completes.""" + from types import SimpleNamespace + + mem_percent = [50.0] # mutable so the mock can be changed mid-scan + + def mock_memory_status(): + return SimpleNamespace(percent=mem_percent[0], available=1_000_000_000) + + scan = bbot_scanner("127.0.0.1", config={"max_mem_percent": 90}) + await scan._prep() + await scan.helpers.dns._mock_dns({"1.1.1.1.in-addr.arpa": {"PTR": ["one.one.one.one"]}}) + monkeypatch.setattr("bbot.core.helpers.misc.memory_status", mock_memory_status) + + # starts in the OK state + assert scan._memory_ok.is_set() + + # simulate high memory — should pause + mem_percent[0] = 95.0 + scan.modules_status(_log=False) + assert not scan._memory_ok.is_set(), "Ingress should be paused when memory exceeds threshold" + assert scan._memory_paused_since is not None + + # memory drops but not below resume threshold (88) — should stay paused + mem_percent[0] = 89.0 + scan.modules_status(_log=False) + assert not scan._memory_ok.is_set(), "Should stay paused until memory drops below resume threshold" + + # memory drops below resume threshold — should resume + mem_percent[0] = 87.0 + scan.modules_status(_log=False) + assert scan._memory_ok.is_set(), "Ingress should resume when memory drops below resume threshold" + assert scan._memory_paused_since is None + + # test safety valve → throttled mode + import time + + mem_percent[0] = 95.0 + scan.modules_status(_log=False) + assert not scan._memory_ok.is_set() + assert not scan._memory_throttled + # fake that we've been paused for a long time + scan._memory_paused_since = time.time() - 200 + scan.modules_status(_log=False) + # should switch to throttled mode: _memory_ok set but _memory_throttled on + assert scan._memory_ok.is_set(), "Should open ingress after max pause duration" + assert scan._memory_throttled, "Should be in throttled mode, not full speed" + # next tick: memory still high → stays throttled (doesn't re-pause) + scan.modules_status(_log=False) + assert scan._memory_ok.is_set(), "Should stay open in throttled mode" + assert scan._memory_throttled, "Should remain throttled while memory is high" + # memory drops → back to full speed + mem_percent[0] = 80.0 + scan.modules_status(_log=False) + assert scan._memory_ok.is_set(), "Should be fully resumed" + assert not scan._memory_throttled, "Throttle should be cleared when memory drops" + + # verify the scan still completes with memory pressure active during the run + mem_percent[0] = 50.0 + scan._memory_ok.set() + scan._memory_throttled = False + scan._memory_paused_since = None + events = [e async for e in scan.async_start()] + assert any(e.type == "IP_ADDRESS" for e in events), "Scan should still produce events" diff --git a/bbot/test/test_step_2/module_tests/test_module_iis_shortnames.py b/bbot/test/test_step_2/module_tests/test_module_iis_shortnames.py index abc2f14719..0821eba42e 100644 --- a/bbot/test/test_step_2/module_tests/test_module_iis_shortnames.py +++ b/bbot/test/test_step_2/module_tests/test_module_iis_shortnames.py @@ -110,7 +110,7 @@ class TestIIS_Shortnames_GatewayError(ModuleTestBase): """Negative test: server returns 502 gateway errors. Should NOT detect IIS shortnames.""" targets = ["http://127.0.0.1:8888"] - modules_overrides = ["httpx", "iis_shortnames"] + modules_overrides = ["http", "iis_shortnames"] async def setup_after_prep(self, module_test): module_test.httpserver.no_handler_status_code = 404 diff --git a/docs/scanning/configuration.md b/docs/scanning/configuration.md index bbc5aa7a22..737307725c 100644 --- a/docs/scanning/configuration.md +++ b/docs/scanning/configuration.md @@ -74,6 +74,9 @@ home: ~/.bbot keep_scans: 20 # Interval for displaying status messages status_frequency: 15 +# Pause event processing when system memory exceeds this percentage. +# Workers finish their current event then stop pulling new ones until memory drops. +max_mem_percent: 90 # Include the raw data of files (i.e. PDFs, web screenshots) as base64 in the event file_blobs: false # Include the raw data of directories (i.e. git repos) as tar.gz base64 in the event