/*
 fahlimit.c
 Created by Kevin Bernhagen, 04 Feb 2004
 This is hereby placed into the Public Domain. There is no warranty whatsoever.
 Latest version should be available at http://calxalot.homeip.net/downloads/fahlimit.c

 Contributors
 kjb  Kevin J. Bernhagen "calxalot" < kjbern @ pacbell . net >
 rph  Richard P. Howell IV

 2004-02-05
 Made much more efficient by directly gathering core pids and signaling them.
 Added command line options and usage.

 2004-02-05
 Linux port courtesy of Richard P. Howell IV

 2004-02-06
 Changed text output to include pids signaled.
 Reverted to strstr in Linux code for name comparison.
 Should be fixed now. I'm a fool.

 2004-02-10  0.6
 Added kill other fahlimit at launch.
 Removed excessive comments from last function.
 Added version string.
 */

#include <stdio.h>
#include <signal.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#ifndef __linux__
# include <errno.h>
# include <sys/sysctl.h>
#else
# include <sys/types.h>
# include <dirent.h>
#endif

static char *version = "fahlimit version 0.6";

static int f_version = 0;
static int f_daemon = 0;
static int f_persist = 0;
static int f_fallback = 0;
static int f_debug = 0;
static int running = 1;
static int quit_signal = 0;
static int paused = 0;
static int need_refresh = 1;
static int percent = 70;
static unsigned run_interval = 7;
static unsigned pause_interval = 3;
static unsigned start_delay = 10;

#define MAX_CORES 64
static pid_t core_pids[MAX_CORES];
static int core_pid_count = 0;

static void usage()
{
    printf("%s\n", version);
    printf("usage: fahlimit [options]\n");
    printf("options:\n");
    printf("  -daemon       run in background. implies -persist\n");
    printf("  -persist      don't exit if no cores are detected\n");
    printf("  -fallback     if no cores are detected, try fallback of /usr/bin/killall. implies -persist\n");
    printf("  -percent nn   percent time cores should run. valid range is 20-90 (default 70)\n");
    printf("  -version      show version info and quit\n");
    printf("description:\n");
    printf("fahlimit is for reducing the cpu load caused by the folding@home core.\n");
    printf("it works by alternately pausing (SIGSTOP) and continuing (SIGCONT) all cores \"FahCore_\".\n");
    printf("this is useful, for instance, in reducing the operating temperature of a laptop.\n");
    printf("typical usage would be\n");
    printf("  fahlimit -daemon   # to start\n");
    printf("  killall fahlimit   # to stop\n");
    fflush(stdout);
    exit(1);
}

static void signal_handler(int sig)
{
    switch (sig)
    {
        case SIGHUP:
            need_refresh = 1;
            break;
        case SIGINT:
        case SIGTERM:
            quit_signal = sig;
            running = 0;
            break;
    }
}

static int GetAllPIDsForProcessNamesWithPrefix(const char* name,
                                               pid_t pids[],
                                               const unsigned int maxpids,
                                               int *pidsfound,
                                               int *SysctlError);

static void refresh_core_pids()
{
    int error;
    error = GetAllPIDsForProcessNamesWithPrefix("FahCore_", core_pids, MAX_CORES, &core_pid_count, NULL);
    if (error == 0)
    {
        need_refresh = 0;
    }
}

static void signal_cores(int sig)
{
    int i;
    pid_t pid;
    if (need_refresh)
        refresh_core_pids();
    if (core_pid_count == 0 && f_fallback)
    {
        if (!f_daemon)
        {
            printf("core count is zero. trying fallback killall...\n");
            fflush(stdout);
        }
        switch (sig)
        {
            case SIGSTOP:
                system("/usr/bin/killall -STOP -m '^FahCore_'");
                break;
            case SIGCONT:
                system("/usr/bin/killall -CONT -m '^FahCore_'");
                break;
        }
    }
    else if (core_pid_count == 0)
    {
        if (!f_daemon)
        {
            printf("no cores found");
        }
    }
    else for (i=0; i < core_pid_count; i++)
    {
        pid = core_pids[i];
        if (pid > 0)
        {
            kill(pid, sig);
            if (!f_daemon)
                printf("%i%s", (int)pid, ((i < core_pid_count-1) ? ", " : ""));
        }
    }
}

static void continue_cores()
{
    if (!f_daemon)
    {
        printf("continuing...");
        fflush(stdout);
    }
    signal_cores(SIGCONT);
    if (!f_daemon)
    {
        printf("\n");
        fflush(stdout);
    }
    paused = 0;
}

static void pause_cores()
{
    if (!f_daemon)
    {
        printf("pausing...");
        fflush(stdout);
    }
    paused = 1;
    signal_cores(SIGSTOP);
    if (!f_daemon)
    {
        printf("\n");
        fflush(stdout);
    }
}

static void parse_arguments(int argc, const char * argv[])
{
    int i;
    const char *p;
    for (i=1; i < argc; i++)
    {
        if (argv[i][0] == '-')
        {
            p = argv[i];
            if (strcmp(p, "-daemon") == 0)
            {
                f_daemon = 1;
                f_persist = 1;
            }
            else if (strcmp(p, "-persist") == 0)
            {
                f_persist = 1;
            }
            else if (strcmp(p, "-debug") == 0)
            {
                f_debug = 1;
            }
            else if (strcmp(p, "-version") == 0)
            {
                f_version = 1;
            }
            else if (strcmp(p, "-fallback") == 0)
            {
                f_fallback = 1;
                f_persist = 1;
            }
            else if (strcmp(p, "-percent") == 0)
            {
                p = argv[++i];
                if (p)
                {
                    percent = atoi(p);
                    if (percent < 20 || percent > 90)
                        usage();
                    // calc run and pause times
                    if (percent % 10 == 0)
                    {
                        run_interval = percent / 10;
                        pause_interval = 10 - run_interval;
                    }
                    else
                    {
                        run_interval = (percent + 2.5) / 5.0;
                        pause_interval = 20 - run_interval;
                        if (run_interval == pause_interval)
                            run_interval = pause_interval = 5;
                    }
                }
                else
                    usage();
            }
            else
                usage();
        }
        else
        {
            usage();
        }
    }
}

int main (int argc, const char * argv[])
{
    int i;
    parse_arguments(argc, argv);
    if (f_version)
    {
        printf("%s\n", version);
        exit(0);
    }
    if (f_debug)
    {
        printf("f_daemon %i\n", f_daemon);
        printf("f_persist %i\n", f_persist);
        printf("f_fallback %i\n", f_fallback);
        printf("percent %i\n", percent);
        printf("run_interval %u\n", run_interval);
        printf("pause_interval %u\n", pause_interval);
        printf("start_delay %u\n", start_delay);
        refresh_core_pids();
        printf("pid count: %i\npids: ", core_pid_count);
        for (i=0; i < core_pid_count; i++)
            printf(" %i", core_pids[i]);
        printf("\n");
        fflush(stdout);
        exit(0);
    }
    signal(SIGHUP, signal_handler);
    signal(SIGINT, signal_handler);
    // kill all others...
    signal(SIGTERM, SIG_IGN);
    system("/usr/bin/killall -TERM fahlimit");
    signal(SIGTERM, signal_handler);
    if (f_daemon)
    {
        printf("becoming a daemon...use \"killall fahlimit\" to stop.\n");
        fflush(stdout);
        daemon(0, 0);
    }
    else
    {
        printf("entering run loop...\n");
        fflush(stdout);
    }
    while (running)
    {
        need_refresh = 1;
        if (core_pid_count == 0 && running)
            sleep(start_delay);
        for (i=0; running && (i < 10); i++)
        {
            sleep(run_interval);
            pause_cores();
            sleep(pause_interval);
            continue_cores();
            if (core_pid_count == 0 && !f_persist)
                running = 0;
        }
    }
    if (!f_daemon)
    {
        printf("\n");
        if (quit_signal)
            printf("fahlimit exiting due to signal %i\n", quit_signal);
        else if (core_pid_count == 0)
            printf("fahlimit exiting because no cores were detected\n");
        else
            printf("fahlimit exiting for unknown reason\n");
        fflush(stdout);
    }
    return 0;
}

// remainder is hacked-up from apple example PIDFromBSDProcessName

enum {
    kSuccess = 0,
    kCouldNotFindRequestedProcess = -1,
    kInvalidArgumentsError = -2,
    kErrorGettingSizeOfBufferRequired = -3,
    kUnableToAllocateMemoryForBuffer = -4,
    kPIDBufferOverrunError = -5
};

int GetAllPIDsForProcessNamesWithPrefix(const char* ProcessName,
                                        pid_t ArrayOfReturnedPIDs[],
                                        const unsigned int NumberOfPossiblePIDsInArray,
                                        int* NumberOfMatchesFound,
                                        int* SysctlError)
{

#ifndef __linux__

    int mib[6] = {0,0,0,0,0,0}; //used for sysctl call.
    int SuccessfullyGotProcessInformation;
    size_t sizeOfBufferRequired = 0; //set to zero to start with.
    int error = 0;
    long NumberOfRunningProcesses = 0;
    unsigned int Counter = 0;
    struct kinfo_proc* BSDProcessInformationStructure = NULL;
    pid_t CurrentExaminedProcessPID = 0;
    char* CurrentExaminedProcessName = NULL;
    int nameLength = MIN(strlen(ProcessName), MAXCOMLEN);

#else   /* __linux__ */

    DIR *ProcDir;
    struct dirent *ProcDirEnt;
    unsigned int ProcDirEntPID;
    char pbuf[256];
    FILE *CommandLineFile;
    char zc;
    int n;
    int nameLength = strlen(ProcessName);

#endif  /* __linux__ */

    if ((ProcessName == NULL) || (*ProcessName == '\0')) //need valid process name
    {
        return(kInvalidArgumentsError);
    }
    if (ArrayOfReturnedPIDs == NULL) //need an actual array
    {
        return(kInvalidArgumentsError);
    }
    if (NumberOfPossiblePIDsInArray <= 0)
    {
        //length of the array must be larger than zero.
        return(kInvalidArgumentsError);
    }
    if (NumberOfMatchesFound == NULL) //need an integer for return.
    {
        return(kInvalidArgumentsError);
    }
    //initializing PID array so all values are zero
    memset(ArrayOfReturnedPIDs, 0, NumberOfPossiblePIDsInArray * sizeof(pid_t));
    *NumberOfMatchesFound = 0; //no matches found yet
    if (SysctlError != NULL) //only set sysctlError if it is present
        *SysctlError = 0;

#ifndef __linux__
    
    mib[0] = CTL_KERN;
    mib[1] = KERN_PROC;
    mib[2] = KERN_PROC_ALL;

    SuccessfullyGotProcessInformation = FALSE;

    while (SuccessfullyGotProcessInformation == FALSE)
    {
        error = sysctl(mib, 3, NULL, &sizeOfBufferRequired, NULL, NULL);

        if (error != 0)
        {
            if (SysctlError != NULL)
            {
                *SysctlError = errno;  //we only set this variable if the pre-allocated variable is given
            }
            return(kErrorGettingSizeOfBufferRequired);
        }

        BSDProcessInformationStructure = (struct kinfo_proc*) malloc(sizeOfBufferRequired);

        if (BSDProcessInformationStructure == NULL)
        {
            if (SysctlError != NULL)
            {
                *SysctlError = ENOMEM;  //we only set this variable if the pre-allocated variable is given
            }
            return(kUnableToAllocateMemoryForBuffer); //unrecoverable error (no memory available) so give up
        }

        error = sysctl(mib, 3, BSDProcessInformationStructure, &sizeOfBufferRequired, NULL, NULL);
        
        if (error == 0)
        {
            SuccessfullyGotProcessInformation = TRUE;
        }
        else
        {
            free(BSDProcessInformationStructure);
        }
    }
    
    NumberOfRunningProcesses = sizeOfBufferRequired / sizeof(struct kinfo_proc);

    for (Counter = 0 ; Counter < NumberOfRunningProcesses ; Counter++)
    {
        CurrentExaminedProcessPID = BSDProcessInformationStructure[Counter].kp_proc.p_pid;
        CurrentExaminedProcessName = BSDProcessInformationStructure[Counter].kp_proc.p_comm;

        if ((CurrentExaminedProcessPID > 0) //Valid PID
            && ((strncmp(CurrentExaminedProcessName, ProcessName, nameLength) == 0))) // name matches
        {
            // --- Got a match add it to the array if possible --- //
            if ((*NumberOfMatchesFound + 1) > NumberOfPossiblePIDsInArray)
            {
                //if we overran the array buffer passed we release the allocated buffer give an error.
                free(BSDProcessInformationStructure);
                return(kPIDBufferOverrunError);
            }            
            //adding the value to the array.
            ArrayOfReturnedPIDs[*NumberOfMatchesFound] = CurrentExaminedProcessPID;
            //incrementing our number of matches found.
            *NumberOfMatchesFound = *NumberOfMatchesFound + 1;
        }
    }

    free(BSDProcessInformationStructure);

#else   /* __linux__ */
    
    // --- Open the </proc> directory and look at the saved command lines --- //

    /* We look for (pseudo)files called </proc/NNN/cmdline>, where NNN is a decimal integer
        * corresponding to the PID number.  The data in that file is a NUL-separated list of
        * command arguments, beginning with arg[0], the name of the program.
        */

    if ((ProcDir = opendir("/proc")) != NULL)
    {
        while ((ProcDirEnt = readdir(ProcDir)) != NULL)
        {
            if (sscanf(ProcDirEnt->d_name, "%u%c", &ProcDirEntPID, &zc) != 1)
                continue;            /* NNN is not a number */
            sprintf(pbuf, "/proc/%u/cmdline", ProcDirEntPID);
            if ((CommandLineFile = fopen(pbuf,"r")) == NULL)
                continue;            /* Strange ??  It should exist */
            n = fread(pbuf, 1, sizeof(pbuf) - 1, CommandLineFile);
            fclose(CommandLineFile);
            if (n < nameLength)
                continue;            /* Name too short, can't match */
            pbuf[n] = '\0';
            if (strstr(pbuf, ProcessName) != NULL)
            {
				// --- Got a match; add it to the array if possible --- //
                if (*NumberOfMatchesFound >= NumberOfPossiblePIDsInArray)
                {
					//if we overran the array buffer passed, clean up and give an error.
                    closedir(ProcDir);
                    return(kPIDBufferOverrunError);
                }
                
				//adding the value to the array.
                ArrayOfReturnedPIDs[*NumberOfMatchesFound] = (pid_t) ProcDirEntPID;
                
				//incrementing our number of matches found.
                ++*NumberOfMatchesFound;
            }
        }
        closedir(ProcDir);
    }

#endif    /* __linux__ */

    if (*NumberOfMatchesFound == 0)
    {
        //didn't find any matches return error.
        return(kCouldNotFindRequestedProcess);
    }
    else
    {
        //found matches return success.
        return(kSuccess);
    }
}