diff --git a/src/runtime/protocol.ts b/src/runtime/protocol.ts index 91635ce..c822329 100644 --- a/src/runtime/protocol.ts +++ b/src/runtime/protocol.ts @@ -17,6 +17,7 @@ export function buildPayload(req: any, subPath: string): string { export function readVsockResponse( socket: Socket, timeout: number, + onStreamChunk?: (chunk: any) => void ): Promise<{ type: string; data: any; error?: string }> { return new Promise((resolve, reject) => { let buffer = ""; @@ -41,6 +42,14 @@ export function readVsockResponse( onData = (chunk: Buffer) => { buffer += chunk.toString(); + if (buffer.length > 10 * 1024 * 1024) { + clearTimeout(timer); + cleanup(); + socket.destroy(); + reject(new Error("Response too large")); + return; + } + let index; while ((index = buffer.indexOf("\n")) >= 0) { @@ -53,6 +62,11 @@ export function readVsockResponse( try { const msg = JSON.parse(line); + if (msg.type === "stream") { + onStreamChunk?.(msg); + continue; + } + if (msg.type === "response" || msg.type === "error") { clearTimeout(timer); cleanup(); diff --git a/src/runtime/transport.ts b/src/runtime/transport.ts index f853189..e4a3c47 100644 --- a/src/runtime/transport.ts +++ b/src/runtime/transport.ts @@ -90,3 +90,20 @@ export async function sendRequest(subPath: string, req: any, res: any, vm: Vm) { res.status(statusCode).json(msg.data ?? { error: msg.error }); } } + +export async function sendMessage( + vm: Vm, + message: Record, + onStreamChunk?: (chunk: any) => void, + timeout: number = 60000, +): Promise { + const socket = await getVmSocket(vm); + socket.write(JSON.stringify(message) + "\n"); + + transportLogger.debug( + { vmId: vm.id, messageType: message.type, messageId: message.id }, + "message sent to VM" + ); + + return readVsockResponse(socket, timeout, onStreamChunk); +} diff --git a/src/types/types.ts b/src/types/types.ts index a5ce129..cdb539d 100644 --- a/src/types/types.ts +++ b/src/types/types.ts @@ -34,3 +34,59 @@ export interface RuntimeFunction { vms: Vm[]; readyVms: Set; } + +export interface ExecuteMessage { + type: "execute"; + id: string; + command: string; + args: string[]; + cwd?: string; + env?: Record; + timeout?: number; +} + +export interface WriteFileMessage { + type: "write_file"; + id: string; + path: string; + content: string; + mode?: number; +} + +export interface ReadFileMessage { + type: "read_file"; + id: string; + path: string; +} + +export interface ListFilesMessage { + type: "list_files"; + id: string; + path?: string; + recursive?: boolean; +} + +export interface CancelMessage { + type: "cancel"; + id: string; +} + +export interface StreamMessage { + type: "stream"; + id: string; + stream: "stdout" | "stderr"; + data: string; +} + +export interface ResponseMessage { + type: "response"; + id: string; + data: any; +} + +export interface ErrorMessage { + type: "error"; + id: string; + error: string; + code?: number; +}