#include "usuals.h"
#include "httpauthd.h"
#include "md5.h"
#include "sha1.h"
#include "bd.h"

#include <sys/time.h>

/* Mysql library */
#include <mysql.h>
#include <errmsg.h>

/* -------------------------------------------------------------------------------
 * Structures
 */

#define DB_PW_CLEAR      0
#define DB_PW_SHA1       1
#define DB_PW_MD5        2
#define DB_PW_CRYPT      3

/* Our hanler context */
typedef struct mysql_context
{
    /* Base Handler ------------------------------------------------------ */
    bd_context_t bd;

    /* Settings ---------------------------------------------------------- */
    const char* host;         /* The connection host or path */
    unsigned int port;        /* The connection port */
    const char* user;         /* The pgsql user name */
    const char* password;     /* The pgsql password */
    const char* database;     /* The database name */
    const char* query;        /* The query */
    const char* pw_column;    /* The database query to retrieve a password */
    int pw_type;              /* The type of password encoded in database */
    const char* ha1_column;   /* The database query to retrieve a ha1 */

    int mysql_max;            /* Number of open connections allowed */
    int mysql_timeout;        /* Maximum amount of time to dedicate to a query */
    int use_unix_socket;

    /* Context ----------------------------------------------------------- */
    MYSQL** pool;            /* Pool of available connections */
    int pool_mark;          /* Amount of connections allocated */
}
mysql_context_t;

/* Forward declarations for callbacks */
static int validate_digest(ha_request_t* rq, const char* user, digest_context_t* dg);
static int validate_basic(ha_request_t* rq, const char* user, const char* password);
static void escape_mysql(const ha_request_t* rq, ha_buffer_t* buf, const char* value);

/* The defaults for the context */
static const mysql_context_t mysql_defaults =
{
    BD_CALLBACKS(validate_digest,
        validate_basic, escape_mysql),
    NULL,               /* host */
    0,                  /* port */
    NULL,               /* user */
    NULL,               /* password */
    NULL,               /* database */
    NULL,               /* query */
    NULL,               /* pw_attr */
    DB_PW_CLEAR,        /* pw_type */
    NULL,               /* ha1_attr */
    10,                 /* mysql_max */
    30,                 /* mysql_timeout */
    0,                  /* use_unix_socket */
    NULL,               /* pool */
    0                   /* pool_mark */
};


/* -------------------------------------------------------------------------------
 * Internal Functions
 */

static void escape_mysql(const ha_request_t* rq, ha_buffer_t* buf, const char* value)
{
    size_t len;
    char* t;

    ASSERT(value);

    len = strlen(value);

    t = (char*)malloc((len * 2) + 1);
    if(t != NULL)
    {
        mysql_escape_string(t, value, len);
        ha_bufcpy(buf, t);
        free(t);
    }
}

static int dec_mysql_binary(const ha_request_t* rq, const char* enc,
                            unsigned char* bytes, size_t len)
{
    size_t enclen;
    void* d;

    ASSERT(rq && enc && bytes && len);
    enclen = strlen(enc);

    /* Hex encoded */
    if(enclen == (len * 2))
    {
        d = ha_bufdechex(rq->buf, enc, &enclen);

        if(d && enclen == len)
        {
            ha_messagex(rq, LOG_DEBUG, "found value in hex encoded format");
            memcpy(bytes, d, len);
            return HA_OK;
        }
    }

    /* TODO: Does mysql have raw binary encoding? */

    /* B64 Encoded */
    enclen = strlen(enc);
    d = ha_bufdec64(rq->buf, enc, &enclen);

    if(d && len == enclen)
    {
        ha_messagex(rq, LOG_DEBUG, "found value in b64 encoded format");
        memcpy(bytes, d, len);
        return HA_OK;
    }

    return ha_buferr(rq->buf) ? HA_CRITERROR : HA_FALSE;
}

static int validate_ha1(ha_request_t* rq, mysql_context_t* ctx, const char* user,
                        const char* clearpw, const char* dbpw)
{
    unsigned char dbha1[MD5_LEN];
    unsigned char ha1[MD5_LEN];
    const char* p;

    int r = dec_mysql_binary(rq, dbpw, dbha1, MD5_LEN);

    if(r < 0)
        return r;

    if(r == HA_OK)
    {
        digest_makeha1(ha1, user, rq->context->realm, clearpw);
        if(memcmp(ha1, dbha1, MD5_LEN) == 0)
            return HA_OK;
    }

    return HA_FALSE;
}

static int validate_password(ha_request_t* rq, mysql_context_t* ctx, const char* user,
                             const char* clearpw, const char* dbpw)
{
    unsigned char buf[SHA1_LEN];
    const char* p;
    int r;

    ASSERT(SHA1_LEN > MD5_LEN);

    switch(ctx->pw_type)
    {

    /* Clear text */
    case DB_PW_CLEAR:

        if(strcmp(clearpw, dbpw) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching clear text password");
            return HA_OK;
        }

        break;

    /* Crypt pw */
    case DB_PW_CRYPT:

        ha_lock(NULL);
            p = (const char*)crypt(clearpw, dbpw);
        ha_unlock(NULL);

        if(p && strcmp(dbpw, p) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching crypt password");
            return HA_OK;
        }

        break;

    /* MD5 */
    case DB_PW_MD5:

        r = dec_mysql_binary(rq, dbpw, buf, MD5_LEN);
        if(r < 0) return r;

        if(r == HA_OK && md5_strcmp(buf, clearpw) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching md5 password");
            return HA_OK;
        }

        break;


    /* SHA1 */
    case DB_PW_SHA1:

        r = dec_mysql_binary(rq, dbpw, buf, SHA1_LEN);
        if(r < 0) return r;

        if(r == HA_OK && sha1_strcmp(buf, clearpw) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching sha1 password");
            return HA_OK;
        }

        break;

    default:
        ASSERT(0 && "invalid password type");
        break;
    };

    return HA_FALSE;
}

static MYSQL* get_mysql_connection(const ha_request_t* rq, mysql_context_t* ctx)
{
    MYSQL* my;
    MYSQL* r;
    int i;

    ASSERT(ctx);

    for(i = 0; i < ctx->mysql_max; i++)
    {
        /* An open connection in the pool */
        if(ctx->pool[i])
        {
            ha_messagex(rq, LOG_DEBUG, "using cached mysql connection");
            my = ctx->pool[i];
            ctx->pool[i] = NULL;
            return my;
        }
    }

    if(ctx->pool_mark >= ctx->mysql_max)
    {
        ha_messagex(rq, LOG_ERR, "too many open mysql connections");
        return NULL;
    }

    my = mysql_init(NULL);
    if(!my)
    {
        ha_messagex(rq, LOG_ERR, "out of memory");
        return NULL;
    }

    mysql_options(my, MYSQL_OPT_CONNECT_TIMEOUT, (char*)&(ctx->mysql_timeout));
    /* mysql_options(my, MYSQL_OPT_READ_TIMEOUT, (char*)&(ctx->mysql_timeout));
    mysql_options(my, MYSQL_OPT_WRITE_TIMEOUT, (char*)&(ctx->mysql_timeout)); */

    /* Apparently mysql_real_connect is not thread safe :( */
    ha_lock(NULL);
        r = mysql_real_connect(my, ctx->use_unix_socket ? NULL : ctx->host,
                               ctx->user, ctx->password, ctx->database, ctx->port,
                               ctx->use_unix_socket ? ctx->host : NULL, 0);
    ha_unlock(NULL);

    if(!r)
    {
        ha_messagex(rq, LOG_ERR, "error opening mysql connection: %s", mysql_error(my));
        mysql_close(my);
        return NULL;
    }

    ctx->pool_mark++;
    ha_messagex(rq, LOG_DEBUG, "opened new mysql connection (total %d)", ctx->pool_mark);
    return my;
}

static void discard_mysql_connection(const ha_request_t* rq, mysql_context_t* ctx, MYSQL* my)
{
    mysql_close(my);
    ctx->pool_mark--;
    ha_messagex(rq, LOG_DEBUG, "discarding mysql connection (total %d)", ctx->pool_mark);
}

static void save_mysql_connection(const ha_request_t* rq, mysql_context_t* ctx, MYSQL* my)
{
    int i, e;

    ASSERT(ctx);

    if(!my)
        return;

    switch(mysql_errno(my))
    {
    case CR_SOCKET_CREATE_ERROR:
    case CR_CONNECTION_ERROR:
    case CR_CONN_HOST_ERROR:
    case CR_IPSOCK_ERROR:
    case CR_UNKNOWN_HOST:
    case CR_SERVER_GONE_ERROR:
    case CR_VERSION_ERROR:
    case CR_WRONG_HOST_INFO:
    case CR_LOCALHOST_CONNECTION:
    case CR_TCP_CONNECTION:
    case CR_SERVER_HANDSHAKE_ERR:
    case CR_SERVER_LOST:
    case CR_COMMANDS_OUT_OF_SYNC:
        break;

    /* Make sure it's worth saving */
    default:
        for(i = 0; i < ctx->mysql_max; i++)
        {
            /* An open connection in the pool */
            if(!ctx->pool[i])
            {
                ha_messagex(rq, LOG_DEBUG, "caching mysql connection for later use");
                ctx->pool[i] = my;
                my = NULL;
                break;
            }
        }
        break;
    };

    if(my != NULL)
        discard_mysql_connection(rq, ctx, my);
}

static int resolve_column(MYSQL_RES* res, const char* column)
{
    int i, fields;

    if(column)
    {
        fields = mysql_num_fields(res);
        for(i = 0; i < fields; i++)
        {
            if(strcasecmp(column, mysql_fetch_field_direct(res, i)->name) == 0)
                return i;
        }
    }

    return -1;
}

static int retrieve_user_rows(ha_request_t* rq, mysql_context_t* ctx,
                              const char* user, MYSQL_RES** results)
{
    MYSQL* my = NULL;
    MYSQL_RES* res = NULL;
    const char* query;
    int ret = HA_OK;

    ASSERT(rq && ctx && user && results);
    *results = NULL;

    my = get_mysql_connection(rq, ctx);
    if(!my)
        RETURN(HA_FAILED);

    ASSERT(ctx->query);

    /* The map can have %u and %r to denote user and realm */
    query = bd_substitute(rq, user, ctx->query);
    if(!query)
        RETURN(HA_CRITERROR);

    ha_messagex(rq, LOG_DEBUG, "executing query: %s", query);
    if(mysql_query(my, query) != 0)
    {
        ha_messagex(rq, LOG_ERR, "error querying database: %s", mysql_error(my));
        RETURN(HA_FAILED);
    }

    res = mysql_store_result(my);
    if(!res)
    {
        if(mysql_field_count(my) == 0)
            ha_messagex(rq, LOG_ERR, "mysql query didn't return results: %s", query);
        else
            ha_messagex(rq, LOG_ERR, "error querying database: %s", mysql_error(my));

        RETURN(HA_FAILED);
    }

    if(mysql_num_rows(res) == 0)
    {
        ha_messagex(rq, LOG_WARNING, "login failed. couldn't find user: %s", user);
        RETURN(HA_FALSE);
    }

    ha_messagex(rq, LOG_DEBUG, "received %d result rows", mysql_num_rows(res));
    *results = res;
    res = NULL;

finally:

    if(res != NULL)
        mysql_free_result(res);

    /* TODO: Look into what happens if we free a mysql connection
     * before processing results: */
    if(my != NULL)
        save_mysql_connection(rq, ctx, my);

    return ret;
}

static int validate_digest(ha_request_t* rq, const char* user, digest_context_t* dg)
{
    mysql_context_t* ctx = (mysql_context_t*)rq->context->ctx_data;
    MYSQL_RES* res = NULL;
    MYSQL_ROW row;
    int ret = HA_FALSE;
    int pw_column = -1;
    int ha1_column = -1;
    int r, i, foundany = 0;
    const char* v;

    ASSERT(rq && user && dg);

    ret = retrieve_user_rows(rq, ctx, user, &res);
    if(ret != HA_OK)
        RETURN(ret);

    ASSERT(res);

    pw_column = resolve_column(res, ctx->pw_column);
    ha1_column = resolve_column(res, ctx->ha1_column);

    if(pw_column == -1 && ha1_column == -1)
    {
        if(mysql_num_fields(res) > 1)
            ha_messagex(rq, LOG_WARNING, "query returned more than 1 column, using first as password");

        pw_column = 0;
    }

    while((row = mysql_fetch_row(res)) != NULL)
    {
        if(pw_column != -1 && ctx->pw_type == DB_PW_CLEAR)
        {
            v = *(row + pw_column);
            if(v != NULL)
            {
                foundany = 1;

                digest_makeha1(dg->ha1, user, rq->context->realm, v);
                ha_messagex(rq, LOG_DEBUG, "testing clear text password for digest auth");

                /* Run the actual check */
                ret = digest_complete_check(dg, rq->buf);

                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }

        if(ha1_column != -1)
        {
            v = *(row + ha1_column);
            if(v != NULL)
            {
                ret = dec_mysql_binary(rq, v, dg->ha1, MD5_LEN);
                if(ret < 0)
                    RETURN(ret)
                else if(ret == HA_FALSE)
                    continue;

                foundany = 1;

                /* Run the actual check */
                ret = digest_complete_check(dg, rq->buf);

                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }
    }

    if(!foundany)
        ha_messagex(rq, LOG_WARNING, "no clear password or ha1 present for user: %s", user);

finally:
    if(res)
        mysql_free_result(res);

    return ret;
}

static int validate_basic(ha_request_t* rq, const char* user, const char* password)
{
    mysql_context_t* ctx = (mysql_context_t*)rq->context->ctx_data;
    MYSQL_RES* res = NULL;
    MYSQL_ROW row;
    int ret = HA_FALSE;
    int pw_column = -1;
    int ha1_column = -1;
    int i, foundany = 0;
    const char* v;

    ASSERT(rq && user && password);

    ret = retrieve_user_rows(rq, ctx, user, &res);
    if(ret != HA_OK)
        RETURN(ret);

    ASSERT(res);

    pw_column = resolve_column(res, ctx->pw_column);
    ha1_column = resolve_column(res, ctx->ha1_column);

    if(pw_column == -1 && ha1_column == -1)
    {
        if(mysql_num_fields(res) > 1)
            ha_messagex(rq, LOG_WARNING, "query returned more than 1 column, using first as password");
        pw_column = 0;
    }


    while((row = mysql_fetch_row(res)) != NULL)
    {
        if(pw_column != -1)
        {
            v = *(row + pw_column);
            if(v != NULL)

            {
                foundany = 1;
                ret = validate_password(rq, ctx, user, password, v);
                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }

        if(ha1_column != -1)
        {
            v = *(row + ha1_column);
            if(v != NULL)
            {
                foundany = 1;
                ret = validate_ha1(rq, ctx, user, password, v);
                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }
    }

    if(!foundany)
        ha_messagex(rq, LOG_WARNING, "no password present for user: %s", user);

finally:
    if(res)
        mysql_free_result(res);

    return ret;
}


/* -------------------------------------------------------------------------------
 * Handler Functions
 */

int mysql_config(ha_context_t* context, const char* name, const char* value)
{
    mysql_context_t* ctx = (mysql_context_t*)(context->ctx_data);

    ASSERT(name && value && value[0]);

    if(strcmp(name, "dbserver") == 0)
    {
        ctx->use_unix_socket = (value[0] == '/');
        ctx->host = value;
        return HA_OK;
    }

    if(strcmp(name, "dbport") == 0)
    {
        return ha_confint(name, value, 0, 65535, &(ctx->port));
    }

    if(strcmp(name, "dbuser") == 0)
    {
        ctx->user = value;
        return HA_OK;
    }

    if(strcmp(name, "dbpassword") == 0)
    {
        ctx->password = value;
        return HA_OK;
    }

    if(strcmp(name, "dbdatabase") == 0)
    {
        ctx->database = value;
        return HA_OK;
    }

    if(strcmp(name, "dbquery") == 0)
    {
        ctx->query = value;
        return HA_OK;
    }

    if(strcmp(name, "dbpwcolumn") == 0)
    {
        ctx->pw_column = value;
        return HA_OK;
    }

    if(strcmp(name, "dbpwtype") == 0)
    {
        if(strcasecmp(value, "clear") == 0)
            ctx->pw_type = DB_PW_CLEAR;
        else if(strcasecmp(value, "crypt") == 0)
            ctx->pw_type = DB_PW_CRYPT;
        else if(strcasecmp(value, "md5") == 0)
            ctx->pw_type = DB_PW_MD5;
        else if(strcasecmp(value, "sha1") == 0)
            ctx->pw_type = DB_PW_SHA1;
        else
        {
            ha_messagex(NULL, LOG_ERR, "invalid value for '%s' (must be 'clear', 'crypt', 'md5' or 'sha1')", name);
            return HA_FAILED;
        }

        return HA_OK;
    }

    if(strcmp(name, "dbha1column") == 0)
    {
        ctx->ha1_column = value;
        return HA_OK;
    }

    else if(strcmp(name, "dbmax") == 0)
    {
        return ha_confint(name, value, 1, 256, &(ctx->mysql_max));
    }

    else if(strcmp(name, "dbtimeout") == 0)
    {
        /* TODO: Implement database timeouts */
        return ha_confint(name, value, 0, 86400, &(ctx->mysql_timeout));
    }

    return HA_FALSE;
}

int mysql_initialize(ha_context_t* context)
{
    int r;

    if((r = bd_init(context)) != HA_OK)
        return r;

    /* Context specific initialization */
    if(context)
    {
        mysql_context_t* ctx = (mysql_context_t*)(context->ctx_data);
        ASSERT(ctx);

        /* Check for mandatory configuration */
        if(!ctx->database || !ctx->query)
        {
            ha_messagex(NULL, LOG_ERR, "mysql configuration incomplete. "
                                       "Must have DBDatabase and DBQuery.");
            return HA_FAILED;
        }

        ASSERT(!ctx->pool);
        ASSERT(ctx->mysql_max > 0);

        /*
         * Our connection pool. It's the size of our maximum
         * amount of pending connections as that's the max
         * we'd be able to use at a time anyway.
         */
        ctx->pool = (MYSQL**)malloc(sizeof(MYSQL*) * ctx->mysql_max);
        if(!ctx->pool)
        {
            ha_messagex(NULL, LOG_CRIT, "out of memory");
            return HA_CRITERROR;
        }

        memset(ctx->pool, 0, sizeof(MYSQL*) * ctx->mysql_max);
        ha_messagex(NULL, LOG_INFO, "initialized mysql handler");
    }

    return HA_OK;
}

void mysql_destroy(ha_context_t* context)
{
    if(context)
    {
        /* Note: We don't need to be thread safe here anymore */
        mysql_context_t* ctx = (mysql_context_t*)(context->ctx_data);
        int i;

        ASSERT(ctx);

        if(ctx->pool)
        {
            /* Close any connections we have open */
            for(i = 0; i < ctx->mysql_max; i++)
            {
                if(ctx->pool[i])
                    mysql_close(ctx->pool[i]);
            }

            /* And free the connection pool */
            free(ctx->pool);
        }
    }

    bd_destroy(context);
    ha_messagex(NULL, LOG_INFO, "uninitialized mysql handler");
}



/* -------------------------------------------------------------------------------
 * Handler Definition
 */

ha_handler_t mysql_handler =
{
    "MYSQL",                 /* The type */
    mysql_initialize,        /* Initialization function */
    mysql_destroy,           /* Uninitialization routine */
    mysql_config,            /* Config routine */
    bd_process,              /* Processing routine */
    &mysql_defaults,         /* The context defaults */
    sizeof(mysql_context_t)
};