#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>
#include <openssl/hmac.h>
#include <openssl/rand.h>

#define NGX_HTTP_POW_CHALLENGE_RANDOM_LENGTH 64
#define NGX_HTTP_POW_HMAC_KEY "thisisaverysecretkey"
#define NGX_HTTP_POW_CHALLENGE_COOKIE_KEY "pow-challenge="
#define NGX_HTTP_POW_CHALLENGE_COOKIE_ATTRIBUTES "; Path=/; SameSite=Lax"
#define NGX_HTTP_POW_COOKIE_NAME "pow"

#define NGX_HTTP_POW_UNUSED __attribute__((unused))

static ngx_int_t ngx_http_pow_init(ngx_conf_t *cf);
static ngx_int_t ngx_http_pow_handler(ngx_http_request_t *r);

static ngx_str_t *pow_html;

static ngx_http_module_t  ngx_http_pow_module_ctx = {
    NULL,                   /* preconfiguration */
    ngx_http_pow_init,      /* postconfiguration */
    NULL,                   /* create main configuration */
    NULL,                   /* init main configuration */
    NULL,                   /* create server configuration */
    NULL,                   /* merge server configuration */
    NULL,                   /* create location configuration */
    NULL                    /* merge location configuration */
};

ngx_module_t  ngx_http_pow = {
    NGX_MODULE_V1,
    &ngx_http_pow_module_ctx,	   /* module context */
    NULL,                          /* module directives */
    NGX_HTTP_MODULE,               /* module type */
    NULL,                          /* init master */
    NULL,                          /* init module */
    NULL,                          /* init process */
    NULL,                          /* init thread */
    NULL,                          /* exit thread */
    NULL,                          /* exit process */
    NULL,                          /* exit master */
    NGX_MODULE_V1_PADDING
};

typedef struct __attribute__((packed)) ngx_http_pow_challenge_s {
	uint32_t  version;
	uint32_t  created;
	uint32_t  ttl;
	uint32_t  difficulty;
	u_char    random[NGX_HTTP_POW_CHALLENGE_RANDOM_LENGTH];
	u_char    digest[EVP_MAX_MD_SIZE];
} ngx_http_pow_challenge_t;

typedef struct __attribute__((packed)) ngx_http_pow_s {
	ngx_http_pow_challenge_t  challenge;
	uint32_t                  nonce;
	u_char                    hash[EVP_MAX_MD_SIZE];
} ngx_http_pow_t;

static ngx_http_pow_challenge_t
*ngx_http_pow_create_challenge(ngx_pool_t *pool)
{
	ngx_http_pow_challenge_t *challenge;

	challenge = ngx_pcalloc(pool, sizeof(*challenge));

	challenge->version    = 1;
	challenge->created    = ngx_time();
	challenge->ttl        = 30;
	challenge->difficulty = 3;

	RAND_bytes(challenge->random, NGX_HTTP_POW_CHALLENGE_RANDOM_LENGTH);
	
	HMAC(EVP_sha256(), NGX_HTTP_POW_HMAC_KEY, sizeof(NGX_HTTP_POW_HMAC_KEY),
			(u_char *)challenge, sizeof(*challenge) - EVP_MAX_MD_SIZE,
			challenge->digest, NULL);

	return challenge;
}

static ngx_str_t
*ngx_http_pow_challenge_marshal(ngx_http_pow_challenge_t *challenge,
		ngx_pool_t *pool)
{
	ngx_str_t   challenge_string;
	ngx_str_t  *marshalled;
	size_t      encoded_length;
	u_char     *buf;

	challenge_string.data = (u_char *)challenge;
	challenge_string.len  = sizeof(*challenge);

	encoded_length = ngx_base64_encoded_length(sizeof(*challenge));
	buf            = ngx_pcalloc(pool, encoded_length);
	marshalled     = ngx_pcalloc(pool, sizeof(ngx_str_t));

 	marshalled->len  = encoded_length;
	marshalled->data = buf;

	ngx_encode_base64(marshalled, &challenge_string);

	return marshalled;
}

static ngx_http_pow_challenge_t NGX_HTTP_POW_UNUSED
*ngx_http_pow_challenge_unmarshal(ngx_str_t *marshalled, ngx_pool_t *pool)
{
	ngx_http_pow_challenge_t  *challenge;
	ngx_str_t                 *unmarshalled;

	challenge    = ngx_pcalloc(pool, sizeof(*challenge));
	unmarshalled = ngx_pcalloc(pool, sizeof(ngx_str_t));

	unmarshalled->data = (u_char *)challenge;
	unmarshalled->len  = sizeof(*challenge);

	ngx_decode_base64(unmarshalled, marshalled);

	return challenge;
}

static void
ngx_http_pow_require_pow(ngx_http_request_t *r)
{
	ngx_http_pow_challenge_t  *challenge;
	ngx_table_elt_t           *header;
	ngx_chain_t                out;
	ngx_str_t                 *marshalled;
	ngx_buf_t                 *body_buf;
	u_char                    *p;

	challenge  = ngx_http_pow_create_challenge(r->pool);
	marshalled = ngx_http_pow_challenge_marshal(challenge, r->pool);
	header     = ngx_list_push(&r->headers_out.headers);

	ngx_str_set(&header->key, "Set-Cookie");

	header->hash = 1;

    header->value.len = sizeof(NGX_HTTP_POW_CHALLENGE_COOKIE_KEY) - 1
		              + marshalled->len
				      + sizeof(NGX_HTTP_POW_CHALLENGE_COOKIE_ATTRIBUTES) - 1;

	header->value.data = ngx_pcalloc(r->pool, header->value.len);

	p = header->value.data;

	p = ngx_cpymem(p, NGX_HTTP_POW_CHALLENGE_COOKIE_KEY,
			sizeof(NGX_HTTP_POW_CHALLENGE_COOKIE_KEY) - 1);
    p = ngx_cpymem(p, marshalled->data, marshalled->len);
    p = ngx_cpymem(p, NGX_HTTP_POW_CHALLENGE_COOKIE_ATTRIBUTES,
            sizeof(NGX_HTTP_POW_CHALLENGE_COOKIE_ATTRIBUTES) - 1);

	r->headers_out.status           = NGX_HTTP_UNAUTHORIZED;
	r->headers_out.content_length_n = pow_html->len;

	body_buf = ngx_calloc_buf(r->pool);
	if (body_buf == NULL) {
		/* We done fucked up */
	}

	body_buf->pos           = pow_html->data;
	body_buf->last          = pow_html->data + pow_html->len;
	body_buf->memory        = 1;
	body_buf->last_buf      = 1;
	body_buf->last_in_chain = 1;
	body_buf->sync          = 0;

	out.buf  = body_buf;
	out.next = NULL;

	ngx_http_send_header(r);
	ngx_http_output_filter(r, &out);
	ngx_http_finalize_request(r, NGX_HTTP_UNAUTHORIZED);
}

static ngx_str_t NGX_HTTP_POW_UNUSED
*ngx_http_pow_marshal(ngx_http_pow_t *pow, ngx_pool_t *pool)
{
	ngx_str_t   pow_string;
	ngx_str_t  *marshalled;
	size_t      encoded_length;
	u_char     *buf;

	pow_string.data = (u_char *)pow;
	pow_string.len  = sizeof(*pow);

	encoded_length = ngx_base64_encoded_length(sizeof(*pow));
	buf            = ngx_pcalloc(pool, encoded_length);
	marshalled     = ngx_pcalloc(pool, sizeof(ngx_str_t));

 	marshalled->len  = encoded_length;
	marshalled->data = buf;

	ngx_encode_base64(marshalled, &pow_string);

	return marshalled;
}

static ngx_http_pow_t
*ngx_http_pow_unmarshal(ngx_pool_t *pool, ngx_str_t *marshalled)
{
	ngx_http_pow_t  *pow;
	ngx_str_t       *unmarshalled;

	pow          = ngx_pcalloc(pool, sizeof(*pow));
	unmarshalled = ngx_pcalloc(pool, sizeof(ngx_str_t));

	unmarshalled->data = (u_char *)pow;
	unmarshalled->len  = sizeof(*pow);

	ngx_decode_base64(unmarshalled, marshalled);

	return pow;
}

static bool
ngx_http_pow_check_difficulty(ngx_http_pow_t *pow, uint32_t difficulty)
{
	size_t  i;
	size_t  zero_bytes     = difficulty / 8;
	size_t  remaining_bits = difficulty % 8;

	static u_char lut[] = {
		0b11111111, 0b01111111, 0b00111111, 0b00011111,
		0b00001111, 0b00000111, 0b00000011, 0b00000001
	};

	for (i = 0; i < zero_bytes; i++) {
		if (pow->hash[i] > 0) {
			return false;
		}
	}

	if (pow->hash[zero_bytes + 1] ^ lut[remaining_bits]) {
		return false;
	}

	return true;
}

static bool
ngx_http_pow_check_digest(void *data, size_t len, u_char *digest)
{
	u_char  d[EVP_MAX_MD_SIZE];

	HMAC(EVP_sha256(), NGX_HTTP_POW_HMAC_KEY, sizeof(NGX_HTTP_POW_HMAC_KEY),
			data, len, d, NULL);

	/* CRYPTO_memcmp returns 0 if compared values are equal, else != 0 */
	if (CRYPTO_memcmp(digest, d, EVP_MAX_MD_SIZE)) {
		return false;
	}

	return true;
}

static bool
ngx_http_pow_check_pow(ngx_http_request_t *r)
{
	ngx_http_pow_challenge_t  *challenge;
	ngx_http_pow_t            *pow;
	ngx_str_t                  pow_cookie_value;

	ngx_str_t  pow_cookie_name = ngx_string(NGX_HTTP_POW_COOKIE_NAME);

	/* No PoW cookie is present */
	if (!ngx_http_parse_multi_header_lines(r, r->headers_in.cookie,
				&pow_cookie_name, &pow_cookie_value)) {
		return false;
	}

	pow = ngx_http_pow_unmarshal(r->pool, &pow_cookie_value);
	challenge = &pow->challenge;

	/* Check digest matches */
	if (!ngx_http_pow_check_digest(challenge,
				sizeof(*challenge) - sizeof(challenge->digest),
				challenge->digest)) {
		return false;
	}

	/* Check pow is still valid */
	if (challenge->created + challenge->ttl > ngx_time()) {
		return false;
	}

	/* Check pow difficulty */
	if (!ngx_http_pow_check_difficulty(pow, challenge->difficulty)) {
		return false;
	}

	/* Check pow is correct */
	if (!ngx_http_pow_check_digest(pow, sizeof(*pow) - sizeof(pow->hash),
				pow->hash)) {
		return false;
	}

	return true;
}

static ngx_int_t
ngx_http_pow_handler(ngx_http_request_t *r)
{
	/* PoW is either not present or presetn, but incorrect, require PoW */
	if (!ngx_http_pow_check_pow(r)) {
		ngx_http_pow_require_pow(r);
		return NGX_DONE;
	}

	/* PoW is present and correct, resume processing */
	return NGX_DECLINED;
}

static ngx_str_t
*ngx_http_pow_load_js(ngx_pool_t *pool, ngx_str_t *filename, ngx_log_t *log)
{
	size_t           filesize;
	ngx_str_t       *pow_html = NULL;
	ngx_file_t       file;
	ngx_file_info_t  file_info;

	file.name = *filename;
	file.log  = log;
	file.fd   = ngx_open_file(filename->data, NGX_FILE_RDONLY, NGX_FILE_OPEN, 0);

	if (file.fd == NGX_INVALID_FILE) {
		goto cleanup;
	}

	if (ngx_file_info(filename->data, &file_info) == NGX_FILE_ERROR) {
		goto cleanup;
	}

	filesize = ngx_file_size(&file_info);

	pow_html       = ngx_pcalloc(pool, sizeof(ngx_str_t));
	pow_html->data = ngx_pcalloc(pool, filesize);
	pow_html->len  = ngx_read_file(&file, pow_html->data, filesize, 0);

	if (/* pow_html->len == NGX_ERROR || */ pow_html->len != filesize) {
		pow_html = NULL;
	}

cleanup:

	ngx_close_file(file.fd);

	return pow_html;
}

static ngx_int_t
ngx_http_pow_init(ngx_conf_t *cf)
{
	ngx_str_t pow_html_path =
		ngx_string("/Users/jona/repos/ngx-pow/ngx-http-pow/html/pow.html");

    ngx_http_handler_pt        *h;
    ngx_http_core_main_conf_t  *cmcf;

    cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);

	pow_html = ngx_http_pow_load_js(cf->pool, &pow_html_path, cf->log);
	if (!pow_html) {
		return NGX_ERROR;
	}

    h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
    if (h == NULL) {
        return NGX_ERROR;
    }

    *h = ngx_http_pow_handler;

    return NGX_OK;
}