From e4178886fa2158cde4c5b70792f5634c4ff758e9 Mon Sep 17 00:00:00 2001 From: Shoubhit Dash Date: Fri, 21 Aug 2026 18:49:57 +0530 Subject: [PATCH] fix(core): align websocket network policy (#43875) --- bun.lock | 2 + packages/core/package.json | 2 + packages/core/src/effect/app-node-platform.ts | 4 +- .../core/src/effect/websocket-constructor.ts | 89 +++++++++++++++++++ .../test/effect-app-node-platform.test.ts | 86 ++++++++++++++++++ .../test/session-model-transport-live.test.ts | 4 +- .../core/test/session-model-transport.test.ts | 29 ++++-- 7 files changed, 206 insertions(+), 10 deletions(-) create mode 100644 packages/core/src/effect/websocket-constructor.ts create mode 100644 packages/core/test/effect-app-node-platform.test.ts diff --git a/bun.lock b/bun.lock index 1bd154090f..b521d91559 100644 --- a/bun.lock +++ b/bun.lock @@ -380,6 +380,8 @@ "google-auth-library": "10.5.0", "gray-matter": "4.0.3", "htmlparser2": "8.0.2", + "http-proxy-agent": "7.0.2", + "https-proxy-agent": "7.0.6", "ignore": "7.0.5", "immer": "11.1.4", "jsonc-parser": "3.3.1", diff --git a/packages/core/package.json b/packages/core/package.json index 5f9b924565..70d180aae6 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -133,6 +133,8 @@ "google-auth-library": "10.5.0", "gray-matter": "4.0.3", "htmlparser2": "8.0.2", + "http-proxy-agent": "7.0.2", + "https-proxy-agent": "7.0.6", "immer": "11.1.4", "ignore": "7.0.5", "jsonc-parser": "3.3.1", diff --git a/packages/core/src/effect/app-node-platform.ts b/packages/core/src/effect/app-node-platform.ts index 3ae252c6e7..dcf00887b3 100644 --- a/packages/core/src/effect/app-node-platform.ts +++ b/packages/core/src/effect/app-node-platform.ts @@ -1,8 +1,8 @@ import { LLMClient, RequestExecutor } from "@opencode-ai/ai/route" -import { NodeSocket } from "@effect/platform-node" import { Socket } from "effect/unstable/socket" import { makeGlobalNode } from "@opencode-ai/util/effect/app-node" import { httpClient } from "@opencode-ai/util/effect/app-node-platform" +import { WebSocketConstructor } from "./websocket-constructor.js" export const requestExecutor = makeGlobalNode({ service: RequestExecutor.Service, @@ -14,7 +14,7 @@ export const llmClient = makeGlobalNode({ service: LLMClient.Service, layer: LLM export const webSocketConstructor = makeGlobalNode({ service: Socket.WebSocketConstructor, - layer: NodeSocket.layerWebSocketConstructorWS, + layer: WebSocketConstructor.layer, deps: [], }) diff --git a/packages/core/src/effect/websocket-constructor.ts b/packages/core/src/effect/websocket-constructor.ts new file mode 100644 index 0000000000..6ffc7af74b --- /dev/null +++ b/packages/core/src/effect/websocket-constructor.ts @@ -0,0 +1,89 @@ +import { NodeSocket } from "@effect/platform-node" +import { HttpProxyAgent } from "http-proxy-agent" +import { HttpsProxyAgent } from "https-proxy-agent" +import { Layer } from "effect" +import { Headers } from "effect/unstable/http" +import { Socket } from "effect/unstable/socket" + +interface WebSocketOptions { + readonly headers?: Headers.Headers + readonly protocols?: string | Array +} + +type BunWebSocketConstructor = new ( + url: string, + options: WebSocketOptions & { readonly proxy?: string }, +) => globalThis.WebSocket + +type Environment = Readonly> + +const environmentValue = (environment: Environment, name: string) => + environment[name] ?? environment[name.toLowerCase()] + +const bypassesProxy = (url: URL, value: string | undefined) => { + if (!value) return false + const port = url.port || (url.protocol === "wss:" ? "443" : "80") + return value.split(/[\s,]+/).some((entry) => { + if (!entry) return false + if (entry === "*") return true + const match = entry.match(/^(.+?):(\d+)$/) + if (match?.[2] && match[2] !== port) return false + const host = (match?.[1] ?? entry).toLowerCase().replace(/^\*/, "") + const hostname = url.hostname.toLowerCase() + return host.startsWith(".") ? hostname.endsWith(host) : hostname === host + }) +} + +const proxy = (value: string, environment: Environment = process.env) => { + const url = new URL(value) + if (["127.0.0.1", "localhost", "::1"].includes(url.hostname)) return undefined + if (bypassesProxy(url, environmentValue(environment, "NO_PROXY"))) return undefined + const protocolProxy = url.protocol === "wss:" ? "WSS_PROXY" : "WS_PROXY" + const standardProxy = url.protocol === "wss:" ? "HTTPS_PROXY" : "HTTP_PROXY" + return ( + environmentValue(environment, protocolProxy) ?? + environmentValue(environment, standardProxy) ?? + environmentValue(environment, "ALL_PROXY") + ) +} + +const constructorOptions = (input: string | Array | undefined): WebSocketOptions => { + if (typeof input === "string" || Array.isArray(input)) return { protocols: input } + // AI routes pass handshake options through Effect's browser-shaped constructor. + return (input ?? {}) as WebSocketOptions +} + +const proxyAgent = (url: string, selectedProxy: string | undefined) => { + if (!selectedProxy) return undefined + if (url.startsWith("wss:") || selectedProxy.startsWith("https:")) return new HttpsProxyAgent(selectedProxy) + return new HttpProxyAgent(selectedProxy) +} + +const layer = Layer.succeed(Socket.WebSocketConstructor, (url, input) => { + const config = constructorOptions(input) + const selectedProxy = proxy(url) + // Keep trust on the runtime store so NODE_EXTRA_CA_CERTS remains additive. + if (typeof Bun !== "undefined") { + // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion -- Bun extends the browser constructor with handshake and network options. + const WebSocket = globalThis.WebSocket as unknown as BunWebSocketConstructor + return new WebSocket(url, { + headers: config.headers, + protocols: config.protocols, + ...(selectedProxy ? { proxy: selectedProxy } : {}), + }) + } + + const native = { + headers: config.headers, + agent: proxyAgent(url, selectedProxy), + // Reject redirects before headers can cross an origin boundary; the caller safely falls back to HTTP. + followRedirects: false, + } + const socket = config.protocols + ? new NodeSocket.NodeWS.WebSocket(url, config.protocols, native) + : new NodeSocket.NodeWS.WebSocket(url, native) + // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion -- ws implements the WebSocket surface consumed by the AI transport. + return socket as unknown as globalThis.WebSocket +}) + +export const WebSocketConstructor = { layer, proxy } as const diff --git a/packages/core/test/effect-app-node-platform.test.ts b/packages/core/test/effect-app-node-platform.test.ts new file mode 100644 index 0000000000..0e987a5015 --- /dev/null +++ b/packages/core/test/effect-app-node-platform.test.ts @@ -0,0 +1,86 @@ +import { describe, expect, test } from "bun:test" +import { WebSocketTransport } from "@opencode-ai/ai/route" +import { WebSocketConstructor } from "@opencode-ai/core/effect/websocket-constructor" +import { Effect } from "effect" +import { Headers } from "effect/unstable/http" +import { Socket } from "effect/unstable/socket" + +const makeServer = (fetch: (request: Request, server: Bun.Server) => Response | undefined) => + Effect.acquireRelease( + Effect.sync(() => + Bun.serve({ + hostname: "127.0.0.1", + port: 0, + fetch, + websocket: { message() {} }, + }), + ), + (server) => Effect.promise(() => server.stop(true)), + ) + +describe("WebSocket network policy", () => { + test("uses protocol-specific and standard proxy variables", () => { + expect( + WebSocketConstructor.proxy("wss://provider.test/responses", { + WSS_PROXY: "http://wss-proxy.test", + HTTPS_PROXY: "http://https-proxy.test", + }), + ).toBe("http://wss-proxy.test") + expect( + WebSocketConstructor.proxy("wss://provider.test/responses", { HTTPS_PROXY: "http://https-proxy.test" }), + ).toBe("http://https-proxy.test") + expect(WebSocketConstructor.proxy("ws://provider.test/responses", { HTTP_PROXY: "http://http-proxy.test" })).toBe( + "http://http-proxy.test", + ) + expect(WebSocketConstructor.proxy("ws://provider.test/responses", { ALL_PROXY: "http://all-proxy.test" })).toBe( + "http://all-proxy.test", + ) + }) + + test("respects no-proxy hosts, domains, ports, and wildcards", () => { + const environment = { HTTPS_PROXY: "http://proxy.test" } + expect( + WebSocketConstructor.proxy("wss://provider.test/responses", { ...environment, NO_PROXY: "provider.test" }), + ).toBeUndefined() + expect( + WebSocketConstructor.proxy("wss://api.provider.test/responses", { ...environment, NO_PROXY: ".provider.test" }), + ).toBeUndefined() + expect( + WebSocketConstructor.proxy("wss://provider.test:8443/responses", { + ...environment, + NO_PROXY: "provider.test:443", + }), + ).toBe("http://proxy.test") + expect( + WebSocketConstructor.proxy("wss://provider.test/responses", { ...environment, NO_PROXY: "*" }), + ).toBeUndefined() + }) + + test("rejects redirects without forwarding authorization", async () => { + let destinationRequests = 0 + await Effect.runPromise( + Effect.gen(function* () { + const destination = yield* makeServer((request, server) => { + destinationRequests++ + if (server.upgrade(request)) return undefined + return new Response("upgrade failed", { status: 426 }) + }) + const redirect = yield* makeServer( + () => + new Response(null, { + status: 302, + headers: { location: destination.url.toString().replace(/^http/, "ws") }, + }), + ) + const constructor = yield* Socket.WebSocketConstructor + + yield* WebSocketTransport.open({ + url: redirect.url.toString().replace(/^http/, "ws"), + headers: Headers.fromInput({ authorization: "Bearer secret" }), + }).pipe(Effect.provideService(Socket.WebSocketConstructor, constructor), Effect.flip) + + expect(destinationRequests).toBe(0) + }).pipe(Effect.scoped, Effect.provide(WebSocketConstructor.layer)), + ) + }) +}) diff --git a/packages/core/test/session-model-transport-live.test.ts b/packages/core/test/session-model-transport-live.test.ts index a3574a4f0d..4514bf0b3b 100644 --- a/packages/core/test/session-model-transport-live.test.ts +++ b/packages/core/test/session-model-transport-live.test.ts @@ -1,5 +1,4 @@ import { describe, expect, test } from "bun:test" -import { NodeSocket } from "@effect/platform-node" import { AIError, LLM, Message } from "@opencode-ai/ai" import { LLMClient, @@ -10,6 +9,7 @@ import { } from "@opencode-ai/ai/route" import { configure } from "@opencode-ai/ai/providers/openai" import { SessionModelTransport } from "@opencode-ai/core/session/model-transport" +import { WebSocketConstructor } from "@opencode-ai/core/effect/websocket-constructor" import { Session } from "@opencode-ai/schema/session" import { Effect, Fiber, Layer, Stream } from "effect" import { Headers } from "effect/unstable/http" @@ -48,7 +48,7 @@ const withServer = ( }), ), ) - }).pipe(Effect.scoped, Effect.provide(NodeSocket.layerWebSocketConstructorWS)), + }).pipe(Effect.scoped, Effect.provide(WebSocketConstructor.layer)), ) const collect = (transport: SessionModelTransport.Interface, item: WebSocketChannelExchange) => diff --git a/packages/core/test/session-model-transport.test.ts b/packages/core/test/session-model-transport.test.ts index b4177068d7..2e2bc23ebf 100644 --- a/packages/core/test/session-model-transport.test.ts +++ b/packages/core/test/session-model-transport.test.ts @@ -78,14 +78,15 @@ const collectComplete = ( const automatic = () => { const connections: Array<{ readonly messages: Queue.Queue + readonly headers: Headers.Headers closed: number sent: string[] }> = [] const connector: WebSocketConnector = { - open: () => + open: (input) => Effect.gen(function* () { const messages = yield* Queue.unbounded() - const record = { messages, closed: 0, sent: [] as string[] } + const record = { messages, headers: input.headers, closed: 0, sent: [] as string[] } connections.push(record) const connection: WebSocketConnection = { sendText: (message) => @@ -583,7 +584,7 @@ describe("SessionModelTransport", () => { ) }) - test("rotates when handshake affinity or connection age changes", async () => { + test("rotates when refreshed authorization changes handshake affinity", async () => { const fixture = automatic() await run( @@ -594,10 +595,26 @@ describe("SessionModelTransport", () => { yield* collect(executor, exchange("first", { headers: { authorization: "one" } })) yield* collect(executor, exchange("second", { headers: { authorization: "one" } })) yield* collect(executor, exchange("third", { headers: { authorization: "two" } })) + expect(fixture.connections).toHaveLength(2) + expect(fixture.connections[0]?.closed).toBe(1) + expect(fixture.connections.map((item) => item.headers.authorization)).toEqual(["one", "two"]) + }), + ) + }) + + test("rotates when the connection exceeds its requested age limit", async () => { + const fixture = automatic() + + await run( + fixture.connector, + Effect.gen(function* () { + const transport = yield* SessionModelTransport.Service + const executor = transport.bind(session) + yield* collect(executor, exchange("first")) yield* Effect.sleep("5 millis") - yield* collect(executor, exchange("fourth", { headers: { authorization: "two" }, rotateAfterMs: 1 })) - expect(fixture.connections).toHaveLength(3) - expect(fixture.connections.slice(0, 2).map((item) => item.closed)).toEqual([1, 1]) + yield* collect(executor, exchange("second", { rotateAfterMs: 1 })) + expect(fixture.connections).toHaveLength(2) + expect(fixture.connections[0]?.closed).toBe(1) }), ) })