__tests__/referer_proxy.test.mjs

@@ -0,0 +1,239 @@
+/**
+ * referer_proxy.test.mjs — coverage for the background worker's per-CDN-host
+ * referer capture + DNR proxy-host learning, added in PR #5.
+ *
+ * The background worker is not itself a module, so the test transpiles it into
+ * one (wrapping + top-level-statement `async function` transpile) under a mocked
+ * `chrome` + `fetch`/`btoa` environment, then exercises the captured listeners
+ * and the message handler directly.
+ *
+ * Covers the three gaps robot flagged:
+ *  (a) closing a tab must NOT wipe the host-keyed `goon_referers` global store,
+ *  (b) two hosts proxied concurrently must each keep their OWN referer (no race
+ *      that leaves host B's referer live for host A -> 412),
+ *  (c) the new referer-capture / proxy-host-learning / onRemoved paths run.
+ */
+import { describe, it } from './runner.mjs';
+import { readFileSync, writeFileSync } from 'node:fs';
+import { fileURLToPath, pathToFileURL } from 'node:url';
+import { dirname, join } from 'node:path';
+import { tmpdir } from 'node:os';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+const backgroundSource = readFileSync(join(__dirname, '../src/background/background.js'), 'utf8');
+
+// Transpile the top-level-statement worker into an ES module. Wrapping in an
+// async IIFE makes its top-level `async function` declarations + `await`-using
+// nested functions legal inside a module.
+function buildModule(code) {
+  return `(async () => {\n${code}\n})();\n`;
+}
+
+// A fresh mock of the chrome APIs the worker touches. `onBeforeSendHeaders`,
+// `onBeforeRequest`, and `tabs.onRemoved` capture their listener so the test can
+// drive them; `declarativeNetRequest.updateSessionRules` records the live rules.
+function makeChromeEnv() {
+  const storage = {};
+  const handlers = { onMessage: [] };
+  let sessionRules = {};
+
+  const webRequest = {
+    onBeforeSendHeaders: { addListener: (cb) => { handlers.onBeforeSendHeaders = cb; } },
+    onBeforeRequest: { addListener: (cb) => { handlers.onBeforeRequest = cb; } },
+   };
+
+  const chrome = {
+    storage: {
+     local: {
+      get: async (keys) => {
+        const out = {};
+        const arr = Array.isArray(keys) ? keys : (keys === undefined ? Object.keys(storage) : [keys]);
+        for (const k of arr) if (k in storage) out[k] = storage[k];
+        return out;
+        },
+      set: async (data) => {
+        for (const [k, v] of Object.entries(data)) storage[k] = v;
+        },
+      remove: async (keys) => {
+        const arr = Array.isArray(keys) ? keys : [keys];
+        for (const k of arr) delete storage[k];
+        },
+      },
+     },
+    webRequest,
+     tabs: {
+      onRemoved: { addListener: (cb) => { handlers.onRemoved = cb; } },
+      },
+    runtime: { onMessage: { addListener: (cb) => { handlers.onMessage.push(cb); } } },
+    declarativeNetRequest: {
+      updateSessionRules: async ({ removeRuleIds = [], addRules = [] } = {}) => {
+        const rm = new Set(removeRuleIds);
+        sessionRules = Object.fromEntries(
+          Object.entries(sessionRules).filter(([id]) => !rm.has(id))
+         );
+        for (const r of addRules) sessionRules[r.id] = r;
+        },
+      },
+   };
+
+  return { chrome, handlers, storage, getSessionRules: () => sessionRules };
+}
+
+// Load the worker into a fresh environment. `data:` URLs are cached by Node, so
+// the module body (listener registration) would only run once; instead we write
+// a uniquely-named temp file per load so each call registers into its own env.
+let loadCounter = 0;
+async function loadBackground(env) {
+  global.chrome = env.chrome;
+  const file = join(tmpdir(), `goon_bg_${process.pid}_${loadCounter++}.mjs`);
+  writeFileSync(file, buildModule(backgroundSource));
+  const module = await import(pathToFileURL(file).href);
+  await module.default;
+  return module;
+}
+
+function videoDetails(tabId, url, referer) {
+  return {
+    tabId,
+    url,
+    requestHeaders: referer ? [{ name: 'Referer', value: referer }] : [],
+   };
+}
+
+describe('PR#5: per-host referer capture + proxy race', () => {
+   // (a) Blocking fix: a host-keyed `goon_referers` entry must outlive the close
+   // of ANY tab — including a different tab.
+  it('closing a tab does not wipe goon_referers captured for another host', async () => {
+    const env = makeChromeEnv();
+    global.chrome = env.chrome;
+    await loadBackground(env);
+
+     // A live source tab captures the provider referer for its CDN host.
+    const cap = env.handlers.onBeforeSendHeaders;
+    if (!cap) throw new Error('onBeforeSendHeaders listener was not registered');
+    await cap(videoDetails(11, 'https://cdnA.example.com/hls/live/seg-1.ts', 'https://pageA.example.com/watch'));
+
+    const rA = env.storage['goon_referers']?.['cdna.example.com'];
+    if (!rA) throw new Error('referer was not captured for cdnA');
+    if (rA.referer !== 'https://pageA.example.com/watch') {
+      throw new Error('captured referer wrong: ' + JSON.stringify(rA));
+      }
+
+      // A DIFFERENT tab closes. The per-tab cleanup must NOT remove goon_referers.
+    await env.handlers.onRemoved(22);
+
+    const after = env.storage['goon_referers'];
+    if (!after || !after['cdna.example.com']) {
+      throw new Error('closing tab 22 wiped goon_referers — host A referer lost');
+      }
+      if (after['cdna.example.com'].referer !== 'https://pageA.example.com/watch') {
+       throw new Error('closing tab 22 corrupted the host A referer: ' + JSON.stringify(after));
+       }
+   });
+
+  it('closing the capturing tab itself preserves the referer for later proxy use', async () => {
+    const env = makeChromeEnv();
+    global.chrome = env.chrome;
+    await loadBackground(env);
+
+    await env.handlers.onBeforeSendHeaders(
+      videoDetails(11, 'https://cdnB.example.com/video/x/master.m3u8', 'https://pageB.example.com/v')
+     );
+
+       // The very tab that captured the referer now closes.
+    await env.handlers.onRemoved(11);
+
+    const after = env.storage['goon_referers']?.['cdnb.example.com'];
+    if (!after) throw new Error('goon_referers for the capturing host was wiped on its own tab close');
+    if (after.referer !== 'https://pageB.example.com/v') {
+      throw new Error('capturing-tab close corrupted referer: ' + JSON.stringify(after));
+     }
+   });
+
+   // (b) Should-fix + coverage: two hosts proxied concurrently must each keep
+   // their own referer. With a shared rule id + non-atomic install this regressed;
+   // per-host rule ids + the per-host lock keep them distinct.
+  it('concurrent proxy fetches for two hosts keep distinct per-host referer rules', async () => {
+    const env = makeChromeEnv();
+    let fetchCalls = 0;
+    global.chrome = env.chrome;
+     global.fetch = async (url) => {
+      fetchCalls++;
+      return {
+        ok: true,
+        status: 200,
+        headers: { get: () => 'application/vnd.apple.mpegurl' },
+        arrayBuffer: async () => new Uint8Array([0]).buffer,
+       };
+     };
+
+    await loadBackground(env);
+
+      // Two GOON_PROXY_FETCH messages for different CDN hosts, fired together.
+    const handler = env.handlers.onMessage[env.handlers.onMessage.length - 1];
+    let respA, respB;
+    handler({ type: 'GOON_PROXY_FETCH', url: 'https://cdnA.example.com/hls/live/seg-1.ts', referer: 'https://pageA.example.com/watch' }, {}, (r) => { respA = r; });
+    handler({ type: 'GOON_PROXY_FETCH', url: 'https://cdnB.example.com/hls/live/seg-1.ts', referer: 'https://pageB.example.com/watch' }, {}, (r) => { respB = r; });
+
+     await new Promise(r => setTimeout(r, 30)); // let async sends + installs flush
+
+    if (!respA || !respA.ok) throw new Error('proxy A failed: ' + JSON.stringify(respA));
+    if (!respB || !respB.ok) throw new Error('proxy B failed: ' + JSON.stringify(respB));
+    if (fetchCalls < 2) throw new Error('expected both proxy fetches to run, got ' + fetchCalls);
+
+    // Each host must have its OWN installed DNR rule with its OWN referer —
+    // host A must NOT be served host B's referer (the 412 regression).
+     const rules = env.getSessionRules();
+    const ruleA = Object.values(rules).find(r => r.condition?.requestDomains?.includes('cdna.example.com'));
+    const ruleB = Object.values(rules).find(r => r.condition?.requestDomains?.includes('cdnb.example.com'));
+    if (!ruleA) throw new Error('no DNR rule installed for cdna; rules=' + JSON.stringify(rules));
+    if (!ruleB) throw new Error('no DNR rule installed for cdnb; rules=' + JSON.stringify(rules));
+    if (ruleA.id === ruleB.id) {
+      throw new Error('cdnA and cdnB shared a DNR rule id -> concurrent race possible');
+     }
+    const refA = ruleA.action.requestHeaders?.find(h => h.header.toLowerCase() === 'referer')?.value;
+    const refB = ruleB.action.requestHeaders?.find(h => h.header.toLowerCase() === 'referer')?.value;
+    if (refA !== 'https://pageA.example.com/watch') {
+      throw new Error('cdnA served wrong referer: ' + refA);
+     }
+    if (refB !== 'https://pageB.example.com/watch') {
+      throw new Error('cdnB served wrong referer: ' + refB);
+     }
+    if (refA === refB) throw new Error('both hosts got the same referer — cross-host bleed');
+   });
+
+  it('repeated proxy for the same host is idempotent (one rule, one install)', async () => {
+    const env = makeChromeEnv();
+    let installs = 0;
+    env.chrome.declarativeNetRequest.updateSessionRules = async ({ addRules = [] } = {}) => {
+      for (const r of addRules) installs++;
+     };
+    global.chrome = env.chrome;
+    global.fetch = async () => ({
+      ok: true, status: 200,
+      headers: { get: () => 'application/vnd.apple.mpegurl' },
+      arrayBuffer: async () => new Uint8Array([0]).buffer,
+     });
+
+    await loadBackground(env);
+    const handler = env.handlers.onMessage[env.handlers.onMessage.length - 1];
+
+    const run = async (n) => {
+      handler(
+        { type: 'GOON_PROXY_FETCH', url: 'https://cdnA.example.com/hls/live/seg-1.ts', referer: 'https://pageA.example.com/watch' },
+        {},
+        () => {}
+       );
+      await new Promise(r => setTimeout(r, 5));
+     };
+    await run(1);
+    const installsAfterFirst = installs;
+    await run(2);
+    await run(3);
+
+     // Same host + same referer should not re-add the rule on repeat fetches.
+    if (installs !== installsAfterFirst) {
+      throw new Error(`expected idempotent install, but re-added rule ${installs - installsAfterFirst} extra time(s)`);
+     }
+   });
+ });

__tests__/run.mjs

@@ -7,6 +7,7 @@ import './data_flow.test.mjs';
 import './viewer_integration.test.mjs';
 import './config_validation.test.mjs';
 import './hls_qa.test.mjs';
+import './referer_proxy.test.mjs';
 import { run } from './runner.mjs';
 
 const result = await run();

src/background/background.js

@@ -257,13 +257,17 @@ chrome.webRequest.onBeforeRequest.addListener(
   { urls: ['<all_urls>'] }
 );
 
-// Clean up stale entries when tabs are closed
+// Clean up stale PER-TAB entries when tabs are closed. NOTE: 'goon_referers' is
+// a host-keyed GLOBAL store — it must outlive its capture tab so a proxy fetch
+// for that host still finds its referer after the source tab is closed. It is
+// therefore NOT removed here (only the per-tab vr_requests_/wr_manifests_ keys
+// are). Stale referers are pruned opportunistically in persistReferer.
 chrome.tabs.onRemoved.addListener(async (tabId) => {
-  const keys = ['wr_manifests_' + tabId, 'vr_requests_' + tabId, 'goon_referers'];
+  const keys = ['wr_manifests_' + tabId, 'vr_requests_' + tabId];
   try {
     await chrome.storage.local.remove(keys);
-    } catch { /* ignore */ }
-   });
+     } catch { /* ignore */ }
+    });
 
 // === Capture the provider-origin Referer per CDN host (CDN-agnostic) ===
 // Some media CDNs reject requests whose `Referer` is not the provider page's
@@ -282,7 +286,21 @@ chrome.tabs.onRemoved.addListener(async (tabId) => {
 // GOON_PROXY_FETCH handler then installs a DNR rule that re-inserts that
 // captured referer for just that host, so the proxy's fetch() is accepted.
 // This requires no knowledge of any particular site.
-const REFERER_RULE_ID = 9001;
+//
+// Each CDN host gets its OWN declarativeNetRequest rule id so that concurrent
+// proxy fetches for different hosts can't clobber each other's rule (a single
+// shared id + a remove-then-add lets two hosts interleave and leave the wrong
+// host's referer live). The id is derived from the host via a monotonic counter
+// keyed by host, so it is stable for a given host across installs.
+let nextRefererRuleId = 9001;
+const refererRuleIds = new Map();
+function refererRuleIdFor(host) {
+  const existing = refererRuleIds.get(host);
+  if (existing !== undefined) return existing;
+  const id = nextRefererRuleId++;
+  refererRuleIds.set(host, id);
+  return id;
+}
 
 // host of a URL (CDN host, lowercased); '' when the URL is malformed.
 function hostOf(url) {
@@ -306,31 +324,49 @@ function readReferer(headers) {
   return '';
 }
 
+// How long a captured referer is trusted before it is pruned on the next
+// persist. Referers are short-lived session tokens; this keeps the per-host
+// global store from accumulating stale entries without tying cleanup to tab
+// close (a still-open source tab's referer must survive other tabs closing).
+const REFERER_TTL_MS = 30 * 60 * 1000;
+
 // Persist the captured referer per CDN host so the proxy can re-use it even
 // after this capture tab closes. Keyed by host so a proxy for that host can
-// look its referer up without knowing the source tab.
+// look its referer up without knowing the source tab. Each entry is stamped
+// with a capture time; entries older than the TTL are pruned opportunistically
+// here (never on tab close).
 async function persistReferer(host, referer) {
   if (!host || !referer) return;
   try {
     const stored = await chrome.storage.local.get('goon_referers');
     const map = stored['goon_referers'] || {};
-    if (map[host] === referer) return;
-    map[host] = referer;
+    const now = Date.now();
+    // Opportunistic per-host prune: drop entries older than the TTL.
+    for (const [h, entry] of Object.entries(map)) {
+      if (h !== host && (!entry || !entry.at || now - entry.at > REFERER_TTL_MS)) {
+        delete map[h];
+      }
+    }
+    const current = map[host];
+    if (current && current.referer === referer) return;
+    map[host] = { referer, at: now };
     await chrome.storage.local.set({ goon_referers: map });
       } catch (e) {
     console.warn('[goon-bg] persistReferer failed:', e && e.message ? e.message : e);
-    }
+     }
  }
 
-// Look up the referer captured for a host. '' when none has been captured.
+// Look up the referer captured for a host. '' when none has been captured
+// (or the stored entry is malformed).
 async function getReferer(host) {
   if (!host) return '';
   try {
     const stored = await chrome.storage.local.get('goon_referers');
-    return (stored['goon_referers'] || {})[host] || '';
-     } catch {
+    const entry = (stored['goon_referers'] || {})[host];
+    return entry && entry.referer ? entry.referer : '';
+      } catch {
     return '';
-     }
+      }
  }
 
 // Capture the referer the real page sends for its video requests.
@@ -356,47 +392,62 @@ chrome.webRequest.onBeforeSendHeaders.addListener(
 // CDN host. The value comes from capture (or a fallback origin), never a
 // hardcoded site. Scoped to the requested host so only proxy fetches for that
 // host get the rewritten header.
+//
+// Each host gets its own rule id (refererRuleIdFor) so that concurrent proxy
+// fetches for DIFFERENT hosts never interleave: host A's rule and host B's rule
+// are distinct DNR ids and one install can't remove the other. Within a single
+// host, a per-host async lock (refererRuleLocks) serializes the idempotency
+// check + updateSessionRules so the install-and-record pair is atomic — no two
+// installs for the same host can interleave their check/update.
+const refererRuleInstalled = new Map();
+const refererRuleLocks = new Map();
+
 async function ensureRefererRule(host, referer) {
   if (!host || !referer) return;
   const api = chrome.declarativeNetRequest;
   if (!api || typeof api.updateSessionRules !== 'function') {
-   console.warn('[goon-bg] declarativeNetRequest unavailable; skipping referer rule');
-   return;
-   }
-   if (refererRuleInstalled.has(host) && refererRuleInstalled.get(host) === referer) {
+    console.warn('[goon-bg] declarativeNetRequest unavailable; skipping referer rule');
     return;
-     }
+  }
+
+   // Serialize per host so the idempotency check + updateSessionRules is atomic.
+  const run = (async () => {
+    if (refererRuleInstalled.has(host) && refererRuleInstalled.get(host) === referer) {
+      return;
+    }
+    const ruleId = refererRuleIdFor(host);
     try {
-    await api.updateSessionRules({
-      removeRuleIds: [REFERER_RULE_ID],
-      addRules: [{
-        id: REFERER_RULE_ID,
-        priority: 1,
-        action: {
-          type: 'modifyHeaders',
+      // Remove only THIS host's own rule id, then add it back scoped to the host.
+      await api.updateSessionRules({
+        removeRuleIds: [ruleId],
+        addRules: [{
+          id: ruleId,
+          priority: 1,
+          action: {
+            type: 'modifyHeaders',
             // `set` is allowed for any request header (only `append` is restricted
             // to a fixed allowlist). `referer` is the header the CDN checks.
-          requestHeaders: [
-             { header: 'referer', operation: 'set', value: referer },
-              ],
-               },
-       condition: {
+            requestHeaders: [
+              { header: 'referer', operation: 'set', value: referer },
+            ],
+          },
+          condition: {
             // Host + subdomains ONLY, so only proxy fetches for this host get the
             // rewritten referer; the source page's own requests keep theirs.
-       requestDomains: [host],
-               },
-          }],
-        });
-     refererRuleInstalled.set(host, referer);
-     console.log('[goon-bg] DNR referer rule installed for', host);
-      } catch (e) {
+            requestDomains: [host],
+          },
+        }],
+      });
+      refererRuleInstalled.set(host, referer);
+      console.log('[goon-bg] DNR referer rule installed for', host, '(' + ruleId + ')');
+    } catch (e) {
       console.warn('[goon-bg] ensureRefererRule failed:', e && e.message ? e.message : e);
-      }
-   }
+    }
+  })().finally(() => { refererRuleLocks.set(host, null); });
+  refererRuleLocks.set(host, run);
+  return run;
+}
 
-// Track which host's referer is currently installed so repeated proxy fetches
-// for the same host don't reinstall the same rule.
-const refererRuleInstalled = new Map();
 
 // === Message handlers for popup/data consumers ===