Jyotisha/frontend/tests/identity-session-cookie-cache.test.ts

import assert from "node:assert/strict";
import test from "node:test";
import { betterAuth } from "better-auth";
import { memoryAdapter, type MemoryDB } from "../node_modules/better-auth/node_modules/@better-auth/memory-adapter/dist/index.mjs";
import type { BetterAuthOptions } from "better-auth";

import {
  buildAuthOptions,
  IDENTITY_SESSION_COOKIE_CACHE_MAX_AGE_SECONDS,
} from "../src/modules/identity/auth-factory.ts";
import { FakeEmailOtpSender } from "../src/modules/identity/email/fake-email-otp-sender.ts";
import type { SelfHostedIdentityConfig } from "../src/modules/identity/config.ts";
import {
  readIdentitySession,
  requireIdentityServerSession,
  type IdentitySessionReader,
} from "../src/modules/identity/session.ts";

const config: SelfHostedIdentityConfig = {
  provider: "self-hosted",
  databaseUrl: "postgresql://identity_runtime:test-password@postgres:5432/jyotisha",
  userOrigin: "https://staging.jyotisha.chat",
  adminOrigin: "https://admin.staging.jyotisha.chat",
  userSecret: "user-secret-that-is-at-least-32-bytes-long",
  resendApiKey: "re_test",
  resendFrom: "Jyotisha <login@staging.jyotisha.chat>",
  testOtp: null,
};

function withCacheFlag(value: string | undefined, run: () => void): void {
  const previous = process.env.IDENTITY_SESSION_COOKIE_CACHE;
  if (value === undefined) delete process.env.IDENTITY_SESSION_COOKIE_CACHE;
  else process.env.IDENTITY_SESSION_COOKIE_CACHE = value;
  try {
    run();
  } finally {
    if (previous === undefined) delete process.env.IDENTITY_SESSION_COOKIE_CACHE;
    else process.env.IDENTITY_SESSION_COOKIE_CACHE = previous;
  }
}

test("user auth omits cookie cache unless the server switch is on", () => {
  withCacheFlag(undefined, () => {
    const options = buildAuthOptions({
      config,
      database: memoryAdapter({}),
      emailSender: new FakeEmailOtpSender(),
    });
    assert.equal(options.advanced?.cookiePrefix, "jyotisha-user");
    assert.equal(options.session && "cookieCache" in options.session, false);
  });
  withCacheFlag("on", () => {
    const options = buildAuthOptions({
      config,
      database: memoryAdapter({}),
      emailSender: new FakeEmailOtpSender(),
    });
    assert.deepEqual(options.session?.cookieCache, {
      enabled: true,
      maxAge: IDENTITY_SESSION_COOKIE_CACHE_MAX_AGE_SECONDS,
    });
  });
});

test("sensitive session readers pass disableCookieCache and ordinary readers do not", async () => {
  const queries: Array<{ disableCookieCache?: boolean }> = [];
  const reader: IdentitySessionReader = {
    async getSession(input) {
      queries.push({ disableCookieCache: input.query?.disableCookieCache });
      return {
        session: {
          id: "session-1",
          token: "session-token",
          expiresAt: new Date("2030-01-01T00:00:00.000Z"),
        },
        user: {
          id: "018f4e6d-7a11-7000-8000-000000000001",
          email: "person@example.com",
          emailVerified: true,
          name: "Person",
          image: null,
          role: "user",
        },
      };
    },
  };
  await readIdentitySession(reader, new Headers());
  await requireIdentityServerSession(reader, new Headers(), { disableCookieCache: true });
  assert.deepEqual(queries, [
    { disableCookieCache: undefined },
    { disableCookieCache: true },
  ]);
});

type SetCookieHeaders = Headers & { getSetCookie?: () => string[] };

function cookieHeaderFrom(response: Response, previous = ""): string {
  const headers = response.headers as SetCookieHeaders;
  const setCookies = headers.getSetCookie?.()
    ?? (headers.get("set-cookie")?.split(/,(?=\s*[^;,=\s]+=[^;,]*)/g) ?? []);
  const fresh = setCookies
    .map((value) => value.split(";", 1)[0] ?? "")
    .filter((value) => value.includes("=") && value.slice(value.indexOf("=") + 1).length > 0);
  const jar = new Map<string, string>();
  for (const part of previous.split("; ").filter(Boolean)) {
    jar.set(part.slice(0, part.indexOf("=")), part);
  }
  for (const part of fresh) jar.set(part.slice(0, part.indexOf("=")), part);
  return [...jar.values()].join("; ");
}

function createCacheAuth(enabled: boolean) {
  const db: MemoryDB = {
    users: [],
    sessions: [],
    accounts: [],
    verifications: [],
    two_factors: [],
    otp_rate_limits: [],
  };
  const sessionReads: string[] = [];
  const factory = memoryAdapter(db);
  const database = ((options: BetterAuthOptions) => {
    const adapter = factory(options);
    return {
      ...adapter,
      async findOne<T>(data: { model: string }) {
        if (data.model === "session" || data.model === "sessions") sessionReads.push(data.model);
        return adapter.findOne<T>(data as never);
      },
    };
  }) as BetterAuthOptions["database"];
  const previous = process.env.IDENTITY_SESSION_COOKIE_CACHE;
  process.env.IDENTITY_SESSION_COOKIE_CACHE = enabled ? "on" : "off";
  try {
    const built = buildAuthOptions({
      config,
      database,
      emailSender: new FakeEmailOtpSender(),
    });
    const auth = betterAuth({
      ...built,
      emailAndPassword: {
        ...built.emailAndPassword,
        enabled: true,
        disableSignUp: false,
      },
    });
    return { auth, db, sessionReads };
  } finally {
    if (previous === undefined) delete process.env.IDENTITY_SESSION_COOKIE_CACHE;
    else process.env.IDENTITY_SESSION_COOKIE_CACHE = previous;
  }
}

type CacheAuth = ReturnType<typeof createCacheAuth>["auth"];

async function signIn(auth: CacheAuth): Promise<string> {
  const email = "person@example.com";
  const password = "long-enough-password";
  const origin = config.userOrigin;
  const signedUp = await auth.api.signUpEmail({
    body: { email, password, name: "Person" },
    headers: new Headers({ origin }),
    asResponse: true,
  });
  assert.equal(signedUp.status, 200, await signedUp.clone().text());
  let cookie = cookieHeaderFrom(signedUp);
  if (!cookie.includes("session_token")) {
    const signedIn = await auth.api.signInEmail({
      body: { email, password },
      headers: new Headers({ origin }),
      asResponse: true,
    });
    assert.equal(signedIn.status, 200, await signedIn.clone().text());
    cookie = cookieHeaderFrom(signedIn, cookie);
  }
  return cookie;
}

async function readSession(
  auth: CacheAuth,
  cookie: string,
  disableCookieCache = false,
): Promise<{ cookie: string; status: number; userId: string | null }> {
  const response = await auth.api.getSession({
    headers: new Headers({ cookie, origin: config.userOrigin }),
    query: disableCookieCache ? { disableCookieCache: true } : undefined,
    asResponse: true,
  });
  const body = await response.json() as { user?: { id?: string } } | null;
  return {
    cookie: cookieHeaderFrom(response, cookie),
    status: response.status,
    userId: body?.user?.id ?? null,
  };
}

test("cookie cache serves three ordinary reads from one database read", async () => {
  const { auth, db, sessionReads } = createCacheAuth(true);
  const cookie = await signIn(auth);
  sessionReads.length = 0;
  let current = cookie
    .split("; ")
    .filter((part) => !part.slice(0, part.indexOf("=")).includes("session_data"))
    .join("; ");
  const reads: Array<string | null> = [];
  for (let index = 0; index < 3; index += 1) {
    const result = await readSession(auth, current);
    current = result.cookie;
    reads.push(result.userId);
  }
  assert.deepEqual(reads, [reads[0], reads[0], reads[0]]);
  assert.ok(reads[0]);
  assert.equal(sessionReads.length, 1);
  const cachedCookie = current;

  sessionReads.length = 0;
  for (let index = 0; index < 3; index += 1) {
    const result = await readSession(auth, cachedCookie, true);
    assert.ok(result.userId);
  }
  assert.equal(sessionReads.length, 3);

  const sessions = db.sessions ?? db.session;
  assert.ok(Array.isArray(sessions));
  sessions.length = 0;
  const sensitive = await readSession(auth, cachedCookie, true);
  assert.equal(sensitive.userId, null);
  const ordinary = await readSession(auth, cachedCookie, false);
  assert.equal(ordinary.userId, reads[0]);

  const realNow = Date.now;
  const started = realNow();
  Date.now = () => started + (IDENTITY_SESSION_COOKIE_CACHE_MAX_AGE_SECONDS + 1) * 1000;
  try {
    const expired = await readSession(auth, cachedCookie, false);
    assert.equal(expired.userId, null);
  } finally {
    Date.now = realNow;
  }
});

test("cookie cache off reads the database on every ordinary request", async () => {
  const { auth, sessionReads } = createCacheAuth(false);
  const cookie = await signIn(auth);
  sessionReads.length = 0;
  let current = cookie
    .split("; ")
    .filter((part) => !part.slice(0, part.indexOf("=")).includes("session_data"))
    .join("; ");
  for (let index = 0; index < 3; index += 1) {
    const result = await readSession(auth, current);
    assert.ok(result.userId);
    current = result.cookie;
  }
  assert.equal(sessionReads.length, 3);
});