@specprotected/spec-proxy-service-worker/index.test.ts

Server Worker API implementation for integrating with Spec Proxy from an Edge Worker

import { jest, describe, beforeEach, test, expect, it } from "@jest/globals";
import rewire from "rewire";
import makeServiceWorkerEnv from "service-worker-mock";
import { parse as parseCookies } from "cookie";

import {
  specProxyProcessRequest,
  specProxyProcessResponse,
  specMakeRequestWithFallback,
} from "./index";

// Import the header constant for testing
const HEADER_SPEC_ACTIVITY = "x-atvak-activity-count";

// Workers supply `crypto` globally, but it's not present in our testing environment
// so let's fake it.
import * as crypto from "crypto";

Object.defineProperty(globalThis, "crypto", {
  value: {
    randomUUID: () => crypto.randomUUID(),
  },
});

const TEST_URL = new URL("https://specprotected.com");
const TEST_SPEC_TRAFFIC_URL = new URL(
  "https://specprotected.com/spec_traffic/somewhere",
);

describe("processing", () => {
  beforeEach(() => {
    // temporary workaround until https://github.com/nodejs/node/pull/46615#issuecomment-1604484636
    // lands in node 18.x
    delete global.location;
    delete global.performance;
    delete global.navigator;
    delete global.self;

    Object.assign(global, makeServiceWorkerEnv());

    // blow away the Promise from Fetch API
    global.fetch = (v: any) => v;
    // create a fake version of the ExtendableEvent to use,
    // also doesn't use Promises
    global.simple_event = {
      waitUntil: jest.fn((v: any) => v),
      request: new Request("https://spectrust.com", {
        method: "GET",
        headers: {
          "X-Forwarded-For": "127.0.0.1",
          Host: "spectrust.com",
        },

        // ReadableStream is difficult to mock so can't test mirroring
        // which modifies the body object...
        // we also have to set it to null because service-worker-mock uses
        // a garbage implementation of this when it's not provided and disrupts
        // our use of the tee() function
        body: null,
      }),
    };
  });

  test("disabled proxy is noop", () => {
    expect(
      specProxyProcessRequest(global.simple_event, { disableSpecProxy: true }),
    ).toBe(global.simple_event.request);

    expect(
      specProxyProcessRequest(global.simple_event, { disableSpecProxy: false }),
    ).not.toBe(global.simple_event.request);
  });

  test("traffic split", () => {
    expect(
      specProxyProcessRequest(global.simple_event, { percentageOfIPs: 0 }),
    ).toBe(global.simple_event.request);

    expect(
      specProxyProcessRequest(global.simple_event, { percentageOfIPs: 100 }),
    ).not.toBe(global.simple_event.request);
  });

  test("proxy forks request in listening mode", () => {
    expect(specProxyProcessRequest(global.simple_event)).toEqual(
      global.simple_event.request,
    );

    expect(global.simple_event.waitUntil.mock.calls.length).toEqual(1);
  });

  // Note: libraries that make the request will change the Host header
  // for us. Cloudflare does this automatically, Fastly does this through
  // the use of the `backend` configuration property on the request.
  test("in listening mode, Host header of mirrored request is unchanged", () => {
    expect(specProxyProcessRequest(global.simple_event)).toEqual(
      global.simple_event.request,
    );

    let specProxyRequest = global.simple_event.waitUntil.mock.calls[0][0];
    console.log(specProxyRequest.headers, global.simple_event.request.headers);
    expect(specProxyRequest.headers.get("host")).toEqual(
      global.simple_event.request.headers.get("host"),
    );
  });

  test("proxy replaces request in inline mode", () => {
    let request = specProxyProcessRequest(global.simple_event, {
      inlineMode: true,
    });
    expect(request).not.toEqual(global.simple_event.request);
    expect(request.url).toEqual("https://spectrust.com.spec-internal.com/");

    expect(global.simple_event.waitUntil.mock.calls.length).toEqual(0);
  });

  test("sending traffic to /spec_traffic in inline mode goes to spec proxy", () => {
    let event = global.simple_event;
    event.request.url = "https://spectrust.com/spec_traffic/somewhere";
    let request = specProxyProcessRequest(global.simple_event, {
      inlineMode: true,
    });
    expect(request.url).toEqual(
      "https://spectrust.com.spec-internal.com/spec_traffic/somewhere",
    );

    expect(global.simple_event.waitUntil.mock.calls.length).toEqual(0);
  });

  test("sending traffic to /spec_traffic in mirror mode goes to spec proxy", () => {
    let event = global.simple_event;
    event.request.url = "https://spectrust.com/spec_traffic/somewhere";
    let request = specProxyProcessRequest(global.simple_event, {
      inlineMode: false,
    });
    expect(request.url).toEqual(
      "https://spectrust.com.spec-internal.com/spec_traffic/somewhere",
    );

    expect(global.simple_event.waitUntil.mock.calls.length).toEqual(0);
  });

  test("sending traffic to /spec_traffic when it's disabled forwards traffic in mirror mode", () => {
    let event = global.simple_event;
    event.request.url = "https://spectrust.com/spec_traffic/somewhere";
    let request = specProxyProcessRequest(global.simple_event, {
      disableSpecTraffic: true,
    });
    expect(request).toEqual(global.simple_event.request);

    expect(global.simple_event.waitUntil.mock.calls.length).toEqual(1);
  });

  test("response sets SpecID cookie", () => {
    let response = new Response(null, {
      status: 200,
      headers: {
        "X-Forwarded-For": "127.0.0.1",
        Host: "spectrust.com",
      },
    });

    response = specProxyProcessResponse(
      global.simple_event.request,
      response,
      {},
    );

    let setCookie = parseCookies(response.headers.get("Set-Cookie"));

    expect(setCookie).toHaveProperty("x-spec-id");
    expect(setCookie["Max-Age"]).toEqual("320000000");

    // now assert that we don't assign a new cookie if one exists!
    response = new Response(null, {
      status: 200,
      headers: {
        "X-Forwarded-For": "127.0.0.1",
        Host: "spectrust.com",
      },
    });
    global.simple_event.request.headers.set("Cookie", "x-spec-id=abcd");

    response = specProxyProcessResponse(global.simple_event.request, response);
    let newSetCookie = response.headers.get("Set-Cookie");
    expect(newSetCookie).toBeNull();
  });

  test("response set SpecID doesn't erase other set-cookies", () => {
    let response = new Response(null, {
      status: 200,
      headers: {
        "X-Forwarded-For": "127.0.0.1",
        Host: "spectrust.com",
        "Set-Cookie": "a=cookie",
      },
    });

    response = specProxyProcessResponse(
      global.simple_event.request,
      response,
      {},
    );

    expect(response.headers.get("Set-Cookie").split(",")).toHaveLength(2);
  });

  test("response doesn't set SpecID in inline mode", () => {
    let response = new Response(null, {
      status: 200,
      headers: {
        "X-Forwarded-For": "127.0.0.1",
        Host: "spectrust.com",
      },
    });

    response = specProxyProcessResponse(global.simple_event.request, response, {
      inlineMode: true,
    });

    expect(response.headers.get("Set-Cookie")).toBeNull();
  });

  test("domain override affects specUrl hostname", () => {
    // Test without domain override
    let request = specProxyProcessRequest(global.simple_event, {
      inlineMode: true,
    });
    expect(request.url).toEqual("https://spectrust.com.spec-internal.com/");
    expect(request.headers.get("x-spec-forward-origin")).toEqual(
      "spectrust.com",
    );

    // Test with domain override
    request = specProxyProcessRequest(global.simple_event, {
      inlineMode: true,
      domainOverride: "staging",
    });
    expect(request.url).toEqual("https://staging.spec-internal.com/");
    expect(request.headers.get("x-spec-forward-origin")).toEqual(
      "spectrust.com",
    );

    // Test with empty domain override (should be same as no override)
    request = specProxyProcessRequest(global.simple_event, {
      inlineMode: true,
      domainOverride: "",
    });
    expect(request.url).toEqual("https://spectrust.com.spec-internal.com/");
    expect(request.headers.get("x-spec-forward-origin")).toEqual(
      "spectrust.com",
    );

    // Test with path and query parameters
    global.simple_event.request.url = "https://spectrust.com/path?query=true";
    request = specProxyProcessRequest(global.simple_event, {
      inlineMode: true,
      domainOverride: "prod",
    });
    expect(request.url).toEqual(
      "https://prod.spec-internal.com/path?query=true",
    );
    expect(request.headers.get("x-spec-forward-origin")).toEqual(
      "spectrust.com",
    );
  });
});

describe("specMakeRequestWithFallback", () => {
  let mockFetch: jest.Mock;
  let mockEvent: any;
  let mockRequest: Request;

  beforeEach(() => {
    // Mock fetch globally
    mockFetch = jest.fn();
    global.fetch = mockFetch;

    // Create mock request
    mockRequest = new Request("https://spectrust.com.spec-internal.com/test", {
      method: "GET",
      headers: {
        "x-spec-forward-origin": "spectrust.com",
      },
    });

    // Create mock event
    mockEvent = {
      request: new Request("https://spectrust.com/test", {
        method: "GET",
        headers: {
          Host: "spectrust.com",
        },
      }),
    };
  });

  test("in inline mode, returns response with header removed when activity header is present", async () => {
    const mockResponse = new Response("success", {
      status: 200,
      headers: {
        [HEADER_SPEC_ACTIVITY]: "1",
        "content-type": "text/plain",
      },
    });

    mockFetch.mockResolvedValueOnce(mockResponse);

    const config = { inlineMode: true };
    const result = await specMakeRequestWithFallback(
      mockEvent,
      mockRequest,
      config,
    );

    expect(mockFetch).toHaveBeenCalledTimes(1);
    expect(mockFetch).toHaveBeenCalledWith(mockRequest);
    expect(result.headers.get(HEADER_SPEC_ACTIVITY)).toBeNull();
    expect(result.headers.get("content-type")).toBe("text/plain");
    expect(await result.text()).toBe("success");
  });

  test("in inline mode, falls back to original request when activity header is missing", async () => {
    const mockSpecProxyResponse = new Response("spec proxy response", {
      status: 200,
      headers: {
        "content-type": "text/plain",
      },
    });

    const mockFallbackResponse = new Response("fallback response", {
      status: 200,
      headers: {
        "content-type": "text/html",
      },
    });

    // First call to spec proxy (no activity header)
    mockFetch.mockResolvedValueOnce(mockSpecProxyResponse);
    // Second call to fallback (original request)
    mockFetch.mockResolvedValueOnce(mockFallbackResponse);

    const config = { inlineMode: true };
    const result = await specMakeRequestWithFallback(
      mockEvent,
      mockRequest,
      config,
    );

    expect(mockFetch).toHaveBeenCalledTimes(2);
    expect(mockFetch).toHaveBeenNthCalledWith(1, mockRequest);
    expect(mockFetch).toHaveBeenNthCalledWith(2, mockEvent.request);
    expect(await result.text()).toBe("fallback response");
    expect(result.headers.get("content-type")).toBe("text/html");
  });

  test("in mirror mode, always uses the provided request without fallback", async () => {
    const mockResponse = new Response("mirror response", {
      status: 200,
      headers: {
        "content-type": "application/json",
      },
    });

    mockFetch.mockResolvedValueOnce(mockResponse);

    const config = { inlineMode: false };
    const result = await specMakeRequestWithFallback(
      mockEvent,
      mockRequest,
      config,
    );

    expect(mockFetch).toHaveBeenCalledTimes(1);
    expect(mockFetch).toHaveBeenCalledWith(mockRequest);
    expect(await result.text()).toBe("mirror response");
    expect(result.headers.get("content-type")).toBe("application/json");
  });

  test("in mirror mode, does not fallback even when activity header is missing", async () => {
    const mockResponse = new Response("mirror response no header", {
      status: 200,
      headers: {
        "content-type": "application/json",
      },
    });

    mockFetch.mockResolvedValueOnce(mockResponse);

    const config = { inlineMode: false };
    const result = await specMakeRequestWithFallback(
      mockEvent,
      mockRequest,
      config,
    );

    expect(mockFetch).toHaveBeenCalledTimes(1);
    expect(mockFetch).toHaveBeenCalledWith(mockRequest);
    expect(await result.text()).toBe("mirror response no header");
  });
});

//---------------------------------------------
// Private Function Testing
//---------------------------------------------

// "rewire" the library module to allow us to reach private functions
let spec = rewire("./dist/index.js");
// setup the jest global mock
spec.__set__({ self: global.self });

describe("shouldHandleRequest", () => {
  let shouldHandleRequest = spec.__get__("shouldHandleRequest");

  test.each([
    ["0.0.0.0", undefined, true],
    ["0.0.0.40", undefined, true],
    ["0.0.0.99", undefined, true],
    ["0.0.0.100", undefined, true], // "wraps" to 0
    ["0.0.0.0", 100, true],
    ["0.0.0.40", 100, true],
    ["0.0.0.99", 100, true],
    ["0.0.0.0", 100, true],
    ["0.0.0.40", 50, true],
    ["0.0.0.49", 50, true],
    ["0.0.0.50", 50, false],
    ["0.0.0.60", 50, false],
    ["0.0.0.0", 0, false],
    ["0.0.0.40", 0, false],
    ["0.0.0.99", 0, false],
    ["0.0.0.100", 0, false],
    // some legit-looking IPs that don't exploit our algorithm for testing
    ["24.68.195.11", 50, false],
    ["74.232.255.255", 50, true],
    ["89.2.79.2", 50, false],
    ["67.67.67.67", 50, false],
    ["10.0.0.8", 50, true],
  ])("traffic split [ip=%s, split=%i]", (ip, split, expected) => {
    let headers = new Headers({ "x-forwarded-for": ip });
    expect(
      shouldHandleRequest(TEST_URL, headers, {
        percentageOfIPs: split,
      }),
    ).toEqual(expected);
  });

  test.each([
    ["0.0.0.100", undefined, true], // "wraps" to 0
    ["0.0.0.99", 100, true],
    ["0.0.0.0", 100, true],
    ["0.0.0.49", 50, true],
    ["0.0.0.50", 50, true],
    ["0.0.0.60", 50, true],
    ["0.0.0.0", 0, true],
    ["0.0.0.40", 0, true],
    ["0.0.0.99", 0, true],
    ["0.0.0.100", 0, true],
    // some legit-looking IPs that don't exploit our algorithm for testing
    ["24.68.195.11", 50, true],
    ["74.232.255.255", 50, true],
    ["89.2.79.2", 50, true],
    ["67.67.67.67", 50, true],
    ["10.0.0.8", 50, true],
  ])(
    "traffic split [ip=%s, split=%i], always send to /spec_traffic when enabled",
    (ip, split, expected) => {
      let headers = new Headers({ "x-forwarded-for": ip });
      expect(
        shouldHandleRequest(TEST_SPEC_TRAFFIC_URL, headers, {
          percentageOfIPs: split,
        }),
      ).toEqual(expected);
    },
  );

  test.each([
    ["0.0.0.100", undefined, true], // "wraps" to 0
    ["0.0.0.99", 100, true],
    ["0.0.0.0", 100, true],
    ["0.0.0.49", 50, true],
    ["0.0.0.50", 50, false],
    ["0.0.0.60", 50, false],
    ["0.0.0.0", 0, false],
    ["0.0.0.40", 0, false],
    ["0.0.0.99", 0, false],
    ["0.0.0.100", 0, false],
    // some legit-looking IPs that don't exploit our algorithm for testing
    ["24.68.195.11", 50, false],
    ["74.232.255.255", 50, true],
    ["89.2.79.2", 50, false],
    ["67.67.67.67", 50, false],
    ["10.0.0.8", 50, true],
  ])(
    "traffic split [ip=%s, split=%i], normal behavior when /spec_traffic disabled",
    (ip, split, expected) => {
      let headers = new Headers({ "x-forwarded-for": ip });
      expect(
        shouldHandleRequest(TEST_SPEC_TRAFFIC_URL, headers, {
          percentageOfIPs: split,
          disableSpecTraffic: true,
        }),
      ).toEqual(expected);
    },
  );

  test.each([
    ["X-Forwarded-For", 50, "0.0.0.0", true],
    ["X-Forwarded-For", 50, "0.0.0.60", false],
    ["Not-Checked-Header", 50, "0.0.0.0", false],
    ["Not-Checked-Header", 50, "0.0.0.99", false],
    // if it's 100% split, we expect to handle the request even if we can't find the
    // IP address on this header
    ["X-Forwarded-For", 100, "0.0.0.60", true],
    ["Not-Checked-Header", 100, "0.0.0.99", true],
  ])(
    'ip from header "%s" with %i%% split for ip %s"',
    (header, split, ip, expected) => {
      let headers = new Headers({});
      headers.set(header, ip);
      expect(
        shouldHandleRequest(TEST_URL, headers, {
          percentageOfIPs: split,
        }),
      ).toEqual(expected);
    },
  );

  // Check NaN handling, just in case.
  test.each([
    ["X-Forwarded-For", 50, "wtf", false],
    ["X-Forwarded-For", 99, "srs", false],
    // if it's 100% split, we expect to handle the request even if we can't find the
    // IP address on this header
    ["X-Forwarded-For", 100, "susfam", true],
    ["X-Forwarded-For", 100, ":ghost-shrug:", true],
  ])(
    'ip from header "%s" with %i%% split for ip %s"',
    (header, split, ip, expected) => {
      let headers = new Headers({});
      headers.set(header, ip);
      expect(
        shouldHandleRequest(TEST_URL, headers, {
          percentageOfIPs: split,
        }),
      ).toEqual(expected);
    },
  );
});

describe("extractTopLevelDomain", () => {
  let extractTopLevelDomain = spec.__get__("extractTopLevelDomain");

  test.each([
    ["spec-trust.com", "spec-trust.com"],
    ["somewhere.spec-trust.com", "spec-trust.com"],
    ["somewhere.spec-trust.com.spec-internal.com", "spec-trust.com"],
    ["somewhere.spec-trust.edu", "spec-trust.edu"],
    ["somewhere.spec-trust.edu.spec-internal.com", "spec-trust.edu"],
    ["somewhere.spec-trust.gov", "spec-trust.gov"],
    ["somewhere.spec-trust.gov.spec-internal.com", "spec-trust.gov"],
    ["somewhere.spec-trust.int", "spec-trust.int"],
    ["somewhere.spec-trust.int.spec-internal.com", "spec-trust.int"],
    ["somewhere.spec-trust.net", "spec-trust.net"],
    ["somewhere.spec-trust.net.spec-internal.com", "spec-trust.net"],
    ["somewhere.spec-trust.org", "spec-trust.org"],
    ["somewhere.spec-trust.org.spec-internal.com", "spec-trust.org"],
    ["somewhere.spec-trust.uk", "spec-trust.uk"],
    ["somewhere.spec-trust.uk.spec-internal.com", "spec-trust.uk"],
    ["many.many.many.many.sub-domains.spec-trust.com", "spec-trust.com"],
    ["any.suffixes", "any.suffixes"],
    // if we don't match any known top-level domain suffix, we just get the whole thing
    [
      "subdomains.irrelevant.any.suffixes",
      "subdomains.irrelevant.any.suffixes",
    ],
    ["invalid-top-level-though", "invalid-top-level-though"],
    // invalid top-level domain suffix, but with spec-internal on the end winds up
    // matching to spec-internal
    ["somewhere.fun.to.go.spec-internal.com", "spec-internal.com"],
    ["staging.special.somebody.co.uk.spec-internal.com", "somebody.co.uk"],
    ["staging.somebody.ac.uk.spec-internal.com", "somebody.ac.uk"],
    ["staging.somebody.co.uk.spec-internal.com", "somebody.co.uk"],
    ["staging.somebody.gov.uk.spec-internal.com", "somebody.gov.uk"],
    ["staging.somebody.ltd.uk.spec-internal.com", "somebody.ltd.uk"],
    ["staging.somebody.me.uk.spec-internal.com", "somebody.me.uk"],
    ["staging.somebody.net.uk.spec-internal.com", "somebody.net.uk"],
    ["staging.somebody.nhs.uk.spec-internal.com", "somebody.nhs.uk"],
    ["staging.somebody.plc.uk.spec-internal.com", "somebody.plc.uk"],
    ["staging.somebody.police.uk.spec-internal.com", "somebody.police.uk"],
  ])("extract top level domain from %s", (host, expected) => {
    let headers = new Headers({ Host: host });
    expect(extractTopLevelDomain(headers)).toEqual(expected);
  });
});