test(dns): add regression test for DNS capture on connection abort

Adds a /CONNECTION_ABORT/ handler to the test server that sends headers
then forcibly closes the connection, and a test proving DNS data is still
captured. This guards against reverting from onHeadersReceived to
onCompleted, which would miss aborted requests.

Author: vringar

test/test_dns_instrument.py

@@ -25,3 +25,24 @@ def test_name_resolution(default_params, task_manager_creator):
     # Each redirect hop should record the URL it was associated with
     redirect_urls = [r["redirect_url"] for r in results]
     assert all(url is not None for url in redirect_urls)
+
+
+def test_dns_captured_on_connection_abort(default_params, task_manager_creator):
+    """Regression test: DNS data must be captured even when the connection
+    aborts before completion. This verifies that the extension uses
+    onHeadersReceived (not onCompleted) to record DNS responses."""
+    manager_params, browser_params = default_params
+    for browser_param in browser_params:
+        browser_param.dns_instrument = True
+
+    manager, db = task_manager_creator((manager_params, browser_params))
+    manager.get("http://localhost:8000/CONNECTION_ABORT/")
+    manager.close()
+
+    results = db_utils.query_db(db, "SELECT * FROM dns_responses")
+    assert len(results) > 0, "No DNS responses captured for aborted connection"
+    result = results[0]
+    assert isinstance(result, Row)
+    assert result["used_address"] is not None
+    assert result["addresses"] is not None
+    assert result["hostname"] == "localhost"

test/utilities.py

@@ -83,6 +83,15 @@ def do_GET(self, *args, **kwargs):
             self.end_headers()
             return
 
+        # 2. Abort connection after sending partial response.
+        if self.path.startswith("/CONNECTION_ABORT/"):
+            self.send_response(200)
+            self.send_header("Content-Length", "99999")
+            self.end_headers()
+            self.wfile.write(b"partial")
+            self.wfile.close()
+            return
+
         # Otherwise, return file from disk
         return SimpleHTTPRequestHandler.do_GET(self, *args, **kwargs)