/**
 *    Copyright (C) 2024 Graham Leggett <minfrin@sharp.fm>
 *
 * 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.
 *
 */

/*
 * redwax_unbound - DNS server access routines.
 *
 */

#include <apr_strings.h>

#include "config.h"
#include "redwax-tool.h"

#include "redwax_unbound.h"
#include "redwax_util.h"

#if HAVE_UNBOUND_H

#include "apr_portable.h"
#include "apr_uri.h"

#include <netdb.h>

#include <unbound.h>

module unbound_module;

typedef struct {
    struct ub_ctx *ctx;
    apr_pollfd_t socket_poll;
    redwax_pollfd_t poll_ctx;
    apr_status_t status;
    int dns_work;
    int dns_server_set;
    int dns_trust_anchor_set;
} unbound_config_t;

static apr_status_t cleanup_ub(void *dummy)
{
    struct ub_ctx *ctx = dummy;

    ub_ctx_delete(ctx);

    return APR_SUCCESS;
}

static apr_status_t redwax_unbound_initialise(redwax_tool_t *r)
{
    unbound_config_t *config;

    config = apr_pcalloc(r->pool, sizeof(unbound_config_t));
    redwax_set_module_config(r->per_module, &unbound_module, config);

    /* create an unbound context */
    config->ctx = ub_ctx_create();
    if(!config->ctx) {
        redwax_print_error(r,
                "Could not create unbound context\n");
        return APR_EINIT;
    }

    apr_pool_cleanup_register(r->pool, config->ctx, cleanup_ub,
            apr_pool_cleanup_null);

    return OK;
}

static apr_status_t redwax_unbound_set_dns_server(redwax_tool_t *r, const char *arg)
{
    unbound_config_t *config = redwax_get_module_config(r->per_module, &unbound_module);

    int rv;

    if ((rv = ub_ctx_set_fwd(config->ctx, arg)) != 0) {
        redwax_print_error(r,
                "Could not assign DNS server '%s': %s\n",
                arg, ub_strerror(rv));
        return APR_EINIT;
    }

    config->dns_server_set++;

    return APR_SUCCESS;
}

static apr_status_t redwax_unbound_set_dns_trust_anchor(redwax_tool_t *r, const char *arg)
{
    unbound_config_t *config = redwax_get_module_config(r->per_module, &unbound_module);

    int rv;

    if ((rv = ub_ctx_add_ta_file(config->ctx, arg)) != 0) {
        redwax_print_error(r,
                "Could not read DNS trust anchor '%s': %s (%s)\n",
                arg, ub_strerror(rv), strerror(errno));
        return APR_EINIT;
    }

    config->dns_trust_anchor_set++;

    return APR_SUCCESS;
}

static apr_status_t redwax_unbound_set_tlsa(redwax_tool_t *r, const char *arg)
{
    apr_uri_t uri;

    apr_status_t status;

    if (r->dane_basename) {
        redwax_print_error(r,
                "URI '%s': filter-verify-tlsa can only be specified once\n", arg);
        return APR_EINVAL;
    }

    status = apr_uri_parse(r->pool, arg, &uri);
    if (APR_SUCCESS != status) {
        redwax_print_error(r,
                "Could not parse URI '%s': %pm\n", arg, &status);
        return status;
    }

    if (!uri.hostname) {
        redwax_print_error(r,
                "URI '%s': hostname missing\n", arg);
        return APR_EINVAL;
    }
    else {
        r->dane_basename = uri.hostname;
    }

    if (!strcmp(uri.scheme, "tcp") || !strcmp(uri.scheme, "udp")) {

        if (uri.port) {

            const char **qname;

            const char *tlsa = apr_psprintf(r->pool, "_%d._%s.%s", uri.port, uri.scheme, uri.hostname);

            redwax_dns_t *dns = apr_array_push(r->dns_requests);

            dns->r = r;
            dns->basename = uri.hostname;
            dns->qname = tlsa;
            dns->rrtype = 52 /* TYPE TLSA */;
            dns->rrclass = 1 /* CLASS IN (internet) */;
            dns->cb = rt_run_process_tlsa;
            dns->port = uri.port;

            qname = apr_array_push(r->tlsa_qnames);
            *qname = tlsa;

            return APR_SUCCESS;
        }
        else {
            redwax_print_error(r,
                    "URI '%s': port missing\n", arg);
            return APR_EINVAL;
        }
    }

    if (uri.scheme) {

        int found = 0;

        setservent(1);

        struct servent *ent;

        while ((ent = getservent())) {

            if (!strcmp(uri.scheme, ent->s_name)) {

                const char **qname;

                const char *tlsa = apr_psprintf(r->pool, "_%d._%s.%s",
                        uri.port ? uri.port : ntohs(ent->s_port),
                                ent->s_proto, uri.hostname);

                redwax_dns_t *dns = apr_array_push(r->dns_requests);

                dns->r = r;
                dns->basename = uri.hostname;
                dns->qname = tlsa;
                dns->rrtype = 52 /* TYPE TLSA */;
                dns->rrclass = 1 /* CLASS IN (internet) */;
                dns->cb = rt_run_process_tlsa;
                dns->port = uri.port ? uri.port : ntohs(ent->s_port);

                qname = apr_array_push(r->tlsa_qnames);
                *qname = tlsa;

                found = 1;
            }

        }

        endservent();

        if (!found) {
            redwax_print_error(r,
                    "URI '%s': protocol '%s' not found\n", arg, uri.scheme);
            return APR_EINVAL;
        }
        else {
            return APR_SUCCESS;
        }

    }

    else {
        redwax_print_error(r,
                "URI '%s': scheme missing\n", arg);
        return APR_EINVAL;
    }

}

static apr_status_t redwax_unbound_set_smimea(redwax_tool_t *r, const char *arg)
{

    return APR_ENOTIMPL;
}

static apr_status_t redwax_unbound_set_tls_in(redwax_tool_t *r, const char *arg)
{
    apr_uri_t uri;

    apr_status_t status;

    int protocol = 0;
    int type = 0;

    status = apr_uri_parse(r->pool, arg, &uri);
    if (APR_SUCCESS != status) {
        redwax_print_error(r,
                "Could not parse URI '%s': %pm\n", arg, &status);
        return status;
    }

    if (!uri.hostname) {
        redwax_print_error(r,
                "URI '%s': hostname missing\n", arg);
        return APR_EINVAL;
    }

    /* fixme: Use inet_pton to detect raw IP/IPv6 addresses and skip lookup */

    if (!strcmp(uri.scheme, "tcp")) {
        protocol = APR_PROTO_TCP;
        type = SOCK_STREAM;
    }
    else if (!strcmp(uri.scheme, "udp")) {
        protocol = APR_PROTO_UDP;
        type = SOCK_DGRAM;
    }
    else if (!strcmp(uri.scheme, "sctp")) {
        protocol = APR_PROTO_SCTP;
        type= SOCK_SEQPACKET;
    }

    if (protocol) {

        if (uri.port) {

            redwax_dns_t *dns;
            redwax_tls_t *tls;

            tls = apr_palloc(r->pool, sizeof(redwax_tls_t));

            tls->unparsed_uri = arg;
            tls->uri = uri;
            tls->protocol = protocol;
            tls->type = type;

            dns = apr_array_push(r->dns_requests);

            dns->r = r;
            dns->basename = uri.hostname;
            dns->qname = uri.hostname;
            dns->rrtype = 28 /* TYPE AAAA */;
            dns->rrclass = 1 /* CLASS IN (internet) */;
            dns->cb = rt_run_process_tls_in;
            dns->ctx = tls;
            dns->port = uri.port;

            dns = apr_array_push(r->dns_requests);

            dns->r = r;
            dns->basename = uri.hostname;
            dns->qname = uri.hostname;
            dns->rrtype = 1 /* TYPE A */;
            dns->rrclass = 1 /* CLASS IN (internet) */;
            dns->cb = rt_run_process_tls_in;
            dns->ctx = tls;
            dns->port = uri.port;

            return APR_SUCCESS;
        }
        else {
            redwax_print_error(r,
                    "URI '%s': port missing\n", arg);
            return APR_EINVAL;
        }
    }

    else if (uri.scheme) {

        int found = 0;

        setservent(1);

        struct servent *ent;

        while ((ent = getservent())) {

            if (!strcmp(uri.scheme, ent->s_name)) {

                redwax_dns_t *dns;
                redwax_tls_t *tls;

                tls = apr_palloc(r->pool, sizeof(redwax_tls_t));

                /* we hard code the protocol to TCP, as the services
                 * file contacts lots of bogus tcp/udp entries.
                 */

                tls->unparsed_uri = arg;
                tls->uri = uri;
                tls->protocol = APR_PROTO_TCP;
                tls->type = SOCK_STREAM;

                dns = apr_array_push(r->dns_requests);

                dns->r = r;
                dns->basename = uri.hostname;
                dns->qname = uri.hostname;
                dns->rrtype = 28 /* TYPE AAAA */;
                dns->rrclass = 1 /* CLASS IN (internet) */;
                dns->cb = rt_run_process_tls_in;
                dns->ctx = tls;
                dns->port = uri.port ? uri.port : ntohs(ent->s_port);

                dns = apr_array_push(r->dns_requests);

                dns->r = r;
                dns->basename = uri.hostname;
                dns->qname = uri.hostname;
                dns->rrtype = 1 /* TYPE A */;
                dns->rrclass = 1 /* CLASS IN (internet) */;
                dns->cb = rt_run_process_tls_in;
                dns->ctx = tls;
                dns->port = uri.port ? uri.port : ntohs(ent->s_port);

                found = 1;

                break;
            }

        }

        endservent();

        if (!found) {
            redwax_print_error(r,
                    "URI '%s': protocol '%s' not found\n", arg, uri.scheme);
            return APR_EINVAL;
        }
        else {
            return APR_SUCCESS;
        }

    }

    else {
        redwax_print_error(r,
                "URI '%s': scheme missing\n", arg);
        return APR_EINVAL;
    }

}

static apr_status_t redwax_unbound_filter_read_cb(redwax_tool_t *r, void *baton, apr_pollfd_t *descriptor)
{
    unbound_config_t *config = redwax_get_module_config(r->per_module, &unbound_module);

    apr_status_t status = APR_SUCCESS;
    int rv;

    /* ack readcb work */
    r->poll_work--;

    rv = ub_process(config->ctx);

    if (rv) {
        redwax_print_error(r,
                "Could not process DNS: %s (%s)\n",
                ub_strerror(rv), strerror(errno));
        return APR_EGENERAL;
    }

    if (config->status) {
        return config->status;
    }

    if (config->dns_work) {

        /* line up next event */
        status = apr_pollcb_add(r->poll, &config->socket_poll);

        if (APR_SUCCESS != status) {
            redwax_print_error(r,
                    "unbound: could not add descriptor to poll: %pm\n", &status);
            return status;
        }

        /* req readcb work */
        r->poll_work++;
    }

    return status;
}

void redwax_unbound_filter_result_cb(void *ctx, int rv, struct ub_result* result)
{
    redwax_dns_t *dns = ctx;

    redwax_tool_t *r = dns->r;

    unbound_config_t *config = redwax_get_module_config(r->per_module, &unbound_module);

    /* ack dns lookup work */
    config->dns_work--;

    if (rv) {
        redwax_print_error(r,
                "Could not resolve: %s\n",
                ub_strerror(rv));
        return;
    }

    dns->canonname = result->canonname;
    dns->havedata = result->havedata;
    dns->nxdomain = result->nxdomain;
    dns->secure = result->secure;
    dns->bogus = result->bogus;
    dns->why_bogus = result->why_bogus;

    if (result->data) {

        int i, count;

        for (count = 0; result->data[count]; count++);

        dns->rdata = apr_array_make(r->pool, count, sizeof(redwax_rdata_t));

        for (i = 0; i < count; i++) {

            redwax_rdata_t *rdata = apr_array_push(dns->rdata);

            rdata->data = (unsigned char *)result->data[i];
            rdata->len = result->len[i];

            rt_run_process_dns(r, dns, rdata);

        }

    }

    config->status = dns->cb(r, dns);

    ub_resolve_free(result);
}

static apr_status_t redwax_unbound_filter_write_cb(redwax_tool_t *r, void *baton, apr_pollfd_t *descriptor)
{
    unbound_config_t *config = redwax_get_module_config(r->per_module, &unbound_module);

    redwax_dns_t *dns = apr_array_pop(r->dns_requests);

    apr_status_t status = APR_SUCCESS;
    int rv;

    /* ack writecb work */
    r->poll_work--;

    rv = ub_resolve_async(config->ctx, dns->qname,
            dns->rrtype,
            dns->rrclass,
            (void *)dns, redwax_unbound_filter_result_cb, NULL);

    if (rv) {
        redwax_print_error(r,
                "Could not resolve '%s': %s (%s)\n",
                dns->qname, ub_strerror(rv), strerror(errno));
        return APR_EGENERAL;
    }

    /* req dns lookup work */
    config->dns_work++;

    /* we are done sending requests, poll for responses only */
    if (!r->dns_requests->nelts) {
        config->socket_poll.reqevents = APR_POLLIN;
    }

    /* line up next event */
    status = apr_pollcb_add(r->poll, &config->socket_poll);

    if (APR_SUCCESS != status) {
        redwax_print_error(r,
                "Could not add read descriptor to poll: %pm\n", &status);
        return status;
    }

    /* req readcb work */
    r->poll_work++;

    return status;
}

static apr_status_t redwax_unbound_filter_cb(redwax_tool_t *r, void *baton, apr_pollfd_t *descriptor)
{
    if (descriptor->rtnevents & APR_POLLIN) {
        /* process reads first if available to relieve pressure */
        return redwax_unbound_filter_read_cb(r, baton, descriptor);
    }
    else if (descriptor->rtnevents & APR_POLLOUT) {
        /* do writes if needed */
        return redwax_unbound_filter_write_cb(r, baton, descriptor);
    }
    else {
        /* errors are handled by ub_process() */
        return redwax_unbound_filter_read_cb(r, baton, descriptor);
    }

    return APR_SUCCESS;
}

static apr_status_t redwax_unbound_filter_poll(redwax_tool_t *r)
{
    unbound_config_t *config;
    apr_socket_t *s = NULL;

    apr_os_sock_t fd;
    apr_status_t status;

    if (!r->dns_requests->nelts) {
        return OK;
    }

    config = redwax_get_module_config(r->per_module, &unbound_module);

    /*
     * Apply defaults if no explicit values are set.
     */

    /* read /etc/resolv.conf for DNS proxy settings (from DHCP) */
    if (!config->dns_server_set) {

        int rv;

        if ((rv = ub_ctx_resolvconf(config->ctx, "/etc/resolv.conf"))) {
            redwax_print_error(r,
                    "Could not read /etc/resolv.conf: %s (%s)\n",
                    ub_strerror(rv), strerror(errno));
            return APR_EINIT;
        }
    }

    /* read public keys for DNSSEC verification */
    if (!config->dns_trust_anchor_set) {

        int rv;

        if ((rv = ub_ctx_add_ta_file(config->ctx, REDWAX_DEFAULT_ROOT_KEY))) {
            redwax_print_error(r,
                    "Could not read " REDWAX_DEFAULT_ROOT_KEY ": %s (%s)\n",
                    ub_strerror(rv), strerror(errno));
            return APR_EINIT;
        }
    }

    fd = ub_fd(config->ctx);

    if (-1 == fd) {
        redwax_print_error(r,
                "Could not obtain unbound fd: %s\n",
                strerror(errno));
        return APR_EGENERAL;
    }

    status = apr_os_sock_put(&s, &fd, r->pool);

    if (APR_SUCCESS != status) {
        redwax_print_error(r,
                "Could not convert unbound fd: %pm\n", &status);
        return status;
    }

    /* set up the read/write callback */

    config->poll_ctx.cb = redwax_unbound_filter_cb;

    /* set up read/write descriptor */
    config->socket_poll.desc_type = APR_POLL_SOCKET;
    config->socket_poll.reqevents = APR_POLLIN | APR_POLLOUT;
    config->socket_poll.desc.s = s;
    config->socket_poll.client_data = &config->poll_ctx;

    status = apr_pollcb_add(r->poll, &config->socket_poll);

    if (APR_SUCCESS != status) {
        redwax_print_error(r,
                "Could not add descriptor to poll: %pm\n", &status);
        return status;
    }

    /* req writecb work */
    r->poll_work++;

    return APR_SUCCESS;
}

void redwax_add_default_unbound_hooks(apr_pool_t *pool)
{
    rt_hook_initialise(redwax_unbound_initialise, NULL, NULL, APR_HOOK_MIDDLE);
    rt_hook_set_dns_server(redwax_unbound_set_dns_server, NULL, NULL, APR_HOOK_MIDDLE);
    rt_hook_set_dns_trust_anchor(redwax_unbound_set_dns_trust_anchor, NULL, NULL, APR_HOOK_MIDDLE);
    rt_hook_set_tlsa(redwax_unbound_set_tlsa, NULL, NULL, APR_HOOK_MIDDLE);
    rt_hook_set_smimea(redwax_unbound_set_smimea, NULL, NULL, APR_HOOK_MIDDLE);
    rt_hook_set_tls_in(redwax_unbound_set_tls_in, NULL, NULL, APR_HOOK_MIDDLE);
    rt_hook_filter_poll(redwax_unbound_filter_poll, NULL, NULL, APR_HOOK_MIDDLE);
}

#else

void redwax_add_default_unbound_hooks(apr_pool_t *pool)
{
}

#endif

REDWAX_DECLARE_MODULE(unbound) =
{
    STANDARD_MODULE_STUFF,
    redwax_add_default_unbound_hooks                   /* register hooks */
};