|
1 | 1 | import { ReadWriteConnection } from "../common/connection";
|
2 |
| -import { NewEvalMessage, ServerMessage, EvalDoneMessage, EvalFailedMessage, TypedValue, ClientMessage } from "../proto"; |
| 2 | +import { NewEvalMessage, ServerMessage, EvalDoneMessage, EvalFailedMessage, TypedValue, ClientMessage, NewSessionMessage, TTYDimensions, SessionOutputMessage, CloseSessionInputMessage } from "../proto"; |
3 | 3 | import { Emitter } from "@coder/events";
|
4 | 4 | import { logger, field } from "@coder/logger";
|
5 |
| - |
| 5 | +import { ChildProcess, SpawnOptions, ServerProcess } from "./command"; |
6 | 6 |
|
7 | 7 | export class Client {
|
8 |
| - |
9 | 8 | private evalId: number = 0;
|
10 | 9 | private evalDoneEmitter: Emitter<EvalDoneMessage> = new Emitter();
|
11 | 10 | private evalFailedEmitter: Emitter<EvalFailedMessage> = new Emitter();
|
12 | 11 |
|
| 12 | + private sessionId: number = 0; |
| 13 | + private sessions: Map<number, ServerProcess> = new Map(); |
| 14 | + |
13 | 15 | public constructor(
|
14 | 16 | private readonly connection: ReadWriteConnection,
|
15 | 17 | ) {
|
@@ -86,20 +88,108 @@ export class Client {
|
86 | 88 | if (failedMsg.getId() === id) {
|
87 | 89 | d1.dispose();
|
88 | 90 | d2.dispose();
|
89 |
| - |
| 91 | + |
90 | 92 | rej(failedMsg.getMessage());
|
91 | 93 | }
|
92 | 94 | });
|
93 |
| - |
| 95 | + |
94 | 96 | return prom;
|
95 | 97 | }
|
96 |
| - |
| 98 | + |
| 99 | + /** |
| 100 | + * Spawns a process from a command. _Somewhat_ reflects the "child_process" API. |
| 101 | + * @param command |
| 102 | + * @param args Arguments |
| 103 | + * @param options Options to execute for the command |
| 104 | + */ |
| 105 | + public spawn(command: string, args: string[] = [], options?: SpawnOptions): ChildProcess { |
| 106 | + return this.doSpawn(command, args, options, false); |
| 107 | + } |
| 108 | + |
| 109 | + /** |
| 110 | + * Fork a module. |
| 111 | + * @param modulePath Path of the module |
| 112 | + * @param args Args to add for the module |
| 113 | + * @param options Options to execute |
| 114 | + */ |
| 115 | + public fork(modulePath: string, args: string[] = [], options?: SpawnOptions): ChildProcess { |
| 116 | + return this.doSpawn(modulePath, args, options, true); |
| 117 | + } |
| 118 | + |
| 119 | + private doSpawn(command: string, args: string[] = [], options?: SpawnOptions, isFork: boolean = false): ChildProcess { |
| 120 | + const id = this.sessionId++; |
| 121 | + const newSess = new NewSessionMessage(); |
| 122 | + newSess.setId(id); |
| 123 | + newSess.setCommand(command); |
| 124 | + newSess.setArgsList(args); |
| 125 | + newSess.setIsFork(isFork); |
| 126 | + if (options) { |
| 127 | + if (options.cwd) { |
| 128 | + newSess.setCwd(options.cwd); |
| 129 | + } |
| 130 | + if (options.env) { |
| 131 | + Object.keys(options.env).forEach((envKey) => { |
| 132 | + newSess.getEnvMap().set(envKey, options.env![envKey]); |
| 133 | + }); |
| 134 | + } |
| 135 | + if (options.tty) { |
| 136 | + const tty = new TTYDimensions(); |
| 137 | + tty.setHeight(options.tty.rows); |
| 138 | + tty.setWidth(options.tty.columns); |
| 139 | + newSess.setTtyDimensions(tty); |
| 140 | + } |
| 141 | + } |
| 142 | + const clientMsg = new ClientMessage(); |
| 143 | + clientMsg.setNewSession(newSess); |
| 144 | + this.connection.send(clientMsg.serializeBinary()); |
| 145 | + |
| 146 | + const serverProc = new ServerProcess(this.connection, id, options ? options.tty !== undefined : false); |
| 147 | + serverProc.stdin.on("close", () => { |
| 148 | + console.log("stdin closed"); |
| 149 | + const c = new CloseSessionInputMessage(); |
| 150 | + c.setId(id); |
| 151 | + const cm = new ClientMessage(); |
| 152 | + cm.setCloseSessionInput(c); |
| 153 | + this.connection.send(cm.serializeBinary()); |
| 154 | + }); |
| 155 | + this.sessions.set(id, serverProc); |
| 156 | + return serverProc; |
| 157 | + } |
| 158 | + |
97 | 159 | private handleMessage(message: ServerMessage): void {
|
98 | 160 | if (message.hasEvalDone()) {
|
99 | 161 | this.evalDoneEmitter.emit(message.getEvalDone()!);
|
100 | 162 | } else if (message.hasEvalFailed()) {
|
101 | 163 | this.evalFailedEmitter.emit(message.getEvalFailed()!);
|
| 164 | + } else if (message.hasNewSessionFailure()) { |
| 165 | + const s = this.sessions.get(message.getNewSessionFailure()!.getId()); |
| 166 | + if (!s) { |
| 167 | + return; |
| 168 | + } |
| 169 | + s.emit("error", new Error(message.getNewSessionFailure()!.getMessage())); |
| 170 | + this.sessions.delete(message.getNewSessionFailure()!.getId()); |
| 171 | + } else if (message.hasSessionDone()) { |
| 172 | + const s = this.sessions.get(message.getSessionDone()!.getId()); |
| 173 | + if (!s) { |
| 174 | + return; |
| 175 | + } |
| 176 | + s.emit("exit", message.getSessionDone()!.getExitStatus()); |
| 177 | + this.sessions.delete(message.getSessionDone()!.getId()); |
| 178 | + } else if (message.hasSessionOutput()) { |
| 179 | + const output = message.getSessionOutput()!; |
| 180 | + const s = this.sessions.get(output.getId()); |
| 181 | + if (!s) { |
| 182 | + return; |
| 183 | + } |
| 184 | + const data = new TextDecoder().decode(output.getData_asU8()); |
| 185 | + const stream = output.getFd() === SessionOutputMessage.FD.STDOUT ? s.stdout : s.stderr; |
| 186 | + stream.emit("data", data); |
| 187 | + } else if (message.hasIdentifySession()) { |
| 188 | + const s = this.sessions.get(message.getIdentifySession()!.getId()); |
| 189 | + if (!s) { |
| 190 | + return; |
| 191 | + } |
| 192 | + s.pid = message.getIdentifySession()!.getPid(); |
102 | 193 | }
|
103 | 194 | }
|
104 |
| - |
105 | 195 | }
|
0 commit comments