[CRX] Use DNR instead of webRequest in preserve-referer

webRequestBlocking is unavailable in MV3. Non-blocking webRequest can
still be used to detect the Referer, but we have to use
declarativeNetRequest to change the Referer header as needed.

extensions/chromium/manifest.json

@@ -10,6 +10,7 @@
     "16": "icon16.png"
   },
   "permissions": [
+    "declarativeNetRequestWithHostAccess",
     "webRequest",
     "webRequestBlocking",
     "<all_urls>",

extensions/chromium/pdfHandler.js

@@ -13,7 +13,6 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
-/* globals saveReferer */
 
 "use strict";
 
@@ -139,9 +138,6 @@ chrome.webRequest.onHeadersReceived.addListener(
 
     var viewerUrl = getViewerURL(details.url);
 
-    // Implemented in preserve-referer.js
-    saveReferer(details);
-
     return { redirectUrl: viewerUrl };
   },
   {

extensions/chromium/preserve-referer.js

@@ -13,20 +13,14 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
-/* globals getHeaderFromHeaders */
-/* exported saveReferer */
 
 "use strict";
 /**
  * This file is one part of the Referer persistency implementation. The other
  * part resides in chromecom.js.
  *
- * This file collects request headers for every http(s) request, and temporarily
- * stores the request headers in a dictionary. Upon completion of the request
- * (success or failure), the headers are discarded.
- * pdfHandler.js will call saveReferer(details) when it is about to redirect to
- * the viewer. Upon calling saveReferer, the Referer header is extracted from
- * the request headers and saved.
+ * This file collects Referer headers for every http(s) request, and temporarily
+ * stores the request headers in a dictionary, for REFERRER_IN_MEMORY_TIME ms.
  *
  * When the viewer is opened, it opens a port ("chromecom-referrer"). This port
  * is used to set up the webRequest listeners that stick the Referer headers to
@@ -36,50 +30,38 @@ limitations under the License.
  * See setReferer in chromecom.js for more explanation of this logic.
  */
 
-// Remembers the request headers for every http(s) page request for the duration
-// of the request.
-var g_requestHeaders = {};
 // g_referrers[tabId][frameId] = referrer of PDF frame.
 var g_referrers = {};
+var g_referrerTimers = {};
+// The background script will eventually suspend after 30 seconds of inactivity.
+// This can be delayed when extension events are firing. To prevent the data
+// from being kept in memory for too long, cap the data duration to 5 minutes.
+var REFERRER_IN_MEMORY_TIME = 300000;
 
-(function () {
-  var requestFilter = {
-    urls: ["*://*/*"],
-    types: ["main_frame", "sub_frame"],
-  };
-  chrome.webRequest.onSendHeaders.addListener(
-    function (details) {
-      g_requestHeaders[details.requestId] = details.requestHeaders;
-    },
-    requestFilter,
-    ["requestHeaders", "extraHeaders"]
-  );
-  chrome.webRequest.onBeforeRedirect.addListener(forgetHeaders, requestFilter);
-  chrome.webRequest.onCompleted.addListener(forgetHeaders, requestFilter);
-  chrome.webRequest.onErrorOccurred.addListener(forgetHeaders, requestFilter);
-  function forgetHeaders(details) {
-    delete g_requestHeaders[details.requestId];
-  }
-})();
+var rIsReferer = /^referer$/i;
+chrome.webRequest.onSendHeaders.addListener(
+  function saveReferer(details) {
+    const { tabId, frameId, requestHeaders } = details;
+    g_referrers[tabId] ??= {};
+    g_referrers[tabId][frameId] = requestHeaders.find(h =>
+      rIsReferer.test(h.name)
+    )?.value;
+    forgetReferrerEventually(tabId);
+  },
+  { urls: ["*://*/*"], types: ["main_frame", "sub_frame"] },
+  ["requestHeaders", "extraHeaders"]
+);
 
-/**
- * @param {object} details - onHeadersReceived event data.
- */
-function saveReferer(details) {
-  var referer =
-    g_requestHeaders[details.requestId] &&
-    getHeaderFromHeaders(g_requestHeaders[details.requestId], "referer");
-  referer = (referer && referer.value) || "";
-  if (!g_referrers[details.tabId]) {
-    g_referrers[details.tabId] = {};
+function forgetReferrerEventually(tabId) {
+  if (g_referrerTimers[tabId]) {
+    clearTimeout(g_referrerTimers[tabId]);
   }
-  g_referrers[details.tabId][details.frameId] = referer;
+  g_referrerTimers[tabId] = setTimeout(() => {
+    delete g_referrers[tabId];
+    delete g_referrerTimers[tabId];
+  }, REFERRER_IN_MEMORY_TIME);
 }
 
-chrome.tabs.onRemoved.addListener(function (tabId) {
-  delete g_referrers[tabId];
-});
-
 // This method binds a webRequest event handler which adds the Referer header
 // to matching PDF resource requests (only if the Referer is non-empty). The
 // handler is removed as soon as the PDF viewer frame is unloaded.
@@ -89,8 +71,11 @@ chrome.runtime.onConnect.addListener(function onReceivePort(port) {
   }
   var tabId = port.sender.tab.id;
   var frameId = port.sender.frameId;
+  var dnrRequestId;
 
   // If the PDF is viewed for the first time, then the referer will be set here.
+  // Note: g_referrers could be empty if the background script was suspended by
+  // the browser. In that case, chromecom.js may send us the referer (below).
   var referer = (g_referrers[tabId] && g_referrers[tabId][frameId]) || "";
   port.onMessage.addListener(function (data) {
     // If the viewer was opened directly (without opening a PDF URL first), then
@@ -99,49 +84,49 @@ chrome.runtime.onConnect.addListener(function onReceivePort(port) {
     if (data.referer) {
       referer = data.referer;
     }
-    chrome.webRequest.onBeforeSendHeaders.removeListener(onBeforeSendHeaders);
-    if (referer) {
-      // Only add a blocking request handler if the referer has to be rewritten.
-      chrome.webRequest.onBeforeSendHeaders.addListener(
-        onBeforeSendHeaders,
-        {
-          urls: [data.requestUrl],
-          types: ["xmlhttprequest"],
-          tabId,
-        },
-        ["blocking", "requestHeaders", "extraHeaders"]
-      );
-    }
-    // Acknowledge the message, and include the latest referer for this frame.
-    port.postMessage(referer);
+    dnrRequestId = data.dnrRequestId;
+    setStickyReferrer(dnrRequestId, tabId, data.requestUrl, referer, () => {
+      // Acknowledge the message, and include the latest referer for this frame.
+      port.postMessage(referer);
+    });
   });
 
   // The port is only disconnected when the other end reloads.
   port.onDisconnect.addListener(function () {
-    if (g_referrers[tabId]) {
-      delete g_referrers[tabId][frameId];
-    }
-    chrome.webRequest.onBeforeSendHeaders.removeListener(onBeforeSendHeaders);
+    unsetStickyReferrer(dnrRequestId);
   });
-
-  function onBeforeSendHeaders(details) {
-    if (details.frameId !== frameId) {
-      return undefined;
-    }
-    var headers = details.requestHeaders;
-    var refererHeader = getHeaderFromHeaders(headers, "referer");
-    if (!refererHeader) {
-      refererHeader = { name: "Referer" };
-      headers.push(refererHeader);
-    } else if (
-      refererHeader.value &&
-      refererHeader.value.lastIndexOf("chrome-extension:", 0) !== 0
-    ) {
-      // Sanity check. If the referer is set, and the value is not the URL of
-      // this extension, then the request was not initiated by this extension.
-      return undefined;
-    }
-    refererHeader.value = referer;
-    return { requestHeaders: headers };
-  }
 });
+
+function setStickyReferrer(dnrRequestId, tabId, url, referer, callback) {
+  if (!referer) {
+    unsetStickyReferrer(dnrRequestId);
+    callback();
+    return;
+  }
+  const rule = {
+    id: dnrRequestId,
+    condition: {
+      urlFilter: `|${url}|`,
+      // The viewer and background are presumed to have the same origin:
+      initiatorDomains: [location.hostname], // = chrome.runtime.id.
+      resourceTypes: ["xmlhttprequest"],
+      tabIds: [tabId],
+    },
+    action: {
+      type: "modifyHeaders",
+      requestHeaders: [{ operation: "set", header: "referer", value: referer }],
+    },
+  };
+  chrome.declarativeNetRequest.updateSessionRules(
+    { removeRuleIds: [dnrRequestId], addRules: [rule] },
+    callback
+  );
+}
+
+function unsetStickyReferrer(dnrRequestId) {
+  if (dnrRequestId) {
+    chrome.declarativeNetRequest.updateSessionRules({
+      removeRuleIds: [dnrRequestId],
+    });
+  }
+}

web/chromecom.js

@@ -267,6 +267,7 @@ if (window === top) {
   });
 }
 
+let dnrRequestId;
 // This port is used for several purposes:
 // 1. When disconnected, the background page knows that the frame has unload.
 // 2. When the referrer was saved in history.state.chromecomState, it is sent
@@ -281,6 +282,7 @@ let port;
 // 3. Background -> page: Send latest referer and save to history.
 // 4. Page: Invoke callback.
 function setReferer(url, callback) {
+  dnrRequestId ??= crypto.getRandomValues(new Uint32Array(1))[0] % 0x80000000;
   if (!port) {
     // The background page will accept the port, and keep adding the Referer
     // request header to requests to |url| until the port is disconnected.
@@ -290,6 +292,7 @@ function setReferer(url, callback) {
   port.onMessage.addListener(onMessage);
   // Initiate the information exchange.
   port.postMessage({
+    dnrRequestId,
     referer: window.history.state?.chromecomState,
     requestUrl: url,
   });