piston/cli/commands/execute.js

const fs = require('fs');
const path = require('path');
const chalk = require('chalk');
const WebSocket = require('ws');

const SIGNALS = ["SIGABRT","SIGALRM","SIGBUS","SIGCHLD","SIGCLD","SIGCONT","SIGEMT","SIGFPE","SIGHUP","SIGILL","SIGINFO","SIGINT","SIGIO","SIGIOT","SIGLOST","SIGPIPE","SIGPOLL","SIGPROF","SIGPWR","SIGQUIT","SIGSEGV","SIGSTKFLT","SIGTSTP","SIGSYS","SIGTERM","SIGTRAP","SIGTTIN","SIGTTOU","SIGUNUSED","SIGURG","SIGUSR1","SIGUSR2","SIGVTALRM","SIGXCPU","SIGXFSZ","SIGWINCH"]


exports.command = ['execute <language> <file> [args..]'];
exports.aliases = ['run'];
exports.describe = 'Executes file with the specified runner';

exports.builder = {
    language_version: {
        string: true,
        desc: 'Set the version of the language to use',
        alias: ['l'],
        default: '*'
    },
    stdin: {
        boolean: true,
        desc: 'Read input from stdin and pass to executor',
        alias: ['i']
    },
    run_timeout: {
        alias: ['rt', 'r'],
        number: true,
        desc: 'Milliseconds before killing run process',
        default: 3000
    },
    compile_timeout: {
        alias: ['ct', 'c'],
        number: true,
        desc: 'Milliseconds before killing compile process',
        default: 10000,
    },
    files: {
        alias: ['f'],
        array: true,
        desc: 'Additional files to add',
    },
    interactive: {
        boolean: true,
        alias: ['t'],
        desc: 'Run interactively using WebSocket transport'
    },
    status: {
        boolean: true,
        alias: ['s'],
        desc: 'Output additional status to stderr'
    }
};

async function handle_interactive(files, argv){
    const ws = new WebSocket(argv.pistonUrl.replace("http", "ws") + "/api/v2/connect")

    const log_message = (process.stderr.isTTY && argv.status) ? console.error : ()=>{};

    process.on("exit", ()=>{
        ws.close();
        process.stdin.end();
        process.stdin.destroy();
        process.exit();        
    })

    for(const signal of SIGNALS){
        process.on(signal, ()=>{
            ws.send(JSON.stringify({type: 'signal', signal}))
        }) 
    }

    

    ws.on('open', ()=>{
        const request = {
            type: "init",
            language: argv.language,
            version: argv['language_version'],
            files: files,
            args: argv.args,
            compile_timeout: argv.ct,
            run_timeout: argv.rt
        }

        ws.send(JSON.stringify(request))
        log_message(chalk.white.bold("Connected"))

        process.stdin.resume();

        process.stdin.on("data", (data) => {
            ws.send(JSON.stringify({
                type: "data",
                stream: "stdin",
                data: data.toString()
            }))
        })
    })

    ws.on("close", (code, reason)=>{
        log_message(
            chalk.white.bold("Disconnected: "),
            chalk.white.bold("Reason: "),
            chalk.yellow(`"${reason}"`),
            chalk.white.bold("Code: "),
            chalk.yellow(`"${code}"`),
            )
        process.stdin.pause()
    })

    ws.on('message', function(data){
        const msg = JSON.parse(data);
        
        switch(msg.type){
            case "runtime":
                log_message(chalk.bold.white("Runtime:"), chalk.yellow(`${msg.language} ${msg.version}`))
                break;
            case "stage":
                log_message(chalk.bold.white("Stage:"), chalk.yellow(msg.stage))
                break;
            case "data":
                if(msg.stream == "stdout") process.stdout.write(msg.data)
                else if(msg.stream == "stderr") process.stderr.write(msg.data)
                else log_message(chalk.bold.red(`(${msg.stream}) `), msg.data)
                break;
            case "exit":
                if(msg.signal === null)
                    log_message(
                        chalk.white.bold("Stage"),
                        chalk.yellow(msg.stage),
                        chalk.white.bold("exited with code"),
                        chalk.yellow(msg.code)
                    )
                else
                    log_message(
                        chalk.white.bold("Stage"),
                        chalk.yellow(msg.stage),
                        chalk.white.bold("exited with signal"),
                        chalk.yellow(msg.signal)
                    )
                    break;
            default:
                log_message(chalk.red.bold("Unknown message:"), msg)
        }
    })

}

async function run_non_interactively(files, argv) {
    

    const stdin = (argv.stdin && await new Promise((resolve, _) => {
        let data = '';
        process.stdin.on('data', d => data += d);
        process.stdin.on('end', _ => resolve(data));
    })) || '';

    const request = {
        language: argv.language,
        version: argv['language_version'],
        files: files,
        args: argv.args,
        stdin,
        compile_timeout: argv.ct,
        run_timeout: argv.rt
    };

    let { data: response } = await argv.axios.post('/api/v2/execute', request);

    const step = (name, ctx) => {
        console.log(chalk.bold(`== ${name} ==`));

        if (ctx.stdout) {
            console.log(chalk.bold(`STDOUT`))
            console.log(ctx.stdout.replace(/\n/g,'\n    '))
        }

        if (ctx.stderr) {
            console.log(chalk.bold(`STDERR`))
            console.log(ctx.stderr.replace(/\n/g,'\n    '))
        }

        if (ctx.code) {
            console.log(
                chalk.bold(`Exit Code:`),
                chalk.bold[ctx.code > 0 ? 'red' : 'green'](ctx.code)
            );
        }

        if (ctx.signal) {
            console.log(
                chalk.bold(`Signal:`),
                chalk.bold.yellow(ctx.signal)
            );
        }
    }

    if (response.compile) {
        step('Compile', response.compile);
    }

    step('Run', response.run);
}

exports.handler = async (argv) => {
    const files = [...(argv.files || []),argv.file]
        .map(file_path => {
            return {
                name: path.basename(file_path),
                content: fs.readFileSync(file_path).toString()
            };
        });

    if(argv.interactive) await handle_interactive(files, argv);
    else await run_non_interactively(files, argv);
}