tango-tfe/platform/src/key_keeper.cpp

#include "key_keeper.h"
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include "MESA/MESA_prof_load.h"
#include "tfe_rpc.h"
#include "event2/http.h"
#include "cjson/cJSON.h"
#define HTABLE_MAX_KEY_LEN 256
#define KEYRING_EXSITED 0
#define KEYRING_NOT_EXSITED -1

struct key_keeper
{
    char mode[TFE_STRING_MAX];
	char ca_path[TFE_STRING_MAX];
	char untrusted_ca_path[TFE_STRING_MAX];
	char cert_store_host[TFE_STRING_MAX];
	unsigned int cert_store_port;
	unsigned int hash_slot_size;
	unsigned int hash_expire_seconds;
    MESA_htable_handle htable;
    void* logger;
};

enum KEY_KEEPER_MODE{
	NORMAL = 0,
	DEBUG,
};

struct keyring_private
{
	struct keyring head;
	pthread_mutex_t mutex;
	size_t references;
};

struct key_keeper_promise_ctx
{
	void* logger;
	MESA_htable_handle htable;
	uchar* key;
	struct future* f_certstore_rpc;
	unsigned int key_len;
};

static void key_keeper_promise_free_ctx(void* ctx)
{
	struct key_keeper_promise_ctx* _ctx = (struct key_keeper_promise_ctx*)ctx;
	free(_ctx->key);
	_ctx->key = NULL;
	free(_ctx);
}

static struct keyring_private* keyring_new(void)
{
	struct keyring_private *kyr;
	if (!(kyr = (struct keyring_private *)ALLOC(struct keyring_private, 1)))
	{
		return NULL;
	}
	if (pthread_mutex_init(&kyr->mutex, NULL)) {
		free(kyr);
		return NULL;
	}
	kyr->references = 1;
	return kyr;
}

/*
 * Passed OpenSSL objects are owned by cert_t; refcount will not be
 * incremented, stack will not be duplicated.
 */
static struct keyring* keyring_new3(EVP_PKEY *key, X509 *cert, STACK_OF(X509) *chain)
{
	struct keyring_private* kyr=NULL;
	kyr=keyring_new();

	kyr->references = 1;
	(kyr->head).key = key;
	(kyr->head).cert = cert;
	(kyr->head).chain = chain;
	if(key)
	{
		ssl_key_refcount_inc(key);
	}
	if(cert)
	{
		ssl_x509_refcount_inc(cert);
	}
	if(chain){
        int i = 0;
		for (i = 0; i < sk_X509_num(chain); i++)
		{
			ssl_x509_refcount_inc(sk_X509_value(chain, i));
		}
	}
	return &(kyr->head);
}

// Increment reference count.
static void keyring_ref_inc(struct keyring_private* kyr)
{
	pthread_mutex_lock(&kyr->mutex);
	kyr->references++;
	pthread_mutex_unlock(&kyr->mutex);
}

/*
 * Thread-safe setter functions; they copy the value (refcounts are inc'd).
 */
static void keyring_set_key(struct keyring_private* kyr, EVP_PKEY *key)
{
	pthread_mutex_lock(&kyr->mutex);
	if ((kyr->head).key)
	{
		EVP_PKEY_free((kyr->head).key);
	}
	(kyr->head).key = key;
	if (key)
	{
		ssl_key_refcount_inc((kyr->head).key);
	}
	pthread_mutex_unlock(&kyr->mutex);
}

static void keyring_set_cert(struct keyring_private* kry, X509 *cert)
{
	pthread_mutex_lock(&kry->mutex);
	if ((kry->head).cert)
	{
		X509_free((kry->head).cert);
	}
	(kry->head).cert = cert;
	if (cert)
	{
		ssl_x509_refcount_inc((kry->head).cert);
	}
	pthread_mutex_unlock(&kry->mutex);
}

static void keyring_set_chain(struct keyring_private* kyr, STACK_OF(X509) *chain)
{
	pthread_mutex_lock(&kyr->mutex);
	if ((kyr->head).chain)
	{
		sk_X509_pop_free((kyr->head).chain, X509_free);
	}
	if (chain)
	{
		(kyr->head).chain = sk_X509_dup(chain);
		int i = 0;
        for (i = 0; i < sk_X509_num((kyr->head).chain); i++)
		{
			ssl_x509_refcount_inc(sk_X509_value((kyr->head).chain, i));
		}
	} else
	{
		(kyr->head).chain = NULL;
	}
	pthread_mutex_unlock(&kyr->mutex);
}

/*
 * Free keyring including internal objects.
 */
void key_keeper_free_keyring(struct keyring *kyr)
{
	struct keyring_private* _kyr = (struct keyring_private*)kyr;
	pthread_mutex_lock(&_kyr->mutex);
	_kyr->references--;
	if (_kyr->references>0)
	{
		pthread_mutex_unlock(&_kyr->mutex);
		return;
	}
	pthread_mutex_unlock(&_kyr->mutex);
	pthread_mutex_destroy(&_kyr->mutex);
	if ((_kyr->head).key)
	{
		EVP_PKEY_free((_kyr->head).key);
		(_kyr->head).key=NULL;
	}
	if ((_kyr->head).cert)
	{
		X509_free((_kyr->head).cert);
		(_kyr->head).cert=NULL;
	}
	if ((_kyr->head).chain)
	{
		sk_X509_pop_free((_kyr->head).chain, X509_free);
		(_kyr->head).chain=NULL;
	}
	free(_kyr);
}


static X509* transform_cert_to_x509(const char* str)
{
	BIO *bio;
	X509 *cert;
	bio = BIO_new(BIO_s_mem());
	BIO_write(bio, (const void*)str, strnlen(str, TFE_STRING_MAX));
	cert = PEM_read_bio_X509(bio, NULL, NULL, NULL);
	BIO_free_all(bio);
	return cert;
}

static EVP_PKEY* transform_key_to_EVP(const char* str)
{
	BIO *mem;
	mem = BIO_new_mem_buf(str, -1);
	EVP_PKEY* key = PEM_read_bio_PrivateKey(mem, NULL, NULL, 0);
	BIO_free(mem);
	return key;
}


static void err_out(X509* cert, EVP_PKEY* key, STACK_OF(X509)* chain)
{
	if(cert)
	{
		X509_free(cert);
	}
	if(key)
	{
		EVP_PKEY_free(key);
	}
	if(chain)
	{
		sk_X509_pop_free(chain, X509_free);
	}
	return;
}

static struct keyring_private* get_keyring_from_response(const char* data)
{
	X509* cert = NULL;
	EVP_PKEY* key = NULL;
	STACK_OF(X509)* chain = NULL;
	if(data == NULL)
	{
		return NULL;
	}
	cJSON* data_json = cJSON_Parse(data);
	if(data_json == NULL)
	{
		return NULL;
	}
	cJSON* cert_json = NULL;
	cJSON* key_json = NULL;
	cJSON* chain_json = NULL;
	cert_json = cJSON_GetObjectItemCaseSensitive(data_json, "CERTIFICATE");
	key_json = cJSON_GetObjectItemCaseSensitive(data_json, "PRIVATE_KEY");
	chain_json = cJSON_GetObjectItemCaseSensitive(data_json, "CERTIFICATE_CHAIN");
	if (cert_json && cert_json->valuestring != NULL)
	{
		cert = transform_cert_to_x509(cert_json->valuestring);
	}
	if(cert == NULL)
	{
		err_out(cert, key, chain);
		return NULL;
	}
	if (key_json && key_json->valuestring != NULL)
	{
		key = transform_key_to_EVP(key_json->valuestring);
	}
	if(key == NULL)
	{
		err_out(cert, key, chain);
		return NULL;
	}
	if(chain_json == NULL)
	{
		err_out(cert, key, chain);
		return NULL;
	}
	cJSON* chain_cert_json = NULL;
	chain = sk_X509_new_null();
	cJSON_ArrayForEach(chain_cert_json, chain_json)
	{
		X509* chain_cert = NULL;
		if (chain_cert_json && chain_cert_json->valuestring != NULL)
		{
			chain_cert = transform_cert_to_x509(chain_cert_json->valuestring);
		}
		if(chain_cert == NULL)
		{
			err_out(cert, key, chain);
			return NULL;
		}
		sk_X509_push(chain, chain_cert);
	}
	struct keyring_private* _kyr= keyring_new();
	keyring_set_cert(_kyr, cert);
	keyring_set_key(_kyr, key);
	keyring_set_chain(_kyr, chain);
	X509_free(cert);
	EVP_PKEY_free(key);
	sk_X509_pop_free(chain, X509_free);
	return _kyr;
}

static long keyring_local_cache_query_cb(void * data, const uchar * key, uint size, void * user_arg)
{
	struct keyring_private* kyr=(struct keyring_private*)data;
	if(kyr == NULL)
	{
		return KEYRING_NOT_EXSITED;
	}
	else
	{
		struct promise* p = (struct promise*)user_arg;
		promise_success(p, data);
		return KEYRING_EXSITED;
	}
}


static struct keyring_private* generate_x509_keyring(X509* origin_cert, int keyring_id, const char* filename)
{
	X509* ca = ssl_x509_load(filename);
	EVP_PKEY* cakey = ssl_key_load(filename);
	EVP_PKEY* forge_key = ssl_key_genrsa(1024);
	X509* forge_cert = ssl_x509_forge(ca, cakey, origin_cert, forge_key, NULL, NULL);
	STACK_OF(X509)* chain = sk_X509_new_null();
	sk_X509_push(chain, ca);
	sk_X509_push(chain, forge_cert);
	struct keyring_private* _kyr= keyring_new();
	keyring_set_key(_kyr, forge_key);
	keyring_set_cert(_kyr, forge_cert);
	keyring_set_chain(_kyr, chain);
	X509_free(ca);
	EVP_PKEY_free(cakey);
	X509_free(forge_cert);
	EVP_PKEY_free(forge_key);
	sk_X509_pop_free(chain, X509_free);
	return _kyr;
}

static void certstore_rpc_on_succ(void* result, void* user)
{
	//printf("call certstore_rpc_on_succ\n");
	struct promise * p = (struct promise *) user;
	struct key_keeper_promise_ctx* ctx = (struct key_keeper_promise_ctx*)promise_get_ctx(p);
	MESA_htable_handle htable= ctx->htable;
	const uchar* key = ctx->key;
	unsigned int key_len = ctx->key_len;
	struct tfe_rpc_response_result* response = tfe_rpc_release(result);
	int status_code = response->status_code;
	const char* status_msg = response->status_msg;
	char* data = response->data;
	size_t len = response->len;
	if(status_code == HTTP_OK)
	{
		*(data+len) = '\0';
		struct keyring_private* kyr= get_keyring_from_response(data);
		keyring_ref_inc(kyr);
		int ret = MESA_htable_add(htable, key, key_len, (void*)kyr);
		if(ret<0)
		{
			key_keeper_free_keyring((struct keyring*)kyr);
		}
		promise_success(p, (void*)kyr);
		key_keeper_free_keyring((struct keyring*)kyr);
	}
	else
	{
		promise_failed(p, FUTURE_ERROR_EXCEPTION, status_msg);
	}
	future_destroy(ctx->f_certstore_rpc);
	//promise_dettach_ctx(p);
	//ctx_destroy_cb((void*)ctx);
}

static void certstore_rpc_on_fail(enum e_future_error err, const char * what, void * user)
{
	struct promise * p = (struct promise *) user;
	promise_failed(p, err, what);
	struct key_keeper_promise_ctx* ctx= (struct key_keeper_promise_ctx*)promise_get_ctx(p);
	future_destroy(ctx->f_certstore_rpc);
	//promise_dettach_ctx(p);
	//ctx_destroy_cb((void*)ctx);
}

/*
 * Certificate, including private key and keyring chain.
 */
static int __wrapper_MESA_htable_set_opt(MESA_htable_handle table, enum MESA_htable_opt opt_type, unsigned int value)
{
    int ret = MESA_htable_set_opt(table, opt_type, &value, (int)(sizeof(value)));
    assert(ret == 0);
    return ret;
}

static MESA_htable_handle create_hash_table(unsigned int slot_size, unsigned int expire_seconds)
{
	int ret = 0;
	unsigned max_num = slot_size * 4;
	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);
    //TODO: how to do overide?
	//ret = __wrapper_MESA_htable_set_opt(htable, MHO_CBFUN_DATA_FREE,
	//	(void *)key_keeper_free_serialized);
	//ret = __wrapper_MESA_htable_set_opt(htable, MHO_CBFUN_DATA_EXPIRE_NOTIFY,
	//	(void *)key_keeper_verify_cb);
	ret = MESA_htable_mature(htable);
	assert(ret == 0);
	return htable;
}

struct key_keeper* key_keeper_init(const char * profile, const char* section, void* logger)
{
	struct key_keeper* keeper = ALLOC(struct key_keeper, 1);
	keeper->logger = logger;
	MESA_load_profile_string_def(profile, section, "mode", keeper->mode, sizeof(keeper->mode), "debug");
	MESA_load_profile_string_def(profile, section, "ca_path", keeper->ca_path, sizeof(keeper->ca_path), "./conf/mesalab-ca.pem");
	MESA_load_profile_string_def(profile, section, "untrusted_ca_path", keeper->untrusted_ca_path, sizeof(keeper->untrusted_ca_path), "./conf/mesalab-ca.pem");
	MESA_load_profile_string_def(profile, section, "cert_store_host", keeper->cert_store_host, sizeof(keeper->cert_store_host), "xxxxx");
	MESA_load_profile_uint_def(profile, section, "cert_store_port", &(keeper->cert_store_port), 80);
	MESA_load_profile_uint_def(profile, section, "hash_slot_size", &(keeper->hash_slot_size), 16);
	MESA_load_profile_uint_def(profile, section, "hash_expire_seconds", &(keeper->hash_expire_seconds), 16);
	keeper->htable = create_hash_table(keeper->hash_slot_size, keeper->hash_expire_seconds);
	return keeper;
}

void key_keeper_destroy(struct key_keeper *keeper)
{
	MESA_htable_destroy(keeper->htable, NULL);
	free(keeper);
	keeper = NULL;
	return;
}

struct keyring* key_keeper_release_keyring(future_result_t* result)
{
	struct keyring_private* kyr=(struct keyring_private*)result;
	keyring_ref_inc(kyr);
	return &(kyr->head);
}


static uchar* get_key_by_cert(X509* cert, int keyring_id, unsigned int* len)
{
	char* cert_fgr = NULL;
	cert_fgr = ssl_x509_fingerprint(cert, 0);
	char* key = (char*)malloc(HTABLE_MAX_KEY_LEN);
	*key = '\0';
	if(cert == NULL)
	{
		return NULL;
	}
	snprintf(key, HTABLE_MAX_KEY_LEN, "%d:", keyring_id);
	strncat(key, cert_fgr, HTABLE_MAX_KEY_LEN);
	*len = strnlen(key, HTABLE_MAX_KEY_LEN);
	free(cert_fgr);
	return (uchar*)key;
}


void key_keeper_async_ask(struct future * f, struct key_keeper * keeper, const char* sni, int keyring_id, X509 * origin_cert, int is_cert_valid, struct event_base * evbase)
{
	struct promise* p = future_to_promise(f);
	struct key_keeper_promise_ctx* ctx = ALLOC(struct key_keeper_promise_ctx, 1);
	unsigned int len = 0;
	uchar* key = get_key_by_cert(origin_cert, keyring_id, &len);
	ctx->logger = keeper->logger;
	ctx->htable = keeper->htable;
	ctx->key = key;
	ctx->key_len = len;
	promise_set_ctx(p, (void*)ctx, key_keeper_promise_free_ctx);
	long int cb_rtn = 0;
	MESA_htable_search_cb(ctx->htable, (const unsigned char*)(ctx->key), ctx->key_len, keyring_local_cache_query_cb, p, &cb_rtn);
	if(cb_rtn == KEYRING_EXSITED)
	{
		printf("KEYRING_EXSITED\n");
		return;
	}
	int mode = 0;
	if(strncmp(keeper->mode, "debug", TFE_STRING_MAX) == 0)
	{
		mode = 1;
	}
	switch(mode){
		case NORMAL:
        {
    		struct future* f_certstore_rpc = future_create("tfe_rpc", certstore_rpc_on_succ, certstore_rpc_on_fail, p);
			ctx->f_certstore_rpc = f_certstore_rpc;
			char url[TFE_STRING_MAX];
			char _sni[TFE_STRING_MAX] = "www.baidu.com";
			if(sni)
			{
				strncpy(_sni, sni, TFE_STRING_MAX);
			}
			snprintf(url, TFE_STRING_MAX, "http://%s:%d/ca?host=%s&flag=1&valid=1&keyring_id=%d", keeper->cert_store_host, keeper->cert_store_port, _sni, keyring_id);
			printf("url is %s\n", url);
    		tfe_rpc_async_ask(f_certstore_rpc, url, GET, DONE_CB, NULL, 0, evbase);
			break;
        }
		case DEBUG:
        {
			//TOOD: generate X509 cert
			char* filename = NULL;
			if(is_cert_valid == 1)
			{
				filename = keeper->ca_path;
			}
			else
			{
				filename = keeper->untrusted_ca_path;
			}
			struct keyring_private* kyr = generate_x509_keyring(origin_cert, keyring_id, filename);
			if(kyr)
			{
				keyring_ref_inc(kyr);
				int ret = MESA_htable_add(ctx->htable, ctx->key, ctx->key_len, (void*)kyr);
				if(ret < 0)
				{
					key_keeper_free_keyring((struct keyring*)kyr);
				}
				promise_success(p, (void*)kyr);
				key_keeper_free_keyring((struct keyring*)kyr);
			}
			else
			{
				promise_failed(p, FUTURE_ERROR_EXCEPTION, "generate X509 cert failed");
			}
			break;
        }
	}
	return;
}