judge: sandbox: add stream interface (#725)

Co-authored-by: undefined <i@undefined.moe>

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')();

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<SandboxResult[]> {
         return superagent.post(`${url}/run`).send(req).then((res) => res.body);
@@ -29,6 +104,9 @@ const client = new Proxy({
     config(): Promise<Record<string, any>> {
         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');

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<string, CopyInFile>;
-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<string, CopyInFile>;
+    copyIn?: Record<string, CopyInFile | Symlink>;
 
     /**
      * 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;
+}

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);
+    });
+};