#include "ssl_internal.h"

#include <gnutls/x509.h>

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

// XXX: remove
#include "log.h"
#include <assert.h>


/**
 * Cast a ssl_client to a sock_fd.
 */
#define SSL_CLIENT_FD(client_ptr) (&(client_ptr)->base_tcp.base_trans.base_fd)

/**
 * Cast a ssl_client to a sock_stream.
 */
#define SSL_CLIENT_TRANSPORT(client_ptr) (&(client_ptr)->base_tcp.base_trans.base_fd.base)



/**
 * Enable the TCP events based on the session's gnutls_record_get_direction().
 */
static err_t ssl_client_ev_enable (struct ssl_client *client, error_t *err)
{
    int ret;
    short mask;

    // gnutls_record_get_direction tells us what I/O operation gnutls would have required for the last
    // operation, so we can use that to determine what events to register
    switch ((ret = gnutls_record_get_direction(client->session))) {
        case 0: 
            // read more data
            mask = TRANSPORT_READ;
            break;
        
        case 1:
            // write buffer full
            mask = TRANSPORT_WRITE;
            break;
        
        default:
            // random error
            RETURN_SET_ERROR_EXTRA(err, ERR_GNUTLS_RECORD_GET_DIRECTION, ret);
    }
    
    // do the enabling
    if ((ERROR_CODE(err) = transport_fd_enable(SSL_CLIENT_FD(client), mask)))
        return ERROR_CODE(err);
    

    return SUCCESS;
}

/**
 * Translate a set of gnutls_certificate_status_t values to a constant error message
 */
static const char* ssl_client_verify_error (unsigned int status)
{
    if (status & GNUTLS_CERT_REVOKED)
        return "certificate was revoked";

    else if (status & GNUTLS_CERT_INVALID) {
        if (status & GNUTLS_CERT_SIGNER_NOT_FOUND)
            return "certificate signer was not found";

        else if (status & GNUTLS_CERT_SIGNER_NOT_CA)
            return "certificate signer is not a Certificate Authority";

        else if (status & GNUTLS_CERT_INSECURE_ALGORITHM)
            return "certificate signed using an insecure algorithm";

        else
            return "certificate could not be verified";

    } else
        return "unknown error";

}

/**
 * Perform the certificate validation procedure on the peer cert.
 *
 * Based on the GnuTLS examples/ex-rfc2818.c
 */
static err_t ssl_client_verify (struct ssl_client *client, error_t *err)
{
    unsigned int status;
    const gnutls_datum_t *cert_list;
    unsigned int cert_list_size;
    gnutls_x509_crt_t cert = NULL;
    time_t t, now;

    // init
    RESET_ERROR(err);
    now = time(NULL);
    
    // inspect the peer's cert chain using the installed trusted CAs
    if ((ERROR_EXTRA(err) = gnutls_certificate_verify_peers2(client->session, &status)))
        JUMP_SET_ERROR(err, ERR_GNUTLS_CERT_VERIFY_PEERS2);

    // verify errors?
    if (status)
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, ssl_client_verify_error(status));
    
    // import the main cert
    assert(gnutls_certificate_type_get(client->session) == GNUTLS_CRT_X509);

    if ((ERROR_EXTRA(err) = gnutls_x509_crt_init(&cert)))
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, "gnutls_x509_crt_init");

    if ((cert_list = gnutls_certificate_get_peers(client->session, &cert_list_size)) == NULL)
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, "gnutls_certificate_get_peers");

    if (!cert_list_size)
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, "cert_list_size");

    if ((ERROR_EXTRA(err) = gnutls_x509_crt_import(cert, &cert_list[0], GNUTLS_X509_FMT_DER)))
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, "gnutls_x509_crt_import");
    
    // check expire/activate... not sure if we need to do this
    if ((t = gnutls_x509_crt_get_expiration_time(cert)) == ((time_t) -1) || t < now)
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, "gnutls_x509_crt_get_expiration_time");

    if ((t = gnutls_x509_crt_get_activation_time(cert)) == ((time_t) -1) || t > now)
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, "gnutls_x509_crt_get_activation_time");
    
    // check hostname
    if (!gnutls_x509_crt_check_hostname(cert, client->hostname))
        JUMP_SET_ERROR_STR(err, ERR_GNUTLS_CERT_VERIFY, "gnutls_x509_crt_check_hostname");

error:
    // cleanup
    if (cert)
        gnutls_x509_crt_deinit(cert);
    
    // should be SUCCESS
    return ERROR_CODE(err);    
}


/**
 * Our handshake driver. This will execute the next gnutls_handshake step, handling E_AGAIN.
 *
 * This updates the ssl_client::handshake state internally, as used by ssl_client_event_handler.
 *
 * If the client is marked as verify, this will perform the verification, returning on any errors, and then unset the
 * verify flag - this ensures that the peer cert is only verified once per connection...
 *
 * @return >0 for finished handshake, 0 for handshake-in-progress, -err_t for errors.
 */
static int ssl_client_handshake (struct ssl_client *client, error_t *err)
{
    int ret;

    // perform the handshake
    if ((ret = gnutls_handshake(client->session)) < 0 && ret != GNUTLS_E_AGAIN)
        JUMP_SET_ERROR_EXTRA(err, ERR_GNUTLS_HANDSHAKE, ret);
    
    // complete?
    if (ret == 0) {
        // update state
        client->handshake = false;

        // verify?
        if (client->verify) {
            // perform the validation
            if (ssl_client_verify(client, err))
                goto error;
            
            // unmark
            client->verify = false;
        }

        // handshake done
        return 1;

    } else {
        // set state, isn't really needed every time, but easier this way
        client->handshake = true;

        // re-enable the event for the next iteration
        return ssl_client_ev_enable(client, err);
    }

error:
    return -ERROR_CODE(err);    
}

/**
 * Our transport_fd event handler. Drive the handshake if that's current, otherwise, invoke user callbacks.
 */
static void ssl_client_on_event (struct transport_fd *fd, short what, void *arg)
{
    struct ssl_client *client = arg;
    error_t err;

    (void) fd;

    // XXX: timeouts
    (void) what;

    // are we in the handshake cycle?
    if (client->handshake) {
        RESET_ERROR(&err);

        // perform the next handshake step
        // this returns zero when the handshake is not yet done, errors/completion then trigger the else-if-else below
        if (ssl_client_handshake(client, &err) == 0) {
            // handshake continues
        
        } else if (!SSL_CLIENT_TRANSPORT(client)->connected) {
            // the async connect+handshake process has completed
            // invoke the user connect callback directly with appropriate error
            transport_connected(SSL_CLIENT_TRANSPORT(client), ERROR_CODE(&err) ? &err : NULL, true);

        } else {
            // in-connection re-handshake completed
            if (ERROR_CODE(&err))
                // the re-handshake failed, so this transport is dead
                transport_error(SSL_CLIENT_TRANSPORT(client), &err);
        
            else
                // re-handshake completed, so continue with the transport_callbacks
                transport_invoke(SSL_CLIENT_TRANSPORT(client), what);
        }

    } else {
        // normal transport operation
        // gnutls might be able to proceed now, so invoke user callbacks
        transport_invoke(SSL_CLIENT_TRANSPORT(client), what);
    }
}

static err_t ssl_client__read (transport_t *transport, void *buf, size_t *len, error_t *err)
{
    struct ssl_client *client = transport_check(transport, &ssl_client_type);
    int ret;
    
    // read gnutls record
    do {
        ret = gnutls_record_recv(client->session, buf, *len);

    } while (ret == GNUTLS_E_INTERRUPTED);
    
    // errors
    // XXX: E_REHANDSHAKE?
    if (ret < 0 && ret != GNUTLS_E_AGAIN)
        RETURN_SET_ERROR_EXTRA(err, ERR_GNUTLS_RECORD_RECV, ret);
    
    else if (ret == 0)
        return SET_ERROR(err, ERR_EOF);


    // EAGAIN?
    if (ret < 0) {
        *len = 0;

    } else {
        // updated length
        *len = ret;

    }

    return SUCCESS;
}

static err_t ssl_client__write (transport_t *transport, const void *buf, size_t *len, error_t *err)
{
    struct ssl_client *client = transport_check(transport, &ssl_client_type);
    int ret;
 
    // read gnutls record
    do {
        ret = gnutls_record_send(client->session, buf, *len);
   
    } while (ret == GNUTLS_E_INTERRUPTED);

    // errors
    if (ret < 0 && ret != GNUTLS_E_AGAIN)
        RETURN_SET_ERROR_EXTRA(err, ERR_GNUTLS_RECORD_SEND, ret);
    
    else if (ret == 0)
        return SET_ERROR(err, ERR_WRITE_EOF);


    // eagain?
    if (ret < 0) {
        *len = 0;

    } else {
        // updated length
        *len = ret;
    }

    return SUCCESS;
}

void ssl_client_deinit (struct ssl_client *client)
{
    // close the session rudely
    gnutls_deinit(client->session);
    client->session = NULL;
 
    // terminate the TCP transport
    tcp_client_deinit(&client->base_tcp);
   
    if (client->cred) {
        // drop the cred ref
        ssl_client_cred_put(client->cred);

        client->cred = NULL;
    }

    // free
    free(client->hostname);
    client->hostname = NULL;
}


static void ssl_client__deinit (transport_t *transport)
{
    struct ssl_client *client = transport_check(transport, &ssl_client_type);
    
    // die
    ssl_client_deinit(client);
}

/**
 * Our tcp_client-invoked connect handler
 */
static void ssl_client__connected (transport_t *transport, const error_t *tcp_err)
{
    struct ssl_client *client = transport_check(transport, &ssl_client_type);
    error_t err;

    // trap errors to let the user handle them directly
    if (tcp_err)
        JUMP_SET_ERROR_INFO(&err, tcp_err);
    
    // bind default transport functions (recv/send) to use the TCP fd
    gnutls_transport_set_ptr(client->session, (gnutls_transport_ptr_t) (long int) SSL_CLIENT_FD(client)->fd);

    // add ourselves as the event handler
    if ((ERROR_CODE(&err) = transport_fd_setup(SSL_CLIENT_FD(client), ssl_client_on_event, client)))
        goto error;

    // start handshake
    if (ssl_client_handshake(client, &err))
        // this should complete with SUCCESS if it returns >0
        goto error;

    // ok, so we wait...
    return;

error:
    // tell the user
    transport_connected(transport, &err, true);
}

struct transport_type ssl_client_type = {
    .base_type = {
        .parent     = &tcp_client_type.base_type,
    },
    .methods = {
        .read       = ssl_client__read,
        .write      = ssl_client__write,
        .deinit     = ssl_client__deinit,
        ._connected = ssl_client__connected,
    },
};



static void ssl_client_destroy (struct ssl_client *client)
{
    ssl_client_deinit(client);

    free(client);
}

err_t ssl_connect (const struct transport_info *info, transport_t **transport_ptr, 
        const char *hostname, const char *service,
        struct ssl_client_cred *cred,
        error_t *err
    )
{
    struct ssl_client *client = NULL;

    // alloc
    if ((client = calloc(1, sizeof(*client))) == NULL)
        return SET_ERROR(err, ERR_CALLOC);

    // initialize base
    transport_init(SSL_CLIENT_TRANSPORT(client), &ssl_client_type, info);

    if (!cred) {
        // default credentials
        cred = &ssl_client_cred_anon;
    
    } else {
        // take a ref
        client->cred = cred;
        cred->refcount++;
    };

    // do verify?
    if (cred->verify)
        client->verify = true;

    // init
    if ((client->hostname = strdup(hostname)) == NULL)
        JUMP_SET_ERROR(err, ERR_STRDUP);

    // initialize TCP
    tcp_client_init(&client->base_tcp);

    // initialize client session
    if ((ERROR_EXTRA(err) = gnutls_init(&client->session, GNUTLS_CLIENT)) < 0)
        JUMP_SET_ERROR(err, ERR_GNUTLS_INIT);

    // ...default priority stuff
    if ((ERROR_EXTRA(err) = gnutls_set_default_priority(client->session)))
        JUMP_SET_ERROR(err, ERR_GNUTLS_SET_DEFAULT_PRIORITY);

    // XXX: silly hack for OpenSSL interop
    gnutls_dh_set_prime_bits(client->session, 512);

    // bind credentials
    if ((ERROR_EXTRA(err) = gnutls_credentials_set(client->session, GNUTLS_CRD_CERTIFICATE, cred->x509)))
        JUMP_SET_ERROR(err, ERR_GNUTLS_CRED_SET);

    // TCP connect
    if (tcp_client_connect_async(&client->base_tcp, hostname, service, err))
        goto error;

    // done, wait for the connect to complete
    *transport_ptr = SSL_CLIENT_TRANSPORT(client);

    return SUCCESS;

error:
    // cleanup
    ssl_client_destroy(client);

    return ERROR_CODE(err);    
}