OmniRoute/src/lib/db/upstreamProxy.ts

/** Upstream proxy config persistence for upstream_proxy_config table. */
import { getDbInstance } from "./core";
import {
  isCloudMetadataHost,
  isPrivateHost as isPrivateNetworkHost,
  mappedIpv4Host,
} from "@/shared/network/outboundUrlGuard";
import { ipVersion, normalizeHost } from "@/shared/network/privateHost";

/** Which embedded proxy handles the retry leg when mode === "fallback". */
export type FallbackBackend = "cliproxyapi" | "dario";

interface UpstreamProxyConfig {
  id: number;
  providerId: string;
  mode: string;
  cliproxyapiModelMapping: Record<string, unknown> | null;
  nativePriority: number;
  cliproxyapiPriority: number;
  enabled: boolean;
  family: string;
  // #dario: retry-leg backend for mode="fallback" ("cliproxyapi" default).
  fallbackBackend: FallbackBackend;
  createdAt: string;
  updatedAt: string;
}

interface UpstreamProxyRow {
  id: unknown;
  provider_id: unknown;
  mode: unknown;
  cliproxyapi_model_mapping: unknown;
  native_priority: unknown;
  cliproxyapi_priority: unknown;
  enabled: unknown;
  family: unknown;
  fallback_backend: unknown;
  created_at: unknown;
  updated_at: unknown;
}

function toRecord(value: unknown): Record<string, unknown> {
  return value && typeof value === "object" ? (value as Record<string, unknown>) : {};
}

const LOOPBACK_HOSTNAMES = new Set(["localhost", "127.0.0.1", "::1"]);

/** IPv4 multicast (224.0.0.0/4) — kept from this module's original rule set. */
function isMulticastIpv4(host: string): boolean {
  const first = Number.parseInt(host.split(".")[0], 10);
  return ipVersion(host) === 4 && first >= 224 && first <= 239;
}

/**
 * Reject a proxy target that is private or cloud-metadata, judging the ADDRESS
 * rather than its spelling.
 *
 * This module used to carry its own prefix regexes, which matched only the
 * dotted form: `http://169.254.169.254` was refused while
 * `http://[::ffff:169.254.169.254]` — the same address, serialised by WHATWG
 * URL as `::ffff:a9fe:a9fe` — was accepted, as were `::ffff:10.0.0.5`,
 * `fd00::/8`, `fe80::/10` and CGNAT `100.64.0.0/10`. #10843 fixed exactly that
 * class in the shared guard; routing this copy through the same helpers keeps
 * the two from drifting apart again.
 *
 * The deliberate exception stays: CLIProxyAPI runs on localhost:8317, so
 * loopback is allowed — and now so is its mapped spelling, for the same
 * address-not-spelling reason.
 */
function isPrivateHost(hostname: string): boolean {
  const normalized = normalizeHost(hostname);
  const asIpv4 = mappedIpv4Host(normalized) ?? normalized;

  if (LOOPBACK_HOSTNAMES.has(normalized) || LOOPBACK_HOSTNAMES.has(asIpv4)) return false;

  return (
    isCloudMetadataHost(normalized) || isPrivateNetworkHost(normalized) || isMulticastIpv4(asIpv4)
  );
}

export function validateProxyUrl(
  url: string
): { valid: true; url: string } | { valid: false; error: string } {
  try {
    const parsed = new URL(url);
    if (!["http:", "https:"].includes(parsed.protocol)) {
      return {
        valid: false,
        error: `Unsupported protocol "${parsed.protocol}" — use http or https`,
      };
    }
    if (isPrivateHost(parsed.hostname)) {
      return {
        valid: false,
        error: `Proxy URL cannot point to private/internal address "${parsed.hostname}"`,
      };
    }
    return { valid: true, url };
  } catch {
    return { valid: false, error: `Invalid URL: "${url}"` };
  }
}

/** Normalize an arbitrary stored/user value to a valid FallbackBackend. */
function normalizeFallbackBackend(value: unknown): FallbackBackend {
  return value === "dario" ? "dario" : "cliproxyapi";
}

function rowToConfig(record: Record<string, unknown>): UpstreamProxyConfig {
  let mapping: Record<string, unknown> | null = null;
  if (record.cliproxyapi_model_mapping && typeof record.cliproxyapi_model_mapping === "string") {
    try {
      mapping = JSON.parse(record.cliproxyapi_model_mapping);
    } catch {
      mapping = null;
    }
  }
  return {
    id: record.id as number,
    providerId: record.provider_id as string,
    mode: record.mode as string,
    cliproxyapiModelMapping: mapping,
    nativePriority: record.native_priority as number,
    cliproxyapiPriority: record.cliproxyapi_priority as number,
    enabled: record.enabled === 1 || record.enabled === true,
    family: typeof record.family === "string" ? record.family : "auto",
    fallbackBackend: normalizeFallbackBackend(record.fallback_backend),
    createdAt: record.created_at as string,
    updatedAt: record.updated_at as string,
  };
}

export async function getUpstreamProxyConfigs() {
  const db = getDbInstance();
  const rows = db
    .prepare("SELECT * FROM upstream_proxy_config ORDER BY provider_id")
    .all() as UpstreamProxyRow[];
  return rows.map((row) => rowToConfig(toRecord(row)));
}

export async function getUpstreamProxyConfig(providerId: string) {
  const db = getDbInstance();
  const row = db
    .prepare("SELECT * FROM upstream_proxy_config WHERE provider_id = ?")
    .get(providerId) as UpstreamProxyRow | undefined;
  if (!row) return null;
  return rowToConfig(toRecord(row));
}

export async function upsertUpstreamProxyConfig(data: {
  providerId: string;
  mode?: string;
  cliproxyapiModelMapping?: Record<string, unknown> | null;
  nativePriority?: number;
  cliproxyapiPriority?: number;
  enabled?: boolean;
  family?: string;
  fallbackBackend?: FallbackBackend;
}) {
  const db = getDbInstance();
  const mode = data.mode ?? "native";
  const cliproxyapiModelMapping =
    data.cliproxyapiModelMapping !== undefined
      ? JSON.stringify(data.cliproxyapiModelMapping)
      : null;
  const nativePriority = data.nativePriority ?? 1;
  const cliproxyapiPriority = data.cliproxyapiPriority ?? 2;
  const enabled = data.enabled !== false ? 1 : 0;
  const family = data.family ?? "auto";
  const fallbackBackend = normalizeFallbackBackend(data.fallbackBackend);

  db.prepare(
    `INSERT INTO upstream_proxy_config
     (provider_id, mode, cliproxyapi_model_mapping, native_priority, cliproxyapi_priority, enabled, family, fallback_backend, created_at, updated_at)
     VALUES (?, ?, ?, ?, ?, ?, ?, ?, datetime('now'), datetime('now'))
     ON CONFLICT(provider_id) DO UPDATE SET
       mode = excluded.mode,
       cliproxyapi_model_mapping = excluded.cliproxyapi_model_mapping,
       native_priority = excluded.native_priority,
       cliproxyapi_priority = excluded.cliproxyapi_priority,
       enabled = excluded.enabled,
       family = excluded.family,
       fallback_backend = excluded.fallback_backend,
       updated_at = datetime('now')`
  ).run(
    data.providerId,
    mode,
    cliproxyapiModelMapping,
    nativePriority,
    cliproxyapiPriority,
    enabled,
    family,
    fallbackBackend
  );

  return getUpstreamProxyConfig(data.providerId);
}

export async function updateUpstreamProxyConfig(
  providerId: string,
  updates: Record<string, unknown>
) {
  const db = getDbInstance();
  const current = await getUpstreamProxyConfig(providerId);
  if (!current) {
    throw new Error(`Provider ${providerId} not found`);
  }

  const sets: string[] = ["updated_at = datetime('now')"];
  const params: unknown[] = [];

  if (updates.mode !== undefined) {
    sets.push("mode = ?");
    params.push(updates.mode);
  }
  if (updates.cliproxyapiModelMapping !== undefined) {
    sets.push("cliproxyapi_model_mapping = ?");
    params.push(
      updates.cliproxyapiModelMapping === null
        ? null
        : JSON.stringify(updates.cliproxyapiModelMapping)
    );
  }
  if (updates.nativePriority !== undefined) {
    sets.push("native_priority = ?");
    params.push(updates.nativePriority);
  }
  if (updates.cliproxyapiPriority !== undefined) {
    sets.push("cliproxyapi_priority = ?");
    params.push(updates.cliproxyapiPriority);
  }
  if (updates.enabled !== undefined) {
    sets.push("enabled = ?");
    params.push(updates.enabled === true ? 1 : 0);
  }
  if (updates.family !== undefined) {
    sets.push("family = ?");
    params.push(updates.family);
  }
  if (updates.fallbackBackend !== undefined) {
    sets.push("fallback_backend = ?");
    params.push(normalizeFallbackBackend(updates.fallbackBackend));
  }

  params.push(providerId);
  db.prepare(`UPDATE upstream_proxy_config SET ${sets.join(", ")} WHERE provider_id = ?`).run(
    ...params
  );

  return getUpstreamProxyConfig(providerId);
}

export async function deleteUpstreamProxyConfig(providerId: string) {
  const db = getDbInstance();
  const result = db
    .prepare("DELETE FROM upstream_proxy_config WHERE provider_id = ?")
    .run(providerId);
  return result.changes > 0;
}

export async function getProvidersByMode(mode: string) {
  const db = getDbInstance();
  const rows = db
    .prepare(
      "SELECT * FROM upstream_proxy_config WHERE mode = ? AND enabled = 1 ORDER BY provider_id"
    )
    .all(mode) as UpstreamProxyRow[];
  return rows.map((row) => rowToConfig(toRecord(row)));
}

export async function getFallbackChainForProvider(providerId: string) {
  const config = await getUpstreamProxyConfig(providerId);
  if (!config) return [];

  const chain: { executor: "native" | "cliproxyapi" | "dario"; priority: number }[] = [];

  if (config.enabled) {
    chain.push({ executor: "native", priority: config.nativePriority });
    if (config.mode === "cliproxyapi") {
      chain.push({ executor: "cliproxyapi", priority: config.cliproxyapiPriority });
    } else if (config.mode === "dario") {
      chain.push({ executor: "dario", priority: config.cliproxyapiPriority });
    } else if (config.mode === "fallback") {
      chain.push({ executor: config.fallbackBackend, priority: config.cliproxyapiPriority });
    }
  }

  chain.sort((a, b) => a.priority - b.priority);
  return chain;
}