From 1b394f7442f6ccf4093cba9ea5c8b3380f5a3822 Mon Sep 17 00:00:00 2001 From: Yang Gao <6821729+criyle@users.noreply.github.com> Date: Sat, 23 Mar 2024 06:00:50 +0800 Subject: [PATCH] judge: sandbox: add stream interface (#725) Co-authored-by: undefined --- packages/hydrojudge/bin/hydrojudge.js | 1 + packages/hydrojudge/src/sandbox/client.ts | 80 +++++++++++++++++++- packages/hydrojudge/src/sandbox/interface.ts | 42 +++++++++- packages/hydrojudge/src/terminal.ts | 59 +++++++++++++++ 4 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 packages/hydrojudge/src/terminal.ts diff --git a/packages/hydrojudge/bin/hydrojudge.js b/packages/hydrojudge/bin/hydrojudge.js index 2d8b997e0..a1b21ed1e 100755 --- a/packages/hydrojudge/bin/hydrojudge.js +++ b/packages/hydrojudge/bin/hydrojudge.js @@ -4,4 +4,5 @@ const argv = require('cac')().parse(); require('@hydrooj/register'); if (argv.args[0] === 'cache') require('../src/cache')(); +else if (argv.args[0] === 'terminal') require('../src/terminal')(); else require('../src/daemon')(); diff --git a/packages/hydrojudge/src/sandbox/client.ts b/packages/hydrojudge/src/sandbox/client.ts index a5579a836..af2f15ea7 100644 --- a/packages/hydrojudge/src/sandbox/client.ts +++ b/packages/hydrojudge/src/sandbox/client.ts @@ -1,11 +1,86 @@ +import { EventEmitter } from 'events'; import fs from 'fs'; import superagent from 'superagent'; +import WebSocket from 'ws'; import { pipeRequest } from '@hydrooj/utils'; import { getConfig } from '../config'; -import { SandboxRequest, SandboxResult, SandboxVersion } from './interface'; +import { + Input, Output, Resize, SandboxRequest, SandboxResult, SandboxVersion, +} from './interface'; let url; +export class Stream extends EventEmitter { + private ws: WebSocket; + + constructor(httpUrl: string, req: SandboxRequest) { + super(); + const wsUrl = `${httpUrl.replace('http', 'ws')}/stream`; + this.ws = new WebSocket(wsUrl); + this.ws.onmessage = (e) => { + const data: Buffer = e.data as Buffer; + switch (data[0]) { + case 1: + this.emit('end', JSON.parse(data.toString('utf-8', 1))); + break; + case 2: + this.emit('output', { + index: (data[1] >> 4) & 0xf, + fd: (data[1]) & 0xf, + content: data.subarray(2), + }); + break; + default: + break; + } + }; + this.ws.onclose = (e) => { + this.emit('close', e); + }; + this.ws.onopen = (e) => { + this.request(req); + this.emit('open', e); + }; + this.ws.onerror = (e) => { + this.emit('error', e); + }; + } + + on(type: 'open', handle: (e: WebSocket.Event) => void): this; + on(type: 'close', handle: (e: WebSocket.CloseEvent) => void): this; + on(type: 'error', handle: (e: WebSocket.ErrorEvent) => void): this; + on(type: 'output', handle: (v: Output) => void): this; + on(type: 'end', handle: (v: SandboxResult[]) => void): this; + on(type: 'open' | 'close' | 'error' | 'output' | 'end', handle: (v: any) => void): this { + super.on(type, handle); + return this; + } + + private request(r: SandboxRequest) { + this.ws.send(Buffer.concat([Buffer.of(1), Buffer.from(JSON.stringify(r))])); + } + + resize(r: Resize) { + this.ws.send(Buffer.concat([Buffer.of(2), Buffer.from(JSON.stringify(r))])); + } + + input(i: Input) { + this.ws.send(Buffer.concat([Buffer.of(3, (i.index << 4) | i.fd), i.content])); + } + + cancel() { + this.ws.send(Buffer.of(4)); + } + + close() { + this.ws.close(); + } + + terminate() { + this.ws.terminate(); + } +} + const client = new Proxy({ run(req: SandboxRequest): Promise { return superagent.post(`${url}/run`).send(req).then((res) => res.body); @@ -29,6 +104,9 @@ const client = new Proxy({ config(): Promise> { return superagent.get(`${url}/config`).then((res) => res.body); }, + stream(req: SandboxRequest): Stream { + return new Stream(url, req); + }, }, { get(self, key) { url = getConfig('sandbox_host'); diff --git a/packages/hydrojudge/src/sandbox/interface.ts b/packages/hydrojudge/src/sandbox/interface.ts index e4213fb6b..6013a545c 100644 --- a/packages/hydrojudge/src/sandbox/interface.ts +++ b/packages/hydrojudge/src/sandbox/interface.ts @@ -6,6 +6,7 @@ export interface SandboxVersion { copyOutOptional?: boolean; pipeProxy?: boolean; addressSpaceLimit?: boolean; + stream?: boolean; } export interface LocalFile { @@ -26,9 +27,25 @@ export interface Collector { pipe?: boolean; } +// Symlink creates symlink for copyIn location: v1.6.0+ +export interface Symlink { + symlink: string; +} + +// StreamIn enables the stream input on /stream interface: v1.8.1+ +export interface StreamIn { + streamIn: boolean; +} + +// StreamOut enables the stream out on /stream interface: v1.8.1+ +export interface StreamOut { + streamOut: boolean; +} + export type CopyInFile = LocalFile | MemoryFile | PreparedFile; export type CopyIn = Record; -export type CmdFile = LocalFile | MemoryFile | PreparedFile | Collector | null; +// CmdFile defines file descriptor for the command. null is reserved for group execution and it must be provided via pipeMapping +export type CmdFile = LocalFile | MemoryFile | PreparedFile | Collector | StreamIn | StreamOut | null; /** * Cmd defines a single command to be executed by sandbox server @@ -59,7 +76,7 @@ export interface Cmd { addressSpaceLimit?: boolean; /** files to be copied into sandbox before execution */ - copyIn?: Record; + copyIn?: Record; /** * files to be copied out from sandbox after execution. @@ -147,3 +164,24 @@ export interface SandboxResult { /** contains detailed error if status is FileError */ fileError?: FileError[]; } + +export interface Resize { + index?: number; + fd?: number; + rows: number; + cols: number; + x?: number; + y?: number; +} + +export interface Input { + index?: number; + fd?: number; + content: Buffer; +} + +export interface Output { + index: number; + fd: number; + content: Buffer; +} diff --git a/packages/hydrojudge/src/terminal.ts b/packages/hydrojudge/src/terminal.ts new file mode 100644 index 000000000..14de853d8 --- /dev/null +++ b/packages/hydrojudge/src/terminal.ts @@ -0,0 +1,59 @@ +import log from './log'; +import client from './sandbox/client'; + +export = async function terminal() { + const version = await client.version(); + if (!version.stream) { + log.error('no stream support, please upgrade hydro-sandbox to v1.8.1+'); + return; + } + if (!process.stdin.isTTY) { + log.error('interactive terminal requires available tty'); + return; + } + process.stdin.setRawMode(true); + log.info('Starting /bin/bash, double press Ctrl+C to kill all child processes.'); + const stream = client.stream({ + cmd: [{ + args: ['/bin/bash'], + env: ['PATH=/usr/local/bin:/usr/bin:/bin', 'HOME=/tmp', `TERM=${process.env['TERM']}`], + files: [{ streamIn: true }, { streamOut: true }, { streamOut: true }], + cpuLimit: (20 * 1e9), + clockLimit: (30 * 60 * 1e9), + memoryLimit: (256 << 20), + procLimit: 50, + tty: true, + }], + }); + stream.on('output', (output) => { + process.stdout.write(output.content); + }); + stream.on('open', async () => { + let stop = false; + process.stdin.on('data', (buf) => { + if (buf.length === 1 && buf[0] === 3) { + if (stop) stream.cancel(); + stop = true; + } else stop = false; + stream.input({ content: buf }); + }); + const resize = () => { + stream.resize({ + rows: process.stdout.rows, + cols: process.stdout.columns, + }); + }; + process.stdout.on('resize', resize); + resize(); + }); + stream.on('end', (result) => { + process.stdin.setRawMode(false); + log.info(result); + process.exit(0); + }); + stream.on('close', (e) => { + process.stdin.setRawMode(false); + log.info('close', e.code); + process.exit(0); + }); +};