Move dns_instrument from onCompleted to onHeadersReceived

Extension/src/background/dns-instrument.ts

@@ -2,10 +2,11 @@ import { PendingResponse } from "../lib/pending-response";
 import { DnsResolved } from "../schema";
 import { allTypes } from "./http-instrument";
 import RequestFilter = browser.webRequest.RequestFilter;
+import { WebRequestOnHeadersReceivedDetails } from "../types/browser-web-request-event-details";
 
 export class DnsInstrument {
   private readonly dataReceiver;
-  private onCompleteListener;
+  private onHeadersReceivedListener;
   private pendingResponses: {
     [requestId: number]: PendingResponse;
   } = {};
@@ -28,25 +29,23 @@ export class DnsInstrument {
     /*
      * Attach handlers to event listeners
      */
-    this.onCompleteListener = (
-      details: browser.webRequest._OnCompletedDetails,
-    ) => {
+    this.onHeadersReceivedListener = (details: browser.webRequest._OnHeadersReceivedDetails) => {
       // Ignore requests made by extensions
       if (requestStemsFromExtension(details)) {
         return;
       }
       const pendingResponse = this.getPendingResponse(details.requestId);
-      pendingResponse.resolveOnCompletedEventDetails(details);
+      pendingResponse.resolveOnHeadersReceivedDetails(details);
 
       this.onCompleteDnsHandler(details, crawlID);
     };
 
-    browser.webRequest.onCompleted.addListener(this.onCompleteListener, filter);
+    browser.webRequest.onHeadersReceived.addListener(this.onHeadersReceivedListener, filter);
   }
 
   public cleanup() {
-    if (this.onCompleteListener) {
-      browser.webRequest.onCompleted.removeListener(this.onCompleteListener);
+    if (this.onHeadersReceivedListener) {
+      browser.webRequest.onHeadersReceived.removeListener(this.onHeadersReceivedListener);
     }
   }
 
@@ -71,7 +70,7 @@ export class DnsInstrument {
   }
 
   private async onCompleteDnsHandler(
-    details: browser.webRequest._OnCompletedDetails,
+    details: WebRequestOnHeadersReceivedDetails,
     crawlID,
   ) {
     // Create and populate DnsResolve object

Extension/src/lib/pending-response.ts

@@ -1,23 +1,36 @@
 import { ResponseBodyListener } from "./response-body-listener";
 
 /**
- * Ties together the two separate events that together holds information about both response headers and body
+ * Ties together three separate events that hold information about both response headers and body
  */
 export class PendingResponse {
-  public readonly onBeforeRequestEventDetails: Promise<browser.webRequest._OnBeforeRequestDetails>;
-  public readonly onCompletedEventDetails: Promise<browser.webRequest._OnCompletedDetails>;
+  public readonly onBeforeRequestEventDetails: Promise<
+  browser.webRequest._OnBeforeRequestDetails
+  >;
+  public readonly onHeadersReceivedEventDetails: Promise<
+  browser.webRequest._OnHeadersReceivedDetails
+  >;
+  public readonly onCompletedEventDetails: Promise<
+    browser.webRequest._OnCompletedDetails
+  >;
   public responseBodyListener: ResponseBodyListener;
   public resolveOnBeforeRequestEventDetails: (
     details: browser.webRequest._OnBeforeRequestDetails,
   ) => void;
+  public resolveOnHeadersReceivedDetails: (
+    details: browser.webRequest._OnHeadersReceivedDetails,
+  ) => void;
   public resolveOnCompletedEventDetails: (
     details: browser.webRequest._OnCompletedDetails,
   ) => void;
   constructor() {
     this.onBeforeRequestEventDetails = new Promise((resolve) => {
       this.resolveOnBeforeRequestEventDetails = resolve;
     });
-    this.onCompletedEventDetails = new Promise((resolve) => {
+    this.onHeadersReceivedEventDetails = new Promise(resolve => {
+      this.resolveOnHeadersReceivedDetails = resolve;
+    });
+    this.onCompletedEventDetails = new Promise(resolve => {
       this.resolveOnCompletedEventDetails = resolve;
     });
   }
@@ -29,6 +42,7 @@ export class PendingResponse {
   public resolved() {
     return Promise.all([
       this.onBeforeRequestEventDetails,
+      this.onHeadersReceivedEventDetails,
       this.onCompletedEventDetails,
     ]);
   }

Extension/src/types/browser-web-request-event-details.ts

@@ -1,11 +1,5 @@
 /* eslint-disable no-underscore-dangle */
 
-/**
- * This file contains selected implicit interfaces copied from node_modules/@types/firefox-webext-browser/index.d.ts
- * Defined and exported here in order for our code to be able to reference them explicitly in helper functions
- * and class methods that accept arguments of these types.
- */
-
 export interface FrameAncestor {
   /** The URL that the document was loaded from. */
   url: string;
@@ -18,3 +12,11 @@ export interface WebRequestOnBeforeSendHeadersEventDetails
   /** Contains information for each document in the frame hierarchy up to the top-level document. The first element in the array contains information about the immediate parent of the document being requested, and the last element contains information about the top-level document. If the load is actually for the top-level document, then this array is empty. */
   frameAncestors: FrameAncestor[];
 }
+
+export interface WebRequestOnHeadersReceivedDetails extends browser.webRequest._OnHeadersReceivedDetails {
+  /**
+   * TODO: Does this actually exist?
+   * The server IP address that the request was actually sent to. Note that it may be a literal IPv6 address.
+   */
+   ip?: string;
+}