diff --git a/README.md b/README.md index 57253f1..7b09b51 100644 --- a/README.md +++ b/README.md @@ -220,6 +220,43 @@ importBudgetState(savedState); --- +### Token Cleanup + +#### `createTokenCleanupStore(options)` + +Creates an in-memory cleanup helper for expired token ids and revoked token deny lists. Use it when your application tracks blocked, revoked, or short-lived tokens and needs a safe periodic cleanup pass with audit counts. + +**Parameters:** + +```typescript +interface TokenCleanupStoreOptions { + revokedTokenMaxAgeMs?: number; // How long revoked tokens stay in the deny list + onCleanup?: (result: TokenCleanupResult) => void; +} +``` + +**Example:** + +```javascript +const { createTokenCleanupStore } = require("tokenfirewall"); + +const cleanupStore = createTokenCleanupStore({ + revokedTokenMaxAgeMs: 24 * 60 * 60 * 1000, + onCleanup: ({ total, expired, revoked }) => { + console.log(`Cleaned ${total} tokens (${expired} expired, ${revoked} revoked)`); + } +}); + +cleanupStore.upsert({ + id: "session-token-id", + expiresAt: Date.now() + 15 * 60 * 1000 +}); + +cleanupStore.startAutoCleanup(5 * 60 * 1000); +``` + +--- + ### Interception #### `patchGlobalFetch()` diff --git a/package-lock.json b/package-lock.json index e9cce39..cdd1200 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "tokenfirewall", - "version": "2.0.1", + "version": "2.1.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "tokenfirewall", - "version": "2.0.1", + "version": "2.1.0", "license": "MIT", "devDependencies": { "@types/node": "^20.0.0", diff --git a/src/core/tokenCleanupStore.ts b/src/core/tokenCleanupStore.ts new file mode 100644 index 0000000..b49c72c --- /dev/null +++ b/src/core/tokenCleanupStore.ts @@ -0,0 +1,163 @@ +export type TokenTimestamp = number | Date; + +export interface TokenCleanupRecord { + id: string; + expiresAt: TokenTimestamp; + revokedAt?: TokenTimestamp | null; + metadata?: TMetadata; +} + +export interface TokenCleanupStoreOptions { + revokedTokenMaxAgeMs?: number; + onCleanup?: (result: TokenCleanupResult) => void; +} + +export interface TokenCleanupOptions { + now?: TokenTimestamp; + revokedTokenMaxAgeMs?: number; +} + +export interface TokenCleanupResult { + expired: number; + revoked: number; + total: number; + remaining: number; +} + +const DEFAULT_REVOKED_TOKEN_MAX_AGE_MS = 24 * 60 * 60 * 1000; + +function normalizeTimestamp(value: TokenTimestamp, fieldName: string): number { + const timestamp = value instanceof Date ? value.getTime() : value; + + if (typeof timestamp !== "number" || Number.isNaN(timestamp) || !Number.isFinite(timestamp)) { + throw new Error(`TokenFirewall: ${fieldName} must be a valid timestamp`); + } + + return timestamp; +} + +function normalizeMaxAge(value: number): number { + if (typeof value !== "number" || Number.isNaN(value) || !Number.isFinite(value) || value < 0) { + throw new Error("TokenFirewall: revokedTokenMaxAgeMs must be a non-negative finite number"); + } + + return value; +} + +/** + * In-memory token cleanup helper for applications that track expired or revoked + * token identifiers alongside TokenFirewall. + */ +export class TokenCleanupStore { + private records = new Map>(); + private cleanupTimer: ReturnType | null = null; + private readonly revokedTokenMaxAgeMs: number; + private readonly onCleanup?: (result: TokenCleanupResult) => void; + + constructor(options: TokenCleanupStoreOptions = {}) { + this.revokedTokenMaxAgeMs = normalizeMaxAge( + options.revokedTokenMaxAgeMs ?? DEFAULT_REVOKED_TOKEN_MAX_AGE_MS, + ); + this.onCleanup = options.onCleanup; + } + + public upsert(record: TokenCleanupRecord): void { + if (!record.id || typeof record.id !== "string" || record.id.trim() === "") { + throw new Error("TokenFirewall: token cleanup record id must be a non-empty string"); + } + + normalizeTimestamp(record.expiresAt, "expiresAt"); + + if (record.revokedAt !== undefined && record.revokedAt !== null) { + normalizeTimestamp(record.revokedAt, "revokedAt"); + } + + this.records.set(record.id, { ...record }); + } + + public delete(id: string): boolean { + return this.records.delete(id); + } + + public get(id: string): TokenCleanupRecord | undefined { + const record = this.records.get(id); + return record ? { ...record } : undefined; + } + + public list(): TokenCleanupRecord[] { + return Array.from(this.records.values(), (record) => ({ ...record })); + } + + public size(): number { + return this.records.size; + } + + public cleanup(options: TokenCleanupOptions = {}): TokenCleanupResult { + const now = normalizeTimestamp(options.now ?? Date.now(), "now"); + const revokedTokenMaxAgeMs = normalizeMaxAge( + options.revokedTokenMaxAgeMs ?? this.revokedTokenMaxAgeMs, + ); + let expired = 0; + let revoked = 0; + + for (const [id, record] of this.records.entries()) { + const expiresAt = normalizeTimestamp(record.expiresAt, "expiresAt"); + const revokedAt = + record.revokedAt === undefined || record.revokedAt === null + ? null + : normalizeTimestamp(record.revokedAt, "revokedAt"); + + if (expiresAt <= now) { + this.records.delete(id); + expired++; + continue; + } + + if (revokedAt !== null && now - revokedAt >= revokedTokenMaxAgeMs) { + this.records.delete(id); + revoked++; + } + } + + const result = { + expired, + revoked, + total: expired + revoked, + remaining: this.records.size, + }; + + if (result.total > 0) { + this.onCleanup?.(result); + } + + return result; + } + + public startAutoCleanup(intervalMs: number): void { + if (typeof intervalMs !== "number" || Number.isNaN(intervalMs) || !Number.isFinite(intervalMs) || intervalMs <= 0) { + throw new Error("TokenFirewall: cleanup interval must be a positive finite number"); + } + + this.stopAutoCleanup(); + this.cleanupTimer = setInterval(() => { + this.cleanup(); + }, intervalMs); + + if (typeof this.cleanupTimer.unref === "function") { + this.cleanupTimer.unref(); + } + } + + public stopAutoCleanup(): void { + if (this.cleanupTimer) { + clearInterval(this.cleanupTimer); + this.cleanupTimer = null; + } + } +} + +export function createTokenCleanupStore( + options?: TokenCleanupStoreOptions, +): TokenCleanupStore { + return new TokenCleanupStore(options); +} diff --git a/src/index.ts b/src/index.ts index c769e99..25a9884 100644 --- a/src/index.ts +++ b/src/index.ts @@ -9,6 +9,7 @@ import { contextRegistry } from "./introspection/contextRegistry"; import { ModelRouter } from "./router/modelRouter"; import { ModelRouterOptions, ApiKeyConfig } from "./router/types"; import { apiKeyManager } from "./router/apiKeyManager"; +import { createTokenCleanupStore, TokenCleanupStore } from "./core/tokenCleanupStore"; let globalBudgetManager: BudgetManager | null = null; let globalModelRouter: ModelRouter | null = null; @@ -192,6 +193,13 @@ export function importBudgetState(state: { totalSpent: number }): void { globalBudgetManager.importState(state); } +/** + * Create an in-memory token cleanup store for expired and revoked token ids. + * Applications can pair this with their own persistence layer or run the + * built-in auto cleanup interval for lightweight token deny lists. + */ +export { createTokenCleanupStore, TokenCleanupStore }; + /** * Create and configure a model router for automatic retries and fallbacks * @param options - Router configuration options @@ -240,6 +248,14 @@ export type { ListModelsOptions, } from "./core/types"; +export type { + TokenCleanupRecord, + TokenCleanupStoreOptions, + TokenCleanupOptions, + TokenCleanupResult, + TokenTimestamp, +} from "./core/tokenCleanupStore"; + export type { ModelInfo as ModelInfoType, ListModelsOptions as ListModelsOptionsType } from "./introspection/modelLister"; export type { diff --git a/tests/token-cleanup.test.js b/tests/token-cleanup.test.js new file mode 100644 index 0000000..c5e0844 --- /dev/null +++ b/tests/token-cleanup.test.js @@ -0,0 +1,100 @@ +/** + * Token cleanup store tests. + * + * Run: npm run build && node tests/token-cleanup.test.js + */ + +const { + createTokenCleanupStore, + TokenCleanupStore, +} = require("../dist/index.js"); + +const results = { total: 0, passed: 0, failed: 0 }; + +function assert(name, condition) { + results.total += 1; + if (condition) { + results.passed += 1; + console.log(`✓ ${name}`); + return; + } + + results.failed += 1; + console.error(`✗ ${name}`); +} + +function assertThrows(name, fn) { + try { + fn(); + assert(name, false); + } catch { + assert(name, true); + } +} + +function run() { + const now = Date.parse("2026-05-29T00:00:00.000Z"); + const store = createTokenCleanupStore({ + revokedTokenMaxAgeMs: 60_000, + }); + + store.upsert({ id: "active", expiresAt: now + 10_000 }); + store.upsert({ id: "expired", expiresAt: now - 1 }); + store.upsert({ + id: "fresh-revoked", + expiresAt: now + 10_000, + revokedAt: now - 30_000, + }); + store.upsert({ + id: "old-revoked", + expiresAt: now + 10_000, + revokedAt: now - 60_000, + }); + + const cleanup = store.cleanup({ now }); + + assert("removes one expired token", cleanup.expired === 1); + assert("removes one stale revoked token", cleanup.revoked === 1); + assert("reports total removed tokens", cleanup.total === 2); + assert("keeps active and freshly revoked tokens", cleanup.remaining === 2); + assert("active token is still present", Boolean(store.get("active"))); + assert("fresh revoked token is still present", Boolean(store.get("fresh-revoked"))); + assert("expired token is removed", !store.get("expired")); + assert("old revoked token is removed", !store.get("old-revoked")); + + const secondCleanup = store.cleanup({ now }); + assert("cleanup is idempotent", secondCleanup.total === 0); + + let callbackResult = null; + const callbackStore = new TokenCleanupStore({ + onCleanup: (result) => { + callbackResult = result; + }, + }); + callbackStore.upsert({ id: "callback-expired", expiresAt: now - 1 }); + callbackStore.cleanup({ now }); + assert("cleanup callback receives audit counts", callbackResult?.expired === 1); + + assertThrows("rejects empty ids", () => { + store.upsert({ id: "", expiresAt: now }); + }); + assertThrows("rejects invalid expiration timestamps", () => { + store.upsert({ id: "bad-expiry", expiresAt: Number.NaN }); + }); + assertThrows("rejects invalid cleanup intervals", () => { + store.startAutoCleanup(0); + }); + + store.startAutoCleanup(60_000); + store.stopAutoCleanup(); + assert("auto cleanup can start and stop", true); + + if (results.failed > 0) { + console.error(`${results.failed}/${results.total} token cleanup tests failed`); + process.exit(1); + } + + console.log(`${results.passed}/${results.total} token cleanup tests passed`); +} + +run();