#include <assert.h>

#include <ssl_sess_cache.h>
#include <ssl_utils.h>
#include <tfe_utils.h>

#include <MESA/MESA_htable.h>
#include <MESA/field_stat2.h>

#define    SESS_CACHE_NOT_FOUND            -1
#define    SESS_CACHE_FOUND                0
#define SESS_CACHE_UPDATE_OLD        1
#define SESS_CACHE_ADD_NEW            2
#define SESS_CACHE_INVALID            3

struct asn1_sess
{
	unsigned char * buff;
	size_t size;
};

struct sess_set_args
{
	MESA_htable_handle hash;
	struct asn1_sess * new_sess;
};

struct sess_cache
{
	enum tfe_conn_dir served_for;
	MESA_htable_handle hash;
	long long hit_cnt, miss_cnt, del_err;
};

static void ssl_sess_free_serialized(void * data)
{
	struct asn1_sess * p = (struct asn1_sess *) data;
	free(p->buff);
	p->size = 0;
	free(p);
	return;
}

static struct asn1_sess * ssl_sess_serialize(SSL_SESSION * sess)
{
	struct asn1_sess * result = ALLOC(struct asn1_sess, 1);

	int __i2d_size = i2d_SSL_SESSION(sess, NULL);
	result->size = (size_t) __i2d_size;
	assert(__i2d_size > 0);

	/* When using i2d_SSL_SESSION(), the memory location pointed to by pp must be large enough to
	 * hold the binary representation of the session. There is no known limit on the size of the
	 * created ASN1 representation, so the necessary amount of space should be obtained by first
	 * calling i2d_SSL_SESSION() with pp=NULL, and obtain the size needed,
	 * then allocate the memory and call i2d_SSL_SESSION() again.*/

	result->buff = ALLOC(unsigned char, result->size);
	i2d_SSL_SESSION(sess, &(result->buff));
	return result;
}

static SSL_SESSION * ssl_sess_deserialize(const struct asn1_sess * asn1)
{
	SSL_SESSION * sess = NULL;
	d2i_SSL_SESSION(&sess, (const unsigned char **) &(asn1->buff), (long) asn1->size); /* increments asn1 */
	return sess;
}

static int ssl_sess_verify_cb(void * data, int eliminate_type)
{
	const struct asn1_sess * asn1 = (struct asn1_sess *) data;
	if (eliminate_type == ELIMINATE_TYPE_NUM)
	{
		return 1;    //direct expired.
	}

	SSL_SESSION * sess = ssl_sess_deserialize(asn1);
	int ret = ssl_session_is_valid(sess);
	SSL_SESSION_free(sess);

	if (ret == 0)
	{
		return 1;    //should be expired (deleted).
	}
	else
	{
		return 0;
	}
}

static long sess_cache_get_cb(void * data, const uchar * key, uint size, void * user_arg)
{
	SSL_SESSION * sess = NULL;
	int is_valid = 0;
	if (data == NULL)
	{
		return SESS_CACHE_NOT_FOUND;
	}

	const struct asn1_sess * asn1 = (struct asn1_sess *) data;
	sess = ssl_sess_deserialize(asn1);
	is_valid = ssl_session_is_valid(sess);

	if (is_valid == 0)
	{
		SSL_SESSION_free(sess);
		return SESS_CACHE_INVALID;
	}
	else
	{
		*(SSL_SESSION **) user_arg = sess;
		return SESS_CACHE_FOUND;
	}
}

static long sess_cache_set_cb(void * data, const uchar * key, uint size, void * user_arg)
{
	struct sess_set_args * args = (struct sess_set_args *) user_arg;
	struct asn1_sess * new_asn1 = args->new_sess;
	struct asn1_sess * cur_asn1 = (struct asn1_sess *) data;

	int ret = 0;
	if (cur_asn1 != NULL)
	{
		free(cur_asn1->buff);
		cur_asn1->size = new_asn1->size;
		cur_asn1->buff = ALLOC(unsigned char, cur_asn1->size);
		memcpy(cur_asn1->buff, new_asn1->buff, cur_asn1->size);
		return SESS_CACHE_UPDATE_OLD;
	}
	else
	{
		ret = MESA_htable_add(args->hash, key, size, new_asn1);
		assert(ret >= 0);
		return SESS_CACHE_ADD_NEW;
	}
}

static size_t upsess_mk_key(struct sockaddr * addr, socklen_t addrlen, const char * sni, unsigned char ** key_buf)
{
	size_t key_size = 0;
	unsigned char * tmp = NULL;
	size_t tmp_size;
	short port;
	size_t snilen;

	switch (addr->sa_family)
	{
		case AF_INET:
			tmp = (unsigned char *)&((struct sockaddr_in *) addr)->sin_addr;
			tmp_size = sizeof(struct in_addr);
			port = ((struct sockaddr_in *) addr)->sin_port;
			break;
		case AF_INET6:
			tmp = (unsigned char *)&((struct sockaddr_in6 *) addr)->sin6_addr;
			tmp_size = sizeof(struct in6_addr);
			port = ((struct sockaddr_in6 *) addr)->sin6_port;
			break;
		default:
			//should never happens.
			assert(0);
			break;
	}

	snilen = sni ? strlen(sni) : 0;
	key_size = tmp_size + sizeof(port) + snilen;

	*key_buf = ALLOC(unsigned char, key_size);
	unsigned char * p = *key_buf;

	memcpy(p, tmp, tmp_size);
	p += tmp_size;
	memcpy(p, (char *) &port, sizeof(port));
	p += sizeof(port);
	return key_size;
}

void up_session_set(struct sess_cache * cache, struct sockaddr * addr, socklen_t addr_len, const char * sni,
	SSL_SESSION * sess)
{
	unsigned char * key = NULL;
	int ret = 0;
	size_t key_size = 0;
	long cb_ret = 0;
	void * no_use = NULL;
	assert(cache->served_for == CONN_DIR_UPSTREAM);
	key_size = upsess_mk_key(addr, addr_len, sni, &key);

	struct asn1_sess * asn1 = NULL;
	asn1 = ssl_sess_serialize(sess);

	struct sess_set_args set_args{.hash = cache->hash, .new_sess = asn1};
	no_use = MESA_htable_search_cb(cache->hash, key, key_size, sess_cache_set_cb, &set_args, &cb_ret);
	if (cb_ret == SESS_CACHE_UPDATE_OLD)
	{
		ssl_sess_free_serialized(asn1);
	}
	free(key);
	return;
}

SSL_SESSION * up_session_get(struct sess_cache * cache, struct sockaddr * addr, socklen_t addr_len, const char * sni)
{
	SSL_SESSION * sess = NULL;
	void * no_use = NULL;
	long cb_ret = 0;

	size_t key_size = 0;
	assert(cache->served_for == CONN_DIR_UPSTREAM);

	unsigned char * key = NULL;
	key_size = upsess_mk_key(addr, addr_len, sni, &key);
	no_use = MESA_htable_search_cb(cache->hash, key, key_size, sess_cache_get_cb, &sess, &cb_ret);
	free(key);
	key = NULL;
	if (cb_ret == 1)
	{
		ATOMIC_INC(&(cache->hit_cnt));
		return sess;
	}
	else
	{
		ATOMIC_INC(&(cache->miss_cnt));
		return NULL;
	}
}

void down_session_set(struct sess_cache * cache, const SSL_SESSION * sess)
{
	unsigned int idlen = 0;
	struct asn1_sess * asn1 = NULL;
	long cb_ret = 0;
	void * no_use = NULL;
	int ret = 0;
	assert(cache->served_for == CONN_DIR_DOWNSTREAM);
	asn1 = ssl_sess_serialize((SSL_SESSION *) sess);

	/*
	* SSL_SESSION_get_id() returns a pointer to the internal session id value for the session s. 
	* The length of the id in bytes is stored in *idlen. The length may be 0. 
	* The caller should not free the returned pointer directly.
	*/

	const unsigned char * id = SSL_SESSION_get_id(sess, &idlen);
	struct sess_set_args set_args{.hash = cache->hash, .new_sess = asn1};

	no_use = MESA_htable_search_cb(cache->hash, id, (unsigned int) idlen, sess_cache_set_cb, &set_args, &cb_ret);
	if (cb_ret == SESS_CACHE_UPDATE_OLD)
	{
		ssl_sess_free_serialized(asn1);
	}
	return;
}

SSL_SESSION * down_session_get(struct sess_cache * cache, const unsigned char * id, int idlen)
{
	SSL_SESSION * sess = NULL;
	void * no_use = NULL;
	long cb_ret = 0;
	assert(cache->served_for == CONN_DIR_DOWNSTREAM);
	no_use = MESA_htable_search_cb(cache->hash, id, (unsigned int) idlen, sess_cache_get_cb, &sess, &cb_ret);
	if (cb_ret == 1)
	{
		ATOMIC_INC(&(cache->hit_cnt));
		return sess;
	}
	else
	{
		ATOMIC_INC(&(cache->miss_cnt));
		return NULL;
	}
}

void down_session_del(struct sess_cache * cache, const SSL_SESSION * sess)
{
	assert(cache->served_for == CONN_DIR_DOWNSTREAM);
	unsigned int len = 0;
	const unsigned char * id = SSL_SESSION_get_id(sess, &len);
	int ret = MESA_htable_del(cache->hash, id, len, NULL);
	if (ret != MESA_HTABLE_RET_OK)
	{
		ATOMIC_INC(&(cache->del_err));
	}
	return;
}

int __wrapper_MESA_htable_set_opt(MESA_htable_handle table, enum MESA_htable_opt opt_type, unsigned value)
{
	int ret = MESA_htable_set_opt(table, opt_type, &value, (int)(sizeof(value)));
	assert(ret == 0);
	return ret;
}

int __wrapper_MESA_htable_set_opt(MESA_htable_handle table, enum MESA_htable_opt opt_type, void * val, size_t len)
{
	int ret = MESA_htable_set_opt(table, opt_type, val, (int)len);
	assert(ret == 0);
	return ret;
}

struct sess_cache * ssl_sess_cache_create(unsigned int slot_size, unsigned int expire_seconds, enum tfe_conn_dir served)
{
	struct sess_cache * cache = ALLOC(struct sess_cache, 1);
	unsigned max_num = slot_size * 4;
	int ret = 0;

	MESA_htable_handle htable = MESA_htable_born();
	ret = __wrapper_MESA_htable_set_opt(htable, MHO_SCREEN_PRINT_CTRL, 0);
	ret = __wrapper_MESA_htable_set_opt(htable, MHO_THREAD_SAFE, 1);

	ret = __wrapper_MESA_htable_set_opt(htable, MHO_MUTEX_NUM, 16);
	ret = __wrapper_MESA_htable_set_opt(htable, MHO_HASH_SLOT_SIZE, slot_size);
	ret = __wrapper_MESA_htable_set_opt(htable, MHO_HASH_MAX_ELEMENT_NUM, max_num);
	ret = __wrapper_MESA_htable_set_opt(htable, MHO_EXPIRE_TIME, expire_seconds);

	ret = __wrapper_MESA_htable_set_opt(htable, MHO_ELIMIMINATE_TYPE,
		HASH_ELIMINATE_ALGO_FIFO);
	ret = __wrapper_MESA_htable_set_opt(htable, MHO_CBFUN_DATA_FREE,
		(void *)ssl_sess_free_serialized, sizeof(&ssl_sess_free_serialized));
	ret = __wrapper_MESA_htable_set_opt(htable, MHO_CBFUN_DATA_EXPIRE_NOTIFY,
		(void *)ssl_sess_verify_cb, sizeof(&ssl_sess_verify_cb));

	ret = MESA_htable_mature(htable);
	assert(ret == 0);

	cache->hash = htable;
	cache->served_for = served;
	return cache;
}
void ssl_sess_cache_destroy(struct sess_cache * cache)
{
	MESA_htable_destroy(cache->hash, NULL);
	cache->hash = NULL;
	free(cache);
	return;
}