Ticket #954: mod_traffic.c

#include <ctype.h>
#include <stdlib.h>
#include <string.h>

#include "base.h"
#include "log.h"
#include "buffer.h"

#include "plugin.h"

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif


/**
 * this is a traffic control plugin derived from mod_skeleton.c for a lighttpd plugin
 * do not use this in production environment! It's my first lighty plugin! :)
 *
 * Bjoern Kalkbrenner <terminar@cyberphoria.org>
 */

/* plugin config for all request/connections */

typedef struct {
        array *match;
} plugin_config;

typedef struct {
        PLUGIN_DATA;

        buffer *match_buf;

        plugin_config **config_storage;

        plugin_config conf;
} plugin_data;

typedef struct {
        int proceeded;
        int handled;
} handler_ctx;

static handler_ctx * handler_ctx_init() {
        handler_ctx * hctx;


        hctx = calloc(1, sizeof(*hctx));

        TRACE("handler_ctx_init for %i",hctx);

        hctx->handled = 0;
        hctx->proceeded = 0;

        return hctx;
}

static void handler_ctx_free(handler_ctx *hctx) 
{
    if (hctx)
    {
        TRACE("handler_ctx_free for %i",hctx);

        free(hctx);
    }
    
}

/* init the plugin data */
INIT_FUNC(mod_traffic_init) {
        plugin_data *p;
        
        UNUSED(srv);

        p = calloc(1, sizeof(*p));

        p->match_buf = buffer_init();

        return p;
}

/* detroy the plugin data */
FREE_FUNC(mod_traffic_free) {
        plugin_data *p = p_d;

        UNUSED(srv);

        if (!p) return HANDLER_GO_ON;

        if (p->config_storage) {
                size_t i;

                for (i = 0; i < srv->config_context->used; i++) {
                        plugin_config *s = p->config_storage[i];

                        if (!s) continue;

                        array_free(s->match);

                        free(s);
                }
                free(p->config_storage);
        }

        buffer_free(p->match_buf);

        free(p);

        return HANDLER_GO_ON;
}

/* handle plugin config and check values */

SETDEFAULTS_FUNC(mod_traffic_set_defaults) {
        plugin_data *p = p_d;
        size_t i = 0;

        config_values_t cv[] = {
                { "traffic.array",             NULL, T_CONFIG_ARRAY, T_CONFIG_SCOPE_CONNECTION },       /* 0 */
                { NULL,                         NULL, T_CONFIG_UNSET, T_CONFIG_SCOPE_UNSET }
        };

        if (!p) return HANDLER_ERROR;

        p->config_storage = calloc(1, srv->config_context->used * sizeof(specific_config *));

        for (i = 0; i < srv->config_context->used; i++) {
                plugin_config *s;

                s = calloc(1, sizeof(plugin_config));
                s->match    = array_init();

                cv[0].destination = s->match;

                p->config_storage[i] = s;

                if (0 != config_insert_values_global(srv, ((data_config *)srv->config_context->data[i])->value, cv)) {
                        return HANDLER_ERROR;
                }
        }

        return HANDLER_GO_ON;
}

static int mod_traffic_patch_connection(server *srv, connection *con, plugin_data *p) {
        size_t i, j;
        plugin_config *s = p->config_storage[0];

        PATCH_OPTION(match);

        /* skip the first, the global context */
        for (i = 1; i < srv->config_context->used; i++) {
                data_config *dc = (data_config *)srv->config_context->data[i];
                s = p->config_storage[i];

                /* condition didn't match */
                if (!config_check_cond(srv, con, dc)) continue;

                /* merge config */
                for (j = 0; j < dc->value->used; j++) {
                        data_unset *du = dc->value->data[j];

                        if (buffer_is_equal_string(du->key, CONST_STR_LEN("traffic.array"))) {
                                PATCH_OPTION(match);
                        }
                }
        }

        return 0;
}

URIHANDLER_FUNC(mod_traffic_uri_handler) {
        plugin_data *p = p_d;
        int s_len;
        size_t k;

        UNUSED(srv);

        if (con->uri.path->used == 0) return HANDLER_GO_ON;

        mod_traffic_patch_connection(srv, con, p);

        s_len = con->uri.path->used - 1;

        for (k = 0; k < p->conf.match->used; k++) {
                data_string *ds = (data_string *)p->conf.match->data[k];
                int ct_len = ds->value->used - 1;

                if (ct_len > s_len) continue;
                if (ds->value->used == 0) continue;

                if (0 == strncmp(con->uri.path->ptr + s_len - ct_len, ds->value->ptr, ct_len)) {
                        con->http_status = 403;

                        return HANDLER_FINISHED;
                }
        }

        /* not found */
        return HANDLER_GO_ON;
}


URIHANDLER_FUNC(mod_traffic_start_backend)
{
    UNUSED(srv);
    UNUSED(con);
    UNUSED(p_d);

    plugin_data *p = p_d;

    handler_ctx *hctx = con->plugin_ctx[p->id];
    if (hctx)
    {
        TRACE("%s","resetting handled, proceeded");
        hctx->proceeded = 0;
        hctx->handled = 0;
    }
    
    
    TRACE("%s","mod_traffic_start_backend");
    return HANDLER_GO_ON;
}

URIHANDLER_FUNC(mod_traffic_send_request_content)
{
    UNUSED(srv);
    UNUSED(con);
    UNUSED(p_d);
    
    TRACE("%s","mod_traffic_send_request_content");
    

    
    return HANDLER_GO_ON;
}

URIHANDLER_FUNC(mod_traffic_response_header)
{
    UNUSED(srv);
    UNUSED(con);
    UNUSED(p_d);
    
    TRACE("%s","mod_traffic_response_header");
    return HANDLER_GO_ON;
}


URIHANDLER_FUNC(mod_traffic_read_response_content)
{

    TRACE("%s","mod_traffic_read_response_content");
                                
    /* not sure if this can go into mod_traffic_send_request_content
     * we use it here to proceed headers which came from mod_proxy_core/fastcgi/php
     * does this also work in mod_traffic_send_request_content even if connections.c doesn't set connection_state to previous state?
     * if it doesn't, mod_traffic_send_request_content is called before the X-LIGHTTPD-KBytes-per-second is set
     * 
     * can the plugin be detached from the current processing loop for specific connections if it is proceeded before?
     */
    plugin_data *p = p_d;
    handler_ctx *hctx;
    
    if (con->plugin_ctx[p->id])
    {
        hctx = con->plugin_ctx[p->id];
    } else
    {
        hctx = handler_ctx_init();
        con->plugin_ctx[p->id] = hctx;
    }

    if (hctx && !hctx->proceeded || (hctx->proceeded && hctx->handled) )
    {
        if (!hctx->proceeded)
        {
            data_string *ds;
        
            if (NULL != (ds = (data_string *)array_get_element(con->response.headers, "X-LIGHTTPD-KBytes-per-second"))) {
                TRACE("Found my header, setting speed: %s",ds->value->ptr);
                con->conf.kbytes_per_second = atol(ds->value->ptr);
                hctx->handled = 1;          
            }
        
            hctx->proceeded = 1;
        }
        
        if (hctx->handled)
        {
            /* this is for processing updates later for other stuff, tracing only at the moment */
            int t_diff;
            /* we don't like div by zero */
            if (0 == (t_diff = srv->cur_ts - con->connection_start)) t_diff = 1;
            int speed = con->bytes_written / t_diff / 1024;

            TRACE("currspeed: %i",speed);
            
        }

    }
        
    return HANDLER_GO_ON;
}

URIHANDLER_FUNC(mod_traffic_connection_reset)
{
    TRACE("%s","traffic_connection_reset");
    
    plugin_data *p = p_d;
    UNUSED(srv);
    
    if (con->plugin_ctx[p->id]) 
    {
        handler_ctx_free(con->plugin_ctx[p->id]);
        con->plugin_ctx[p->id] = NULL;
    }

    return HANDLER_GO_ON;
}

/* this function is called at dlopen() time and inits the callbacks */
int mod_traffic_plugin_init(plugin *p) {
        p->version     = LIGHTTPD_VERSION_ID;
        p->name        = buffer_init_string("traffic");

        p->init        = mod_traffic_init;
        p->handle_uri_clean  = mod_traffic_uri_handler;
        p->set_defaults  = mod_traffic_set_defaults;
        p->cleanup     = mod_traffic_free;

        p->handle_start_backend = mod_traffic_start_backend;
        p->handle_send_request_content = mod_traffic_send_request_content;
        p->handle_response_header = mod_traffic_response_header;
        p->handle_read_response_content = mod_traffic_read_response_content;
        
        p->connection_reset = mod_traffic_connection_reset;
//      p->handle_connection_close = mod_traffic_connection_close;

        p->data        = NULL;

        return 0;
}