#include "h2o.h"

#include "h2o/configurator.h"

static h2o_header_t *find_header(h2o_headers_t *headers, h2o_iovec_t *name)

{

    ssize_t index;

    if (h2o_iovec_is_token(name)) {

        index = h2o_find_header(headers, (void *)name, -1);

    } else {

        index = h2o_find_header_by_str(headers, name->base, name->len, -1);

    }

    if (index == -1)

        return NULL;

    return headers->entries + index;

}

static int is_in_list(const char *base, size_t len, h2o_headers_command_t *cmd)

{

    size_t i;

    h2o_iovec_t name = h2o_iovec_init(base, len);

    for (i = 0; i != cmd->num_args; ++i) {

        if (h2o_iovec_is_token(cmd->args[i].name)) {

            if (cmd->args[i].name->base == name.base) {

                return 1;

            }

        } else {

            if (h2o_memis(cmd->args[i].name->base, cmd->args[i].name->len, name.base, name.len))

                return 1;

        }

    }

    return 0;

}

static void filter_cookie(h2o_mem_pool_t *pool, char **base, size_t *len, h2o_headers_command_t *cmd)

{

    h2o_iovec_t iter = h2o_iovec_init(*base, *len), token_value;

    const char *token;

    size_t token_len;

    char dst[*len * 2];

    size_t dst_len = 0;

    do {

        if ((token = h2o_next_token(&iter, ';', ';', &token_len, &token_value)) == NULL)

            break;

        int found = is_in_list(token, token_len, cmd);

        if ((cmd->cmd == H2O_HEADERS_CMD_COOKIE_UNSETUNLESS && found) || (cmd->cmd == H2O_HEADERS_CMD_COOKIE_UNSET && !found)) {

            if (dst_len != 0) {

                memcpy(dst + dst_len, H2O_STRLIT("; "));

                dst_len += 2;

            }

            memcpy(dst + dst_len, token, token_len);

            dst_len += token_len;

            if (token_value.len > 0) {

                memcpy(dst + dst_len, H2O_STRLIT("="));

                dst_len++;

                memcpy(dst + dst_len, token_value.base, token_value.len);

                dst_len += token_value.len;

            }

        }

    } while (1);

    if (dst_len > *len)

        *base = h2o_mem_alloc_pool(pool, *dst, dst_len);

    memcpy(*base, dst, dst_len);

    *len = dst_len;

}

static void cookie_cmd(h2o_mem_pool_t *pool, h2o_headers_t *headers, h2o_headers_command_t *cmd)

{

    ssize_t header_index;

    for (header_index = -1; (header_index = h2o_find_header(headers, H2O_TOKEN_COOKIE, header_index)) != -1;) {

        h2o_header_t *header = headers->entries + header_index;

        filter_cookie(pool, &header->value.base, &header->value.len, cmd);

        if (header->value.len == 0)

            h2o_delete_header(headers, header_index);

    }

}

static void remove_header_unless(h2o_headers_t *headers, h2o_headers_command_t *cmd)

{

    size_t src, dst = 0;

    for (src = 0; src != headers->size; ++src) {

        if (!is_in_list(headers->entries[src].name->base, headers->entries[src].name->len, cmd))

            continue;

        /* not matched */

        if (dst != src)

            headers->entries[dst] = headers->entries[src];

        ++dst;

    }

    headers->size = dst;

}

static void remove_header(h2o_headers_t *headers, h2o_headers_command_t *cmd)

{

    size_t src, dst = 0;

    for (src = 0; src != headers->size; ++src) {

        if (h2o_iovec_is_token(cmd->args[0].name)) {

            if (headers->entries[src].name == cmd->args[0].name)

                continue;

        } else {

            if (h2o_memis(headers->entries[src].name->base, headers->entries[src].name->len, cmd->args[0].name->base,

                          cmd->args[0].name->len))

                continue;

        }

        /* not matched */

        if (dst != src)

            headers->entries[dst] = headers->entries[src];

        ++dst;

    }

    headers->size = dst;

}

static void dispose_h2o_headers_command(void *_cmds)

{

    h2o_headers_command_t *cmds = _cmds;

    size_t i;

    for (i = 0; cmds[i].cmd != H2O_HEADERS_CMD_NULL; ++i)

        free(cmds[i].args);

}

void h2o_headers_append_command(h2o_headers_command_t **cmds, int cmd, h2o_headers_command_arg_t *args, size_t num_args,

                                h2o_headers_command_when_t when)

{

    h2o_headers_command_t *new_cmds;

    size_t i, cnt;

    if (*cmds != NULL) {

        for (cnt = 0; (*cmds)[cnt].cmd != H2O_HEADERS_CMD_NULL; ++cnt)

            ;

    } else {

        cnt = 0;

    }

    new_cmds = h2o_mem_alloc_shared(NULL, (cnt + 2) * sizeof(*new_cmds), dispose_h2o_headers_command);

    if (*cmds != NULL)

        memcpy(new_cmds, *cmds, cnt * sizeof(*new_cmds));

    new_cmds[cnt] = (h2o_headers_command_t){};

    new_cmds[cnt].cmd = cmd;

    new_cmds[cnt].when = when;

    new_cmds[cnt].args = h2o_mem_alloc(sizeof(*new_cmds->args) * num_args);

    for (i = 0; i < num_args; i++)

        new_cmds[cnt].args[i] = args[i];

    new_cmds[cnt].num_args = num_args;

    new_cmds[cnt + 1] = (h2o_headers_command_t){H2O_HEADERS_CMD_NULL};

    if (*cmds != NULL) {

        (*cmds)[0] = (h2o_headers_command_t){H2O_HEADERS_CMD_NULL};

        h2o_mem_release_shared(*cmds);

    }

    *cmds = new_cmds;

}

void h2o_rewrite_headers(h2o_mem_pool_t *pool, h2o_headers_t *headers, h2o_headers_command_t *cmd)

{

    h2o_header_t *target;

    switch (cmd->cmd) {

    case H2O_HEADERS_CMD_ADD:

        goto AddHeader;

    case H2O_HEADERS_CMD_APPEND:

        assert(cmd->num_args == 1);

        if ((target = find_header(headers, cmd->args[0].name)) == NULL)

            goto AddHeader;

        goto AppendToken;

    case H2O_HEADERS_CMD_MERGE:

        assert(cmd->num_args == 1);

        if ((target = find_header(headers, cmd->args[0].name)) == NULL)

            goto AddHeader;

        if (h2o_contains_token(target->value.base, target->value.len, cmd->args[0].value.base, cmd->args[0].value.len, ','))

            return;

        goto AppendToken;

    case H2O_HEADERS_CMD_SET:

        remove_header(headers, cmd);

        goto AddHeader;

    case H2O_HEADERS_CMD_SETIFEMPTY:

        assert(cmd->num_args == 1);

        if (find_header(headers, cmd->args[0].name) != NULL)

            return;

        goto AddHeader;

    case H2O_HEADERS_CMD_UNSET:

        remove_header(headers, cmd);

        return;

    case H2O_HEADERS_CMD_UNSETUNLESS:

        remove_header_unless(headers, cmd);

        return;

    case H2O_HEADERS_CMD_COOKIE_UNSET:

    case H2O_HEADERS_CMD_COOKIE_UNSETUNLESS:

        cookie_cmd(pool, headers, cmd);

        return;

    }

    assert(!"FIXME");

    return;

AddHeader:

    assert(cmd->num_args == 1);

    if (h2o_iovec_is_token(cmd->args[0].name)) {

        h2o_add_header(pool, headers, (void *)cmd->args[0].name, NULL, cmd->args[0].value.base, cmd->args[0].value.len);

    } else {

        h2o_add_header_by_str(pool, headers, cmd->args[0].name->base, cmd->args[0].name->len, 0, NULL, cmd->args[0].value.base,

                              cmd->args[0].value.len);

    }

    return;

AppendToken:

    assert(cmd->num_args == 1);

    if (target->value.len != 0) {

        h2o_iovec_t v;

        v.len = target->value.len + 2 + cmd->args[0].value.len;

        v.base = h2o_mem_alloc_pool(pool, char, v.len);

        memcpy(v.base, target->value.base, target->value.len);

        v.base[target->value.len] = ',';

        v.base[target->value.len + 1] = ' ';

        memcpy(v.base + target->value.len + 2, cmd->args[0].value.base, cmd->args[0].value.len);

        target->value = v;

    } else {

        target->value = cmd->args[0].value;

    }

    return;

}