cinny/src/app/state/sessions.ts

// import { atom } from 'jotai';
// import {
//   atomWithLocalStorage,
//   getLocalStorageItem,
//   setLocalStorageItem,
// } from './utils/atomWithLocalStorage';

// OIDC/next-gen-auth (MSC3861) metadata needed to refresh + revoke tokens. Since
// matrix-js-sdk 43 only `clientId` is used (the SDK finds the auth server via
// /auth_metadata); `idTokenClaims` is kept for sessions saved by older builds.
export type OidcSessionMeta = {
  issuer: string;
  clientId: string;
  redirectUri: string;
  idTokenClaims?: Record<string, unknown>;
};

export type Session = {
  baseUrl: string;
  userId: string;
  deviceId: string;
  accessToken: string;
  expiresInMs?: number;
  refreshToken?: string;
  fallbackSdkStores?: boolean;
  oidc?: OidcSessionMeta;
};

// Legacy per-field localStorage keys. Kept for dual-write (see below) so a
// rollback to an older build that only understands these keys still works.
const LEGACY_KEYS = {
  accessToken: 'cinny_access_token',
  deviceId: 'cinny_device_id',
  userId: 'cinny_user_id',
  baseUrl: 'cinny_hs_base_url',
} as const;

// OIDC-only localStorage keys (absent for password/legacy-SSO sessions).
const OIDC_KEYS = {
  refreshToken: 'cinny_refresh_token',
  expiresAt: 'cinny_expires_at',
  issuer: 'cinny_oidc_issuer',
  clientId: 'cinny_oidc_client_id',
  redirectUri: 'cinny_oidc_redirect_uri',
  idTokenClaims: 'cinny_oidc_id_token_claims',
} as const;

// Single-key atomic session blob. The whole session is serialised and written
// in ONE `setItem`, so a reader can never observe a torn/partial session the
// way the multi-key legacy layout could. Bumping the schema means bumping the
// `_v1` suffix.
const SESSION_BLOB_KEY = 'cinny_session_v1';

// [Gitea #105] Desktop, step 2: the login tokens live in the OS keychain.
// localStorage then holds a SHELL blob (everything but the tokens, marked
// `tokensIn: 'keychain'`), and the full session is kept in memory, filled at
// boot (bootKeychainSession) before the app renders, so every reader of
// getFallbackSession stays synchronous. A new session is written to
// localStorage in full first; the tokens are stripped only once the keychain
// write was read back (stripTokensFromStorage), so at every instant one durable
// store holds the newest tokens (a token refresh invalidates the old ones).
let keychainMode = false;
let memorySession: PersistedSession | undefined;

type ShellSession = Omit<PersistedSession, 'accessToken' | 'refreshToken'> & {
  tokensIn: 'keychain';
};

// The exact shape stored inside SESSION_BLOB_KEY. Note it stores an ABSOLUTE
// `expiresAt` (ms since epoch) rather than a relative lifetime — identical to
// the legacy `cinny_expires_at` semantics — so reads stay drift-free.
type PersistedSession = {
  accessToken: string;
  deviceId: string;
  userId: string;
  baseUrl: string;
  refreshToken?: string;
  expiresAt?: number;
  oidc?: OidcSessionMeta;
};

// Build the persisted shape from the public setFallbackSession arguments. This
// is the single source of truth written to BOTH the blob and the legacy keys.
const buildPersisted = (
  accessToken: string,
  deviceId: string,
  userId: string,
  baseUrl: string,
  extra?: FallbackSessionExtra,
): PersistedSession => {
  const persisted: PersistedSession = { accessToken, deviceId, userId, baseUrl };
  if (extra?.refreshToken) persisted.refreshToken = extra.refreshToken;
  // Store ABSOLUTE expiry to avoid drift across reloads.
  if (typeof extra?.expiresInMs === 'number') persisted.expiresAt = Date.now() + extra.expiresInMs;
  if (extra?.oidc) persisted.oidc = extra.oidc;
  return persisted;
};

// Convert a persisted shape into the public Session returned to callers. Keeps
// behaviour identical to the original getFallbackSession assembly: derives the
// REMAINING lifetime from the absolute expiry, and only surfaces `oidc` when the
// three required OIDC fields are present.
const sessionFromPersisted = (p: PersistedSession): Session => {
  const session: Session = {
    baseUrl: p.baseUrl,
    userId: p.userId,
    deviceId: p.deviceId,
    accessToken: p.accessToken,
    fallbackSdkStores: true,
  };

  if (p.refreshToken) session.refreshToken = p.refreshToken;

  if (typeof p.expiresAt === 'number' && Number.isFinite(p.expiresAt)) {
    // Expose the REMAINING lifetime (clamped at 0); the SDK refreshes on 401.
    session.expiresInMs = Math.max(0, p.expiresAt - Date.now());
  }

  if (p.oidc && p.oidc.issuer && p.oidc.clientId && p.oidc.redirectUri) {
    session.oidc = {
      issuer: p.oidc.issuer,
      clientId: p.oidc.clientId,
      redirectUri: p.oidc.redirectUri,
      idTokenClaims: p.oidc.idTokenClaims,
    };
  }

  return session;
};

// Read the atomic blob. Returns undefined when absent, unparseable, or missing
// any of the four required credential fields — callers then fall back to the
// legacy keys.
const readSessionBlob = (): PersistedSession | undefined => {
  const raw = localStorage.getItem(SESSION_BLOB_KEY);
  if (!raw) return undefined;
  let parsed: unknown;
  try {
    parsed = JSON.parse(raw);
  } catch {
    // Corrupt JSON — treat as absent and let the legacy path take over.
    return undefined;
  }
  if (!parsed || typeof parsed !== 'object') return undefined;
  const p = parsed as Partial<PersistedSession>;
  if (
    typeof p.accessToken !== 'string' ||
    typeof p.deviceId !== 'string' ||
    typeof p.userId !== 'string' ||
    typeof p.baseUrl !== 'string'
  ) {
    // Partial/corrupt blob — fall back to legacy assembly.
    return undefined;
  }
  return p as PersistedSession;
};

// Assemble a session from the legacy per-field keys, or undefined when the four
// required keys are not all present. Used for transparent migration from builds
// that predate the atomic blob.
const readLegacyKeys = (): PersistedSession | undefined => {
  const baseUrl = localStorage.getItem(LEGACY_KEYS.baseUrl);
  const userId = localStorage.getItem(LEGACY_KEYS.userId);
  const deviceId = localStorage.getItem(LEGACY_KEYS.deviceId);
  const accessToken = localStorage.getItem(LEGACY_KEYS.accessToken);

  if (!(baseUrl && userId && deviceId && accessToken)) return undefined;

  const persisted: PersistedSession = { accessToken, deviceId, userId, baseUrl };

  const refreshToken = localStorage.getItem(OIDC_KEYS.refreshToken);
  if (refreshToken) persisted.refreshToken = refreshToken;

  const expiresAtRaw = localStorage.getItem(OIDC_KEYS.expiresAt);
  if (expiresAtRaw) {
    const expiresAt = Number(expiresAtRaw);
    if (Number.isFinite(expiresAt)) persisted.expiresAt = expiresAt;
  }

  const issuer = localStorage.getItem(OIDC_KEYS.issuer);
  const clientId = localStorage.getItem(OIDC_KEYS.clientId);
  const redirectUri = localStorage.getItem(OIDC_KEYS.redirectUri);
  if (issuer && clientId && redirectUri) {
    let idTokenClaims: Record<string, unknown> | undefined;
    const claimsRaw = localStorage.getItem(OIDC_KEYS.idTokenClaims);
    if (claimsRaw) {
      try {
        idTokenClaims = JSON.parse(claimsRaw);
      } catch {
        /* corrupt claims — ignore, the refresher will re-validate on use */
      }
    }
    persisted.oidc = { issuer, clientId, redirectUri, idTokenClaims };
  }

  return persisted;
};

// Write the legacy per-field keys (dual-write half). Mirrors the original
// setFallbackSession body so a rollback to an older build keeps working.
const writeLegacyKeys = (p: PersistedSession): void => {
  localStorage.setItem(LEGACY_KEYS.accessToken, p.accessToken);
  localStorage.setItem(LEGACY_KEYS.deviceId, p.deviceId);
  localStorage.setItem(LEGACY_KEYS.userId, p.userId);
  localStorage.setItem(LEGACY_KEYS.baseUrl, p.baseUrl);

  // OIDC fields — written only when present; otherwise cleared so a password
  // session never carries stale OIDC state.
  if (p.refreshToken) localStorage.setItem(OIDC_KEYS.refreshToken, p.refreshToken);
  else localStorage.removeItem(OIDC_KEYS.refreshToken);

  if (typeof p.expiresAt === 'number')
    localStorage.setItem(OIDC_KEYS.expiresAt, String(p.expiresAt));
  else localStorage.removeItem(OIDC_KEYS.expiresAt);

  if (p.oidc) {
    localStorage.setItem(OIDC_KEYS.issuer, p.oidc.issuer);
    localStorage.setItem(OIDC_KEYS.clientId, p.oidc.clientId);
    localStorage.setItem(OIDC_KEYS.redirectUri, p.oidc.redirectUri);
    if (p.oidc.idTokenClaims) {
      localStorage.setItem(OIDC_KEYS.idTokenClaims, JSON.stringify(p.oidc.idTokenClaims));
    } else localStorage.removeItem(OIDC_KEYS.idTokenClaims);
  } else {
    localStorage.removeItem(OIDC_KEYS.issuer);
    localStorage.removeItem(OIDC_KEYS.clientId);
    localStorage.removeItem(OIDC_KEYS.redirectUri);
    localStorage.removeItem(OIDC_KEYS.idTokenClaims);
  }
};

export type FallbackSessionExtra = {
  refreshToken?: string;
  expiresInMs?: number;
  oidc?: OidcSessionMeta;
};

export type Sessions = Session[];
export type SessionStoreName = {
  sync: string;
  crypto: string;
};

/**
 * Migration code for old session
 */
// const FALLBACK_STORE_NAME: SessionStoreName = {
//   sync: 'web-sync-store',
//   crypto: 'crypto-store',
// } as const;

// [Gitea #105] Listeners told about every session write in THIS tab (login,
// token rotation) and removal (logout), e.g. the desktop keychain mirror.
// A throwing listener never breaks the write.
type PersistListener = (session: Session | null) => void;
const persistListeners = new Set<PersistListener>();
const notifyPersisted = (session: Session | null): void => {
  persistListeners.forEach((cb) => {
    try {
      cb(session);
    } catch {
      /* listener errors must not affect login/logout */
    }
  });
};
export const onSessionPersisted = (cb: PersistListener): (() => void) => {
  persistListeners.add(cb);
  return () => {
    persistListeners.delete(cb);
  };
};

// Persist the session. Writes the atomic blob FIRST (so the consistent,
// never-torn copy is established before the multi-key legacy write), then
// dual-writes the legacy keys for rollback safety. Signature is unchanged —
// callers (login/register/OIDC callback/token refresher) are untouched.
export function setFallbackSession(
  accessToken: string,
  deviceId: string,
  userId: string,
  baseUrl: string,
  extra?: FallbackSessionExtra,
) {
  const persisted = buildPersisted(accessToken, deviceId, userId, baseUrl, extra);
  if (keychainMode) memorySession = persisted;
  // ONE setItem — the blob can never be observed half-written.
  localStorage.setItem(SESSION_BLOB_KEY, JSON.stringify(persisted));
  // Dual-write the legacy keys (removal of this half is a future release).
  writeLegacyKeys(persisted);
  notifyPersisted(sessionFromPersisted(persisted));
}

// Clear BOTH the atomic blob and every legacy key so no reader (blob-preferring
// or legacy-fallback) can resurrect a logged-out session.
export const removeFallbackSession = () => {
  memorySession = undefined;
  localStorage.removeItem(SESSION_BLOB_KEY);
  Object.values(LEGACY_KEYS).forEach((key) => localStorage.removeItem(key));
  Object.values(OIDC_KEYS).forEach((key) => localStorage.removeItem(key));
  notifyPersisted(null);
};

// Read the session, preferring the atomic blob. If the blob is absent or
// corrupt/partial we transparently assemble from the legacy keys (migration);
// the next setFallbackSession then persists the blob. When both exist the blob
// wins by construction.
// The stored session with its tokens, choosing blob vs legacy as below.
const readStoredPersisted = (): PersistedSession | undefined => {
  const blob = readSessionBlob();
  const legacy = readLegacyKeys();
  if (
    blob &&
    legacy &&
    typeof legacy.expiresAt === 'number' &&
    (typeof blob.expiresAt !== 'number' || legacy.expiresAt > blob.expiresAt)
  ) {
    return legacy;
  }
  return blob ?? legacy;
};

export const getFallbackSession = (): Session | undefined => {
  if (keychainMode) return memorySession ? sessionFromPersisted(memorySession) : undefined;
  const blob = readSessionBlob();
  const legacy = readLegacyKeys();
  // Prefer the atomic blob, EXCEPT when the legacy keys carry a later expiry: a
  // pre-blob build's token refresh writes only the legacy keys, so a
  // downgrade→upgrade can leave a stale blob newer than fresh legacy keys →
  // booting on a dead token. Whichever has the later expiresAt wins.
  let persisted = blob ?? legacy;
  if (
    blob &&
    legacy &&
    typeof legacy.expiresAt === 'number' &&
    (typeof blob.expiresAt !== 'number' || legacy.expiresAt > blob.expiresAt)
  ) {
    persisted = legacy;
  }
  if (!persisted) return undefined;
  return sessionFromPersisted(persisted);
};
// --- [Gitea #105] desktop keychain session -----------------------------------

export type KeychainSessionTokens = {
  userId: string;
  deviceId: string;
  accessToken: string;
  refreshToken?: string;
};

/** The shell blob written while the tokens are in the keychain, if any. */
const readShell = (): Omit<ShellSession, 'tokensIn'> | undefined => {
  const raw = localStorage.getItem(SESSION_BLOB_KEY);
  if (!raw) return undefined;
  try {
    const p = JSON.parse(raw) as Partial<ShellSession>;
    if (
      p?.tokensIn !== 'keychain' ||
      typeof p.deviceId !== 'string' ||
      typeof p.userId !== 'string' ||
      typeof p.baseUrl !== 'string'
    )
      return undefined;
    const { tokensIn: _drop, ...rest } = p as ShellSession;
    return rest;
  } catch {
    return undefined;
  }
};

/** A session in localStorage that still carries its tokens (blob or legacy). */
export const readStoredSessionWithTokens = (): Session | undefined => {
  const persisted = readStoredPersisted();
  return persisted ? sessionFromPersisted(persisted) : undefined;
};

/** Whether localStorage holds a token-less shell (tokens are in the keychain). */
export const hasKeychainShell = (): boolean => readShell() !== undefined;

export type KeychainBoot =
  /** Use the keychain: memory holds `session` (or none = signed out). */
  | { mode: 'keychain'; session: Session | undefined }
  /** A shell says the tokens are in the keychain, but they can't be read. */
  | { mode: 'missing'; userId: string };

/**
 * Decide the boot session from the keychain's tokens and localStorage, and
 * switch this tab to keychain mode. `tokens` is what the keychain returned
 * (null = no entry). Pure apart from reading localStorage and setting the
 * in-memory state.
 */
export const enterKeychainMode = (tokens: KeychainSessionTokens | null): KeychainBoot => {
  const withTokens = readStoredPersisted();
  const shell = readShell();
  keychainMode = true;
  if (withTokens) {
    // Pre-step-2 install, or a write not yet confirmed: localStorage has the
    // full session and stays authoritative until the mirror verifies it.
    memorySession = withTokens;
    return { mode: 'keychain', session: sessionFromPersisted(withTokens) };
  }
  if (shell) {
    if (tokens && tokens.userId === shell.userId && tokens.deviceId === shell.deviceId) {
      memorySession = {
        ...shell,
        accessToken: tokens.accessToken,
        ...(tokens.refreshToken ? { refreshToken: tokens.refreshToken } : {}),
      };
      return { mode: 'keychain', session: sessionFromPersisted(memorySession) };
    }
    // Never treat this as a logout: the shell stays, the caller offers Retry.
    memorySession = undefined;
    return { mode: 'missing', userId: shell.userId };
  }
  memorySession = undefined;
  return { mode: 'keychain', session: undefined };
};

/** Leave keychain mode (the keychain failed): localStorage is authoritative again. */
export const leaveKeychainMode = (): void => {
  keychainMode = false;
  memorySession = undefined;
};

/**
 * The keychain now verifiably holds `tokens`: if they're still the current
 * session's, drop the tokens from localStorage (shell blob, no legacy keys).
 */
export const stripTokensFromStorage = (tokens: KeychainSessionTokens): boolean => {
  const m = memorySession;
  if (
    !keychainMode ||
    !m ||
    m.userId !== tokens.userId ||
    m.deviceId !== tokens.deviceId ||
    m.accessToken !== tokens.accessToken ||
    (m.refreshToken ?? undefined) !== (tokens.refreshToken ?? undefined)
  )
    return false;
  const { accessToken: _a, refreshToken: _r, ...rest } = m;
  const shell: ShellSession = { ...rest, tokensIn: 'keychain' };
  localStorage.setItem(SESSION_BLOB_KEY, JSON.stringify(shell));
  Object.values(LEGACY_KEYS).forEach((key) => localStorage.removeItem(key));
  Object.values(OIDC_KEYS).forEach((key) => localStorage.removeItem(key));
  return true;
};

/**
 * End of migration code for old session
 */

// Session keys whose cross-tab change indicates a login/logout/token-rotation
// in another tab. localStorage.clear() dispatches a storage event with a null
// key, which we also treat as a session change.
const SESSION_STORAGE_KEYS = new Set<string>([
  SESSION_BLOB_KEY,
  ...Object.values(LEGACY_KEYS),
  ...Object.values(OIDC_KEYS),
]);

/**
 * Subscribe to session changes made in OTHER tabs/windows. The browser only
 * dispatches `storage` events to tabs that did NOT perform the write, so this
 * is inherently guarded against reacting to our own same-tab writes — no
 * echo-suppression needed. The callback receives the freshly-read session, or
 * `null` when the session was removed (logout in another tab, or a full
 * localStorage.clear()). Returns an unsubscribe function.
 */
export const subscribeSessionChanges = (
  callback: (session: Session | null) => void,
): (() => void) => {
  const handleStorage = (evt: StorageEvent) => {
    // A null key means localStorage.clear(); otherwise only react to our keys.
    if (evt.key !== null && !SESSION_STORAGE_KEYS.has(evt.key)) return;
    callback(getFallbackSession() ?? null);
  };
  window.addEventListener('storage', handleStorage);
  return () => {
    window.removeEventListener('storage', handleStorage);
  };
};

// export const getSessionStoreName = (session: Session): SessionStoreName => {
//   if (session.fallbackSdkStores) {
//     return FALLBACK_STORE_NAME;
//   }

//   return {
//     sync: `sync${session.userId}`,
//     crypto: `crypto${session.userId}`,
//   };
// };

// export const MATRIX_SESSIONS_KEY = 'matrixSessions';
// const baseSessionsAtom = atomWithLocalStorage<Sessions>(
//   MATRIX_SESSIONS_KEY,
//   (key) => {
//     const defaultSessions: Sessions = [];
//     const sessions = getLocalStorageItem(key, defaultSessions);

//     // Before multi account support session was stored
//     // as multiple item in local storage.
//     // So we need these migration code.
//     const fallbackSession = getFallbackSession();
//     if (fallbackSession) {
//       removeFallbackSession();
//       sessions.push(fallbackSession);
//       setLocalStorageItem(key, sessions);
//     }
//     return sessions;
//   },
//   (key, value) => {
//     setLocalStorageItem(key, value);
//   }
// );

// export type SessionsAction =
//   | {
//       type: 'PUT';
//       session: Session;
//     }
//   | {
//       type: 'DELETE';
//       session: Session;
//     };

// export const sessionsAtom = atom<Sessions, [SessionsAction], undefined>(
//   (get) => get(baseSessionsAtom),
//   (get, set, action) => {
//     if (action.type === 'PUT') {
//       const sessions = [...get(baseSessionsAtom)];
//       const sessionIndex = sessions.findIndex(
//         (session) => session.userId === action.session.userId
//       );
//       if (sessionIndex === -1) {
//         sessions.push(action.session);
//       } else {
//         sessions.splice(sessionIndex, 1, action.session);
//       }
//       set(baseSessionsAtom, sessions);
//       return;
//     }
//     if (action.type === 'DELETE') {
//       const sessions = get(baseSessionsAtom).filter(
//         (session) => session.userId !== action.session.userId
//       );
//       set(baseSessionsAtom, sessions);
//     }
//   }
// );