feat(harness): secure session model configuration

2026-09-07 15:55:30 +08:00
parent 85739f58ed
commit f9ea4ce429
9 changed files with 511 additions and 1 deletions
@@ -3,6 +3,10 @@ export * from "./approvals/in-memory-approval-broker.js";
 export * from "./events/event-store.js";
 export * from "./events/in-memory-event-store.js";
 export * from "./errors/harness-error.js";
+export * from "./models/base-url-policy.js";
+export * from "./models/in-memory-model-connection-store.js";
+export * from "./models/model-connection-store.js";
+export * from "./models/pi-model-runtime-factory.js";
 export * from "./security/redact.js";
 export * from "./skills/skill-catalog.js";
 export * from "./skills/skill-event-inference.js";
@@ -0,0 +1,157 @@
+import { lookup } from "node:dns/promises";
+import { BlockList, isIP } from "node:net";
+import { HarnessError } from "../errors/harness-error.js";
+
+export type ResolvedAddress = { address: string; family: 4 | 6 };
+export type AddressResolver = (hostname: string) => Promise<ResolvedAddress[]>;
+export type BaseUrlPolicyOptions = {
+  allowLocalhostHttp?: boolean;
+  allowPrivateNetwork?: boolean;
+};
+export type ValidatedBaseUrl = {
+  normalizedUrl: string;
+  addresses: string[];
+  redirectPolicy: "same-origin";
+};
+
+const alwaysBlocked = createAlwaysBlockedList();
+const privateBlocked = createPrivateBlockedList();
+const sensitiveQueryNames = new Set([
+  "apikey",
+  "authorization",
+  "token",
+  "secret",
+  "password",
+  "credential",
+]);
+
+export async function validateBaseUrl(
+  value: string,
+  options: BaseUrlPolicyOptions = {},
+  resolver: AddressResolver = defaultResolver,
+): Promise<ValidatedBaseUrl> {
+  let url: URL;
+  try {
+    url = new URL(value);
+  } catch {
+    throw notAllowed("Base URL is malformed");
+  }
+
+  const hostname = stripIpv6Brackets(url.hostname);
+  const localhostHttp = url.protocol === "http:" && isLocalhostName(hostname) && options.allowLocalhostHttp === true;
+  if (url.protocol !== "https:" && !localhostHttp) {
+    throw notAllowed("Base URL must use HTTPS");
+  }
+  if (url.username !== "" || url.password !== "") {
+    throw notAllowed("Base URL must not contain user information");
+  }
+  if (url.hash !== "") throw notAllowed("Base URL must not contain a fragment");
+  for (const key of url.searchParams.keys()) {
+    const normalized = key.toLowerCase().replaceAll(/[^a-z0-9]/g, "");
+    if ([...sensitiveQueryNames].some((name) => normalized === name || normalized.endsWith(name))) {
+      throw notAllowed("Base URL must not contain credentials in its query");
+    }
+  }
+  if (hostname === "") throw notAllowed("Base URL must contain a hostname");
+
+  const addresses = isIP(hostname) === 0
+    ? await resolveAddresses(hostname, resolver)
+    : [{ address: hostname, family: isIP(hostname) as 4 | 6 }];
+  const uniqueAddresses = [...new Set(addresses.map(({ address }) => stripIpv6Brackets(address)))];
+  if (uniqueAddresses.length === 0) throw notAllowed("Base URL hostname did not resolve");
+  for (const address of uniqueAddresses) {
+    assertAllowedAddress(address, options.allowPrivateNetwork === true);
+  }
+
+  return {
+    normalizedUrl: url.toString().replace(/\/$/, url.pathname === "/" ? "" : "/"),
+    addresses: uniqueAddresses,
+    redirectPolicy: "same-origin",
+  };
+}
+
+async function defaultResolver(hostname: string): Promise<ResolvedAddress[]> {
+  const results = await lookup(hostname, { all: true, verbatim: true });
+  return results.flatMap(({ address, family }) =>
+    family === 4 || family === 6 ? [{ address, family }] : [],
+  );
+}
+
+async function resolveAddresses(hostname: string, resolver: AddressResolver): Promise<ResolvedAddress[]> {
+  try {
+    return await resolver(hostname);
+  } catch {
+    throw notAllowed("Base URL hostname could not be resolved");
+  }
+}
+
+function assertAllowedAddress(address: string, allowPrivateNetwork: boolean): void {
+  const familyNumber = isIP(address);
+  if (familyNumber === 0) throw notAllowed("Resolver returned an invalid IP address");
+  const family = familyNumber === 4 ? "ipv4" : "ipv6";
+  if (alwaysBlocked.check(address, family)) throw notAllowed("Base URL resolves to a restricted address");
+  if (!allowPrivateNetwork && privateBlocked.check(address, family)) {
+    throw notAllowed("Base URL resolves to a private address");
+  }
+}
+
+function isLocalhostName(hostname: string): boolean {
+  const normalized = hostname.toLowerCase().replace(/\.$/, "");
+  return normalized === "localhost" || normalized.endsWith(".localhost");
+}
+
+function stripIpv6Brackets(value: string): string {
+  return value.startsWith("[") && value.endsWith("]") ? value.slice(1, -1) : value;
+}
+
+function notAllowed(message: string): HarnessError {
+  return new HarnessError("BASE_URL_NOT_ALLOWED", message);
+}
+
+function createAlwaysBlockedList(): BlockList {
+  const list = new BlockList();
+  addSubnets(list, "ipv4", [
+    ["0.0.0.0", 8],
+    ["169.254.0.0", 16],
+    ["224.0.0.0", 4],
+    ["240.0.0.0", 4],
+  ]);
+  addSubnets(list, "ipv6", [
+    ["::", 128],
+    ["::ffff:0.0.0.0", 104],
+    ["::ffff:169.254.0.0", 112],
+    ["::ffff:224.0.0.0", 100],
+    ["::ffff:240.0.0.0", 100],
+    ["fe80::", 10],
+    ["ff00::", 8],
+  ]);
+  list.addAddress("fd00:ec2::254", "ipv6");
+  return list;
+}
+
+function createPrivateBlockedList(): BlockList {
+  const list = new BlockList();
+  addSubnets(list, "ipv4", [
+    ["10.0.0.0", 8],
+    ["100.64.0.0", 10],
+    ["127.0.0.0", 8],
+    ["172.16.0.0", 12],
+    ["192.168.0.0", 16],
+    ["198.18.0.0", 15],
+  ]);
+  addSubnets(list, "ipv6", [
+    ["::1", 128],
+    ["::ffff:10.0.0.0", 104],
+    ["::ffff:100.64.0.0", 106],
+    ["::ffff:127.0.0.0", 104],
+    ["::ffff:172.16.0.0", 108],
+    ["::ffff:192.168.0.0", 112],
+    ["::ffff:198.18.0.0", 111],
+    ["fc00::", 7],
+  ]);
+  return list;
+}
+
+function addSubnets(list: BlockList, family: "ipv4" | "ipv6", entries: Array<[string, number]>): void {
+  for (const [network, prefix] of entries) list.addSubnet(network, prefix, family);
+}
@@ -0,0 +1,115 @@
+import {
+  ModelConnectionInputSchema,
+  RedactedModelConnectionSchema,
+  type ModelConnectionInput,
+  type RedactedModelConnection,
+} from "@agent-studio/shared";
+import { HarnessError } from "../errors/harness-error.js";
+import { validateBaseUrl, type ValidatedBaseUrl } from "./base-url-policy.js";
+import type { ModelConnectionStore, ResolvedModelConnection } from "./model-connection-store.js";
+
+type ConnectionMetadata = Omit<ModelConnectionInput, "apiKey">;
+type UrlValidator = (url: string) => Promise<ValidatedBaseUrl>;
+
+export type InMemoryModelConnectionStoreOptions = {
+  environment?: ModelConnectionInput;
+  validateUrl?: UrlValidator;
+};
+
+export class InMemoryModelConnectionStore implements ModelConnectionStore {
+  readonly #metadata = new Map<string, ConnectionMetadata>();
+  readonly #credentials = new Map<string, string>();
+  readonly #environmentMetadata: ConnectionMetadata | undefined;
+  readonly #environmentKey: string | undefined;
+  readonly #validateUrl: UrlValidator;
+
+  public constructor(options: InMemoryModelConnectionStoreOptions = {}) {
+    if (options.environment === undefined) {
+      this.#environmentMetadata = undefined;
+      this.#environmentKey = undefined;
+    } else {
+      const parsed = ModelConnectionInputSchema.parse(options.environment);
+      assertProviderApi(parsed);
+      const { apiKey, ...metadata } = parsed;
+      this.#environmentMetadata = metadata;
+      this.#environmentKey = apiKey;
+    }
+    this.#validateUrl = options.validateUrl ?? (async (url) => validateBaseUrl(url));
+  }
+
+  public async configure(sessionId: string, input: ModelConnectionInput): Promise<void> {
+    const parsed = ModelConnectionInputSchema.parse(input);
+    assertProviderApi(parsed);
+    let normalized = parsed;
+    if (parsed.baseUrl !== undefined) {
+      const result = await this.#validateUrl(parsed.baseUrl);
+      normalized = { ...parsed, baseUrl: result.normalizedUrl };
+    }
+    const { apiKey, ...metadata } = normalized;
+    this.#metadata.set(sessionId, metadata);
+    this.#credentials.delete(sessionId);
+    if (apiKey !== undefined) this.#credentials.set(sessionId, apiKey);
+  }
+
+  public getRedacted(sessionId: string): RedactedModelConnection | undefined {
+    const resolved = this.#resolveOptional(sessionId);
+    if (resolved === undefined) return undefined;
+    const key = resolved.apiKey;
+    const safeConnection = { ...resolved.connection };
+    delete safeConnection.apiKey;
+    return RedactedModelConnectionSchema.parse({
+      ...safeConnection,
+      configured: true,
+      source: resolved.source,
+      ...(key === undefined ? {} : { keyHint: keyHint(key) }),
+    });
+  }
+
+  public resolveConnection(sessionId: string): ResolvedModelConnection {
+    const resolved = this.#resolveOptional(sessionId);
+    if (resolved !== undefined) return resolved;
+    throw new HarnessError("MODEL_NOT_CONFIGURED", `No model is configured for session ${sessionId}`);
+  }
+
+  public clearCredentials(sessionId: string): void {
+    this.#credentials.delete(sessionId);
+    this.#metadata.delete(sessionId);
+  }
+
+  #resolveOptional(sessionId: string): ResolvedModelConnection | undefined {
+    const metadata = this.#metadata.get(sessionId);
+    if (metadata !== undefined) {
+      const apiKey = this.#credentials.get(sessionId);
+      return {
+        connection: { ...metadata, ...(apiKey === undefined ? {} : { apiKey }) },
+        ...(apiKey === undefined ? {} : { apiKey }),
+        source: "session",
+      };
+    }
+    if (this.#environmentMetadata === undefined) return undefined;
+    return {
+      connection: {
+        ...this.#environmentMetadata,
+        ...(this.#environmentKey === undefined ? {} : { apiKey: this.#environmentKey }),
+      },
+      ...(this.#environmentKey === undefined ? {} : { apiKey: this.#environmentKey }),
+      source: "environment",
+    };
+  }
+}
+
+function keyHint(apiKey: string): string {
+  return `••••${apiKey.slice(-4)}`;
+}
+
+function assertProviderApi(input: ModelConnectionInput): void {
+  const mismatch =
+    (input.providerId === "anthropic" && input.api !== "anthropic-messages") ||
+    (input.providerId === "openai" && input.api === "anthropic-messages");
+  if (mismatch) {
+    throw new HarnessError(
+      "MODEL_API_MISMATCH",
+      `Provider ${input.providerId} is incompatible with API ${input.api}`,
+    );
+  }
+}
@@ -0,0 +1,14 @@
+import type { ModelConnectionInput, RedactedModelConnection } from "@agent-studio/shared";
+
+export type ResolvedModelConnection = {
+  connection: ModelConnectionInput;
+  apiKey?: string;
+  source: "session" | "environment";
+};
+
+export interface ModelConnectionStore {
+  configure(sessionId: string, input: ModelConnectionInput): Promise<void>;
+  getRedacted(sessionId: string): RedactedModelConnection | undefined;
+  resolveConnection(sessionId: string): ResolvedModelConnection;
+  clearCredentials(sessionId: string): void;
+}
@@ -0,0 +1,40 @@
+import { InMemoryCredentialStore } from "@earendil-works/pi-ai";
+import { ModelRuntime } from "@earendil-works/pi-coding-agent";
+import type { ResolvedModelConnection } from "./model-connection-store.js";
+
+export async function createPiModelRuntime(
+  resolved: ResolvedModelConnection,
+  signal?: AbortSignal,
+): Promise<ModelRuntime> {
+  const credentials = new InMemoryCredentialStore();
+  const runtime = await ModelRuntime.create({
+    credentials,
+    modelsPath: null,
+    allowModelNetwork: false,
+    refreshOnCreate: false,
+    ...(signal === undefined ? {} : { signal }),
+  });
+  const { connection } = resolved;
+  runtime.registerProvider(connection.providerId, {
+    name: connection.displayName ?? connection.providerId,
+    api: connection.api,
+    ...(connection.baseUrl === undefined ? {} : { baseUrl: connection.baseUrl }),
+    models: [
+      {
+        id: connection.modelId,
+        name: connection.displayName ?? connection.modelId,
+        api: connection.api,
+        ...(connection.baseUrl === undefined ? {} : { baseUrl: connection.baseUrl }),
+        reasoning: false,
+        input: connection.supportsImages ? ["text", "image"] : ["text"],
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+        contextWindow: connection.contextWindow ?? 32_000,
+        maxTokens: connection.maxTokens ?? 4_096,
+      },
+    ],
+  });
+  if (resolved.apiKey !== undefined) {
+    await runtime.setRuntimeApiKey(connection.providerId, resolved.apiKey, signal === undefined ? {} : { signal });
+  }
+  return runtime;
+}