From 418a349d88a2928a034f22da8dbacbf788bc95a9 Mon Sep 17 00:00:00 2001 From: 0xMarcio Date: Thu, 22 Oct 2026 21:48:45 +0200 Subject: [PATCH] Keep archive response failures isolated and respect batch limits --- scripts/archive_pocs.py | 14 +++++++++----- tests/test_archive_pocs.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 5 deletions(-) diff --git a/scripts/archive_pocs.py b/scripts/archive_pocs.py index 508ef3a6fe1..7e236133a01 100644 --- a/scripts/archive_pocs.py +++ b/scripts/archive_pocs.py @@ -12,6 +12,7 @@ import re import tempfile import threading import time +import zlib from collections import Counter, defaultdict from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED from datetime import datetime, timedelta, timezone @@ -177,8 +178,11 @@ def response_bytes(response, limit: int) -> bytes: if encoding in {'', 'identity'}: body = response.read(limit + 1) elif encoding == 'gzip': - with gzip.GzipFile(fileobj=response) as stream: - body = stream.read(limit + 1) + try: + with gzip.GzipFile(fileobj=response) as stream: + body = stream.read(limit + 1) + except (gzip.BadGzipFile, EOFError, zlib.error): + raise ArchiveError('invalid-compressed-response') from None else: raise ArchiveError('unsupported-response-encoding') if len(body) > limit: @@ -723,7 +727,7 @@ def lookup_pages(targets: dict, rows: dict, client: Client, directory: Path, lim stopped = True fatal |= problem.fatal control.write_text(json.dumps({'retry_after': after(problem.cooldown or 300), 'reason': str(problem)}) + '\n') - if getattr(problem, 'operation', '') == 'replay' or str(problem) == 'service-transport-failures': + if getattr(problem, 'operation', '') == 'replay': (directory / 'control.json').write_text(control.read_text()) print(f'Archive lookup paused: {problem}', flush=True) print(f'{completed}: {row["status"]} {row["original"]}', flush=True) @@ -857,7 +861,7 @@ def main() -> int: break time.sleep(5) continue - elif selected and hosts: + elif selected and hosts and completed < args.limit and room: time.sleep(5) continue else: @@ -884,7 +888,7 @@ def main() -> int: control.write_text(json.dumps({'retry_after': rows[key]['next_check_at'], 'reason': str(exc)}) + '\n') print(f'Archive service requested backoff during {client.operation}: {exc}', flush=True) break - if str(exc) in {'response-too-large', 'unsupported-response-encoding'}: + if str(exc) in {'response-too-large', 'unsupported-response-encoding', 'invalid-compressed-response'}: rows[key].pop('candidate', None) rows[key].pop('verify_attempts', None) rows[key]['next_check_at'] = after(7 * 86400) diff --git a/tests/test_archive_pocs.py b/tests/test_archive_pocs.py index 4d1be10bf5a..7c656875a09 100644 --- a/tests/test_archive_pocs.py +++ b/tests/test_archive_pocs.py @@ -408,6 +408,32 @@ class ArchiveTests(unittest.TestCase): self.assertEqual(client.failure(str(i % 3), archive.ArchiveError('transport-ConnectionResetError')).cooldown, 0) self.assertEqual(client.failure('third', archive.ArchiveError('replay-URLError')).cooldown, 900) + def test_availability_breaker_does_not_pause_capture_api(self): + for operation in ('availability', 'replay'): + with self.subTest(operation=operation), tempfile.TemporaryDirectory() as directory: + self.client.operation = operation + self.client.lookup.side_effect = archive.ArchiveError('service-transport-failures', cooldown=900) + archive.lookup_pages({URL: self.target}, {}, self.client, Path(directory), 1, 10) + self.assertTrue((Path(directory) / 'lookup-control.json').exists()) + self.assertEqual((Path(directory) / 'control.json').exists(), operation == 'replay') + + def test_capture_limit_exits_without_waiting_for_an_unrelated_host_cooldown(self): + blocked = 'https://example.com/poc' + targets = {u: {**self.target, 'original': u} for u in (URL, URL + '-other', blocked)} + def process(key, target, rows, client, *, capture): + rows[key] = {**target, 'status': 'not-found'} + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) + archive.save_records({blocked: {**targets[blocked], 'host_retry_after': archive.after(86400)}}, {blocked}, path) + with patch.object(archive, 'inventory', return_value=targets), \ + patch.object(archive, 'credentials', return_value=('dummy', 'dummy')), \ + patch.object(archive.Client, 'capacity', return_value=2), \ + patch.object(archive, 'process', side_effect=process) as visit, \ + patch.object(archive.time, 'sleep', side_effect=AssertionError('Unexpected idle wait')), \ + patch.object(sys, 'argv', ['archive_pocs', '--capture', '--directory', str(path), '--limit', '1']): + self.assertEqual(archive.main(), 0) + self.assertEqual(visit.call_count, 1) + def test_capture_loop_serves_new_targets_and_keeps_candidate_backoff(self): pending = URL + '-pending' target = {**self.target, 'original': pending} @@ -469,6 +495,14 @@ class ArchiveTests(unittest.TestCase): with self.assertRaisesRegex(archive.ArchiveError, 'too-large'): archive.response_bytes(response, 1000) + def test_corrupt_compressed_responses_are_target_errors(self): + for body in (b'not-gzip', gzip.compress(b'poc')[:-8], b'\x1f\x8b\x08\x00' + b'\x00' * 6 + b'\x07'): + with self.subTest(body=body): + response = io.BytesIO(body) + response.headers = {'Content-Encoding': 'gzip'} + with self.assertRaisesRegex(archive.ArchiveError, 'invalid-compressed-response'): + archive.response_bytes(response, 1000) + if __name__ == '__main__': unittest.main()