#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;
	int version;
};

struct sess_set_args
{
	MESA_htable_handle hash;
	struct asn1_sess * new_sess;
};
struct sess_get_args
{
	SSL_SESSION *sess;
	int version;	
};

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, int version)
{
	struct asn1_sess * result = ALLOC(struct asn1_sess, 1);

	int i = i2d_SSL_SESSION(sess, NULL);
	UNUSED	int j=0;
	result->size = (size_t) i;
	unsigned char* temp=NULL;
	assert(i > 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.
	 *Note that this will advance the value contained in *pp so it is necessary to save a copy of the original allocation.*/
	result->buff = temp = ALLOC(unsigned char, result->size);
	
	j=i2d_SSL_SESSION(sess, &(temp));
	assert(i == j);
	assert(result->buff + i == temp);
	result->version=version;
	return result;
}

static SSL_SESSION * ssl_sess_deserialize(const struct asn1_sess * asn1)
{
	SSL_SESSION * sess = NULL;
	
	const unsigned char *p=asn1->buff;
	/* i2d_SSL_SESSION increments the pointer pointed to by p to point one byte after the saved data 
	 * We save the pointer first.*/
	sess=d2i_SSL_SESSION(NULL, &(p), (long) asn1->size); /* increments asn1 */
	assert(sess!=NULL);
	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)
{
	struct sess_get_args *result=(struct sess_get_args *)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
	{
		result->sess=sess;
		result->version=asn1->version;
		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;

	UNUSED 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);
		cur_asn1->version=new_asn1->version;
		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 * res, socklen_t addrlen, const char * sni, unsigned char ** key_buf)
{
	size_t key_size = 0;
	unsigned short port=0;
		
	char *s = NULL;
	switch(res->sa_family) {
		case AF_INET: {
			struct sockaddr_in *addr_in = (struct sockaddr_in *)res;
			s = (char*) malloc(INET_ADDRSTRLEN);
			inet_ntop(AF_INET, &(addr_in->sin_addr), s, INET_ADDRSTRLEN);
			port = ntohs(addr_in->sin_port);
			break;
		}
		case AF_INET6: {
			struct sockaddr_in6 *addr_in6 = (struct sockaddr_in6 *)res;
			s = (char*) malloc(INET6_ADDRSTRLEN);
			inet_ntop(AF_INET6, &(addr_in6->sin6_addr), s, INET6_ADDRSTRLEN);
			port = ntohs(addr_in6->sin6_port);
			break;
		}
		default:
			break;
	}
	
	key_size=asprintf((char**)key_buf,"%s:%u:%s",s, port, sni?sni:"null");
	free(s);
	return key_size;
}

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

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

	struct sess_set_args set_args={.hash = cache->hash, .new_sess = asn1};
	MESA_htable_search_cb(cache->hash, key, key_size, sess_cache_set_cb, &set_args, &cb_ret);
//	printf("%s %s\n", __FUNCTION__, key);
	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, int min_ver, int max_ver)
{	
	struct sess_get_args args={NULL, 0};
	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);
	MESA_htable_search_cb(cache->hash, key, key_size, sess_cache_get_cb, &args, &cb_ret);
//	printf("%s %s\n", __FUNCTION__, key);
	free(key);
	key = NULL;
	if (cb_ret == SESS_CACHE_FOUND && args.version>=min_ver && args.version<=max_ver)
	{
		ATOMIC_INC(&(cache->hit_cnt));
		return args.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;
	assert(cache->served_for == CONN_DIR_DOWNSTREAM);
	asn1 = ssl_sess_serialize((SSL_SESSION *) sess, 0);

	/*
	* 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};

	MESA_htable_search_cb(cache->hash, id, 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)
{
	struct sess_get_args result;
	memset(&result, 0, sizeof(result));
	long cb_ret = 0;
	assert(cache->served_for == CONN_DIR_DOWNSTREAM);
	MESA_htable_search_cb(cache->hash, id, (unsigned int) idlen, sess_cache_get_cb, &result, &cb_ret);
	if (cb_ret == SESS_CACHE_FOUND)
	{
		ATOMIC_INC(&(cache->hit_cnt));
		return result.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;
}
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;
	UNUSED int ret = 0;
	int opt_val=0;
	
	MESA_htable_handle htable = MESA_htable_born();
	opt_val=0;
	ret = MESA_htable_set_opt(htable, MHO_SCREEN_PRINT_CTRL, &opt_val, sizeof(opt_val));
	opt_val=1;
	ret = MESA_htable_set_opt(htable, MHO_THREAD_SAFE, &opt_val, sizeof(opt_val));
	opt_val=16;
	ret = MESA_htable_set_opt(htable, MHO_MUTEX_NUM, &opt_val, sizeof(opt_val));
	ret = MESA_htable_set_opt(htable, MHO_HASH_SLOT_SIZE, &slot_size, sizeof(slot_size));
	ret = MESA_htable_set_opt(htable, MHO_HASH_MAX_ELEMENT_NUM, &max_num, sizeof(max_num));
	ret = MESA_htable_set_opt(htable, MHO_EXPIRE_TIME, &expire_seconds, sizeof(expire_seconds));

	opt_val=HASH_ELIMINATE_ALGO_FIFO;
	ret = MESA_htable_set_opt(htable, MHO_ELIMIMINATE_TYPE,
		&opt_val, sizeof(int));
	ret = MESA_htable_set_opt(htable, MHO_CBFUN_DATA_FREE,
		(void *)ssl_sess_free_serialized, sizeof(&ssl_sess_free_serialized));
	ret = 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;
}
void ssl_sess_cache_stat(struct sess_cache * cache, long long *size, long long *n_query, long long* n_hit)
{
	*size=MESA_htable_get_elem_num(cache->hash);
	*n_hit=cache->hit_cnt;
	*n_query=cache->hit_cnt+cache->miss_cnt;
	return;
}