mod_cgroupmin.c

/* 
 **  mod_cgroup.c -- Apache sample cgroup module
 **  [Autogenerated via ``apxs -n cgroup -g'']
 **
 ** Stripped down version of https://github.com/MatthewIfe/mod_cgroup.git without libcgroup dependency
 ** Build with: apxs -c mod_cgroupmin.c -o mod_cgroupmin.la
 */

#define ACTIVE_ON 1
#define ACTIVE_OFF 0

#define CGROUP_OK                       0
#define CGROUP_ERROR                    -1
#define CGROUP_MIGRATION_NOTALLOWED     -2

#include "httpd.h"
#include "http_config.h"
#include "http_log.h"
#include "http_protocol.h"
#include "http_request.h"
#include "ap_config.h"
#include "apr_strings.h"

#include <sys/types.h>
#include <unistd.h>
#include <string.h>

typedef struct cgroup cgroup;
typedef struct cgroup_config cgroup_config;

static int inside_default_cgroup = 0;

struct cgroup_config {
    char default_cgroup[256];
    char cgroup[256];
    int relinquish;
};

module AP_MODULE_DECLARE_DATA cgroupmin_module;

// check if cgroup process limit is reached
// returns 0 if if limit is not reached
static int cgroup_procs_not_allowed(apr_pool_t *pool, server_rec *r, char *cgroup) {
    apr_file_t *cg;
    apr_status_t rc;
    const char *cg_filename;
    char buf[1024];
    unsigned long pid_cur, pid_max;

    cg_filename = apr_pstrcat(pool, "/sys/fs/cgroup/", cgroup, "/pids.max", NULL);
    if ((rc = apr_file_open(&cg, cg_filename, APR_FOPEN_READ, APR_FPROT_OS_DEFAULT, pool)) != APR_SUCCESS) {
        return CGROUP_ERROR;
    }
    rc = apr_file_read_full(cg, buf, sizeof(buf), NULL);
    if (rc != APR_SUCCESS && rc != APR_EOF) {
        apr_file_close(cg);
        return CGROUP_ERROR;
    }
    apr_file_close(cg);
    if (!strncmp("max", buf, 3)) {
        return CGROUP_OK;
    }
    pid_max = apr_atoi64(buf);

    cg_filename = apr_pstrcat(pool, "/sys/fs/cgroup/", cgroup, "/pids.current", NULL);
    if ((rc = apr_file_open(&cg, cg_filename, APR_FOPEN_READ, APR_FPROT_OS_DEFAULT, pool)) != APR_SUCCESS) {
        return CGROUP_ERROR;
    }
    rc = apr_file_read_full(cg, buf, sizeof(buf), NULL);
    if (rc != APR_SUCCESS && rc != APR_EOF) {
        apr_file_close(cg);
        return CGROUP_ERROR;
    }
    apr_file_close(cg);
    pid_cur = apr_atoi64(buf);

    ap_log_error(APLOG_MARK, APLOG_DEBUG, errno, r, "Migration to cgroup: %s limit check: pids.current (%lu) >= pids.max (%lu).", cgroup, pid_cur, pid_max);
    if (pid_cur >= pid_max) {
        ap_log_error(APLOG_MARK, APLOG_ERR, errno, r, "Migration to cgroup: %s not allowed. pids.current (%lu) >= pids.max (%lu). Limit reached.", cgroup, pid_cur, pid_max);
        return CGROUP_MIGRATION_NOTALLOWED;
    }

    return CGROUP_OK;
}


static int cgroup_attach_task(apr_pool_t *pool, server_rec *r, char *cgroup) {
    apr_file_t *cg;
    apr_status_t rc;
    const char *cg_filename;
    const char *pid;

#if 0
    {
        apr_file_t *hf;
        const char *hat_filename;
        hat_filename = apr_psprintf(pool, "/proc/%d/attr/current", getpid());

        if ((rc = apr_file_open(&hf, hat_filename, APR_FOPEN_READ, APR_FPROT_OS_DEFAULT, pool)) == APR_SUCCESS) {
            char buf[1024];
            apr_file_read_full(hf, buf, sizeof(buf), NULL);
        }
    }
#endif

    cg_filename = apr_pstrcat(pool, "/sys/fs/cgroup/", cgroup, "/cgroup.procs", NULL);
    if ((rc = apr_file_open(&cg, cg_filename, APR_FOPEN_WRITE, APR_FPROT_OS_DEFAULT, pool)) != APR_SUCCESS) {
        ap_log_error(APLOG_MARK, APLOG_ERR, errno, r, "Could not migrate to cgroup: %s. Opening cgroup.procs file failed.", cgroup);
        return rc;
    }

    pid = apr_psprintf(pool, "%d", getpid());
    if (!pid) {
        ap_log_error(APLOG_MARK, APLOG_ERR, errno, r, "Could not migrate to cgroup: %s. pid buffer allocation failed.", cgroup);
        return FALSE;
    }
    if ((rc = apr_file_write_full(cg, pid, strlen(pid), NULL)) != APR_SUCCESS) {
        ap_log_error(APLOG_MARK, APLOG_ERR, errno, r, "Could not migrate to cgroup: %s. Write to cgroup.procs failed.", cgroup);
        return rc;
    }
    return apr_file_close(cg);
}

static const char* cgroup_vhost(cmd_parms *cmd, void *mconfig, const char *v1) {
    cgroup_config *dirconfig = (cgroup_config *) mconfig;
    cgroup_config *cgconf = ap_get_module_config(cmd->server->module_config, &cgroupmin_module);
    apr_cpystrn(cgconf->cgroup, v1, 256);
    apr_cpystrn(dirconfig->cgroup, v1, 256);
    return NULL;
}

static const char* cgroup_default(cmd_parms *cmd, void *mconfig, const char *v1) {
    cgroup_config *dirconfig = (cgroup_config *) mconfig;
    cgroup_config *cgconf = ap_get_module_config(cmd->server->module_config, &cgroupmin_module);
    apr_cpystrn(cgconf->default_cgroup, v1, 256);
    apr_cpystrn(dirconfig->default_cgroup, v1, 256);
    return NULL;
}

static const char* cgroup_relinquish(cmd_parms *cmd, void *mconfig, int arg) {
    cgroup_config *dirconfig = (cgroup_config *) mconfig;
    cgroup_config *cgconf = ap_get_module_config(cmd->server->module_config, &cgroupmin_module);
    cgconf->relinquish = ACTIVE_ON;
    dirconfig->relinquish = ACTIVE_ON;
    if (arg == 0) {
        cgconf->relinquish = ACTIVE_OFF;
        dirconfig->relinquish = ACTIVE_OFF;
    }
    return NULL;
}

static void cgroup_child_init(apr_pool_t *pool, server_rec *server)
{
    cgroup_config *cgconf = ap_get_module_config(server->module_config, &cgroupmin_module);
    ap_log_error(APLOG_MARK, APLOG_DEBUG, 0, NULL, "Using default cgroup: %s (handler)", cgconf->default_cgroup);
    if (cgroup_attach_task(pool, server, cgconf->default_cgroup) == 0)
        inside_default_cgroup = 1;

}

static int cgroup_enter(request_rec *r)
{
    /* We only change cgroup for the main request, not subrequests */
    if (r->main)
        return OK;

    if (inside_default_cgroup) {
        cgroup_config *cgconf = ap_get_module_config(r->per_dir_config, &cgroupmin_module);


        // check if cgroup process limit is reached
        if (cgroup_procs_not_allowed(r->pool, r->server, cgconf->cgroup) == CGROUP_MIGRATION_NOTALLOWED) {
            return HTTP_TOO_MANY_REQUESTS;
        }

        ap_log_error(APLOG_MARK, APLOG_DEBUG, 0, NULL, "Using cgroup: %s (handler)", cgconf->cgroup);
        cgroup_attach_task(r->pool, r->server, cgconf->cgroup);
        inside_default_cgroup = 0;
    }

    return OK;
}


static int cgroup_exit(request_rec *r) {
    cgroup_config *cgconf = ap_get_module_config(r->per_dir_config, &cgroupmin_module);

    if (cgconf->relinquish == ACTIVE_OFF) {
        return DECLINED;
    }

    ap_log_error(APLOG_MARK, APLOG_DEBUG, 0, NULL, "Moving back to cgroup %s", cgconf->default_cgroup);
    if (cgroup_attach_task(r->pool, r->server, cgconf->default_cgroup) == 0)
        inside_default_cgroup = 0;

    return OK;
}

static void cgroup_register_hooks(apr_pool_t *p) {
    static const char * const runBeforeUs[] = { "mod_appamor.c", NULL };
    static const char * const runAfterUs[] = { "mod_apparmor.c", NULL };
    ap_hook_child_init(cgroup_child_init, NULL, runAfterUs, APR_HOOK_MIDDLE);
    ap_hook_access_checker(cgroup_enter, NULL, runAfterUs, APR_HOOK_FIRST);
    ap_hook_log_transaction(cgroup_exit, runBeforeUs, NULL, APR_HOOK_LAST+1);
}

static void *cgroup_create_server(apr_pool_t *pool, server_rec *server) {
    cgroup_config *cgconf = apr_pcalloc(pool, sizeof(cgroup_config));
    apr_cpystrn(cgconf->default_cgroup, "/", 256);
    apr_cpystrn(cgconf->cgroup, "/", 256);
    cgconf->relinquish = ACTIVE_ON;
    return cgconf;
}

static void *cgroup_merge_server(apr_pool_t *pool, void *server1_conf, void *server2_conf) {
    cgroup_config *s1 = (cgroup_config *) server1_conf;
    cgroup_config *s2 = (cgroup_config *) server2_conf;
    cgroup_config *merged = apr_pcalloc(pool, sizeof(cgroup_config));

    apr_cpystrn(merged->default_cgroup, s1->default_cgroup, 256);
    apr_cpystrn(merged->cgroup, s2->cgroup, 256);
    merged->relinquish = s1->relinquish;

    return (void *) merged;
}

static void *cgroup_create_dir(apr_pool_t *pool, char *dirspec) {
    char *dname;
    dname = dirspec;
    cgroup_config *cgconf = apr_pcalloc(pool, sizeof(cgroup_config));
    apr_cpystrn(cgconf->default_cgroup, "/", 256);
    apr_cpystrn(cgconf->cgroup, "/", 256);
    cgconf->relinquish = ACTIVE_ON;

    return cgconf;
}

static void *cgroup_merge_dir(apr_pool_t *pool, void *parent_conf, void *child_conf) {
    cgroup_config *s1 = (cgroup_config *) parent_conf;
    cgroup_config *s2 = (cgroup_config *) child_conf;
    cgroup_config *merged = apr_pcalloc(pool, sizeof(cgroup_config));

    apr_cpystrn(merged->default_cgroup, s1->default_cgroup, 256);
    apr_cpystrn(merged->cgroup, s2->cgroup, 256);
    merged->relinquish = s1->relinquish;
    ap_log_perror(APLOG_MARK, APLOG_DEBUG, errno, pool, "s1: %s s2: %s m: %s", s1->cgroup, s2->cgroup, merged->cgroup);

    return (void *) merged;
}

static const command_rec cgroup_cmds[] = {
    AP_INIT_TAKE1("cgroup",
            cgroup_vhost, NULL, RSRC_CONF|ACCESS_CONF,
            "The cgroup you want to allocate the vhost to"),
    AP_INIT_TAKE1("defaultcgroup", 
            cgroup_default, NULL, RSRC_CONF,
            "The default cgroup apache should reside in"),
    AP_INIT_FLAG("relinquishcgroup",
            cgroup_relinquish, NULL, RSRC_CONF,
            "Whether to switch out of the cgroup once finishing a request. Takes 'on' or 'off'"),
    {NULL}
};

/* Dispatch list for API hooks */

module AP_MODULE_DECLARE_DATA cgroupmin_module = {
    STANDARD20_MODULE_STUFF, 
    cgroup_create_dir,                  /* create per-dir    config structures */
    cgroup_merge_dir,                  /* merge  per-dir    config structures */
    cgroup_create_server,     /* create per-server config structures */
    cgroup_merge_server,                  /* merge  per-server config structures */
    cgroup_cmds,                  /* table of config file commands       */
    cgroup_register_hooks  /* register hooks                      */
};


// # vim: set ts=4 sw=4 sts=4 et :