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
@@ -0,0 +1,70 @@
+import { describe, expect, it } from "vitest";
+import * as harness from "../src/index.js";
+
+type Resolver = (hostname: string) => Promise<Array<{ address: string; family: 4 | 6 }>>;
+type Validate = (
+  url: string,
+  options?: { allowLocalhostHttp?: boolean; allowPrivateNetwork?: boolean },
+  resolver?: Resolver,
+) => Promise<{ normalizedUrl: string; addresses: string[]; redirectPolicy: "same-origin" }>;
+
+function validate(): Validate {
+  const value = Reflect.get(harness, "validateBaseUrl") as Validate | undefined;
+  expect(value, "validateBaseUrl must be exported").toBeDefined();
+  return value!;
+}
+
+const publicResolver: Resolver = () => Promise.resolve([{ address: "93.184.216.34", family: 4 }]);
+
+describe("validateBaseUrl", () => {
+  it("normalizes an HTTPS URL and pins all validated addresses", async () => {
+    const result = await validate()("https://api.example.com/v1", {}, () => Promise.resolve([
+      { address: "2606:2800:220:1:248:1893:25c8:1946", family: 6 },
+      { address: "93.184.216.34", family: 4 },
+    ]));
+
+    expect(result).toEqual({
+      normalizedUrl: "https://api.example.com/v1",
+      addresses: ["2606:2800:220:1:248:1893:25c8:1946", "93.184.216.34"],
+      redirectPolicy: "same-origin",
+    });
+  });
+
+  it.each([
+    "http://169.254.169.254/latest/meta-data",
+    "https://127.0.0.1/v1",
+    "file:///etc/passwd",
+    "https://user:pass@example.com/v1",
+    "https://example.com/v1#secret",
+    "https://example.com/v1?api_key=secret",
+  ])("rejects restricted base URL %s", async (url) => {
+    await expect(validate()(url, {}, publicResolver)).rejects.toMatchObject({
+      code: "BASE_URL_NOT_ALLOWED",
+    });
+  });
+
+  it.each(["10.0.0.1", "100.64.0.1", "169.254.10.2", "172.16.0.1", "192.168.1.1", "::1", "::ffff:127.0.0.1", "fc00::1", "fe80::1"])(
+    "rejects a restricted resolved address %s",
+    async (address) => {
+      const family = address.includes(":") ? 6 : 4;
+      await expect(
+        validate()("https://api.example.com", {}, () => Promise.resolve([{ address, family }])),
+      ).rejects.toMatchObject({ code: "BASE_URL_NOT_ALLOWED" });
+    },
+  );
+
+  it("permits explicit localhost HTTP only in development mode", async () => {
+    const result = await validate()("http://localhost:11434/v1", { allowLocalhostHttp: true, allowPrivateNetwork: true }, () => Promise.resolve([
+      { address: "127.0.0.1", family: 4 },
+    ]));
+    expect(result.normalizedUrl).toBe("http://localhost:11434/v1");
+  });
+
+  it("never permits the metadata address even when private networks are enabled", async () => {
+    await expect(
+      validate()("https://metadata.internal", { allowPrivateNetwork: true }, () => Promise.resolve([
+        { address: "169.254.169.254", family: 4 },
+      ])),
+    ).rejects.toMatchObject({ code: "BASE_URL_NOT_ALLOWED" });
+  });
+});
@@ -0,0 +1,106 @@
+import { describe, expect, it } from "vitest";
+import * as harness from "../src/index.js";
+
+type Input = {
+  providerId: string;
+  api: "openai-responses" | "openai-completions" | "anthropic-messages";
+  apiKey?: string;
+  baseUrl?: string;
+  modelId: string;
+  displayName?: string;
+  supportsImages: boolean;
+  contextWindow?: number;
+  maxTokens?: number;
+};
+type Resolved = { connection: Input; apiKey?: string; source: "session" | "environment" };
+type Store = {
+  configure(sessionId: string, input: Input): Promise<void>;
+  getRedacted(sessionId: string): Record<string, unknown>;
+  resolveConnection(sessionId: string): Resolved;
+  clearCredentials(sessionId: string): void;
+};
+
+function input(overrides: Partial<Input> = {}): Input {
+  return {
+    providerId: "training-openai",
+    api: "openai-responses",
+    apiKey: "sk-test-secret",
+    baseUrl: "https://api.example.com/v1",
+    modelId: "training-model",
+    displayName: "Training Model",
+    supportsImages: true,
+    contextWindow: 32_000,
+    maxTokens: 4_096,
+    ...overrides,
+  };
+}
+
+function createStore(environment?: Input): Store {
+  const Constructor = Reflect.get(harness, "InMemoryModelConnectionStore") as unknown as
+    | (new (options?: Record<string, unknown>) => Store)
+    | undefined;
+  expect(Constructor, "InMemoryModelConnectionStore must be exported").toBeDefined();
+  return new Constructor!({
+    ...(environment === undefined ? {} : { environment }),
+    validateUrl: (url: string) => Promise.resolve({ normalizedUrl: url, addresses: ["93.184.216.34"], redirectPolicy: "same-origin" }),
+  });
+}
+
+describe("InMemoryModelConnectionStore", () => {
+  it("never returns the raw key", async () => {
+    const store = createStore();
+    await store.configure("ses_1", input());
+
+    const json = JSON.stringify(store.getRedacted("ses_1"));
+    expect(json).not.toContain("sk-test-secret");
+    expect(store.getRedacted("ses_1").keyHint).toBe("••••cret");
+  });
+
+  it("uses session config before environment and falls back after an idempotent clear", async () => {
+    const store = createStore(input({ providerId: "env-provider", apiKey: "env-secret", modelId: "env-model" }));
+    await store.configure("ses_1", input({ apiKey: "session-secret" }));
+    expect(store.resolveConnection("ses_1")).toMatchObject({ source: "session", apiKey: "session-secret" });
+
+    store.clearCredentials("ses_1");
+    store.clearCredentials("ses_1");
+    expect(store.resolveConnection("ses_1")).toMatchObject({ source: "environment", apiKey: "env-secret" });
+  });
+
+  it("rotates keys without exposing the old credential", async () => {
+    const store = createStore();
+    await store.configure("ses_1", input({ apiKey: "first-secret" }));
+    await store.configure("ses_1", input({ apiKey: "second-secret" }));
+
+    expect(store.resolveConnection("ses_1").apiKey).toBe("second-secret");
+    expect(JSON.stringify(store.getRedacted("ses_1"))).not.toContain("first-secret");
+  });
+
+  it("fails explicitly when no connection is configured", () => {
+    const store = createStore();
+    expect(() => store.resolveConnection("ses_missing")).toThrowError(
+      expect.objectContaining({ code: "MODEL_NOT_CONFIGURED" }),
+    );
+  });
+
+  it.each([
+    input({ providerId: "anthropic", api: "openai-responses" }),
+    input({ providerId: "openai", api: "anthropic-messages" }),
+  ])("rejects a built-in provider/API mismatch", async (invalid) => {
+    await expect(createStore().configure("ses_1", invalid)).rejects.toMatchObject({
+      code: "MODEL_API_MISMATCH",
+    });
+  });
+
+  it("creates a Pi runtime with only the configured model and in-memory credentials", async () => {
+    const createRuntime = Reflect.get(harness, "createPiModelRuntime") as
+      | ((resolved: Resolved) => Promise<{ getModel(providerId: string, modelId: string): unknown; getAuth(providerId: string): Promise<unknown> }>)
+      | undefined;
+    expect(createRuntime, "createPiModelRuntime must be exported").toBeDefined();
+    const resolved = { connection: input(), apiKey: "runtime-secret", source: "session" as const };
+
+    const runtime = await createRuntime!(resolved);
+
+    expect(runtime.getModel("training-openai", "training-model")).toBeDefined();
+    expect(await runtime.getAuth("training-openai")).toMatchObject({ auth: { apiKey: "runtime-secret" } });
+  });
+});