#include <stdexcept>

#include <linux/seccomp.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <seccomp.h>

#include <QIPC/ipc.hpp>
#include <dlhandle.hpp>


struct Limits {
    size_t max_mem = 4 * 1000; // 4 KB
    bool enable_seccomp = true;
    bool close_stdio = true;

    Limits() {
        auto HSS_MAX_MEM = getenv("HSS_MAX_MEM");
        auto HSS_NO_SECCOMP = getenv("HSS_NO_SECCOMP");
        auto HSS_KEEP_STDIO = getenv("HSS_KEEP_STDIO");
        if (HSS_MAX_MEM) {
            max_mem = std::stoul(HSS_MAX_MEM);
        }
        if (HSS_NO_SECCOMP) {
            enable_seccomp = false;
        }
        if (HSS_KEEP_STDIO) {
            close_stdio = false;
        }
    }
};

void enable_limits(const Limits& limits) {
    // rlimit
    if (limits.max_mem) {
        rlimit memRLimit{limits.max_mem, limits.max_mem};
        if (setrlimit(RLIMIT_AS, &memRLimit) < 0) {
            throw std::runtime_error("Error setting ressource limits");
        }
    }
    // Seccomp
    if (limits.enable_seccomp) {
        scmp_filter_ctx ctx = seccomp_init(SCMP_ACT_ERRNO(EOPNOTSUPP));
        for (const auto sysc : {
             SCMP_SYS(mmap), SCMP_SYS(mmap2), SCMP_SYS(munmap), SCMP_SYS(mremap), SCMP_SYS(brk),
             SCMP_SYS(write), SCMP_SYS(read), SCMP_SYS(close),
             SCMP_SYS(exit), SCMP_SYS(exit_group)
                                }) {
            seccomp_rule_add(ctx, SCMP_ACT_ALLOW, sysc, 0);
        }
        if (seccomp_load(ctx) < 0) {
            throw std::runtime_error("Error setting up seccomp");
        }
    }
    if (limits.close_stdio) {
        // Close stdio
        close(STDIN_FILENO);
        close(STDOUT_FILENO);
        close(STDERR_FILENO);
    }
}

int main(int argc, char **argv) {
    // Get IPC
    if (argc != 6) {
        abort();
    }
    QBiIPC ipc({QIPC(std::stoi(argv[1]), std::stoi(argv[2])), QIPC(std::stoi(argv[3]), std::stoi(argv[4]))});
    // Launch
    Dlhandle dl(argv[5], RTLD_NOW | RTLD_LOCAL);
    auto entry = dl.get<void*(QBiIPC&)>("entry");
    enable_limits(Limits());
    entry(ipc);
    exit(EXIT_SUCCESS);
}