Ticket #1658: mod_cookie_balancer.c

/* 
 * Author: Kirill S. Vlasov (kirakuna@yahoo.com)
 * Date: 7 May 2008
 */

#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

typedef struct {
        buffer *prefix; // one symbol prefix for fcgi
        buffer *keyword;
        int count;      // must be power of two
} plugin_config;

typedef struct {
        size_t foo;
} handler_ctx;

typedef struct {
        PLUGIN_DATA;
        buffer *match_buf;
        plugin_config **config_storage;
        plugin_config conf;
} plugin_data;

static handler_ctx * handler_ctx_init() {
        handler_ctx * hctx;
        hctx = calloc(1, sizeof(*hctx));
        return hctx;
}

static void handler_ctx_free(handler_ctx *hctx) {
        free(hctx);
}

INIT_FUNC(mod_cookie_balancer_init) {
        plugin_data *p;

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

        p->match_buf = buffer_init();

        return p;
}

FREE_FUNC(mod_cookie_balancer_free) {
        plugin_data *p = p_d;

        UNUSED(srv);

        if (!p) return HANDLER_GO_ON;

        buffer_free(p->match_buf);
        if (p->config_storage) {
                size_t i;
                for (i = 0; i < srv->config_context->used; i++) {
                        plugin_config *s = p->config_storage[i];
                        free(s);
                }
                free(p->config_storage);
        }

        free(p);

        return HANDLER_GO_ON;
}


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

        config_values_t cv[] = {
                { "cb.count",NULL, T_CONFIG_SHORT, T_CONFIG_SCOPE_CONNECTION },
                { "cb.prefix",NULL, T_CONFIG_STRING, T_CONFIG_SCOPE_CONNECTION },
                { "cb.keyword",NULL, T_CONFIG_STRING, T_CONFIG_SCOPE_CONNECTION },
                { 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->prefix   = buffer_init();
                s->keyword = buffer_init();
                s->count = 0;

                cv[0].destination = &(s->count);
                cv[1].destination = s->prefix;
                cv[2].destination = s->keyword;

                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;
}

#define PATCH(x) \
        p->conf.x = s->x;
static int mod_cookie_balancer_patch_connection(server *srv, connection *con, plugin_data *p) {
        size_t i, j;
        plugin_config *s = p->config_storage[0];

        PATCH(count);
        PATCH(prefix);
        

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

                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("cb.count"))) {
                                PATCH(count);
                        } else if (buffer_is_equal_string(du->key, CONST_STR_LEN("cb.prefix"))) {
                                PATCH(prefix);
                        } else if (buffer_is_equal_string(du->key, CONST_STR_LEN("cb.keyword"))) {
                                PATCH(keyword);
                        }
                }
        }
        char *h = p->conf.prefix->ptr;

        return 0;
}

URIHANDLER_FUNC(mod_cookie_balancer_con_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;
}

URIHANDLER_FUNC(mod_cookie_balancer_uri_handler) {
        plugin_data *p = p_d;
        data_string *ds;

        mod_cookie_balancer_patch_connection(srv, con, p);

        char *ch = con->request.uri->ptr;
        if (*ch == '/') {
                ch++;
                char *substr = strstr(ch, p->conf.prefix->ptr);
                if (substr && strcmp(substr, ch) == 0)
                        return HANDLER_GO_ON;
        }
        ds = (data_string *)array_get_element(con->request.headers, "Cookie");
        if (NULL != ds) {
                char *g;
                if (NULL != (g = strstr(ds->value->ptr, p->conf.keyword->ptr))) {
                        char *nc = g;
                        while (*nc != '=' && *nc != '\t') nc++;
                        if (*nc == '=') {
                                nc++;
                                buffer *tmp = buffer_init();
                                buffer_append_string_buffer(tmp, con->request.uri);
                                buffer_reset(con->request.uri);
                                buffer_copy_string(con->request.uri, "/");
                                buffer_append_string(con->request.uri, p->conf.prefix->ptr);
                                char num[1];
                                sprintf(&num,"%d",(*nc & p->conf.count));
                                buffer_append_string(con->request.uri, &num);
                                buffer_append_string_buffer(con->request.uri, tmp);
                                return HANDLER_COMEBACK;
                        }
                }
        }
        /*
        else {
                if (NULL == (ds = (data_string *)array_get_unused_element(con->response.headers, TYPE_STRING)))
                        ds = data_response_init();
                buffer_copy_string(ds->key, "Set-Cookie");
                buffer_append_string(ds->value, "cookiename=12345671");
                array_insert_unique(con->response.headers, (data_unset *)ds);
        }
        */
        return HANDLER_GO_ON;
}

int mod_cookie_balancer_plugin_init(plugin *p) {
        p->version     = LIGHTTPD_VERSION_ID;
        p->name        = buffer_init_string("cookie_balancer");

        p->init        = mod_cookie_balancer_init;

        p->handle_uri_raw = mod_cookie_balancer_uri_handler;
        p->set_defaults = mod_cookie_balancer_set_defaults;
        p->cleanup     = mod_cookie_balancer_free;
        p->connection_reset = mod_cookie_balancer_con_reset;

        p->data        = NULL;

        return 0;
}