CookieBridge/src/relay/store.ts

import sodium from "sodium-native";
import type { EncryptedCookieBlob } from "../protocol/spec.js";
import { MAX_STORED_COOKIES_PER_DEVICE } from "../protocol/spec.js";

type CookieKey = string; // "domain|name|path"

function blobKey(blob: Pick<EncryptedCookieBlob, "domain" | "cookieName" | "path">): CookieKey {
  return `${blob.domain}|${blob.cookieName}|${blob.path}`;
}

/**
 * In-memory encrypted cookie storage.
 * The server stores opaque ciphertext — it cannot read cookie values.
 * Keyed by (deviceId, domain, cookieName, path).
 */
export class CookieBlobStore {
  // deviceId -> (cookieKey -> blob)
  private store = new Map<string, Map<CookieKey, EncryptedCookieBlob>>();

  /** Upsert an encrypted cookie blob. LWW by lamportTs. */
  upsert(blob: Omit<EncryptedCookieBlob, "id" | "updatedAt">): EncryptedCookieBlob {
    let deviceMap = this.store.get(blob.deviceId);
    if (!deviceMap) {
      deviceMap = new Map();
      this.store.set(blob.deviceId, deviceMap);
    }

    const key = blobKey(blob);
    const existing = deviceMap.get(key);

    if (existing && existing.lamportTs >= blob.lamportTs) {
      return existing; // stale update
    }

    const stored: EncryptedCookieBlob = {
      ...blob,
      id: existing?.id ?? generateId(),
      updatedAt: new Date().toISOString(),
    };

    deviceMap.set(key, stored);

    // Enforce per-device limit
    if (deviceMap.size > MAX_STORED_COOKIES_PER_DEVICE) {
      // Remove oldest by updatedAt
      let oldest: { key: CookieKey; time: string } | null = null;
      for (const [k, v] of deviceMap) {
        if (!oldest || v.updatedAt < oldest.time) {
          oldest = { key: k, time: v.updatedAt };
        }
      }
      if (oldest) deviceMap.delete(oldest.key);
    }

    return stored;
  }

  /** Delete a cookie blob. */
  delete(deviceId: string, domain: string, cookieName: string, path: string): boolean {
    const deviceMap = this.store.get(deviceId);
    if (!deviceMap) return false;
    return deviceMap.delete(blobKey({ domain, cookieName, path }));
  }

  /** Get all blobs for a device, optionally filtered by domain. */
  getByDevice(deviceId: string, domain?: string): EncryptedCookieBlob[] {
    const deviceMap = this.store.get(deviceId);
    if (!deviceMap) return [];
    const blobs = Array.from(deviceMap.values());
    if (domain) return blobs.filter((b) => b.domain === domain);
    return blobs;
  }

  /** Get blobs across all paired devices for a set of deviceIds. */
  getByDevices(deviceIds: string[], domain?: string): EncryptedCookieBlob[] {
    const result: EncryptedCookieBlob[] = [];
    for (const id of deviceIds) {
      result.push(...this.getByDevice(id, domain));
    }
    return result;
  }

  /** Get all stored blobs across all devices. */
  getAll(): EncryptedCookieBlob[] {
    const result: EncryptedCookieBlob[] = [];
    for (const deviceMap of this.store.values()) {
      result.push(...deviceMap.values());
    }
    return result;
  }

  /** Get a single blob by its ID. */
  getById(id: string): EncryptedCookieBlob | null {
    for (const deviceMap of this.store.values()) {
      for (const blob of deviceMap.values()) {
        if (blob.id === id) return blob;
      }
    }
    return null;
  }

  /** Delete a blob by its ID. Returns true if found and deleted. */
  deleteById(id: string): boolean {
    for (const deviceMap of this.store.values()) {
      for (const [key, blob] of deviceMap) {
        if (blob.id === id) {
          deviceMap.delete(key);
          return true;
        }
      }
    }
    return false;
  }

  /** Get all blobs updated after a given timestamp (for polling). */
  getUpdatedSince(deviceIds: string[], since: string): EncryptedCookieBlob[] {
    const result: EncryptedCookieBlob[] = [];
    for (const id of deviceIds) {
      const deviceMap = this.store.get(id);
      if (!deviceMap) continue;
      for (const blob of deviceMap.values()) {
        if (blob.updatedAt > since) {
          result.push(blob);
        }
      }
    }
    return result;
  }
}

function generateId(): string {
  const buf = Buffer.alloc(16);
  sodium.randombytes_buf(buf);
  return buf.toString("hex");
}