stack/packages/shared/src/sessions.test.ts

import { describe, expect, it } from "vitest";
import { InternalSession } from "./sessions";

/**
 * Builds a decodable (unsigned) access-token JWT with a valid payload. `refreshTokenId` controls the
 * `refresh_token_id` claim (the session identifier); `iatOffsetSeconds` lets two tokens for the same session
 * differ as strings while sharing a `refresh_token_id`.
 */
function createAccessTokenString(refreshTokenId: string, options?: { iatOffsetSeconds?: number, sub?: string }): string {
  const encode = (value: unknown) => Buffer.from(JSON.stringify(value)).toString("base64url");
  const nowSeconds = Math.floor(Date.now() / 1000) + (options?.iatOffsetSeconds ?? 0);
  return [
    encode({ alg: "none", typ: "JWT" }),
    encode({
      sub: options?.sub ?? "user-id",
      exp: nowSeconds + 60,
      iat: nowSeconds,
      iss: "https://api.example.test",
      aud: "project-id",
      project_id: "project-id",
      branch_id: "main",
      refresh_token_id: refreshTokenId,
      role: "authenticated",
      name: null,
      email: null,
      email_verified: false,
      selected_team_id: null,
      signed_up_at: nowSeconds,
      is_anonymous: false,
      is_restricted: false,
      restricted_reason: null,
      requires_totp_mfa: false,
    }),
    "",
  ].join(".");
}

function createAccessOnlySession(accessToken: string): InternalSession {
  return new InternalSession({
    refreshAccessTokenCallback: async () => null,
    refreshToken: null,
    accessToken,
  });
}

const currentToken = (session: InternalSession) => session.getAccessTokenIfNotExpiredYet(20_000, null)?.token;

describe("InternalSession.calculateSessionKey", () => {
  it("keys by the refresh token when one is present (ignoring any access token)", () => {
    expect(InternalSession.calculateSessionKey({ refreshToken: "rt-abc" })).toBe("refresh-rt-abc");
    expect(InternalSession.calculateSessionKey({ refreshToken: "rt-abc", accessToken: createAccessTokenString("rtid-1") }))
      .toBe("refresh-rt-abc");
  });

  it("returns not-logged-in when neither token is present", () => {
    expect(InternalSession.calculateSessionKey({ refreshToken: null })).toBe("not-logged-in");
    expect(InternalSession.calculateSessionKey({ refreshToken: null, accessToken: null })).toBe("not-logged-in");
  });

  it("keys an access-only session by its refresh_token_id", () => {
    expect(InternalSession.calculateSessionKey({ refreshToken: null, accessToken: createAccessTokenString("rtid-1") }))
      .toBe("access-session-rtid-1");
  });

  it("is stable across re-minted access tokens for the same session (the regression this fixes)", () => {
    const first = createAccessTokenString("rtid-1", { iatOffsetSeconds: 0 });
    const second = createAccessTokenString("rtid-1", { iatOffsetSeconds: 1 });
    expect(second).not.toBe(first);
    expect(InternalSession.calculateSessionKey({ refreshToken: null, accessToken: second }))
      .toBe(InternalSession.calculateSessionKey({ refreshToken: null, accessToken: first }));
  });

  it("distinguishes access-only sessions with different refresh_token_ids", () => {
    expect(InternalSession.calculateSessionKey({ refreshToken: null, accessToken: createAccessTokenString("rtid-1") }))
      .not.toBe(InternalSession.calculateSessionKey({ refreshToken: null, accessToken: createAccessTokenString("rtid-2") }));
  });

  it("falls back to the raw token when the access token can't be decoded", () => {
    expect(InternalSession.calculateSessionKey({ refreshToken: null, accessToken: "not-a-jwt" })).toBe("access-not-a-jwt");
  });
});

describe("InternalSession#updateAccessToken", () => {
  it("installs a fresh token for the same access-only session in place", () => {
    const initial = createAccessTokenString("rtid-1", { iatOffsetSeconds: 0 });
    const refreshed = createAccessTokenString("rtid-1", { iatOffsetSeconds: 1 });
    const session = createAccessOnlySession(initial);

    session.updateAccessToken({ accessToken: refreshed, refreshToken: null });
    expect(currentToken(session)).toBe(refreshed);
    // identity is unchanged — same session key, same object
    expect(session.sessionKey).toBe("access-session-rtid-1");
  });

  it("rejects a token pair belonging to a different access-only session", () => {
    const initial = createAccessTokenString("rtid-1");
    const foreign = createAccessTokenString("rtid-2", { sub: "other-user" });
    const session = createAccessOnlySession(initial);

    session.updateAccessToken({ accessToken: foreign, refreshToken: null });
    expect(currentToken(session)).toBe(initial);
  });

  it("is a no-op for an unchanged, null, or undecodable token", () => {
    const initial = createAccessTokenString("rtid-1");
    const session = createAccessOnlySession(initial);

    session.updateAccessToken({ accessToken: initial, refreshToken: null });
    session.updateAccessToken({ accessToken: null, refreshToken: null });
    session.updateAccessToken({ accessToken: "not-a-jwt", refreshToken: null });
    expect(currentToken(session)).toBe(initial);
  });

  it("never revives an invalidated session", () => {
    const session = createAccessOnlySession(createAccessTokenString("rtid-1"));
    session.markInvalid();

    session.updateAccessToken({ accessToken: createAccessTokenString("rtid-1", { iatOffsetSeconds: 1 }), refreshToken: null });
    expect(session.isKnownToBeInvalid()).toBe(true);
    expect(currentToken(session)).toBeUndefined();
  });

  it("updates a refresh-token-backed session's access token in place when the refresh token matches", () => {
    const session = new InternalSession({
      refreshAccessTokenCallback: async () => null,
      refreshToken: "rt-abc",
      accessToken: createAccessTokenString("rtid-1"),
    });
    const refreshed = createAccessTokenString("rtid-2", { iatOffsetSeconds: 1 });

    session.updateAccessToken({ accessToken: refreshed, refreshToken: "rt-abc" });
    expect(currentToken(session)).toBe(refreshed);
    expect(session.sessionKey).toBe("refresh-rt-abc");
  });

  it("rejects a token pair carrying a different refresh token for a refresh-backed session", () => {
    const initial = createAccessTokenString("rtid-1");
    const session = new InternalSession({
      refreshAccessTokenCallback: async () => null,
      refreshToken: "rt-abc",
      accessToken: initial,
    });

    session.updateAccessToken({ accessToken: createAccessTokenString("rtid-2"), refreshToken: "rt-other" });
    expect(currentToken(session)).toBe(initial);
  });
});