makelore/electron/services/works-square-session.ts

import { NIANCODE_AUTH_CONFIG } from '../api/auth-config';
import { proxyAwareFetch } from '../utils/proxy-fetch';
import { logger } from '../utils/logger';
import { WORKS_SQUARE_SESSION_IDLE_TIMEOUT_MS } from '../../shared/auth-session';

const TOKEN_REFRESH_SKEW_MS = 30_000;
const SESSION_STORE_SCHEMA_VERSION = 1;
export { WORKS_SQUARE_SESSION_IDLE_TIMEOUT_MS } from '../../shared/auth-session';

export type WorksSquareSessionInput = {
  accessToken: string;
  refreshToken?: string | null;
  tokenType?: string | null;
  expiresAt?: number | null;
  lastActiveAt?: number | null;
};

export type WorksSquareTokenPayload = {
  access_token?: unknown;
  refresh_token?: unknown;
  token_type?: unknown;
  expires_in?: unknown;
};

export type WorksSquareSessionSnapshot = {
  accessToken: string;
  tokenType: string;
  expiresAt: number | null;
  lastActiveAt: number;
  canRefresh: boolean;
};

type StoredWorksSquareSession = {
  accessToken: string;
  refreshToken: string | null;
  tokenType: string;
  expiresAt: number | null;
  lastActiveAt: number;
};

export interface WorksSquareSessionPersistence {
  discardedInvalidRecord?: boolean;
  load(): Promise<WorksSquareSessionInput | null>;
  save(session: WorksSquareSessionInput | null): Promise<void>;
}

export type WorksSquareSessionChangeReason = 'changed' | 'terminal';
export type WorksSquareSessionListener = (
  session: WorksSquareSessionSnapshot | null,
  reason: WorksSquareSessionChangeReason,
  previousSession?: WorksSquareSessionSnapshot | null,
) => void;
export type WorksSquareSessionRestoreStatus = 'ready' | 'unavailable';

type RefreshFlight = {
  generation: number;
  promise: Promise<string | null>;
};

let currentSession: StoredWorksSquareSession | null = null;
let credentialGeneration = 0;
let refreshFlight: RefreshFlight | null = null;
let sessionPersistence: WorksSquareSessionPersistence | null = null;
let sessionPersistenceFactory: (() => Promise<WorksSquareSessionPersistence | null>) | null = null;
let persistenceQueue: Promise<void> = Promise.resolve();
let credentialPersistenceBarrier: Promise<void> | null = null;
let restoreStatus: WorksSquareSessionRestoreStatus = 'ready';
let startupRuntimeCleanupRequired = false;
const sessionListeners = new Set<WorksSquareSessionListener>();

function asString(value: unknown): string | null {
  return typeof value === 'string' && value.trim() ? value.trim() : null;
}

function expiresAtFromExpiresIn(expiresIn: unknown, nowMs = Date.now()): number | null {
  const seconds = typeof expiresIn === 'number'
    ? expiresIn
    : (typeof expiresIn === 'string' ? Number(expiresIn) : NaN);
  return Number.isFinite(seconds) && seconds > 0 ? nowMs + seconds * 1000 : null;
}

function createBasicAuthHeader(clientId: string, clientSecret: string): string {
  return `Basic ${Buffer.from(`${clientId}:${clientSecret}`).toString('base64')}`;
}

function toPublicSnapshot(
  session: StoredWorksSquareSession | null,
): WorksSquareSessionSnapshot | null {
  if (!session) return null;
  return {
    accessToken: session.accessToken,
    tokenType: session.tokenType,
    expiresAt: session.expiresAt,
    lastActiveAt: session.lastActiveAt,
    canRefresh: Boolean(session.refreshToken),
  };
}

function toPersistenceInput(
  session: StoredWorksSquareSession | null,
): WorksSquareSessionInput | null {
  return session ? { ...session } : null;
}

function normalizeSession(
  input: WorksSquareSessionInput,
  nowMs = Date.now(),
  requireLastActiveAt = false,
): StoredWorksSquareSession | null {
  const accessToken = asString(input.accessToken);
  if (!accessToken) return null;
  if (
    requireLastActiveAt
    && (typeof input.lastActiveAt !== 'number' || !Number.isFinite(input.lastActiveAt))
  ) {
    return null;
  }
  return {
    accessToken,
    refreshToken: asString(input.refreshToken),
    tokenType: asString(input.tokenType) ?? 'Bearer',
    expiresAt: typeof input.expiresAt === 'number' && Number.isFinite(input.expiresAt)
      ? input.expiresAt
      : null,
    lastActiveAt: typeof input.lastActiveAt === 'number' && Number.isFinite(input.lastActiveAt)
      ? Math.min(input.lastActiveAt, nowMs)
      : nowMs,
  };
}

function parseTokenPayload(
  payload: WorksSquareTokenPayload,
  fallbackRefreshToken: string | null,
  nowMs = Date.now(),
  lastActiveAt = nowMs,
): StoredWorksSquareSession {
  const accessToken = asString(payload.access_token);
  if (!accessToken) {
    throw new Error('Auth refresh response did not include access_token');
  }

  return {
    accessToken,
    refreshToken: asString(payload.refresh_token) ?? fallbackRefreshToken,
    tokenType: asString(payload.token_type) ?? 'Bearer',
    expiresAt: expiresAtFromExpiresIn(payload.expires_in, nowMs),
    lastActiveAt,
  };
}

function isSessionIdle(session: StoredWorksSquareSession, nowMs: number): boolean {
  return nowMs - session.lastActiveAt >= WORKS_SQUARE_SESSION_IDLE_TIMEOUT_MS;
}

function notifySessionListeners(
  reason: WorksSquareSessionChangeReason = 'changed',
  previousSession: WorksSquareSessionSnapshot | null = null,
): void {
  const snapshot = toPublicSnapshot(currentSession);
  for (const listener of sessionListeners) {
    try {
      listener(
        snapshot ? { ...snapshot } : null,
        reason,
        previousSession ? { ...previousSession } : null,
      );
    } catch {
      // Session observers must not affect authentication state changes.
    }
  }
}

function queuePersistence(session: StoredWorksSquareSession | null): Promise<void> {
  const persistence = sessionPersistence;
  if (!persistence) return Promise.resolve();
  const snapshot = toPersistenceInput(session);
  const operation = persistenceQueue
    .catch(() => undefined)
    .then(() => persistence.save(snapshot));
  persistenceQueue = operation;
  return operation;
}

function replaceCredentialSession(
  session: StoredWorksSquareSession | null,
  notify = true,
): {
    generation: number;
    persisted: Promise<void>;
    previousSession: WorksSquareSessionSnapshot | null;
  } {
  const previousSession = toPublicSnapshot(currentSession);
  currentSession = session;
  credentialGeneration += 1;
  refreshFlight = null;
  const generation = credentialGeneration;
  const persisted = queuePersistence(session);
  if (sessionPersistence) {
    credentialPersistenceBarrier = persisted;
    void persisted.then(
      () => {
        if (credentialPersistenceBarrier === persisted) credentialPersistenceBarrier = null;
      },
      () => {
        if (credentialPersistenceBarrier === persisted) credentialPersistenceBarrier = null;
      },
    );
  }
  if (notify) notifySessionListeners('changed', previousSession);
  return { generation, persisted, previousSession };
}

async function createElectronSessionPersistence(): Promise<WorksSquareSessionPersistence | null> {
  const [{ default: Store }, { safeStorage }] = await Promise.all([
    import('electron-store'),
    import('electron'),
  ]);
  const usesUnprotectedLinuxBackend = process.platform === 'linux'
    && safeStorage.getSelectedStorageBackend() === 'basic_text';
  if (!safeStorage.isEncryptionAvailable() || usesUnprotectedLinuxBackend) {
    logger.warn('[works-square-session] OS credential encryption is unavailable; session restore disabled');
    return null;
  }

  type EncryptedSessionRecord = {
    version: number;
    authBase: string;
    ciphertext: string;
  };
  const storeOptions = {
    name: 'works-square-session',
    configFileMode: 0o600,
  } as const;
  let discardedInvalidRecord = false;
  let store: Store<{ record?: EncryptedSessionRecord }>;
  try {
    store = new Store<{ record?: EncryptedSessionRecord }>(storeOptions);
  } catch (error) {
    logger.warn('[works-square-session] Invalid session store config; discarding it', error);
    store = new Store<{ record?: EncryptedSessionRecord }>({
      ...storeOptions,
      clearInvalidConfig: true,
    });
    discardedInvalidRecord = true;
  }

  return {
    discardedInvalidRecord,
    async load() {
      const record = store.get('record');
      if (!record) return null;
      if (
        record.version !== SESSION_STORE_SCHEMA_VERSION
        || record.authBase !== NIANCODE_AUTH_CONFIG.gatewayAuthUrl
      ) {
        store.delete('record');
        return null;
      }
      return JSON.parse(
        safeStorage.decryptString(Buffer.from(record.ciphertext, 'base64')),
      ) as WorksSquareSessionInput;
    },
    async save(session) {
      if (!session) {
        store.delete('record');
        return;
      }
      const encrypted = safeStorage.encryptString(JSON.stringify(session));
      store.set('record', {
        version: SESSION_STORE_SCHEMA_VERSION,
        authBase: NIANCODE_AUTH_CONFIG.gatewayAuthUrl,
        ciphertext: encrypted.toString('base64'),
      });
    },
  };
}

export async function initializeWorksSquareSession(
  options: {
    persistence?: WorksSquareSessionPersistence;
    nowMs?: number;
  } = {},
): Promise<WorksSquareSessionSnapshot | null> {
  persistenceQueue = Promise.resolve();
  credentialPersistenceBarrier = null;
  startupRuntimeCleanupRequired = false;
  sessionPersistenceFactory = options.persistence ? null : createElectronSessionPersistence;
  try {
    sessionPersistence = options.persistence
      ?? await sessionPersistenceFactory!();
  } catch (error) {
    logger.warn('[works-square-session] Failed to initialize secure session persistence', error);
    sessionPersistence = null;
    restoreStatus = 'unavailable';
    startupRuntimeCleanupRequired = true;
    currentSession = null;
    credentialGeneration += 1;
    refreshFlight = null;
    return null;
  }
  startupRuntimeCleanupRequired = Boolean(sessionPersistence?.discardedInvalidRecord);

  if (!sessionPersistence) {
    restoreStatus = 'ready';
    currentSession = null;
    credentialGeneration += 1;
    refreshFlight = null;
    return null;
  }

  let restored: StoredWorksSquareSession | null;
  let hadPersistedSession: boolean;
  try {
    const persisted = await sessionPersistence.load();
    hadPersistedSession = Boolean(persisted);
    restored = persisted
      ? normalizeSession(persisted, options.nowMs ?? Date.now(), true)
      : null;
  } catch (error) {
    logger.warn('[works-square-session] Failed to restore the persisted session', error);
    restoreStatus = 'unavailable';
    startupRuntimeCleanupRequired = true;
    currentSession = null;
    credentialGeneration += 1;
    refreshFlight = null;
    return null;
  }

  const nowMs = options.nowMs ?? Date.now();
  if (restored && !isSessionIdle(restored, nowMs)) {
    restoreStatus = 'ready';
    currentSession = restored;
    credentialGeneration += 1;
    refreshFlight = null;
    return toPublicSnapshot(restored);
  }

  if (!hadPersistedSession) {
    restoreStatus = 'ready';
    currentSession = null;
    credentialGeneration += 1;
    refreshFlight = null;
    return null;
  }

  startupRuntimeCleanupRequired = hadPersistedSession;
  restoreStatus = 'ready';
  const cleared = replaceCredentialSession(null, false);
  try {
    await cleared.persisted;
  } catch (error) {
    restoreStatus = 'unavailable';
    logger.warn('[works-square-session] Failed to clear an invalid persisted session', error);
  }
  return null;
}

async function readResponsePayload(response: Response): Promise<unknown> {
  const text = await response.text();
  if (!text.trim()) return null;
  try {
    return JSON.parse(text) as unknown;
  } catch {
    return text;
  }
}

export function storeWorksSquareSession(input: WorksSquareSessionInput): void {
  const session = normalizeSession(input);
  if (!session || isSessionIdle(session, Date.now())) {
    clearWorksSquareSession();
    return;
  }
  replaceCredentialSession(session);
}

export async function commitWorksSquareSession(
  input: WorksSquareSessionInput,
): Promise<WorksSquareSessionSnapshot | null> {
  const session = normalizeSession(input);
  if (!session || isSessionIdle(session, Date.now())) {
    await clearWorksSquareSessionPersisted('terminal');
    return null;
  }
  const update = replaceCredentialSession(session, false);
  try {
    await update.persisted;
  } catch (error) {
    await failClosedAfterPersistenceError(update.generation, error);
    throw error;
  }
  restoreStatus = 'ready';
  if (credentialGeneration !== update.generation) return null;
  notifySessionListeners('changed', update.previousSession);
  return getWorksSquareSessionSnapshot();
}

export function storeWorksSquareSessionFromTokenPayload(
  payload: WorksSquareTokenPayload,
  fallbackRefreshToken: string | null = null,
  nowMs = Date.now(),
  lastActiveAt = nowMs,
): WorksSquareSessionSnapshot {
  const session = parseTokenPayload(payload, fallbackRefreshToken, nowMs, lastActiveAt);
  replaceCredentialSession(session);
  return toPublicSnapshot(session)!;
}

export async function commitWorksSquareSessionFromTokenPayload(
  payload: WorksSquareTokenPayload,
  fallbackRefreshToken: string | null = null,
  nowMs = Date.now(),
  lastActiveAt = nowMs,
): Promise<WorksSquareSessionSnapshot> {
  const session = parseTokenPayload(payload, fallbackRefreshToken, nowMs, lastActiveAt);
  const update = replaceCredentialSession(session, false);
  try {
    await update.persisted;
  } catch (error) {
    await failClosedAfterPersistenceError(update.generation, error);
    throw error;
  }
  restoreStatus = 'ready';
  if (credentialGeneration !== update.generation) {
    throw new Error('Session changed while credentials were being persisted');
  }
  notifySessionListeners('changed', update.previousSession);
  return toPublicSnapshot(session)!;
}

export function clearWorksSquareSession(): void {
  replaceCredentialSession(null);
}

async function clearWorksSquareSessionPersisted(
  reason: WorksSquareSessionChangeReason = 'changed',
): Promise<void> {
  const update = replaceCredentialSession(null, false);
  if (credentialGeneration === update.generation) {
    notifySessionListeners(reason, update.previousSession);
  }
  await update.persisted;
}

export async function flushWorksSquareSessionPersistence(): Promise<void> {
  await persistenceQueue;
}

export function subscribeWorksSquareSession(listener: WorksSquareSessionListener): () => void {
  sessionListeners.add(listener);
  return () => {
    sessionListeners.delete(listener);
  };
}

export function getWorksSquareSessionSnapshot(): WorksSquareSessionSnapshot | null {
  const snapshot = toPublicSnapshot(currentSession);
  return snapshot ? { ...snapshot } : null;
}

export function getWorksSquareSessionRestoreStatus(): WorksSquareSessionRestoreStatus {
  return restoreStatus;
}

export function consumeWorksSquareStartupRuntimeCleanupRequired(): boolean {
  const required = startupRuntimeCleanupRequired;
  startupRuntimeCleanupRequired = false;
  return required;
}

export async function retryWorksSquareSessionRestore(
  nowMs = Date.now(),
): Promise<WorksSquareSessionSnapshot | null> {
  if (restoreStatus !== 'unavailable') {
    return getWorksSquareSessionSnapshot();
  }

  if (!sessionPersistence && sessionPersistenceFactory) {
    try {
      sessionPersistence = await sessionPersistenceFactory();
      startupRuntimeCleanupRequired ||= Boolean(sessionPersistence?.discardedInvalidRecord);
    } catch (error) {
      logger.warn('[works-square-session] Secure session persistence is still unavailable', error);
      return null;
    }
  }
  if (!sessionPersistence) {
    restoreStatus = 'ready';
    return null;
  }

  let restored: StoredWorksSquareSession | null;
  let hadPersistedSession: boolean;
  try {
    const persisted = await sessionPersistence.load();
    hadPersistedSession = Boolean(persisted);
    restored = persisted ? normalizeSession(persisted, nowMs, true) : null;
  } catch (error) {
    logger.warn('[works-square-session] Persisted session restore is still unavailable', error);
    startupRuntimeCleanupRequired = true;
    return null;
  }

  if (restored && !isSessionIdle(restored, nowMs)) {
    restoreStatus = 'ready';
    currentSession = restored;
    credentialGeneration += 1;
    refreshFlight = null;
    notifySessionListeners();
    return getWorksSquareSessionSnapshot();
  }

  if (!hadPersistedSession) {
    restoreStatus = 'ready';
    return null;
  }

  restoreStatus = 'ready';
  try {
    await clearWorksSquareSessionPersisted('terminal');
    restoreStatus = 'ready';
  } catch (error) {
    restoreStatus = 'unavailable';
    logger.warn('[works-square-session] Failed to clear an unusable persisted session', error);
  }
  return null;
}

export async function discardUnrestorableWorksSquareSession(): Promise<boolean> {
  if (restoreStatus !== 'unavailable') return true;
  const persistence = sessionPersistence;
  if (!persistence) return false;

  const cleared = replaceCredentialSession(null, false);
  try {
    await cleared.persisted;
  } catch (error) {
    restoreStatus = 'unavailable';
    logger.warn('[works-square-session] Failed to discard the unrestorable session', error);
    return false;
  }
  restoreStatus = 'ready';
  startupRuntimeCleanupRequired = true;
  if (credentialGeneration === cleared.generation) {
    notifySessionListeners('terminal', cleared.previousSession);
  }
  return true;
}

export async function markWorksSquareSessionActive(
  nowMs = Date.now(),
): Promise<WorksSquareSessionSnapshot | null> {
  const session = currentSession;
  if (!session) return null;
  if (isSessionIdle(session, nowMs)) {
    await clearWorksSquareSessionPersisted('terminal');
    return null;
  }

  const generation = credentialGeneration;
  const previousLastActiveAt = session.lastActiveAt;
  currentSession = { ...session, lastActiveAt: nowMs };
  try {
    await queuePersistence(currentSession);
  } catch (error) {
    if (
      credentialGeneration === generation
      && currentSession?.lastActiveAt === nowMs
    ) {
      currentSession = { ...currentSession, lastActiveAt: previousLastActiveAt };
    }
    throw error;
  }
  return getWorksSquareSessionSnapshot();
}

async function failClosedAfterPersistenceError(
  generation: number,
  error: unknown,
): Promise<void> {
  logger.error('[works-square-session] Failed to persist rotated credentials; clearing session', error);
  if (credentialGeneration !== generation) return;
  const cleared = replaceCredentialSession(null, false);
  if (credentialGeneration === cleared.generation) {
    notifySessionListeners('terminal', cleared.previousSession);
  }
  try {
    await cleared.persisted;
  } catch (clearError) {
    logger.error('[works-square-session] Failed to persist the fail-closed session clear', clearError);
  }
}

function sessionStillMatches(
  session: StoredWorksSquareSession,
  generation: number,
): boolean {
  return credentialGeneration === generation
    && currentSession?.accessToken === session.accessToken
    && currentSession.refreshToken === session.refreshToken;
}

async function waitForCredentialPersistence(): Promise<boolean> {
  const barrier = credentialPersistenceBarrier;
  if (!barrier) return true;
  try {
    await barrier;
    return true;
  } catch {
    return false;
  }
}

async function refreshWorksSquareSession(
  session: StoredWorksSquareSession,
  generation: number,
  options: { fetchImpl?: typeof fetch; nowMs?: number } = {},
): Promise<string | null> {
  if (!session.refreshToken) return null;

  const fetchImpl = options.fetchImpl ?? proxyAwareFetch;
  const nowMs = options.nowMs ?? Date.now();
  const body = new URLSearchParams({
    grant_type: 'refresh_token',
    refresh_token: session.refreshToken,
  });
  const response = await fetchImpl(`${NIANCODE_AUTH_CONFIG.gatewayAuthUrl.replace(/\/+$/, '')}/oauth2/token`, {
    method: 'POST',
    headers: {
      Authorization: createBasicAuthHeader(
        NIANCODE_AUTH_CONFIG.clientId,
        NIANCODE_AUTH_CONFIG.clientSecret,
      ),
      'Content-Type': 'application/x-www-form-urlencoded',
    },
    body,
  });

  if (!response.ok) {
    logger.warn('[works-square-session] Refresh failed', { status: response.status });
    if ((response.status === 400 || response.status === 401) && sessionStillMatches(session, generation)) {
      await clearWorksSquareSessionPersisted('terminal');
    }
    return null;
  }

  const payload = await readResponsePayload(response);
  if (!payload || typeof payload !== 'object' || Array.isArray(payload)) {
    logger.warn('[works-square-session] Refresh returned an invalid payload');
    return null;
  }

  if (!sessionStillMatches(session, generation)) {
    return null;
  }

  const nextSession = parseTokenPayload(
    payload as WorksSquareTokenPayload,
    session.refreshToken,
    nowMs,
    currentSession!.lastActiveAt,
  );
  const update = replaceCredentialSession(nextSession, false);
  try {
    await update.persisted;
  } catch (error) {
    await failClosedAfterPersistenceError(update.generation, error);
    return null;
  }
  if (credentialGeneration !== update.generation) return null;
  notifySessionListeners('changed', update.previousSession);
  return currentSession?.accessToken ?? null;
}

export async function getValidWorksSquareAccessToken(
  options: {
    fetchImpl?: typeof fetch;
    nowMs?: number;
    forceRefresh?: boolean;
  } = {},
): Promise<string | null> {
  if (restoreStatus === 'unavailable') {
    await retryWorksSquareSessionRestore(options.nowMs ?? Date.now());
    if (restoreStatus === 'unavailable') return null;
  }
  if (credentialPersistenceBarrier && !await waitForCredentialPersistence()) return null;
  const session = currentSession;
  if (!session) return null;

  const nowMs = options.nowMs ?? Date.now();
  if (isSessionIdle(session, nowMs)) {
    await clearWorksSquareSessionPersisted('terminal');
    return null;
  }
  const nearExpiry = Boolean(
    session.expiresAt
    && session.expiresAt <= nowMs + TOKEN_REFRESH_SKEW_MS,
  );
  if (!nearExpiry && !options.forceRefresh) return session.accessToken;
  if (!session.refreshToken) {
    await clearWorksSquareSessionPersisted('terminal');
    return null;
  }

  const generation = credentialGeneration;
  if (refreshFlight?.generation === generation) {
    return await refreshFlight.promise;
  }

  const flight: RefreshFlight = {
    generation,
    promise: Promise.resolve(null),
  };
  flight.promise = refreshWorksSquareSession(session, generation, options).finally(() => {
    if (refreshFlight === flight) refreshFlight = null;
  });
  refreshFlight = flight;
  return await flight.promise;
}

export function resetWorksSquareSessionForTests(): void {
  currentSession = null;
  credentialGeneration += 1;
  refreshFlight = null;
  sessionPersistence = null;
  sessionPersistenceFactory = null;
  persistenceQueue = Promise.resolve();
  credentialPersistenceBarrier = null;
  restoreStatus = 'ready';
  startupRuntimeCleanupRequired = false;
  sessionListeners.clear();
}