Design a session manager

export interface DeviceCtx { userAgent: string }

export interface Session {
  idHash: string;               // SHA-256 of the opaque id; the raw id exists only in the cookie
  userId: string;
  createdAt: number;            // ms since epoch
  lastSeenAt: number;
  absoluteExpiresAt: number;    // fixed at login; carried, never extended
  elevatedUntil?: number;       // step-up auth window
  device: { uaHash: string };
  rotated: boolean;             // a successor exists; this one is in its grace window
  revokedAt?: number;           // effective from this instant (future = grace)
  version: number;              // for compare-and-set
}

export type DenyReason = "malformed" | "unknown" | "revoked" | "idle_timeout" | "absolute_timeout" | "device_mismatch";
export type Validation = { ok: true; session: Session; elevated: boolean } | { ok: false; reason: DenyReason };

export interface SessionStore {
  get(idHash: string): Promise<Session | null>;
  /** Fails if idHash already exists. Should set a store TTL at absoluteExpiresAt. */
  create(s: Session): Promise<void>;
  /** Compare-and-set: writes only if the stored version equals expectedVersion. */
  replace(idHash: string, expectedVersion: number, next: Session): Promise<boolean>;
  /** Revokes every session of the user not already revoked by at. Returns how many. */
  revokeAllForUser(userId: string, at: number): Promise<number>;
}

export interface SessionConfig { idleMs: number; absoluteMs: number; elevatedMs: number; rotationGraceMs: number; touchEveryMs: number }

export class SessionError extends Error {
  constructor(readonly reason: DenyReason | "already_rotated" | "contention") { super(reason); this.name = "SessionError"; }
}

import { createHash, randomBytes } from "node:crypto";

const ID_RE = /^[A-Za-z0-9_-]{43}$/;                        // 32 random bytes, base64url
const sha256 = (s: string) => createHash("sha256").update(s).digest("base64url");
const deviceOf = (ctx: DeviceCtx) => ({ uaHash: sha256(ctx.userAgent) });
type UpdateResult = "written" | "missing" | "skipped";

export class SessionManager {
  constructor(private readonly store: SessionStore, private readonly cfg: SessionConfig, private readonly clock: () => number = Date.now) {
    const { idleMs, absoluteMs, elevatedMs, rotationGraceMs, touchEveryMs } = cfg;
    if (![idleMs, absoluteMs, elevatedMs, rotationGraceMs, touchEveryMs].every((n) => Number.isFinite(n) && n > 0)) throw new Error("timeouts must be positive");
    if (idleMs >= absoluteMs) throw new Error("idle timeout must be shorter than the absolute timeout");
    if (touchEveryMs >= idleMs) throw new Error("touchEveryMs must be shorter than idleMs");
    if (rotationGraceMs > 60_000) throw new Error("rotation grace covers in-flight requests: keep it under a minute");
  }

  async login(userId: string, ctx: DeviceCtx, preLoginId?: string): Promise<{ id: string }> {
    const now = this.clock();
    if (preLoginId !== undefined && ID_RE.test(preLoginId)) {
      await this.update(sha256(preLoginId), (s) => ({ ...s, revokedAt: now }));   // fixation: a pre-auth id is never upgraded
    }
    return this.mint({ userId, createdAt: now, absoluteExpiresAt: now + this.cfg.absoluteMs, device: deviceOf(ctx) }, now);
  }

  async validate(rawId: string, ctx: DeviceCtx): Promise<Validation> {
    if (!ID_RE.test(rawId)) return { ok: false, reason: "malformed" };
    const s = await this.store.get(sha256(rawId));
    if (!s) return { ok: false, reason: "unknown" };
    const now = this.clock();
    if (s.revokedAt !== undefined && now >= s.revokedAt) return { ok: false, reason: "revoked" };
    if (now >= s.absoluteExpiresAt) return { ok: false, reason: "absolute_timeout" };
    if (now - s.lastSeenAt >= this.cfg.idleMs) return { ok: false, reason: "idle_timeout" };
    if (s.device.uaHash !== deviceOf(ctx).uaHash) {
      await this.update(s.idHash, (cur) => ({ ...cur, revokedAt: now }));      // likely a stolen cookie: kill it
      return { ok: false, reason: "device_mismatch" };
    }
    if (now - s.lastSeenAt >= this.cfg.touchEveryMs) {
      // CAS, not a blind put: if a revoke landed since our read, this write loses instead of resurrecting the session
      await this.store.replace(s.idHash, s.version, { ...s, lastSeenAt: now, version: s.version + 1 });
    }
    return { ok: true, session: s, elevated: (s.elevatedUntil ?? 0) > now };
  }

  async rotate(rawId: string, ctx: DeviceCtx, opts: { elevate: boolean }): Promise<{ id: string }> {
    const v = await this.validate(rawId, ctx);
    if (!v.ok) throw new SessionError(v.reason);
    const now = this.clock();
    const old = v.session;
    const res = await this.update(old.idHash, (s) =>
      s.rotated ? null : { ...s, rotated: true, elevatedUntil: undefined, revokedAt: Math.min(s.revokedAt ?? Infinity, now + this.cfg.rotationGraceMs) });
    if (res === "skipped") throw new SessionError("already_rotated");
    if (res === "missing") throw new SessionError("unknown");
    return this.mint({
      userId: old.userId, createdAt: old.createdAt, device: old.device,
      absoluteExpiresAt: old.absoluteExpiresAt,                                    // rotation never extends the lifetime
      elevatedUntil: opts.elevate ? Math.min(now + this.cfg.elevatedMs, old.absoluteExpiresAt) : old.elevatedUntil,
    }, now);
  }

  async logout(rawId: string): Promise<void> {
    if (ID_RE.test(rawId)) await this.update(sha256(rawId), (s) => ({ ...s, revokedAt: this.clock() }));
  }

  revokeAll(userId: string): Promise<number> { return this.store.revokeAllForUser(userId, this.clock()); }

  private async mint(base: Pick<Session, "userId" | "createdAt" | "absoluteExpiresAt" | "device" | "elevatedUntil">, now: number): Promise<{ id: string }> {
    const id = randomBytes(32).toString("base64url");
    await this.store.create({ ...base, idHash: sha256(id), lastSeenAt: now, rotated: false, version: 0 });
    return { id };
  }

  /** Read-modify-CAS with a bounded retry. f returns null to skip the write. */
  private async update(idHash: string, f: (s: Session) => Session | null): Promise<UpdateResult> {
    for (let attempt = 0; attempt < 3; attempt++) {
      const cur = await this.store.get(idHash);
      if (!cur) return "missing";
      const next = f(cur);
      if (!next) return "skipped";
      if (await this.store.replace(idHash, cur.version, { ...next, version: cur.version + 1 })) return "written";
    }
    throw new SessionError("contention");
  }
}

export class MemorySessionStore implements SessionStore {
  private readonly rows = new Map<string, Session>();
  async get(k: string) { return this.rows.get(k) ?? null; }
  async create(s: Session) {
    if (this.rows.has(s.idHash)) throw new Error("duplicate session id");
    this.rows.set(s.idHash, s);
  }
  async replace(k: string, expectedVersion: number, next: Session) {
    const cur = this.rows.get(k);
    if (!cur || cur.version !== expectedVersion) return false;
    this.rows.set(k, next);
    return true;
  }
  async revokeAllForUser(userId: string, at: number) {
    let n = 0;
    for (const [k, s] of this.rows) {
      if (s.userId !== userId || (s.revokedAt !== undefined && s.revokedAt <= at)) continue;   // sessions in grace are cut too
      this.rows.set(k, { ...s, revokedAt: at, version: s.version + 1 });
      n++;
    }
    return n;
  }
}

import { describe, it, expect } from "vitest";

let t = 0;
const cfg = { idleMs: 30 * 60_000, absoluteMs: 12 * 3_600_000, elevatedMs: 10 * 60_000, rotationGraceMs: 10_000, touchEveryMs: 60_000 };
const make = () => new SessionManager(new MemorySessionStore(), cfg, () => t);
const ua = { userAgent: "Mozilla/5.0 test" };

describe("SessionManager", () => {
  it("ends a quiet session at the idle timeout", async () => {
    t = 0; const m = make(); const { id } = await m.login("u1", ua);
    t = cfg.idleMs;
    expect(await m.validate(id, ua)).toEqual({ ok: false, reason: "idle_timeout" });
  });
  it("neither activity nor rotation extends the absolute timeout", async () => {
    t = 0; const m = make(); let { id } = await m.login("u1", ua);
    for (let i = 1; i < 36; i++) {                      // every 20 minutes for 11h40m
      t = i * 20 * 60_000;
      if (i === 18) ({ id } = await m.rotate(id, ua, { elevate: false }));
      expect((await m.validate(id, ua)).ok).toBe(true);
    }
    t = cfg.absoluteMs;
    expect(await m.validate(id, ua)).toEqual({ ok: false, reason: "absolute_timeout" });
  });
  it("never upgrades a pre-login id", async () => {
    t = 0; const m = make();
    const anon = (await m.login("anon", ua)).id;
    const { id } = await m.login("u1", ua, anon);
    expect(id).not.toBe(anon);
    expect(await m.validate(anon, ua)).toEqual({ ok: false, reason: "revoked" });
  });
  it("keeps the old id briefly after elevation, never elevated, and only rotates once", async () => {
    t = 0; const m = make(); const { id: old } = await m.login("u1", ua);
    const { id: fresh } = await m.rotate(old, ua, { elevate: true });
    expect(await m.validate(fresh, ua)).toMatchObject({ ok: true, elevated: true });
    expect(await m.validate(old, ua)).toMatchObject({ ok: true, elevated: false });
    await expect(m.rotate(old, ua, { elevate: true })).rejects.toThrow("already_rotated");
    t = cfg.rotationGraceMs;
    expect(await m.validate(old, ua)).toEqual({ ok: false, reason: "revoked" });
  });
  it("kills a cookie replayed from another client", async () => {
    t = 0; const m = make(); const { id } = await m.login("u1", ua);
    expect(await m.validate(id, { userAgent: "curl/8.0" })).toEqual({ ok: false, reason: "device_mismatch" });
    expect(await m.validate(id, ua)).toEqual({ ok: false, reason: "revoked" });
  });
});