src/shared/security/forceReferrer.ts

/**
 * IDB is used to persist rather than using a `Map`, because SWs are suicidal
 */

type SiteDirective = "same-origin" | "same-site" | "cross-site" | "none";
interface RedirectTracker {
	originalReferrer: string;
	mostRestrictiveSite: SiteDirective;
	referrerPolicy: string;
	chainStarted: number;
}
interface ReferrerPolicyData {
	policy: string;
	referrer: string;
}

// Persist the redirect trackers for an hour
const TRACKER_EXPIRY = 60 * 60 * 1000;
const SITE_HIERARCHY: Record<SiteDirective, number> = {
	none: 0,
	"same-origin": 1,
	"same-site": 2,
	"cross-site": 3,
};

/**
 * Gets a connection to the IndexedDB database
 *
 * @returns Promise that resolves to the database connection
 */
async function getDB(): Promise<IDBDatabase> {
	const request = indexedDB.open("$scramjet", 1);

	return new Promise((resolve, reject) => {
		request.onerror = () => reject(request.error);
		request.onsuccess = () => resolve(request.result);
	});
}

/**
 * Retrieves a redirect tracker for a given URL
 *
 * @param url The URL to look up
 * @returns Redirect tracker if found, or `null`
 */
async function getTracker(url: string): Promise<RedirectTracker | null> {
	const db = await getDB();
	const tx = db.transaction("redirectTrackers", "readonly");
	const store = tx.objectStore("redirectTrackers");

	return new Promise((resolve) => {
		const request = store.get(url);
		request.onsuccess = () => resolve(request.result || null);
		request.onerror = () => resolve(null);
	});
}

/**
 * Store or update a redirect tracker for a given URL
 *
 * @param url URL to store the tracker for
 * @param tracker Redirect tracker data to store
 */
async function setTracker(
	url: string,
	tracker: RedirectTracker
): Promise<void> {
	const db = await getDB();
	const tx = db.transaction("redirectTrackers", "readwrite");
	const store = tx.objectStore("redirectTrackers");

	return new Promise((resolve, reject) => {
		const request = store.put(tracker, url);
		request.onsuccess = () => resolve();
		request.onerror = () => reject(request.error);
	});
}

/**
 * Delete a redirect tracker for a given URL
 *
 * @param url URL whose tracker should be deleted
 */
async function deleteTracker(url: string): Promise<void> {
	const db = await getDB();
	const tx = db.transaction("redirectTrackers", "readwrite");
	const store = tx.objectStore("redirectTrackers");

	return new Promise((resolve, reject) => {
		const request = store.delete(url);
		request.onsuccess = () => resolve();
		request.onerror = () => reject(request.error);
	});
}

/**
 * Initialize tracking for a new request that might redirect
 *
 * @param requestUrl URL of the request being made
 * @param referrer Referrer URL of the request, or `null`
 * @param initialSite Initial Sec-Fetch-Site directive
 */
export async function initializeTracker(
	requestUrl: string,
	referrer: string | null,
	initialSite: string
): Promise<void> {
	const existing = await getTracker(requestUrl);
	if (existing) {
		return;
	}

	await setTracker(requestUrl, {
		originalReferrer: referrer || "",
		mostRestrictiveSite: initialSite as SiteDirective,
		referrerPolicy: "",
		chainStarted: Date.now(),
	});
}

/**
 * Update tracker when a redirect is encountered
 *
 * @param originalUrl URL that is redirecting
 * @param redirectUrl URL being redirected to
 * @param newReferrerPolicy Referrer Policy from the redirect response
 */
export async function updateTracker(
	originalUrl: string,
	redirectUrl: string,
	newReferrerPolicy?: string
): Promise<void> {
	const tracker = await getTracker(originalUrl);
	if (!tracker) return;

	await deleteTracker(originalUrl);
	if (newReferrerPolicy) {
		tracker.referrerPolicy = newReferrerPolicy;
	}
	await setTracker(redirectUrl, tracker);
}

/**
 * Get most restrictive site value for a request
 *
 * @param requestUrl The URL of the current request
 * @param currentSite The current `Sec-Fetch-Site` directive for this request
 * @returns Most restrictive `Sec-Fetch-Site` directive from the redirect chain
 */
export async function getMostRestrictiveSite(
	requestUrl: string,
	currentSite: string
): Promise<string> {
	const tracker = await getTracker(requestUrl);
	if (!tracker) return currentSite;

	const trackedValue = SITE_HIERARCHY[tracker.mostRestrictiveSite];
	const currentValue = SITE_HIERARCHY[currentSite as SiteDirective] ?? 0;

	if (currentValue > trackedValue) {
		tracker.mostRestrictiveSite = currentSite as SiteDirective;
		await setTracker(requestUrl, tracker);

		return currentSite;
	}

	return tracker.mostRestrictiveSite;
}

/**
 * Clean up tracker after request completes
 * @param requestUrl URL of the completed request
 */
export async function cleanTracker(requestUrl: string): Promise<void> {
	await deleteTracker(requestUrl);
}

/**
 * Clean up expired trackers
 */
export async function cleanExpiredTrackers(): Promise<void> {
	const now = Date.now();
	const db = await getDB();
	const tx = db.transaction("redirectTrackers", "readwrite");
	const store = tx.objectStore("redirectTrackers");

	const request = store.openCursor();

	request.onsuccess = (event) => {
		const cursor = (event.target as IDBRequest).result;
		if (cursor) {
			const tracker = cursor.value as RedirectTracker;
			if (now - tracker.chainStarted > TRACKER_EXPIRY) {
				cursor.delete();
			}
			cursor.continue();
		}
	};
}

/**
 * Store referrer policy for a URL
 *
 * @param url URL to store the policy for
 * @param policy Referrer policy to store
 * @param referrer The referrer URL that set this policy
 */
export async function storeReferrerPolicy(
	url: string,
	policy: string,
	referrer: string
): Promise<void> {
	const db = await getDB();
	const tx = db.transaction("referrerPolicies", "readwrite");
	const store = tx.objectStore("referrerPolicies");

	const data: ReferrerPolicyData = { policy, referrer };

	return new Promise((resolve, reject) => {
		const request = store.put(data, url);
		request.onsuccess = () => resolve();
		request.onerror = () => reject(request.error);
	});
}

/**
 * Get referrer policy data for a URL
 *
 * @param url URL to get the policy for
 * @returns Referrer policy data if found, or `null`
 */
export async function getReferrerPolicy(
	url: string
): Promise<ReferrerPolicyData | null> {
	const db = await getDB();
	const tx = db.transaction("referrerPolicies", "readonly");
	const store = tx.objectStore("referrerPolicies");

	return new Promise((resolve) => {
		const request = store.get(url);
		request.onsuccess = () => resolve(request.result || null);
		request.onerror = () => resolve(null);
	});
}