WRNexusJS/packages/authz/src/store.ts

import type { AuthzScope, SubjectAssignments } from "./types.ts";

export type GrantEffect = "allow" | "deny";

export interface PermissionStore {
  assignmentsFor(subjectId: string, scope?: AuthzScope): Promise<SubjectAssignments>;
  assignRole(subjectId: string, role: string, scope?: AuthzScope): Promise<void>;
  revokeRole(subjectId: string, role: string, scope?: AuthzScope): Promise<void>;
  grant(
    subjectId: string,
    permission: string,
    effect: GrantEffect,
    scope?: AuthzScope,
  ): Promise<void>;
  revokeGrant(subjectId: string, permission: string, scope?: AuthzScope): Promise<void>;
  listSubjects(scope?: AuthzScope): Promise<string[]>;
}

/** Global assignments are stored under the empty-string scope key. */
export function scopeKey(scope?: AuthzScope): string {
  return scope?.tenantId ?? "";
}

interface Row {
  subjectId: string;
  scope: string;
}
interface RoleRow extends Row {
  role: string;
}
interface GrantRow extends Row {
  permission: string;
  effect: GrantEffect;
}

export function memoryPermissionStore(): PermissionStore {
  const roles: RoleRow[] = [];
  const grants: GrantRow[] = [];

  // A request inside tenant t sees global assignments plus t's own.
  const visible = (row: Row, key: string) => row.scope === "" || row.scope === key;

  return {
    async assignmentsFor(subjectId, scope) {
      const key = scopeKey(scope);
      const mine = (row: Row) => row.subjectId === subjectId && visible(row, key);
      const matched = grants.filter(mine);
      return {
        roles: roles.filter(mine).map((row) => row.role),
        grants: matched.filter((row) => row.effect === "allow").map((row) => row.permission),
        denies: matched.filter((row) => row.effect === "deny").map((row) => row.permission),
      };
    },
    async assignRole(subjectId, role, scope) {
      const key = scopeKey(scope);
      if (roles.some((r) => r.subjectId === subjectId && r.scope === key && r.role === role))
        return;
      roles.push({ subjectId, scope: key, role });
    },
    async revokeRole(subjectId, role, scope) {
      const key = scopeKey(scope);
      const at = roles.findIndex(
        (r) => r.subjectId === subjectId && r.scope === key && r.role === role,
      );
      if (at !== -1) roles.splice(at, 1);
    },
    async grant(subjectId, permission, effect, scope) {
      const key = scopeKey(scope);
      const at = grants.findIndex(
        (g) => g.subjectId === subjectId && g.scope === key && g.permission === permission,
      );
      if (at !== -1) grants.splice(at, 1);
      grants.push({ subjectId, scope: key, permission, effect });
    },
    async revokeGrant(subjectId, permission, scope) {
      const key = scopeKey(scope);
      const at = grants.findIndex(
        (g) => g.subjectId === subjectId && g.scope === key && g.permission === permission,
      );
      if (at !== -1) grants.splice(at, 1);
    },
    async listSubjects(scope) {
      const key = scopeKey(scope);
      const ids = new Set<string>();
      for (const row of roles) if (row.scope === key) ids.add(row.subjectId);
      for (const row of grants) if (row.scope === key) ids.add(row.subjectId);
      return [...ids];
    },
  };
}

export interface CachedPermissionStore extends PermissionStore {
  /** Drop one subject. Call after changing roles out of band. */
  invalidate(subjectId: string, scope?: AuthzScope): void;
  invalidateAll(): void;
  /** Cached entry count, for tests and diagnostics. */
  size(): number;
}

export interface CacheOptions {
  ttlMs?: number;
  max?: number;
}

/**
 * Caches assignment reads. Writes through this decorator invalidate the
 * affected subject immediately; changes made directly against the inner store
 * need an explicit `invalidate()` call rather than waiting out the TTL.
 */
export function cachedPermissionStore(
  inner: PermissionStore,
  options: CacheOptions = {},
): CachedPermissionStore {
  const ttlMs = options.ttlMs ?? 5_000;
  const max = options.max ?? 1_000;
  const entries = new Map<string, { at: number; value: SubjectAssignments; subjectId: string }>();
  const bySubject = new Map<string, Set<string>>();

  const cacheKey = (subjectId: string, scope?: AuthzScope) =>
    JSON.stringify([scopeKey(scope), subjectId]);
  const drop = (subjectId: string, scope?: AuthzScope) => {
    // A global write changes what every tenant sees for that subject.
    if (scopeKey(scope) === "") {
      const keys = bySubject.get(subjectId);
      if (keys) for (const key of keys) entries.delete(key);
      bySubject.delete(subjectId);
      return;
    }
    const key = cacheKey(subjectId, scope);
    entries.delete(key);
    bySubject.get(subjectId)?.delete(key);
  };

  return {
    async assignmentsFor(subjectId, scope) {
      const key = cacheKey(subjectId, scope);
      const hit = entries.get(key);
      if (hit && Date.now() - hit.at < ttlMs) return hit.value;
      const value = await inner.assignmentsFor(subjectId, scope);
      if (entries.size >= max) {
        const oldestKey = entries.keys().next().value!;
        const oldest = entries.get(oldestKey);
        entries.delete(oldestKey);
        if (oldest) bySubject.get(oldest.subjectId)?.delete(oldestKey);
      }
      entries.set(key, { at: Date.now(), value, subjectId });
      let keys = bySubject.get(subjectId);
      if (!keys) {
        keys = new Set();
        bySubject.set(subjectId, keys);
      }
      keys.add(key);
      return value;
    },
    async assignRole(subjectId, role, scope) {
      await inner.assignRole(subjectId, role, scope);
      drop(subjectId, scope);
    },
    async revokeRole(subjectId, role, scope) {
      await inner.revokeRole(subjectId, role, scope);
      drop(subjectId, scope);
    },
    async grant(subjectId, permission, effect, scope) {
      await inner.grant(subjectId, permission, effect, scope);
      drop(subjectId, scope);
    },
    async revokeGrant(subjectId, permission, scope) {
      await inner.revokeGrant(subjectId, permission, scope);
      drop(subjectId, scope);
    },
    listSubjects: (scope) => inner.listSubjects(scope),
    invalidate: drop,
    invalidateAll: () => {
      entries.clear();
      bySubject.clear();
    },
    size: () => entries.size,
  };
}