makelore/electron/services/preview-data-session.ts

import { randomBytes, timingSafeEqual } from 'node:crypto';
import path from 'node:path';
import type { IncomingMessage } from 'node:http';
import type {
  AgentBrowserLifecycleEvent,
} from '../agent-browser/module';
import type { CodingProjectService } from '../coding-projects/project-service';
import {
  getWorksSquareAccountBinding,
  getWorksSquareSessionSnapshot,
  isCurrentWorksSquareAccountBinding,
  subscribeWorksSquareSession,
  type WorksSquareAccountBinding,
  type WorksSquareSessionListener,
} from './works-square-session';
import { getPort } from '../utils/config';

export const PREVIEW_DATA_ROUTE_ROOT = '/api/runtime/data/v1';
export const PREVIEW_DATA_CONTRACT_VERSION = 1 as const;
export const PREVIEW_DATA_MAX_REQUEST_BYTES = 98_304;
export const PREVIEW_DATA_MAX_RESPONSE_BYTES = 1_310_720;
export const PREVIEW_DATA_RATE = 5;
export const PREVIEW_DATA_BUCKET_CAPACITY = 30;

const PREVIEW_DATA_METHODS = new Set(['GET', 'PUT', 'DELETE', 'OPTIONS']);
const PREVIEW_DATA_HEADERS = new Set(['authorization', 'content-type', 'if-match']);
const PREVIEW_DATA_TRANSPORT_HEADERS = new Set([
  'accept',
  'accept-encoding',
  'accept-language',
  'cache-control',
  'connection',
  'content-length',
  'host',
  'origin',
  'priority',
  'pragma',
  'referer',
  'sec-ch-ua',
  'sec-ch-ua-mobile',
  'sec-ch-ua-platform',
  'sec-fetch-dest',
  'sec-fetch-mode',
  'sec-fetch-site',
  'sec-fetch-user',
  'transfer-encoding',
  'upgrade-insecure-requests',
  'user-agent',
]);

export type PreviewDataSessionInvalidationReason =
  | 'preview_closed'
  | 'project_deactivated'
  | 'project_identity_changed'
  | 'session_cleared'
  | 'account_changed'
  | 'cross_origin_navigation'
  | 'browser_generation_replaced'
  | 'browser_detached'
  | 'browser_crashed'
  | 'main_shutdown'
  | 'replaced'
  | 'manual';

export type PreviewDataSessionInvalidationListener = (
  reason: PreviewDataSessionInvalidationReason,
) => void;

export type PreviewDataSessionOpenInput = {
  projectPath: string;
  origin: string;
  browserGeneration: number;
};

export type PreviewDataSessionSnapshot = Readonly<{
  projectPath: string;
  projectId: string;
  origin: string;
  browserGeneration: number;
  createdAt: number;
}>;

export type PreviewDataInjectionValue = Readonly<{
  endpoint: string;
  token: string;
  contractVersion: typeof PREVIEW_DATA_CONTRACT_VERSION;
}>;

export type PreviewDataAuthorizationSuccess = Readonly<{
  ok: true;
  method: string;
  session: PreviewDataSessionSnapshot;
}>;

export type PreviewDataAuthorizationFailure = Readonly<{
  ok: false;
  status: number;
  code: string;
  message: string;
  retryable: boolean;
  retryAfterSeconds?: number;
}>;

export type PreviewDataAuthorizationResult =
  | PreviewDataAuthorizationSuccess
  | PreviewDataAuthorizationFailure;

export class PreviewDataSessionError extends Error {
  constructor(
    readonly code: 'invalid_origin' | 'invalid_generation' | 'runtime_unavailable',
    message: string,
  ) {
    super(message);
    this.name = 'PreviewDataSessionError';
  }
}

export type PreviewDataSessionManagerOptions = {
  projects: Pick<CodingProjectService, 'requireActiveRealProjectWithIdentity'>;
  hostPort?: number;
  now?: () => number;
  randomBytes?: (size: number) => Buffer;
  subscribeWorksSquareSession?: typeof subscribeWorksSquareSession;
  getWorksSquareSessionSnapshot?: typeof getWorksSquareSessionSnapshot;
  getWorksSquareAccountBinding?: typeof getWorksSquareAccountBinding;
  isCurrentWorksSquareAccountBinding?: typeof isCurrentWorksSquareAccountBinding;
};

type InternalSession = {
  snapshot: PreviewDataSessionSnapshot;
  token: string;
  tokenBytes: Buffer;
  tokens: number;
  lastRefillMs: number;
};

function normalizePath(value: string): string {
  const resolved = path.resolve(value);
  return process.platform === 'win32' ? resolved.toLocaleLowerCase('en-US') : resolved;
}

function samePath(left: string, right: string): boolean {
  return normalizePath(left) === normalizePath(right);
}

function normalizeOrigin(value: string): string {
  let parsed: URL;
  try {
    parsed = new URL(value);
  } catch {
    throw new PreviewDataSessionError('invalid_origin', 'Preview Origin is invalid');
  }
  const hostname = parsed.hostname.toLowerCase();
  const isLoopback = hostname === 'localhost'
    || hostname === '127.0.0.1'
    || hostname === '[::1]';
  if (
    !isLoopback
    || (parsed.protocol !== 'http:' && parsed.protocol !== 'https:')
    || parsed.username
    || parsed.password
    || parsed.pathname !== '/'
    || parsed.search
    || parsed.hash
  ) {
    throw new PreviewDataSessionError('invalid_origin', 'Preview Origin is invalid');
  }
  return parsed.origin;
}

function normalizeGeneration(value: number): number {
  if (!Number.isSafeInteger(value) || value < 1) {
    throw new PreviewDataSessionError('invalid_generation', 'Browser generation is invalid');
  }
  return value;
}

function headerValue(req: IncomingMessage, name: string): string | undefined {
  const value = req.headers[name];
  return typeof value === 'string' ? value : undefined;
}

function failure(
  status: number,
  code: string,
  message: string,
  retryable = false,
  retryAfterSeconds?: number,
): PreviewDataAuthorizationFailure {
  return {
    ok: false,
    status,
    code,
    message,
    retryable,
    ...(retryAfterSeconds === undefined ? {} : { retryAfterSeconds }),
  };
}

function requestHeadersAllowed(req: IncomingMessage): boolean {
  for (const [name, value] of Object.entries(req.headers)) {
    if (Array.isArray(value)) return false;
    if (PREVIEW_DATA_HEADERS.has(name) || PREVIEW_DATA_TRANSPORT_HEADERS.has(name)) continue;
    if (name === 'access-control-request-method') {
      if (value !== undefined && !PREVIEW_DATA_METHODS.has(value.toUpperCase())) return false;
      continue;
    }
    if (name === 'access-control-request-headers') {
      if (value === undefined) continue;
      const requested = value.split(',').map((item) => item.trim().toLowerCase()).filter(Boolean);
      if (requested.some((item) => !PREVIEW_DATA_HEADERS.has(item))) return false;
      continue;
    }
    return false;
  }
  return true;
}

function isExactToken(candidate: string, expected: Buffer): boolean {
  if (!candidate || !/^[A-Za-z0-9_-]+$/.test(candidate)) return false;
  const candidateBytes = Buffer.from(candidate, 'base64url');
  return candidateBytes.length === expected.length && timingSafeEqual(candidateBytes, expected);
}

function retryAfterSeconds(tokens: number): number {
  return Math.max(1, Math.ceil((1 - tokens) / PREVIEW_DATA_RATE));
}

export class PreviewDataSessionManager {
  private readonly projects: Pick<CodingProjectService, 'requireActiveRealProjectWithIdentity'>;
  private readonly hostPort: number;
  private readonly now: () => number;
  private readonly createRandomBytes: (size: number) => Buffer;
  private readonly getSessionSnapshot: typeof getWorksSquareSessionSnapshot;
  private readonly getAccountBinding: typeof getWorksSquareAccountBinding;
  private readonly isCurrentAccountBinding: typeof isCurrentWorksSquareAccountBinding;
  private readonly unsubscribeWorksSession: () => void;
  private readonly invalidationListeners = new Set<PreviewDataSessionInvalidationListener>();
  private session: InternalSession | null = null;
  private boundAccountBinding: WorksSquareAccountBinding | null;
  private disposed = false;

  constructor(options: PreviewDataSessionManagerOptions) {
    this.projects = options.projects;
    this.hostPort = options.hostPort ?? getPort('NIANCODE_HOST_API');
    this.now = options.now ?? Date.now;
    this.createRandomBytes = options.randomBytes ?? randomBytes;
    this.getSessionSnapshot = options.getWorksSquareSessionSnapshot ?? getWorksSquareSessionSnapshot;
    this.getAccountBinding = options.getWorksSquareAccountBinding ?? getWorksSquareAccountBinding;
    this.isCurrentAccountBinding = options.isCurrentWorksSquareAccountBinding
      ?? isCurrentWorksSquareAccountBinding;
    // The Works observer intentionally has no initial callback. Read the
    // current snapshot once, then rely on the unsubscribe-able observer for
    // subsequent logout/account transitions.
    const initialSession = this.getSessionSnapshot();
    this.boundAccountBinding = initialSession ? this.getAccountBinding() : null;

    const subscribe = options.subscribeWorksSquareSession ?? subscribeWorksSquareSession;
    const listener: WorksSquareSessionListener = (nextSession) => {
      if (!nextSession) {
        this.invalidate('session_cleared');
        return;
      }
      if (!this.boundAccountBinding) {
        this.invalidate('account_changed');
        return;
      }
      if (!this.isCurrentAccountBinding(this.boundAccountBinding)) {
        this.invalidate('account_changed');
      }
    };
    this.unsubscribeWorksSession = subscribe(listener);
  }

  async open(input: PreviewDataSessionOpenInput): Promise<PreviewDataSessionSnapshot> {
    if (this.disposed) {
      throw new PreviewDataSessionError('runtime_unavailable', 'Preview data runtime is unavailable');
    }
    const origin = normalizeOrigin(input.origin);
    const browserGeneration = normalizeGeneration(input.browserGeneration);
    const active = await this.projects.requireActiveRealProjectWithIdentity(input.projectPath);
    if (this.disposed) {
      throw new PreviewDataSessionError('runtime_unavailable', 'Preview data runtime is unavailable');
    }

    const now = this.now();
    const tokenBytes = this.createRandomBytes(32);
    if (tokenBytes.length !== 32) {
      throw new PreviewDataSessionError('runtime_unavailable', 'Preview data runtime is unavailable');
    }
    const token = tokenBytes.toString('base64url');
    this.invalidate('replaced');
    const snapshot: PreviewDataSessionSnapshot = Object.freeze({
      projectPath: active.path,
      projectId: active.projectId,
      origin,
      browserGeneration,
      createdAt: now,
    });
    this.session = {
      snapshot,
      token,
      tokenBytes,
      tokens: PREVIEW_DATA_BUCKET_CAPACITY,
      lastRefillMs: now,
    };
    this.boundAccountBinding = this.getAccountBinding();
    return snapshot;
  }

  getSnapshot(): PreviewDataSessionSnapshot | null {
    return this.session?.snapshot ?? null;
  }

  subscribeInvalidation(listener: PreviewDataSessionInvalidationListener): () => void {
    if (this.disposed) return () => undefined;
    this.invalidationListeners.add(listener);
    return () => {
      this.invalidationListeners.delete(listener);
    };
  }

  getInjectionValue(hostPort = this.hostPort): PreviewDataInjectionValue | null {
    const session = this.session;
    if (!session || !Number.isSafeInteger(hostPort) || hostPort < 1 || hostPort > 65_535) {
      return null;
    }
    return Object.freeze({
      endpoint: `http://127.0.0.1:${hostPort}${PREVIEW_DATA_ROUTE_ROOT}`,
      token: session.token,
      contractVersion: PREVIEW_DATA_CONTRACT_VERSION,
    });
  }

  private accountBindingIsCurrent(): boolean {
    const currentSession = this.getSessionSnapshot();
    const currentBinding = this.getAccountBinding();
    if (!this.boundAccountBinding) return !currentSession && !currentBinding;
    return Boolean(currentBinding && this.isCurrentAccountBinding(this.boundAccountBinding));
  }

  corsOrigin(req: IncomingMessage): string | null {
    const origin = headerValue(req, 'origin');
    const session = this.session;
    return session && origin === session.snapshot.origin ? origin : null;
  }

  authorizeRequest(req: IncomingMessage): PreviewDataAuthorizationResult {
    const session = this.session;
    if (!session) return failure(503, 'runtime_unavailable', 'Preview data runtime is unavailable', true);
    if (!this.accountBindingIsCurrent()) {
      this.invalidate('account_changed');
      return failure(503, 'runtime_unavailable', 'Preview data runtime is unavailable', true);
    }

    const origin = headerValue(req, 'origin');
    if (origin !== session.snapshot.origin) {
      return failure(403, 'origin_not_allowed', 'Preview Origin is not allowed');
    }
    if (!requestHeadersAllowed(req)) {
      return failure(400, 'invalid_request', 'Preview request headers are not allowed');
    }

    const method = (req.method ?? '').toUpperCase();
    if (method === 'OPTIONS') {
      return { ok: true, method, session: session.snapshot };
    }

    const authorization = headerValue(req, 'authorization') ?? '';
    if (!authorization.startsWith('Bearer ') || !isExactToken(authorization.slice(7), session.tokenBytes)) {
      return failure(401, 'authentication_required', 'Preview data authorization is required');
    }
    if (!PREVIEW_DATA_METHODS.has(method)) {
      return failure(405, 'method_not_allowed', 'Preview data method is not allowed');
    }

    const now = this.now();
    const elapsed = Math.max(0, now - session.lastRefillMs);
    session.tokens = Math.min(
      PREVIEW_DATA_BUCKET_CAPACITY,
      session.tokens + elapsed / 1_000 * PREVIEW_DATA_RATE,
    );
    session.lastRefillMs = now;
    if (session.tokens < 1) {
      return failure(
        429,
        'rate_limited',
        'Preview data request rate limit exceeded',
        true,
        retryAfterSeconds(session.tokens),
      );
    }
    session.tokens -= 1;
    return { ok: true, method, session: session.snapshot };
  }

  isCurrentSession(snapshot: PreviewDataSessionSnapshot): boolean {
    return this.session?.snapshot === snapshot;
  }

  handleAgentBrowserLifecycle(event: AgentBrowserLifecycleEvent): void {
    const session = this.session;
    if (!session || !samePath(session.snapshot.projectPath, event.projectPath)) return;
    if (event.generation !== session.snapshot.browserGeneration) return;
    switch (event.type) {
      case 'closed':
        this.invalidate('preview_closed');
        break;
      case 'cross-origin-navigation':
        this.invalidate('cross_origin_navigation');
        break;
      case 'generation-replaced':
        this.invalidate('browser_generation_replaced');
        break;
      case 'detached':
        this.invalidate('browser_detached');
        break;
      case 'crashed':
        this.invalidate('browser_crashed');
        break;
      default:
        break;
    }
  }

  invalidate(reason: PreviewDataSessionInvalidationReason = 'manual'): void {
    if (!this.session) return;
    this.session = null;
    for (const listener of this.invalidationListeners) {
      try {
        listener(reason);
      } catch {
        // A browser cleanup observer must not interrupt token invalidation.
      }
    }
  }

  dispose(): void {
    if (this.disposed) return;
    this.disposed = true;
    this.invalidate('main_shutdown');
    this.unsubscribeWorksSession();
    this.invalidationListeners.clear();
  }
}

export function createPreviewDataSessionManager(
  options: PreviewDataSessionManagerOptions,
): PreviewDataSessionManager {
  return new PreviewDataSessionManager(options);
}