@linxiraos/pi-ai 1.0.0

package/src/auth-broker/snapshot-cache.ts

@@ -0,0 +1,200 @@
+/**
+ * AES-GCM encrypted local cache for auth-broker snapshots.
+ *
+ * The cache is defense-in-depth for at-rest snapshots: a copied cache file is
+ * useless without the matching broker bearer token and URL. The token itself is
+ * still the trust boundary; a process that can read both the token and this file
+ * can decrypt the snapshot.
+ */
+import * as fs from "node:fs/promises";
+import * as path from "node:path";
+import { isEnoent, logger } from "@linxiraos/pi-utils";
+import type { SnapshotResponse } from "./types";
+
+const MAGIC = new Uint8Array([0x4f, 0x4d, 0x50, 0x53]); // "OMPS"
+const VERSION = 2;
+const VERSION_OFFSET = MAGIC.byteLength;
+const IV_OFFSET = VERSION_OFFSET + 1;
+const IV_LENGTH = 12;
+const HEADER_LENGTH = IV_OFFSET + IV_LENGTH;
+const AES_ALGORITHM = "AES-GCM";
+const TEXT_ENCODER = new TextEncoder();
+const TEXT_DECODER = new TextDecoder();
+const HEX = "0123456789abcdef";
+
+export interface ReadAuthBrokerSnapshotCacheOptions {
+	path: string;
+	token: string;
+	url: string;
+	ttlMs: number;
+	/** Override clock for deterministic tests. */
+	now?: () => number;
+}
+
+export interface WriteAuthBrokerSnapshotCacheOptions {
+	path: string;
+	token: string;
+	url: string;
+	snapshot: SnapshotResponse;
+}
+
+/**
+ * Cheap structural guard for a decrypted cache payload. The bytes are already
+ * AES-256-GCM authenticated, so this only rejects shape/version drift (a cache
+ * written by a different omp build, or a buggy write) — not tampering. A
+ * mismatch returns null so the caller refetches a fresh snapshot.
+ */
+function isSnapshotResponseShape(v: unknown): v is SnapshotResponse {
+	if (typeof v !== "object" || v === null) return false;
+	const o = v as Record<string, unknown>;
+	return (
+		typeof o.generation === "number" &&
+		typeof o.generatedAt === "number" &&
+		typeof o.serverNowMs === "number" &&
+		typeof o.refresher === "object" &&
+		o.refresher !== null &&
+		Array.isArray(o.credentials)
+	);
+}
+
+export async function readAuthBrokerSnapshotCache(
+	opts: ReadAuthBrokerSnapshotCacheOptions,
+): Promise<SnapshotResponse | null> {
+	if (opts.ttlMs <= 0) return null;
+	let data: Uint8Array;
+	try {
+		data = await fs.readFile(opts.path);
+	} catch (error) {
+		if (isEnoent(error)) return null;
+		throw error;
+	}
+
+	try {
+		const plaintext = await decryptCachePayload(data, opts.token, opts.url);
+		if (!plaintext) return null;
+		const parsed: unknown = JSON.parse(TEXT_DECODER.decode(plaintext));
+		if (!isSnapshotResponseShape(parsed)) {
+			logger.debug("auth-broker snapshot cache schema invalid", { path: opts.path });
+			return null;
+		}
+		const snapshot = parsed;
+		const now = opts.now?.() ?? Date.now();
+		if (now - snapshot.generatedAt > opts.ttlMs) return null;
+		return snapshot;
+	} catch (error) {
+		logger.debug("auth-broker snapshot cache read failed", { path: opts.path, error: String(error) });
+		return null;
+	}
+}
+
+export async function writeAuthBrokerSnapshotCache(opts: WriteAuthBrokerSnapshotCacheOptions): Promise<void> {
+	const payload = await encryptCachePayload(opts.snapshot, opts.token, opts.url);
+	await fs.mkdir(path.dirname(opts.path), { recursive: true });
+	const tmpPath = `${opts.path}.${process.pid}.${randomHex(8)}.tmp`;
+	let removeTemp = false;
+	try {
+		const handle = await fs.open(tmpPath, "wx", 0o600);
+		removeTemp = true;
+		try {
+			await handle.writeFile(payload);
+		} finally {
+			await handle.close();
+		}
+		await fs.chmod(tmpPath, 0o600);
+		await fs.rename(tmpPath, opts.path);
+		removeTemp = false;
+	} finally {
+		if (removeTemp) await fs.rm(tmpPath, { force: true }).catch(() => {});
+	}
+}
+
+async function encryptCachePayload(snapshot: SnapshotResponse, token: string, url: string): Promise<Uint8Array> {
+	const key = await deriveAesKey(token, ["encrypt"]);
+	const iv = new Uint8Array(IV_LENGTH);
+	globalThis.crypto.getRandomValues(iv);
+	const plaintext = TEXT_ENCODER.encode(JSON.stringify(snapshot));
+	const ciphertext = new Uint8Array(
+		await globalThis.crypto.subtle.encrypt(
+			{
+				name: AES_ALGORITHM,
+				iv,
+				additionalData: cacheAdditionalData(url),
+			},
+			key,
+			plaintext,
+		),
+	);
+	const payload = new Uint8Array(HEADER_LENGTH + ciphertext.byteLength);
+	payload.set(MAGIC, 0);
+	payload[VERSION_OFFSET] = VERSION;
+	payload.set(iv, IV_OFFSET);
+	payload.set(ciphertext, HEADER_LENGTH);
+	return payload;
+}
+
+async function decryptCachePayload(data: Uint8Array, token: string, url: string): Promise<Uint8Array | null> {
+	if (data.byteLength <= HEADER_LENGTH) {
+		logger.debug("auth-broker snapshot cache file too short");
+		return null;
+	}
+	for (let i = 0; i < MAGIC.byteLength; i++) {
+		if (data[i] !== MAGIC[i]) {
+			logger.debug("auth-broker snapshot cache magic mismatch");
+			return null;
+		}
+	}
+	if (data[VERSION_OFFSET] !== VERSION) {
+		logger.debug("auth-broker snapshot cache version mismatch", { version: data[VERSION_OFFSET] });
+		return null;
+	}
+	const key = await deriveAesKey(token, ["decrypt"]);
+	const iv = asStrict(data.subarray(IV_OFFSET, HEADER_LENGTH));
+	const ciphertext = asStrict(data.subarray(HEADER_LENGTH));
+	try {
+		return new Uint8Array(
+			await globalThis.crypto.subtle.decrypt(
+				{
+					name: AES_ALGORITHM,
+					iv,
+					additionalData: cacheAdditionalData(url),
+				},
+				key,
+				ciphertext,
+			),
+		);
+	} catch (error) {
+		logger.debug("auth-broker snapshot cache decrypt failed", { error: String(error) });
+		return null;
+	}
+}
+
+function cacheAdditionalData(url: string): Uint8Array<ArrayBuffer> {
+	const urlBytes = TEXT_ENCODER.encode(url);
+	const additionalData = new Uint8Array(IV_OFFSET + urlBytes.byteLength);
+	additionalData.set(MAGIC, 0);
+	additionalData[VERSION_OFFSET] = VERSION;
+	additionalData.set(urlBytes, IV_OFFSET);
+	return additionalData;
+}
+
+async function deriveAesKey(token: string, usages: Array<"encrypt" | "decrypt">): Promise<CryptoKey> {
+	const digest = await globalThis.crypto.subtle.digest("SHA-256", TEXT_ENCODER.encode(token));
+	return globalThis.crypto.subtle.importKey("raw", digest, AES_ALGORITHM, false, usages);
+}
+
+function asStrict(bytes: Uint8Array): Uint8Array<ArrayBuffer> {
+	if (bytes.buffer instanceof ArrayBuffer && bytes.byteOffset === 0 && bytes.byteLength === bytes.buffer.byteLength) {
+		return bytes as Uint8Array<ArrayBuffer>;
+	}
+	const copy = new Uint8Array(bytes.byteLength);
+	copy.set(bytes);
+	return copy;
+}
+
+function randomHex(byteLength: number): string {
+	const bytes = new Uint8Array(byteLength);
+	globalThis.crypto.getRandomValues(bytes);
+	let out = "";
+	for (const byte of bytes) out += HEX[byte >> 4] + HEX[byte & 15];
+	return out;
+}

package/src/auth-broker/types.ts

@@ -0,0 +1,193 @@
+/**
+ * Wire types shared between the auth-broker server and clients.
+ *
+ * The broker holds OAuth refresh tokens and exposes a redacted snapshot;
+ * clients use `access` tokens directly and call back to the broker when a
+ * credential expires or a 401 surfaces on a supposedly-fresh credential.
+ */
+
+import type {
+	AuthCredential,
+	AuthCredentialSnapshot,
+	AuthCredentialSnapshotEntry,
+	DisabledCredentialSummary,
+	StoredCredentialBlock,
+} from "../auth-storage";
+import type { ClientUsageClientSummary, ClientUsageReport, UsageHistoryEntry, UsageReport } from "../usage";
+
+/** GET /v1/healthz response body. */
+export interface HealthzResponse {
+	ok: boolean;
+	version?: string;
+}
+
+export interface RefresherSchedule {
+	enabled: boolean;
+	intervalMs: number;
+	skewMs: number;
+	nextSweepInMs: number;
+}
+
+export type CredentialBlockSnapshot = Omit<StoredCredentialBlock, "credentialId">;
+
+export type SnapshotEntry = AuthCredentialSnapshotEntry & {
+	rotatesInMs: number | null;
+	blocks?: CredentialBlockSnapshot[];
+};
+
+/** GET /v1/snapshot response body. */
+export interface SnapshotResponse extends Omit<AuthCredentialSnapshot, "credentials"> {
+	serverNowMs: number;
+	refresher: RefresherSchedule;
+	credentials: SnapshotEntry[];
+}
+
+/** GET /v1/usage response body — matches the local `AuthStorage.fetchUsageReports` shape. */
+export interface UsageResponse {
+	generatedAt: number;
+	reports: UsageReport[];
+}
+
+/**
+ * GET /v1/usage/history response body. Entries come from the broker host's
+ * durable `usage_history` — the broker performs every upstream usage fetch in
+ * broker deployments, so this is the only complete utilization record.
+ */
+export interface UsageHistoryResponse {
+	generatedAt: number;
+	entries: UsageHistoryEntry[];
+}
+
+/** POST /v1/usage/observed request body — one client's batched observed usage. */
+export type ClientUsageReportRequest = ClientUsageReport;
+
+/** POST /v1/usage/observed response body. */
+export interface ClientUsageReportResponse {
+	ok: boolean;
+}
+
+/** GET /v1/usage/clients response body — per-client token burn aggregates. */
+export interface ClientUsageSummaryResponse {
+	generatedAt: number;
+	clients: ClientUsageClientSummary[];
+}
+
+/** POST /v1/credential/:id/refresh response body. */
+export interface CredentialRefreshResponse {
+	entry: AuthCredentialSnapshotEntry;
+}
+
+/** POST /v1/credential/:id/disable request body. */
+export interface CredentialDisableRequest {
+	cause: string;
+}
+
+/** POST /v1/credential/:id/disable response body. */
+export interface CredentialDisableResponse {
+	ok: boolean;
+}
+
+/** GET /v1/credentials/disabled response body — tombstones of auto-disabled rows. */
+export interface DisabledCredentialsResponse {
+	generatedAt: number;
+	disabled: DisabledCredentialSummary[];
+}
+
+/** POST /v1/credential/:id/block request body. */
+export type CredentialBlockRequest = CredentialBlockSnapshot;
+
+/** POST /v1/credential/:id/block response body. */
+export interface CredentialBlockResponse {
+	ok: boolean;
+}
+
+/** DELETE /v1/credential/:id/blocks response body. */
+export interface CredentialBlocksDeleteResponse {
+	ok: boolean;
+}
+
+/** POST /v1/usage/stale response body. */
+export interface UsageStaleResponse {
+	ok: boolean;
+}
+
+/**
+ * POST /v1/credential request body. The OAuth `refresh` must be the *real*
+ * refresh token (not the sentinel) — the broker is the canonical writer.
+ */
+export interface CredentialUploadRequest {
+	provider: string;
+	credential: AuthCredential;
+}
+
+/** POST /v1/credential response body — redacted snapshot of the provider's rows after upsert. */
+export interface CredentialUploadResponse {
+	entries: AuthCredentialSnapshotEntry[];
+}
+
+/**
+ * SSE event kinds emitted on `GET /v1/snapshot/stream`. The same value is set
+ * as the SSE `event:` name (load-bearing for clients) **and** embedded as a
+ * `kind` field inside the JSON body so a Zod discriminated union can validate
+ * the payload without consulting the line metadata.
+ */
+export type SnapshotStreamEventKind = "snapshot" | "entry" | "removed";
+
+/** Initial frame emitted on connect — the full {@link SnapshotResponse}. */
+export interface SnapshotStreamSnapshotEvent extends SnapshotResponse {
+	kind: "snapshot";
+}
+
+/** Single credential added/changed (upsert or refresh). */
+export interface SnapshotStreamEntryEvent {
+	kind: "entry";
+	generation: number;
+	serverNowMs: number;
+	refresher: RefresherSchedule;
+	entry: SnapshotEntry;
+}
+
+/** Single credential disabled/deleted. */
+export interface SnapshotStreamRemovedEvent {
+	kind: "removed";
+	generation: number;
+	serverNowMs: number;
+	refresher: RefresherSchedule;
+	id: number;
+}
+
+/** Discriminated union of every event the snapshot stream emits. */
+export type SnapshotStreamEvent = SnapshotStreamSnapshotEvent | SnapshotStreamEntryEvent | SnapshotStreamRemovedEvent;
+
+/**
+ * Default bearer-protected route prefix. The broker exposes `/v1/healthz`
+ * unauthenticated for liveness probes; everything else requires a bearer.
+ */
+export const AUTH_BROKER_API_PREFIX = "/v1";
+
+/** Request header used by clients to advertise optional auth-broker protocol features. */
+export const AUTH_BROKER_CAPABILITIES_HEADER = "OMP-Auth-Broker-Capabilities";
+
+/** Client understands independent Codex `chat` and `spark` credential-block scopes. */
+export const AUTH_BROKER_CAPABILITY_CODEX_METER_BLOCK_SCOPES = "codex-meter-block-scopes";
+
+/** Default port when none is configured. Loopback-only, no external exposure. */
+export const DEFAULT_AUTH_BROKER_BIND = "127.0.0.1:8765";
+
+/** Default broker→provider refresh skew. Refresh credentials this close to expiry. */
+export const DEFAULT_REFRESH_SKEW_MS = 5 * 60_000;
+
+/** Default broker refresh-loop cadence. */
+export const DEFAULT_REFRESH_INTERVAL_MS = 60_000;
+
+/** Default freshness window for the encrypted local broker snapshot cache. */
+export const DEFAULT_SNAPSHOT_CACHE_TTL_MS = 60 * 60_000;
+
+/** Keepalive cadence for `GET /v1/snapshot/stream` SSE comments. */
+export const DEFAULT_STREAM_KEEPALIVE_MS = 20_000;
+
+/**
+ * Bun.serve `idleTimeout` (seconds) used by the broker. Default Bun idle
+ * timeout (10s) would close long-lived SSE connections between keepalives.
+ */
+export const DEFAULT_SERVER_IDLE_TIMEOUT_S = 255;