#include <linux/capability.h>
#include <sys/prctl.h>

#include "initHelpers.hpp"

bool NInit::isSudo() {
    return getuid() != geteuid() || !geteuid();
}

void NInit::gainRealTime() {
    const int          minPrio = sched_get_priority_min(SCHED_RR);
    int                old_policy;
    struct sched_param param;

    if (pthread_getschedparam(pthread_self(), &old_policy, &param)) {
        Log::logger->log(Log::WARN, "Failed to get old pthread scheduling priority");
        return;
    }

    param.sched_priority = minPrio;

    if (pthread_setschedparam(pthread_self(), SCHED_RR, &param)) {
        Log::logger->log(Log::WARN, "Failed to change process scheduling strategy");
        return;
    }

    // NixOS-specific fix to prevent all children from inheriting
    // CAP_SYS_NICE due to how the security wrapper works.
    prctl(PR_CAP_AMBIENT, PR_CAP_AMBIENT_LOWER, CAP_SYS_NICE, 0, 0);

    pthread_atfork(nullptr, nullptr, []() {
        const struct sched_param param = {.sched_priority = 0};
        if (pthread_setschedparam(pthread_self(), SCHED_OTHER, &param))
            Log::logger->log(Log::WARN, "Failed to reset process scheduling strategy");
    });
}