GOJ-Runtime/main.c

#define _GNU_SOURCE

#include <stdio.h>
#include <stdbool.h>
#include <string.h>
#include <stdint.h>
#include <time.h>
#include <stdlib.h>
#include <sys/wait.h>
#include <sys/resource.h>
#include <sys/types.h>
#include <unistd.h>
#include <stdarg.h>
#include <errno.h>
#include <wchar.h>

#include "structs.h"
#include "serialization.h"

#define COMPILE_TIMEOUT 120
#define COMPILE_MEMORYLIMIT 1000000 // 1GB

#define RUNNER_UID 1000
#define RUNNER_GID 1000

#ifndef RUNNER_UID
    #error "RUNNER_UID must be set"
#endif

#ifndef RUNNER_GID
    #error "RUNNER_GID must be set"
#endif

// TODO: Better Panics
// Give an OE
_Noreturn void panic(const char* msg, ...) {
    va_list vargs;
    va_start(vargs, msg);

    fprintf(stderr, "panic: ");
    vfprintf(stderr, msg, vargs);

    va_end(vargs);

    exit(1);
}

void execCommand(char* command) {
    setgid(RUNNER_GID);
    setuid(RUNNER_UID);
    char* args[] = {"sh", "-c", command, NULL};
    execvp("sh", args);
    
    // If exec returns, we're beyond cooked so panic
    panic("Failed to exec in process! Errno: %d", errno);
}

void getExecStats(pid_t pid, RunResult* result) {
    struct rusage usage;
    time_t start = time(NULL);

    if (wait4(pid, &result->exitCode, 0, &usage) <= 0) {
        panic("wait4 non-zero! Errno: %d", errno);
    }

    result->exitCode = WEXITSTATUS(result->exitCode);

    result->maxRss = usage.ru_maxrss;

    time_t end = time(NULL);

    time_t executionTime = end - start;

    result->runtime = executionTime;
}

void setExecFlags(RunResult* result, time_t runtimeSeconds, size_t memoryLimit) {
    if (result->exitCode == 69) {
        result->flags |= Nice;
        strcpy(result->feedback, "Nice!");
    }

    bool errorFlag = false;

    if (result->exitCode != 0) {
        result->flags |= ProgramError;
        errorFlag = true;
    }

    if (result->runtime > runtimeSeconds) {
        result->flags |= TimeLimitExceeded;
        errorFlag = true;
    }

    if (result->maxRss > memoryLimit) {
        result->flags |= MemoryLimitExceeded;
        errorFlag = true;
    }

    if (errorFlag) {
        result->flags |= Err;
    } else {
        result->flags |= Ok;
    }
}

void getRunSummary() {
}

/**
 * @brief Runs commands with given constraints
 *
 * @details
 * executeCommand will run the command string with <em>runtimeSeconds</em>
 * as a time limit and <em>memoryLimit
 * </em> as max RSS in KB. If time
 * is exceeded, the function will immediately return error. However,
 * a program that exceeds the memory limit will continue to run but, it's max
 * RSS usage will be checked and will fail if it exceeds the memory limit
 *
 * Command strings are also passed in as an argument to "shell -c" allowing
 * shell syntax to be used. This can be particularly useful for multi-step
 * compilation.
 *
 * @param[in] command           The shell commmand/syntax to run
 * @param[in] runtimeSeconds    Timeout in seconds
 * @param[in] memoryLimit       Max RSS limit in KB
 *
 */
RunResult executeCommand(char* command, time_t runtimeSeconds, size_t memoryLimit) {
    int process = fork();

    RunResult result;
    
    // stop undefined behavior
    // from thinking a string is present
    result.feedback[0] = '\0';

    if (process > 0) {
        getExecStats(process, &result);
    } else if (process == 0) {
        execCommand(command);
    } else {
        panic("Failed to spawn process");
    }

    setExecFlags(&result, runtimeSeconds, memoryLimit);

    result.flags |= ExecutionFinal;

    return result;
}

int main(int argc, char* argv[]) {
    if (getuid() != 0) {
        panic("Must run as root");
    }

    if (argc < 5) {
        panic("Not enough args");
    }

    size_t runMemoryLimit = atol(argv[3]);
    time_t runRuntimeSeconds = atol(argv[4]);

    if (runMemoryLimit <= 0 || runRuntimeSeconds <= 0) {
        panic("invalid args");
    }

    RunResult compileResult = executeCommand(argv[1], COMPILE_TIMEOUT, COMPILE_MEMORYLIMIT);
    RunResult runResult = executeCommand(argv[2], runRuntimeSeconds, runMemoryLimit);

    SerializationResult serialResult = serializeResult(&runResult);

    write(STDOUT_FILENO, serialResult.data, serialResult.size);

    free(serialResult.data);
}