Skip to content
Source
Sign inSign up

CodexManagedTokenBroker.ts

225 lines8.4 KB

1// @effect-diagnostics nodeBuiltinImport:off - Server-owned private pipe/Unix socket, never a public listener.
2import * as NodeFSP from "node:fs/promises";
3import * as NodeNet from "node:net";
4import * as NodeOS from "node:os";
5import * as NodePath from "node:path";
6import * as NodeCrypto from "node:crypto";
7import { EnvironmentId, ProviderInstanceId } from "@t3tools/contracts";
8import * as HostProcess from "@t3tools/shared/HostProcess";
9import * as Context from "effect/Context";
10import * as Effect from "effect/Effect";
11import * as FiberSet from "effect/FiberSet";
12import * as Layer from "effect/Layer";
13import * as Exit from "effect/Exit";
14import * as Scope from "effect/Scope";
15import * as Schema from "effect/Schema";
16import * as ServerEnvironment from "../environment/ServerEnvironment.ts";
17import { makeCodexChatGptAuth } from "./CodexChatGptAuth.ts";
18import * as Protocol from "./CodexManagedTokenProtocol.ts";
19​
20const MAX_CONNECTIONS = 32;
21const MAX_JOBS = 16;
22// Stay below macOS's 104-byte and Linux's 108-byte sockaddr_un limits, including NUL.
23const UNIX_SOCKET_PATH_BYTES = 100;
24const encodeResponse = Schema.encodeSync(Schema.fromJsonString(Protocol.Response));
25class BrokerStartError extends Schema.TaggedError<BrokerStartError>()(
26 "ManagedTokenBrokerStartError",
27 {
28 cause: Schema.Defect(),
29 },
30) {
31 override get message() {
32 return "Could not start the private managed token listener.";
33 }
35export class CodexManagedTokenBroker extends Context.Service<
36 CodexManagedTokenBroker,
37 {
38 readonly endpoint: Effect.Effect<string, BrokerStartError>;
39 readonly awaitIdle: Effect.Effect<void>;
40 }
41>()("t3/provider/CodexManagedTokenBroker") {}
42​
43const io = <A>(f: () => Promise<A>) =>
44 Effect.tryPromise({ try: f, catch: (cause) => new BrokerStartError({ cause }) });
45const acquireListener = Effect.gen(function* () {
46 const identity = yield* ServerEnvironment.ServerEnvironmentIdentity;
47 const environmentId = yield* identity.getEnvironmentId;
48 const platform = yield* HostProcess.Platform;
49 const uid = yield* HostProcess.UserId;
50 const owned = yield* Effect.acquireRelease(
51 io(async () => {
52 if (platform === "win32")
53 return { endpoint: `\\\\.\\pipe\\t3-token-${NodeCrypto.randomUUID()}` };
54 if (uid === undefined) throw new Error("Missing IPC owner");
55 const acquireDirectory = async (candidate: string) => {
56 const parent = await NodeFSP.realpath(candidate);
57 const stat = await NodeFSP.lstat(parent);
58 if (
59 !stat.isDirectory() ||
60 stat.isSymbolicLink() ||
61 !(
62 (stat.uid === uid && (stat.mode & 0o022) === 0) ||
63 (stat.uid === 0 && ((stat.mode & 0o022) === 0 || (stat.mode & 0o1000) !== 0))
64 )
65 )
66 throw new Error("Untrusted IPC parent");
67 const directory = await NodeFSP.mkdtemp(NodePath.join(parent, "t3-token-"));
68 try {
69 await NodeFSP.chmod(directory, 0o700);
70 const created = await NodeFSP.lstat(directory);
71 if (
72 !created.isDirectory() ||
73 created.isSymbolicLink() ||
74 created.uid !== uid ||
75 (created.mode & 0o777) !== 0o700
76 )
77 throw new Error("Unprotected IPC directory");
78 const endpoint = NodePath.join(directory, "s");
79 if (Buffer.byteLength(endpoint) > UNIX_SOCKET_PATH_BYTES)
80 throw new Error("IPC path too long");
81 return { endpoint, directory, created };
82 } catch (error) {
83 await NodeFSP.rmdir(directory);
84 throw error;
85 }
86 };
87 return acquireDirectory(NodeOS.tmpdir()).catch(() => acquireDirectory("/tmp"));
88 }),
89 (owned) =>
90 Effect.promise(async () => {
91 if (!("directory" in owned)) return;
92 // Never recursively remove a replaced directory or another process's files.
93 const current = await NodeFSP.lstat(owned.directory).catch(() => undefined);
94 if (
95 current &&
96 current.dev === owned.created.dev &&
97 current.ino === owned.created.ino &&
98 !current.isSymbolicLink()
99 )
100 await NodeFSP.rmdir(owned.directory).catch(() => undefined);
101 }),
102 );
103 const jobs = yield* FiberSet.make<void>();
104 const runJob =
105 yield* FiberSet.runtime(jobs)<
106 Exclude<
107 Effect.Services<ReturnType<typeof makeCodexChatGptAuth>>,
108 import("effect/Scope").Scope
109 >
110 >();
111 const connections = yield* FiberSet.make<void>();
112 const runConnection = yield* FiberSet.runtime(connections)();
113 const sockets = new Set<NodeNet.Socket>();
114 let closing = false;
115 let admitted = 0;
116 // Keep branded validation here so the helper need not load the contracts graph.
117 const decode = Schema.decodeUnknownEffect(
118 Schema.fromJsonString(
119 Schema.Struct({
120 ...Protocol.Request.fields,
121 instanceId: ProviderInstanceId,
122 environmentId: EnvironmentId,
123 }),
124 ),
125 {
126 onExcessProperty: "error",
127 },
128 );
129 const server = NodeNet.createServer((socket) => {
130 socket.on("error", () => {});
131 if (closing || sockets.size >= MAX_CONNECTIONS) {
132 socket.destroy();
133 return;
134 }
135 sockets.add(socket);
136 socket.once("close", () => sockets.delete(socket));
137 let transferred = false;
138 runConnection(
139 Effect.gen(function* () {
140 // Windows pipes permit read-only peers by default; send no bytes until a complete request.
141 const request = yield* Protocol.readFrame(
142 socket,
143 Protocol.REQUEST_BYTES,
144 Protocol.SOCKET_IDLE_MS,
145 ).pipe(Effect.flatMap(decode));
146 if (closing || admitted >= MAX_JOBS || request.environmentId !== environmentId) return;
147 admitted++;
148 transferred = true;
149 runJob(
150 Effect.scoped(
151 Effect.gen(function* () {
152 const auth = yield* makeCodexChatGptAuth({
153 instanceId: request.instanceId,
154 expectedProfileKey: request.profileKey,
155 });
156 // Retire abandoned permit waits; access saves a started rotation before honoring interruption.
157 const credentials = yield* auth.access.pipe(
158 Effect.timeout(Protocol.COMMAND_TIMEOUT_MS),
159 );
160 yield* Protocol.writeFrame(
161 socket,
162 encodeResponse({ token: credentials.accessToken }),
163 Protocol.RESPONSE_BYTES,
164 );
165 }),
166 ).pipe(
167 Effect.ignoreCause,
168 Effect.ensuring(
169 Effect.sync(() => {
170 admitted--;
171 socket.destroy();
172 }),
173 ),
174 ),
175 );
176 }).pipe(
177 Effect.ignoreCause,
178 Effect.ensuring(
179 Effect.sync(() => {
180 if (!transferred) socket.destroy();
181 }),
182 ),
183 ),
184 );
185 });
186 server.on("error", () => {});
187 // Drain accepted rotations before the FiberSet finalizers interrupt their fibers.
188 yield* Effect.addFinalizer(() =>
189 Effect.gen(function* () {
190 closing = true;
191 for (const socket of sockets) socket.destroy();
192 yield* Effect.promise(() => new Promise<void>((resolve) => server.close(() => resolve())));
193 yield* FiberSet.awaitEmpty(connections);
194 yield* FiberSet.awaitEmpty(jobs);
195 }),
196 );
197 yield* Effect.callback<void, BrokerStartError>((resume) => {
198 const onError = (cause: Error) => resume(Effect.fail(new BrokerStartError({ cause })));
199 server.once("error", onError);
200 try {
201 server.listen(owned.endpoint, () => resume(Effect.void));
202 } catch (cause) {
203 resume(Effect.fail(new BrokerStartError({ cause })));
204 }
205 return Effect.sync(() => server.off("error", onError));
206 });
207 if (platform !== "win32") yield* io(() => NodeFSP.chmod(owned.endpoint, 0o600));
208 return { endpoint: owned.endpoint, awaitIdle: FiberSet.awaitEmpty(jobs) };
209});
210const make = Effect.gen(function* () {
211 const scope = yield* Scope.fork(yield* Effect.scope);
212 const acquired = yield* acquireListener.pipe(Scope.provide(scope), Effect.result);
213 if (acquired._tag === "Failure") {
214 yield* Scope.close(scope, Exit.fail(acquired.failure));
215 return CodexManagedTokenBroker.of({
216 endpoint: Effect.fail(acquired.failure),
217 awaitIdle: Effect.void,
218 });
219 }
220 return CodexManagedTokenBroker.of({
221 endpoint: Effect.succeed(acquired.success.endpoint),
222 awaitIdle: acquired.success.awaitIdle,
223 });
224});
225export const layer = Layer.effect(CodexManagedTokenBroker, make);