OmniRoute/src/lib/oauth/connectionPersistence.ts

/**
 * Shared upsert for OAuth provider connections, used by both the authenticated
 * OAuth route (`device-complete`) and the public Codex device-flow completion
 * endpoint. Mirrors the exchange/poll/poll-callback persistence: normalize the
 * display name, compute expiry, match an existing connection by id or email
 * (+ Codex workspaceId) and update it, else create a new one, then sync to Cloud.
 */
import { timingSafeEqual } from "crypto";
import {
  createProviderConnection,
  updateProviderConnection,
  getProviderConnections,
  isCloudEnabled,
} from "@/models";
import { getConsistentMachineId } from "@/shared/utils/machineId";
import { syncToCloud } from "@/lib/cloudSync";
import {
  antigravityDegradedProjectState,
  antigravityPersistStatus,
  type AntigravityDegradedProjectState,
} from "@/lib/oauth/antigravityProjectGate";

/**
 * Constant-time string comparison to prevent timing-oracle attacks (CWE-208).
 * Handles null/undefined safely and different-length strings.
 */
export function safeEqual(a: string | null | undefined, b: string | null | undefined): boolean {
  if (a == null || b == null) return a === b;
  const ba = Buffer.from(String(a));
  const bb = Buffer.from(String(b));
  if (ba.length !== bb.length) return false;
  return timingSafeEqual(ba, bb);
}

/**
 * #7737: does this existing Codex connection represent the SAME account as the
 * incoming login? Prefer workspaceId when either side has one (Team plans).
 * When NEITHER side has a workspaceId (Personal-plan logins, or two accounts
 * that both lack a Team workspace), a bare email match is not enough to prove
 * it's the same account — two different ChatGPT accounts can share an email
 * alias. Require chatgptUserId to agree; otherwise treat it as a distinct
 * account so the caller falls through to createProviderConnection (which
 * already does this same disambiguation, added under #6706).
 */
function isSameCodexAccount(
  existingProviderData: Record<string, any> | null | undefined,
  incomingProviderData: Record<string, any> | null | undefined
): boolean {
  const incomingWorkspace = incomingProviderData?.workspaceId;
  const existingWorkspace = existingProviderData?.workspaceId;
  if (incomingWorkspace || existingWorkspace) {
    return safeEqual(existingWorkspace, incomingWorkspace);
  }
  const incomingUserId = incomingProviderData?.chatgptUserId;
  const existingUserId = existingProviderData?.chatgptUserId;
  return Boolean(incomingUserId) && safeEqual(existingUserId, incomingUserId);
}

/**
 * Does this existing Claude connection represent the SAME account as the
 * incoming login? One Anthropic identity can belong to several organizations
 * at once — the personal workspace (`organizationType` `claude_max`/
 * `claude_pro`) plus any Team/Enterprise organization (`claude_team`) — and
 * they all authenticate with the same email and the same `accountUUID`, each
 * with its own tokens, plan and rate limits. `organizationUUID` is the only
 * field that separates them, so an email-only match silently overwrote
 * whichever organization had been connected first. Require `organizationUUID`
 * to agree whenever BOTH sides carry it; when either side lacks it (rows
 * stored before Claude started returning it) the legacy bare-email match
 * still applies, so re-authenticating an existing connection keeps updating
 * it in place instead of forking a duplicate.
 */
function isSameClaudeAccount(
  existingProviderData: Record<string, any> | null | undefined,
  incomingProviderData: Record<string, any> | null | undefined
): boolean {
  const incomingOrganization = incomingProviderData?.organizationUUID;
  const existingOrganization = existingProviderData?.organizationUUID;
  if (incomingOrganization && existingOrganization) {
    return safeEqual(existingOrganization, incomingOrganization);
  }
  return true;
}

/**
 * Find the existing OAuth connection (if any) that an incoming token payload
 * should be merged into, shared by every OAuth-completion call site
 * (persistOAuthConnection, and the exchange/poll/poll-callback branches in
 * `src/app/api/oauth/[provider]/[action]/route.ts`). Matches by explicit
 * connectionId first, then by same email + auth type — with Codex requiring
 * workspaceId/chatgptUserId agreement (#7737) and Claude requiring
 * organizationUUID agreement, to avoid silently overwriting a different
 * account that merely shares an email.
 */
export function findExistingOAuthConnectionMatch(
  existing: Array<Record<string, any>>,
  provider: string,
  tokenData: Record<string, any>,
  connectionId?: string
): Record<string, any> | undefined {
  return existing.find((c) => {
    if (c.id && safeEqual(connectionId, c.id)) return true;
    // Email dedup only when the payload actually carries an email. Without this
    // guard `safeEqual(undefined, undefined)` is true, so an email-less payload
    // would false-match the first email-less connection of the provider.
    if (!tokenData.email) return false;
    if (!safeEqual(c.email, tokenData.email) || c.authType !== "oauth") return false;
    if (provider === "codex") {
      return isSameCodexAccount(c.providerSpecificData, tokenData.providerSpecificData);
    }
    if (provider === "claude") {
      return isSameClaudeAccount(c.providerSpecificData, tokenData.providerSpecificData);
    }
    return true;
  });
}

/**
 * Build the create payload for a brand-new OAuth connection.
 *
 * #5326: mirror the freshly computed `expiresAt` into `tokenExpiresAt` at creation
 * time. The dashboard token-health badge prefers `tokenExpiresAt` over `expiresAt`
 * (ConnectionRow.tsx: `connection.tokenExpiresAt || connection.expiresAt`). If
 * `tokenExpiresAt` stays null on a freshly created connection, the badge falls back
 * to the original grant clock and can flash a false amber/"Token Expired" until the
 * first background refresh writes both fields together. All refresh paths already
 * persist `expiresAt` and `tokenExpiresAt` in lockstep
 * (tokenHealthCheck onPersist, tokenRefresh.updateProviderCredentials); this makes
 * creation consistent with them.
 */
export function buildOAuthConnectionCreatePayload(
  provider: string,
  tokenData: Record<string, any>,
  expiresAt: string | null,
  // warning is HTTP-response only (oauth route); persistence copies error fields.
  degradedProject?: AntigravityDegradedProjectState | null
) {
  return {
    provider,
    authType: "oauth" as const,
    ...tokenData,
    expiresAt,
    tokenExpiresAt: expiresAt,
    // #11284: degraded when Cloud Code projectId discovery failed at connect
    // time — the row is saved (refresh token stored, request-time bootstrap
    // can self-heal) but visibly NOT active. Spread AFTER tokenData so stale
    // error fields in the payload cannot leak through.
    ...antigravityPersistStatus(degradedProject),
  };
}

async function syncToCloudIfEnabled(): Promise<void> {
  try {
    const cloudEnabled = await isCloudEnabled();
    if (!cloudEnabled) return;
    const machineId = await getConsistentMachineId();
    await syncToCloud(machineId);
  } catch (error) {
    console.log("Error syncing to cloud after OAuth:", error);
  }
}

export async function persistOAuthConnection(
  provider: string,
  tokenData: any,
  connectionId?: string
) {
  // Normalize: if name is missing, use email or displayName as fallback label.
  if (!tokenData.name && (tokenData.email || tokenData.displayName)) {
    tokenData.name = tokenData.email || tokenData.displayName;
  }

  const expiresAt = tokenData.expiresIn
    ? new Date(Date.now() + tokenData.expiresIn * 1000).toISOString()
    : null;
  const degradedProject = antigravityDegradedProjectState(provider, tokenData);

  let connection: any;
  // A connectionId is an explicit "update THIS connection" signal (token refresh
  // / re-auth of a known connection); honor it even when the payload has no
  // top-level email. Some providers (e.g. GitHub Copilot) keep identity under
  // providerSpecificData, so gating dedup on tokenData.email alone created a
  // duplicate connection on every refresh (#8059).
  if (connectionId || tokenData.email) {
    const existing = await getProviderConnections({ provider });
    const match = findExistingOAuthConnectionMatch(existing, provider, tokenData, connectionId);
    const matchId = typeof match?.id === "string" ? match.id : null;
    if (matchId) {
      connection = await updateProviderConnection(matchId, {
        ...tokenData,
        expiresAt,
        ...antigravityPersistStatus(degradedProject),
        isActive: true,
      });
    }
  }
  if (!connection) {
    connection = await createProviderConnection(
      buildOAuthConnectionCreatePayload(provider, tokenData, expiresAt, degradedProject)
    );
  }

  await syncToCloudIfEnabled();
  return connection;
}