/* Copyright 1999-2007 The Apache Software Foundation or its licensors, as
 * applicable.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * Original Copyright (c) Netmask.IT!® 2006-2007
 *
 * DNS Protocol module for Apache 2.x
 */

#include "mod_dns.h"

#define DNS_HEADER_SIZE 12

/** TODO: Add serialize/unserialize for query + message and then consider
 *  reimplementing network I/O to work with character arrays */

/** Helper network-read routines */
/** WARNING: This will trash the contents of bb! */
static apr_status_t dns_bb_read_byte(apr_bucket_brigade *bb,
                                     ap_filter_t *filters_in,
                                     apr_byte_t *rval)
{
    apr_status_t rv;
    int dlen = 1;

    apr_brigade_cleanup(bb);
    rv = ap_get_brigade(filters_in, bb, AP_MODE_READBYTES, APR_BLOCK_READ, 1);
    if (rv != APR_SUCCESS)
        return rv;

    if ((rv = apr_brigade_flatten(bb, (char *)rval,
                                  &dlen)) != APR_SUCCESS)
        return rv;

    if (dlen != 1)
        return APR_EGENERAL;

    apr_brigade_cleanup(bb);
    return APR_SUCCESS;
}

/** WARNING: This will trash the contents of bb! */
static apr_status_t dns_bb_read_short(apr_bucket_brigade *bb,
                                      ap_filter_t *filters_in,
                                      apr_uint16_t *rval)
{
    apr_status_t rv;
    int dlen = 2;

    apr_brigade_cleanup(bb);
    rv = ap_get_brigade(filters_in, bb, AP_MODE_READBYTES, APR_BLOCK_READ, 2);
    if (rv != APR_SUCCESS)
        return rv;

    if ((rv = apr_brigade_flatten(bb, (char *)rval,
                                  &dlen)) != APR_SUCCESS)
        return rv;

    if (dlen != 2)
        return APR_EGENERAL;

    apr_brigade_cleanup(bb);
    *rval = ntohs(*rval);
    return APR_SUCCESS;
}

/* The socket is hiding at:
if (!csd) {
            csd = ap_get_module_config(c->conn_config, &core_module);
        }
*/

DNS_DECLARE(apr_status_t) dns_query_serialize(dns_query_t *q, char *data,
                                              int *dlen)
{
    apr_byte_t llen;
    char *ptr, *label, *last = NULL;
    apr_uint16_t s;
    apr_status_t rv = APR_SUCCESS;

    /** The length byte takes the same space as a char, so we just need
     *  strlen(cname) of space (cname should have trailing .) */
    *dlen = strlen(q->qname) + 1;
    if (q->qname[*dlen-2] != '.')
        (*dlen)++;

    /* 2 bytes type, 2 bytes class */
    *dlen += 4;

    if (data == NULL)
        return rv;

    /** NAME */
    ptr = data;
    label = apr_strtok(q->qname, ".", &last);
    while (label != NULL) {
        llen = strlen(label);
        if (llen == 0) {
            /** Trailing . - Set NULL and break */
            *ptr = 0;
            ptr++;
            break;
        }
        /** Write length token */
        *ptr = llen;
        ptr++;
        /** Write label */
        memcpy(ptr, label, llen);
        /** Move pointer */
        ptr+=llen;
        /** Advance to next token */
        label = apr_strtok(NULL, ".", &last);
    }
    *ptr = 0;
    ptr++;

    /** TYPE */
    s = htons(q->qtype);
    memcpy(ptr, (const void *)&s, 2);
    ptr += 2;

    /** CLASS */
    s = htons(q->qclass);
    memcpy(ptr, (const void *)&s, 2);
    ptr += 2;

    return rv;
}

DNS_DECLARE(apr_status_t) dns_query_pserialize(dns_query_t *q,
                                               apr_pool_t *pool,
                                               char **data, int *dlen)
{
    apr_status_t rv;
    *dlen = 0;
    if ((rv = dns_query_serialize(q, NULL, dlen)) != APR_SUCCESS)
        return rv;
    *data = apr_palloc(pool, *dlen);
    return dns_query_serialize(q, *data, dlen);
}


DNS_DECLARE(apr_status_t) dns_write_response(dns_message_t *dns,
                                             char **response,
                                             apr_size_t *len)
{
    dns_header_t *h = dns->header;
    dns_query_t *q;
    dns_rr_t *rr;
    apr_uint16_t s;
    apr_uint32_t l;
    char *data, *header;
    char *label, *last = NULL;

    apr_uint16_t flags = 0;


    /** TODO: Move to fixup */
    h->qr = DNS_FLAG_ON;
    /** TODO: Can we be authorative? */
    h->aa = DNS_FLAG_OFF;
    /** TODO: Can we be recursive? Are we already? */
    h->ra = DNS_FLAG_OFF;
    /** Someone else should be setting rcode */

    /** Create flags bitfield */
    flags |= h->qr << 15;
    flags |= h->opcode << 11;
    flags |= h->aa << 10;
    flags |= h->tc << 9;
    flags |= h->rd << 8;
    flags |= h->ra << 7;
    flags |= h->z  << 4;
    flags |= h->rcode;

    /** Prepare response */
    header = data = malloc(DNS_HEADER_SIZE);
    *len = DNS_HEADER_SIZE;

    /** Header */
    s = htons(h->id);
    memcpy(data, &s, 2);
    data += 2;
    s = htons(flags);
    memcpy(data, &s, 2);
    data += 2;
    s = htons(h->qdcount);
    memcpy(data, &s, 2);
    data += 2;
    s = htons(h->ancount);
    memcpy(data, &s, 2);
    data += 2;
    s = htons(h->nscount);
    memcpy(data, &s, 2);
    data += 2;
    s = htons(h->arcount);
    memcpy(data, &s, 2);
    data += 2;

    /** Queries */
    while ((q = apr_array_pop(dns->query)) != NULL) {
        dns_query_serialize(q, NULL, (int *)&l);
        data = malloc(*len + l);
        memcpy(data, header, *len);
        free(header);
        label = data + *len;
        dns_query_serialize(q, label, (int *)&l);
        header = data;
        *len +=l;
    }

    /** Answers */
    while ((rr = apr_array_pop(dns->answer)) != NULL) {
        dns_rr_serialize(rr, NULL, (int *)&l);
        data = malloc(*len + l);
        memcpy(data, header, *len);
        free(header);
        label = data + *len;
        dns_rr_serialize(rr, label, (int *)&l);
        header = data;
        *len +=l;
    }

    /** Authority */
    while ((rr = apr_array_pop(dns->authority)) != NULL) {
        dns_rr_serialize(rr, NULL, (int *)&l);
        data = malloc(*len + l);
        memcpy(data, header, *len);
        free(header);
        label = data + *len;
        dns_rr_serialize(rr, label, (int *)&l);
        header = data;
        *len +=l;
    }

    /** Additional */
    while ((rr = apr_array_pop(dns->additional)) != NULL) {
        dns_rr_serialize(rr, NULL, (int *)&l);
        data = malloc(*len + l);
        memcpy(data, header, *len);
        free(header);
        label = data + *len;
        dns_rr_serialize(rr, label, (int *)&l);
        header = data;
        *len +=l;
    }

    *response = header;
    return APR_SUCCESS;
}

/** On return, q will point to the newly read query, and r to the newly
 *  created request_rec */
DNS_DECLARE(apr_status_t) dns_read_request(dns_message_t *dns,
                                           request_rec **r_in,
                                           dns_query_t **q_in)
{
    apr_bucket_brigade *bb;
    apr_byte_t llen;
    apr_size_t dlen;
    apr_status_t rv;
    char *label;
    request_rec *r;
    dns_query_t *q;

    *r_in = r = dns_create_request(dns);
    if (!(r))
        return APR_EGENERAL;

    *q_in = q = apr_array_push(dns->query);

    /* This is the only read we're going to get, so now's the time to do
     * input filters */
    ap_run_insert_filter(r);
    rv = dns_invoke_filter_init(r->input_filters);
    if (rv != OK) {
        return rv;
    }

    /** Ensure qname has a pointee before we start reading the label */
    q->qname = apr_pstrdup(dns->pool, "");

    bb = apr_brigade_create(dns->pool, dns->conn->bucket_alloc);

    /** Read in the label.  TODO: input_compression filter */
    rv = dns_bb_read_byte(bb, r->input_filters, &llen);
    while (rv == APR_SUCCESS && llen > 0) {
        dns->bytesread++;
        rv = ap_get_brigade(r->input_filters, bb, AP_MODE_READBYTES,
                            APR_BLOCK_READ, llen);
        if (rv != APR_SUCCESS)
            return rv;
        dns->bytesread += llen;
        /** Allocate individual labels from request_rec pool, but qname should
         *  be allocated from dns pool */
        dlen = llen;
        rv = apr_brigade_pflatten(bb, &label, &dlen, r->pool);
        if (rv != APR_SUCCESS)
            return rv;

        if (dlen != llen)
            return APR_EGENERAL;

        apr_cpystrn(label, (const char *)label, dlen + 1);
        q->qname = apr_pstrcat(dns->pool, q->qname, label, ".", NULL);

        rv = dns_bb_read_byte(bb, r->input_filters, &llen);
    }
    dns->bytesread++;
    if (rv != APR_SUCCESS)
        return rv;

    /** Read in qclass / qtype */
    rv = dns_bb_read_short(bb, r->input_filters, &(q->qtype));
    if (rv !=APR_SUCCESS)
        return rv;
    dns->bytesread += 2;
    rv = dns_bb_read_short(bb, r->input_filters, &(q->qclass));
    if (rv !=APR_SUCCESS)
        return rv;
    dns->bytesread += 2;

    /** Populate request_rec */
    r->the_request = apr_pstrcat(r->pool, "QUERY ", q->qname, " ",
                                 dns_get_name_class(q->qclass), " ",
                                 dns_get_name_type(q->qtype), NULL);
    r->protocol = apr_pstrdup(r->pool, "DNS");
    r->proto_num = 1000;
    r->hostname = apr_pstrdup(r->pool, q->qname);
    r->method = apr_pstrdup(r->pool, dns_get_name_type(q->qtype));
    r->method_number = (int)q->qtype;
    r->unparsed_uri = apr_pstrcat(r->pool, "dns://", q->qname, "/",
                                           dns_get_name_class(q->qclass), "/",
                                           dns_get_name_type(q->qtype), NULL);
    r->uri = apr_pstrcat(r->pool, "/", dns_get_name_class(q->qclass), "/",
                                  dns_get_name_type(q->qtype), NULL);
    apr_uri_parse(r->pool, r->unparsed_uri, &(r->parsed_uri));
    /** Set handler so modules like mod_dir don't try to pick this up */
    r->handler = apr_pstrdup(r->pool, DNS_MAGIC_TYPE);
    /** The virtualhost info is worthless, but we want to call fix_hostname
     *  and this is the only way to do it */
    ap_update_vhost_from_headers(r);
    ap_run_post_read_request(r);
    return APR_SUCCESS;
}

/** Create a dns_message_t based on data in bb (via filters_in) */
/** Only read the header at this point */
DNS_DECLARE(dns_message_t *) dns_read_message_header(conn_rec *c)
{
    apr_pool_t *p;
    apr_bucket_brigade *tmp_bb;
    dns_message_t *dns;
    dns_header_t *h;
    apr_byte_t data[DNS_HEADER_SIZE];
    int sd_type;
    int dlen = 0;

    apr_uint16_t flags;
    apr_status_t rv;

    if (apr_pool_create(&p, c->pool) != APR_SUCCESS)
        return NULL;

    dns = apr_pcalloc(p, sizeof(*dns));

    if (!(dns))
        return NULL;

    dns->pool = p;
    dns->conn = c;
    dns->filters_in = c->input_filters;
    dns->filters_out = c->output_filters;

    h = dns->header = apr_pcalloc(dns->pool, sizeof(*h));
    dns->query = apr_array_make(p, 0, sizeof(dns_query_t));
    dns->answer = apr_array_make(p, 0, sizeof(dns_rr_t));
    dns->authority = apr_array_make(p, 0, sizeof(dns_rr_t));
    dns->additional = apr_array_make(p, 0, sizeof(dns_rr_t));

    tmp_bb = apr_brigade_create(p, c->bucket_alloc);
    apr_brigade_cleanup(tmp_bb);

    /** If we have a TCP connection, we'll need to read h->length */
    apr_socket_type_get(ap_get_module_config(c->conn_config, &core_module),
                         &sd_type);
    if (sd_type == SOCK_STREAM) {
        rv = dns_bb_read_short(tmp_bb, dns->filters_in, &(dns->length));
        if (rv != APR_SUCCESS) {
            apr_brigade_destroy(tmp_bb);
            return NULL;
        }
        apr_brigade_cleanup(tmp_bb);
    } else {
        dns->length = 0;
    }

    rv = ap_get_brigade(dns->filters_in, tmp_bb, AP_MODE_READBYTES,
                        APR_BLOCK_READ, DNS_HEADER_SIZE);

    if (rv != APR_SUCCESS) {
        apr_brigade_destroy(tmp_bb);
        return NULL;
    }
    dlen = DNS_HEADER_SIZE;
    if (apr_brigade_flatten(tmp_bb, (char *)data, &dlen) != APR_SUCCESS) {
        apr_brigade_destroy(tmp_bb);
        return NULL;
    }
    if (dlen != DNS_HEADER_SIZE) {
        apr_brigade_destroy(tmp_bb);
        return NULL;
    }

    apr_brigade_destroy(tmp_bb);

    memcpy(&(h->id), data, 2);
    h->id = ntohs(h->id);
    memcpy(&flags, data + 2, 2);
    flags = ntohs(flags);

    /** Seperate flags bitfield */
    h->qr = (flags & 32768) >> 15;
    h->opcode = (flags & 30720) >> 11;
    h->aa = (flags & 1024) >> 10;
    h->tc = (flags & 512) >> 9;
    h->rd = (flags & 256) >> 8;
    h->ra = (flags & 128) >> 7;
    h->z  = (flags & 112) >>4;
    h->rcode = flags & 15;

    memcpy(&(h->qdcount), data + 4, 2);
    h->qdcount = ntohs(h->qdcount);
    /** These should all be 0 anyway now... */
    memcpy(&(h->ancount), data + 6, 2);
    h->ancount = ntohs(h->ancount);
    memcpy(&(h->nscount), data + 8, 2);
    h->nscount = ntohs(h->nscount);
    memcpy(&(h->arcount), data + 10, 2);
    h->arcount = ntohs(h->arcount);

    dns->bytesread = DNS_HEADER_SIZE;

    /** Initialize rcode */
    dns->header->rcode = DNS_SUCCESS;

    return dns;
}

/** Output routines.  We expect our generic buffered output filter to be
 *  in place to make this somewhat efficient */
DNS_DECLARE(apr_status_t) dns_add_rr_answer(request_rec *r, dns_rr_t *rr)
{
    apr_bucket_brigade *bb;
    apr_bucket *b;
    apr_status_t rv;
    apr_size_t len;
    char *data;


    if (rr == NULL)
        return APR_SUCCESS;

    bb = apr_brigade_create(r->pool, r->connection->bucket_alloc);
    dns_rr_pserialize(rr, r->pool, &data, &len);
    b = dns_bucket_rr_answer_create((const char *)data, len,
                                    r->pool, bb->bucket_alloc);
    APR_BRIGADE_INSERT_TAIL(bb, b);
    rv = ap_pass_brigade(r->output_filters, bb);
    apr_brigade_destroy(bb);
    return rv;
}

DNS_DECLARE(apr_status_t) dns_add_rr_authority(request_rec *r, dns_rr_t *rr)
{
    apr_bucket_brigade *bb;
    apr_bucket *b;
    apr_status_t rv;
    apr_size_t len;
    char *data;

    if (rr == NULL)
        return APR_SUCCESS;

    bb = apr_brigade_create(r->pool, r->connection->bucket_alloc);
    dns_rr_pserialize(rr, r->pool, &data, &len);
    b = dns_bucket_rr_answer_create((const char *)data, len,
                                    r->pool, bb->bucket_alloc);
    APR_BRIGADE_INSERT_TAIL(bb, b);
    rv = ap_pass_brigade(r->output_filters, bb);
    apr_brigade_destroy(bb);
    return rv;
}

DNS_DECLARE(apr_status_t) dns_add_rr_additional(request_rec *r, dns_rr_t *rr)
{
    apr_bucket_brigade *bb;
    apr_bucket *b;
    apr_status_t rv;
    apr_size_t len;
    char *data;

    if (rr == NULL)
        return APR_SUCCESS;

    bb = apr_brigade_create(r->pool, r->connection->bucket_alloc);
    dns_rr_pserialize(rr, r->pool, &data, &len);
    b = dns_bucket_rr_answer_create((const char *)data, len,
                                    r->pool, bb->bucket_alloc);
    APR_BRIGADE_INSERT_TAIL(bb, b);
    rv = ap_pass_brigade(r->output_filters, bb);
    apr_brigade_destroy(bb);
    return rv;
}