goon_plugin/__tests__/referer_proxy.test.mjs

/**
 * 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. `onBeforeRequest` and
  // `tabs.onRemoved` capture their listener, and `tabs.get` serves the source
  // tab's URL (the captured referer); `declarativeNetRequest.updateSessionRules`
  // records the live rules.
function makeChromeEnv() {
  const storage = {};
  const handlers = { onMessage: [] };
  const tabUrls = {};
  let sessionRules = {};

  const webRequest = {
    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: {
      get: async (tabId) => tabUrls[tabId] ? { tabId, url: tabUrls[tabId] } : undefined,
      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, tabUrls,
     setTabUrl: (tabId, url) => { tabUrls[tabId] = url; },
     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;
}

// Drive a referer capture the way the worker now does it: the worker reads the
// source tab's URL (via chrome.tabs.get) as the provider referer, so we set the
// tab's URL and fire an onBeforeRequest for a video URL on that tab.
async function capture(env, videoUrl, referer) {
  env.setTabUrl(11, referer);
  await env.handlers.onBeforeRequest({ tabId: 11, url: videoUrl });
}

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.
    await capture(env, 'https://cdnA.example.com/hls/live/seg-1.ts', 'https://pageA.example.com/watch');

    const rA = env.storage['goon_referers']?.['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['example.com']) {
      throw new Error('closing tab 22 wiped goon_referers — host A referer lost');
      }
      if (after['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 capture(env, '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']?.['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 distinct registrable domains keep distinct 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.other.com/hls/live/seg-1.ts', referer: 'https://pageB.other.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('example.com'));
    const ruleB = Object.values(rules).find(r => r.condition?.requestDomains?.includes('other.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.other.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)`);
     }
   });

   // (c) Domain generalization — the core 412 fix. A provider referer captured for
   // ONE edge subdomain must cover a SIBLING subdomain: the DNR rule is scoped to
   // the registrable domain and the store is domain-keyed, so a proxy for ev-h
   // needs no per-subdomain capture.
  it('a proxy for a sibling subdomain reuses the captured referer via the registrable domain', async () => {
    const env = makeChromeEnv();
    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);

       // Capture the provider referer for hv-h.example.net (one edge subdomain).
    await capture(env, 'https://hv-h.example.net/hls/live/seg-1.ts', 'https://provider.example.com/watch');

       // Proxy a SIBLING subdomain ev-h.example.net with NO referer in the message:
     // the SW must resolve it from the domain-keyed store.
    const handler = env.handlers.onMessage[env.handlers.onMessage.length - 1];
    let resp;
    handler({ type: 'GOON_PROXY_FETCH', url: 'https://ev-h.example.net/hls/live/seg-1.ts' }, {}, (r) => { resp = r; });
    await new Promise(r => setTimeout(r, 20));

    if (!resp || !resp.ok) throw new Error('sibling proxy failed: ' + JSON.stringify(resp));
    const rules = env.getSessionRules();
    const rule = Object.values(rules).find(r => r.condition?.requestDomains?.includes('example.net'));
    if (!rule) throw new Error('no domain-scoped DNR rule for example.net; rules=' + JSON.stringify(rules));
    const ref = rule.action.requestHeaders?.find(h => h.header.toLowerCase() === 'referer')?.value;
    if (ref !== 'https://provider.example.com/watch') {
      throw new Error('sibling edge served wrong referer: ' + ref);
       }
     // The rule is scoped to the registrable domain, NOT the exact subdomain.
    if (rule.condition.requestDomains?.includes('ev-h.example.net')) {
      throw new Error('rule scoped to exact subdomain instead of registrable domain');
       }
    });

   // (d) With no captured referer and none in the message, install NO rule: the
   // self-origin fallback (guaranteed-wrong for hotlink protection) is dropped.
  it('no captured referer and no message referer installs no DNR rule', async () => {
    const env = makeChromeEnv();
    let fetchCalls = 0;
    global.chrome = env.chrome;
    global.fetch = async () => {
      fetchCalls++;
      return {
        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];
    let resp;
    handler({ type: 'GOON_PROXY_FETCH', url: 'https://cdnA.example.com/hls/live/seg-1.ts' }, {}, (r) => { resp = r; });
    await new Promise(r => setTimeout(r, 20));

    if (!resp || !resp.ok) throw new Error('proxy fetch failed: ' + JSON.stringify(resp));
    if (fetchCalls < 1) throw new Error('proxy fetch did not run');
    const rules = env.getSessionRules();
    if (Object.keys(rules).length !== 0) {
      throw new Error('installed a rule with no referer (self-origin fallback not dropped): ' + JSON.stringify(rules));
       }
    });

   // (e) GOON_GET_REFERERS (used by scan-time capture) resolves each manifest host
   // via the domain-keyed store: a sibling subdomain gets the provider referer,
   // and an uncaptured host is omitted.
  it('GOON_GET_REFERERS resolves manifest hosts through the domain-keyed store', async () => {
    const env = makeChromeEnv();
    global.chrome = env.chrome;
    await loadBackground(env);

        // Capture for hv-h.example.net only.
    await capture(env, 'https://hv-h.example.net/hls/live/seg-1.ts', 'https://provider.example.com/watch');

    const handler = env.handlers.onMessage[env.handlers.onMessage.length - 1];
    let resp;
    handler(
       { type: 'GOON_GET_REFERERS', hosts: ['ev-h.example.net', 'hv-h.example.net', 'cdnA.example.com'] },
       {},
       (r) => { resp = r; });
    await new Promise(r => setTimeout(r, 15));

    if (!resp || !resp.ok) throw new Error('GOON_GET_REFERERS failed: ' + JSON.stringify(resp));
    const referers = resp.referers || {};
    if (referers['ev-h.example.net'] !== 'https://provider.example.com/watch') {
      throw new Error('sibling subdomain did not resolve the captured referer: ' + JSON.stringify(referers));
       }
    if (referers['hv-h.example.net'] !== 'https://provider.example.com/watch') {
      throw new Error('captured subdomain did not resolve its own referer: ' + JSON.stringify(referers));
       }
    if ('cdnA.example.com' in referers) {
      throw new Error('uncaptured host should be omitted from referers: ' + JSON.stringify(referers));
       }
    });

     // Regression for the live 412 bug: the worker captures the provider referer
     // from the source tab's URL via chrome.tabs.get inside onBeforeRequest. The
     // old path read request headers via onBeforeSendHeaders + extraInfoSpec, which
     // is rejected by some Chromium builds and fails the whole service worker, so it
     // was removed; this asserts the new tab-URL capture is wired up and actually
     // stores a referer.
   it('the onBeforeRequest capture reads the referer from the source tab URL', async () => {
     const env = makeChromeEnv();
     global.chrome = env.chrome;
     await loadBackground(env);

     const cap = env.handlers.onBeforeRequest;
     if (!cap) throw new Error('onBeforeRequest listener was not registered');

     // A video request for cdnA from a tab whose URL IS the provider referer.
     env.setTabUrl(11, 'https://pageA.example.com/watch');
     await cap({ tabId: 11, url: 'https://cdnA.example.com/hls/live/seg-1.ts' });

     const rA = env.storage['goon_referers']?.['example.com'];
     if (!rA) throw new Error('referer was not captured from the source tab URL');
     if (rA.referer !== 'https://pageA.example.com/watch') {
       throw new Error('captured referer wrong: ' + JSON.stringify(rA));
        }
     });
   });