#include "logger.h"
#include <tfe_scan.h>
#include <MESA/stream.h>

extern void increase_redirect_policy_hit_num(void);

#define DOH_CTX_MAGIC_NUM 20200601

#define REQ_METHOD_IS_GET(method) ((method == TFE_HTTP_METHOD_GET) ? 1 : 0)
#define REQ_METHOD_IS_POST(method) ((method == TFE_HTTP_METHOD_POST) ? 1 : 0)

struct dns_str2idx
{
	int index;
	int len;
	char *type;
};

struct dns_str2idx str2index[] = {
	{DNS_TYPE_CNAME, 5, (char *)"CNAME"},
	{DNS_TYPE_MX, 2, (char *)"MX"},
	{DNS_TYPE_A, 1, (char *)"A"},
	{DNS_TYPE_NS, 2, (char *)"NS"},
	{DNS_TYPE_AAAA, 4, (char *)"AAAA"},
	{DNS_TYPE_TXT, 3, (char *)"TXT"},
	{DNS_TYPE_PTR, 3, (char *)"PTR"}};

static struct doh_conf *g_doh_conf = NULL;

static int doh_type2index(char *type)
{
	int i = 0;
	for (i = 0; i < (int)(sizeof(str2index) / sizeof(struct dns_str2idx)); i++)
	{
		if (str2index[i].len == (int)strlen(type) && (strncasecmp(str2index[i].type, type, str2index[i].len)) == 0)
		{
			return str2index[i].index;
		}
	}

	return -1;
}

static void doh_addr_tfe2sapp(const struct tfe_stream_addr *tfe_addr, struct ipaddr *sapp_addr)
{
	if (tfe_addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V4 || tfe_addr->addrtype == TFE_ADDR_IPV4)
	{
		sapp_addr->addrtype = ADDR_TYPE_IPV4;
	}
	else
	{
		sapp_addr->addrtype = ADDR_TYPE_IPV6;
	}
	sapp_addr->paddr = (char *)tfe_addr->paddr;
}

static int doh_get_answer_ttl(cJSON *object)
{
	int min = 0;
	int max = 0;
	cJSON *item = NULL;

	item = cJSON_GetObjectItem(object, "min");
	if (item)
	{
		min = item->valueint;
	}

	item = cJSON_GetObjectItem(object, "max");
	if (item)
	{
		max = item->valueint;
	}

	return (rand() % (max - min + 1) + min);
}

static cJSON *doh_get_answer_records(struct doh_ctx *ctx, cJSON *object, int qtype)
{
	int i = 0;

	cJSON *vsys_id=cJSON_GetObjectItem(object, "vsys_id");
	if(vsys_id && vsys_id->type==cJSON_Number)
	{
		ctx->vsys_id = vsys_id->valueint;
	}
	cJSON *resolution = cJSON_GetObjectItem(object, "resolution");
	int size = cJSON_GetArraySize(resolution);
	for (i = 0; i < size; i++)
	{
		cJSON *item = cJSON_GetArrayItem(resolution, i);
		cJSON *tmp = cJSON_GetObjectItem(item, "qtype");

		if (doh_type2index(tmp->valuestring) == qtype)
		{
			return cJSON_GetObjectItem(item, "answer");
		}
	}

	return NULL;
}

struct doh_action_param
{
	int ref_cnt;
	int action;
	char *message;
	char *position;
	float enforcement_ratio;
	int profile_id;
	int status_code;
	size_t n_rule;
	void *repl_rule;
	size_t e_rule;
	void *elem_rule;
	struct doh_maat_rule_t hit_rule;
	pthread_mutex_t lock;
};

void doh_action_param_free_cb(int table_id, void **ad, long argl, void *argp)
{
	if(*ad==NULL)
	{
		return;
	}
	struct doh_action_param* param=(struct doh_action_param*)*ad;
	pthread_mutex_lock(&(param->lock));
	param->ref_cnt--;
	if(param->ref_cnt>0)
	{
		pthread_mutex_unlock(&(param->lock));
		return;
	}
	pthread_mutex_unlock(&(param->lock));
	pthread_mutex_destroy(&(param->lock));

	if(param->hit_rule.srv_def_large)
		FREE(&(param->hit_rule.srv_def_large))

	FREE(&(param));
	return;
}

static void doh_get_cheat_data(long long p_result, int qtype, struct doh_ctx *ctx, const char *str_stream_info)
{
	int i;
	int answer_size = 0;
	cJSON *items = NULL;
	cJSON *item = NULL;
	cJSON *object = NULL;
	cJSON *answer_array = NULL;
	int table_id=0;

	table_id=maat_get_table_id(g_doh_conf->maat, "PXY_CTRL_COMPILE");
	if(table_id < 0)
	{
		return;
	}
	struct doh_action_param *get_ex_param=(struct doh_action_param *)maat_plugin_table_get_ex_data(g_doh_conf->maat, table_id, (const char *)&p_result);
	if(get_ex_param==NULL)
	{
		return;
	}

	struct doh_maat_rule_t *hit_rule = &(get_ex_param->hit_rule);
	if(hit_rule==NULL || hit_rule->srv_def_large==NULL)
	{
		goto end;
	}
	memcpy(ctx->result, hit_rule, sizeof(struct doh_maat_rule_t));
	TFE_LOG_INFO(g_doh_conf->local_logger, "%s hit %lld %s", str_stream_info, p_result, hit_rule->srv_def_large);

	object = cJSON_Parse(hit_rule->srv_def_large);
	if (object == NULL)
	{
		goto end;
	}

	answer_array = doh_get_answer_records(ctx, object, qtype);
	if (answer_array != NULL)
	{
		answer_size = cJSON_GetArraySize(answer_array);
		ctx->opts = ALLOC(cheat_pkt_opt_t, answer_size);
		ctx->opts_num = answer_size;

		for (i = 0; i < answer_size; i++)
		{
			items = cJSON_GetArrayItem(answer_array, i);

			item = cJSON_GetObjectItem(items, "atype");
			ctx->opts[i].res_type = doh_type2index(item->valuestring);

			item = cJSON_GetObjectItem(items, "ttl");
			ctx->opts[i].ttl = doh_get_answer_ttl(item);

			if (ctx->min_ttl == 0)
			{
				ctx->min_ttl = ctx->opts[i].ttl;
			}
			else
			{
				if (ctx->min_ttl > ctx->opts[i].ttl)
				{
					ctx->min_ttl = ctx->opts[i].ttl;
				}
			}

			item = cJSON_GetObjectItem(items, "value");
			if (item)
			{
				switch (ctx->opts[i].res_type)
				{
				case DNS_TYPE_A:
					ctx->opts[i].res_len = sizeof(unsigned int);
					inet_pton(AF_INET, item->valuestring, (void *)ctx->opts[i].res_info);
					break;
				case DNS_TYPE_AAAA:
					ctx->opts[i].res_len = IPV6_ADDR_LEN;
					inet_pton(AF_INET6, item->valuestring, (void *)ctx->opts[i].res_info);
					break;
				default:
					ctx->opts[i].res_len = (strlen(item->valuestring) > sizeof(ctx->opts[i].res_info) - 1) ? sizeof(ctx->opts[i].res_info) - 1 : strlen(item->valuestring);
					memcpy(ctx->opts[i].res_info, item->valuestring, ctx->opts[i].res_len);
					break;
				}
			}
		}
	}

end:
	if (object)
	{
		cJSON_Delete(object);
		object = NULL;
	}
	if(get_ex_param)
	{
		doh_action_param_free_cb(0, (void**)&get_ex_param, 0, NULL);
	}
}

static long long doh_fetch_rule(long long *result, int result_num)
{
	int i = 0;
	long long p_result = 0;

	for (i = 0; i < result_num && i < MAX_SCAN_RESULT; i++)
	{
		if (p_result == 0)
		{
			p_result = result[i];
			continue;
		}

		if (result[i] > p_result)
		{
			p_result = result[i];
		}
	}

	return p_result;
}

static void doh_maat_scan(const struct tfe_stream *stream, const struct tfe_http_session *session, struct doh_ctx *ctx, char *qname, int qtype)
{
	int hit_cnt = 0;
	int scan_ret = 0;
	int app_id = 8006;
	size_t n_hit_result;
	struct ipaddr sapp_addr;
	long long p_result = 0;
	long long result[MAX_SCAN_RESULT];

	scan_ret = tfe_scan_subscribe_id(stream, result, ctx->scan_mid, hit_cnt, g_doh_conf->local_logger);
	if (scan_ret > 0)
	{
		hit_cnt += scan_ret;
	}

	scan_ret = tfe_scan_ip_location(stream, result, ctx->scan_mid, hit_cnt, g_doh_conf->local_logger, &(ctx->location_server), &(ctx->location_client));
	if (scan_ret > 0)
	{
		hit_cnt += scan_ret;
	}

	scan_ret = tfe_scan_ip_asn(stream, result, ctx->scan_mid, hit_cnt, g_doh_conf->local_logger, &(ctx->asn_server), &(ctx->asn_client));
	if (scan_ret > 0)
	{
		hit_cnt += scan_ret;
	}

	// scan server host
	const char *host = session->req->req_spec.host;
	if (host)
	{
 		scan_ret = maat_scan_string(g_doh_conf->maat, g_doh_conf->tables[TYPE_HOST].id,host, strlen(host),
									result + hit_cnt, MAX_SCAN_RESULT - hit_cnt, &n_hit_result, ctx->scan_mid);
		if (scan_ret == MAAT_SCAN_HIT)
		{
            TFE_LOG_INFO(g_doh_conf->local_logger, "Scan %s, Hit host: %s scan ret: %d policy_id: %lld addr: %s",
                         g_doh_conf->tables[TYPE_HOST].name, host, scan_ret, result[hit_cnt], stream->str_stream_info);
            hit_cnt += n_hit_result;
		}
		else
		{
            TFE_LOG_INFO(g_doh_conf->local_logger, "Scan %s, NO hit host: %s scan ret: %d addr: %s",
                         g_doh_conf->tables[TYPE_HOST].name, host, scan_ret, stream->str_stream_info);
        }

        scan_ret = tfe_scan_fqdn_cat(stream, result, ctx->scan_mid, hit_cnt, g_doh_conf->local_logger, g_doh_conf->tables[TYPE_HOST_CAT].id);
        if( scan_ret > 0)
		{
			hit_cnt += scan_ret;
		}
	}

	// scan addr
	doh_addr_tfe2sapp(stream->addr, &sapp_addr);
	if (sapp_addr.addrtype == ADDR_TYPE_IPV4)
	{
		scan_ret = maat_scan_ipv4(g_doh_conf->maat, g_doh_conf->tables[TYPE_SRC_ADDR].id,sapp_addr.v4->saddr,
								  sapp_addr.v4->source, 6, result+hit_cnt, MAX_SCAN_RESULT-hit_cnt,
								  &n_hit_result, ctx->scan_mid);
		if (n_hit_result == MAAT_SCAN_HIT)
		{
			hit_cnt += n_hit_result;
		}
		scan_ret = maat_scan_ipv4(g_doh_conf->maat, g_doh_conf->tables[TYPE_DST_ADDR].id,sapp_addr.v4->daddr,
								  sapp_addr.v4->dest, 6, result+hit_cnt, MAX_SCAN_RESULT-hit_cnt,
								  &n_hit_result, ctx->scan_mid);

		if(scan_ret == MAAT_SCAN_HIT)
		{
			hit_cnt += n_hit_result;
		}
	}
	if (sapp_addr.addrtype == ADDR_TYPE_IPV6)
	{
		scan_ret = maat_scan_ipv6(g_doh_conf->maat, g_doh_conf->tables[TYPE_SRC_ADDR].id, sapp_addr.v6->saddr,
								 sapp_addr.v6->source, 6, result+hit_cnt, MAX_SCAN_RESULT-hit_cnt,
								 &n_hit_result, ctx->scan_mid);
		if (scan_ret == MAAT_SCAN_HIT)
		{
			hit_cnt += n_hit_result;
		}
		scan_ret = maat_scan_ipv6(g_doh_conf->maat,g_doh_conf->tables[TYPE_DST_ADDR].id, sapp_addr.v6->daddr,
								 sapp_addr.v6->dest, 6, result+hit_cnt, MAX_SCAN_RESULT-hit_cnt,
								 &n_hit_result, ctx->scan_mid);
		if (scan_ret == MAAT_SCAN_HIT)
		{
			hit_cnt += n_hit_result;
		}
	}
	// scan appid
	scan_ret=maat_scan_integer(g_doh_conf->maat, g_doh_conf->tables[TYPE_APPID].id, app_id, result + hit_cnt, MAX_SCAN_RESULT - hit_cnt, &n_hit_result, ctx->scan_mid);
	if (scan_ret == MAAT_SCAN_HIT)
	{
        TFE_LOG_INFO(g_doh_conf->local_logger, "Scan %s, Hit proto: %d scan ret: %d policy_id: %lld addr: %s",
                     g_doh_conf->tables[TYPE_APPID].name, app_id, scan_ret, result[hit_cnt], stream->str_stream_info);
        hit_cnt += n_hit_result;
	}
	else
	{
        TFE_LOG_INFO(g_doh_conf->local_logger, "Scan %s, NO hit proto: %d scan ret: %d addr: %s",
                     g_doh_conf->tables[TYPE_APPID].name, app_id, scan_ret, stream->str_stream_info);
    }

	// scan qname
	scan_ret = maat_scan_string(g_doh_conf->maat, g_doh_conf->tables[TYPE_QNAME].id, qname, strlen(qname),
							    result + hit_cnt, MAX_SCAN_RESULT - hit_cnt, &n_hit_result, ctx->scan_mid);
	if (scan_ret == MAAT_SCAN_HIT)
	{
        TFE_LOG_INFO(g_doh_conf->local_logger, "Scan %s, Hit domain: %s scan ret: %d qtype: %d policy_id: %lld addr: %s",
                     g_doh_conf->tables[TYPE_QNAME].name, qname, scan_ret, qtype, result[hit_cnt], stream->str_stream_info);
        hit_cnt += n_hit_result;
	}
	else
	{
        TFE_LOG_INFO(g_doh_conf->local_logger, "Scan %s, NO hit domain: %s scan ret: %d addr: %s",
                     g_doh_conf->tables[TYPE_QNAME].name, qname, scan_ret, stream->str_stream_info);
    }

	if (hit_cnt)
	{
		p_result = doh_fetch_rule(result, hit_cnt);
		if (p_result != 0)
		{
			ctx->result_num = 1;
			ctx->result = ALLOC(struct doh_maat_rule_t, ctx->result_num);
			doh_get_cheat_data(p_result, qtype, ctx, stream->str_stream_info);
		}
	}
}

static int doh_maat_init(const char *profile, const char *section)
{
	g_doh_conf->maat = (struct maat *)tfe_bussiness_resouce_get(STATIC_MAAT);
	MESA_load_profile_string_def(profile, section, "table_appid", g_doh_conf->tables[TYPE_APPID].name, TFE_STRING_MAX, "TSG_OBJ_APP_ID");
	MESA_load_profile_string_def(profile, section, "table_src_addr", g_doh_conf->tables[TYPE_SRC_ADDR].name, TFE_STRING_MAX, "TSG_SECURITY_SOURCE_ADDR");
	MESA_load_profile_string_def(profile, section, "table_dst_addr", g_doh_conf->tables[TYPE_DST_ADDR].name, TFE_STRING_MAX, "TSG_SECURITY_DESTINATION_ADDR");
	MESA_load_profile_string_def(profile, section, "table_qname", g_doh_conf->tables[TYPE_QNAME].name, TFE_STRING_MAX, "TSG_FIELD_DOH_QNAME");
	MESA_load_profile_string_def(profile, section, "table_host", g_doh_conf->tables[TYPE_HOST].name, TFE_STRING_MAX, "TSG_FIELD_DOH_HOST");
	MESA_load_profile_string_def(profile, section, "table_host_cat", g_doh_conf->tables[TYPE_HOST_CAT].name, TFE_STRING_MAX, "TSG_FIELD_DOH_HOST_CAT");

	for (int i = 0; i < TYPE_MAX; i++)
	{
		g_doh_conf->tables[i].id = maat_get_table_id(g_doh_conf->maat, g_doh_conf->tables[i].name);
		if (g_doh_conf->tables[i].id < 0)
		{
			TFE_LOG_ERROR(g_doh_conf->local_logger, "maat_get_table_id failed, table_name: %s", g_doh_conf->tables[i].name);
			return -1;
		}
	}

	return 0;
}

static void doh_gc_cb(evutil_socket_t fd, short what, void *arg)
{
	int i = 0;

	for (i = 0; i < DOH_STAT_MAX; i++)
	{
		FS_operate(g_doh_conf->fs_handle, g_doh_conf->fs_id[i], 0, FS_OP_SET, ATOMIC_READ(&(g_doh_conf->stat_val[i])));
	}
}

static int doh_field_init()
{
	int i = 0;
	struct timeval gc_delay = {0, 500 * 1000}; //Microseconds, we set 500 miliseconds here.
	const char *spec[DOH_STAT_MAX] = {0};

	spec[STAT_SESSION] = "doh_sess";
	spec[STAT_LOG_NUM] = "doh_log";
	spec[STAT_ACTION_HIJACK] = "doh_hijack";

	for (i = 0; i < DOH_STAT_MAX; i++)
	{
		if (spec[i] != NULL)
		{
			g_doh_conf->fs_id[i] = FS_register(g_doh_conf->fs_handle, FS_STYLE_FIELD, FS_CALC_CURRENT, spec[i]);
		}
	}
	g_doh_conf->gcev = event_new(g_doh_conf->gc_evbase, -1, EV_PERSIST, doh_gc_cb, NULL);
	evtimer_add(g_doh_conf->gcev, &gc_delay);

	return 0;
}

static struct doh_ctx *doh_ctx_new(unsigned int thread_id)
{
	struct doh_ctx *ctx = ALLOC(struct doh_ctx, 1);
	assert(ctx);

	ctx->magic_num = DOH_CTX_MAGIC_NUM;
	ctx->thread_id = (int)thread_id;
	ctx->scan_mid = maat_state_new(g_doh_conf->maat, thread_id);
	ctx->opts_num = 0;
	ctx->opts = NULL;
	ctx->min_ttl = 0;
	ctx->doh_req = NULL;

	return ctx;
}

static void doh_ctx_free(struct doh_ctx *ctx)
{
	assert(ctx->magic_num == DOH_CTX_MAGIC_NUM);

	if (ctx->result)
	{
		free(ctx->result);
		ctx->result = NULL;
	}
	if (ctx->doh_req)
	{
		dns_free(ctx->doh_req);
		ctx->doh_req = NULL;
	}

	if (ctx->opts_num)
	{
		free(ctx->opts);
		ctx->opts = NULL;
	}

	if(ctx->scan_mid)
	{
		maat_state_free(ctx->scan_mid);
		ctx->scan_mid = NULL;
	}

	if (ctx->http_req_body)
	{
		evbuffer_free(ctx->http_req_body);
		ctx->http_req_body = NULL;
	}

	if (ctx->asn_client)
	{
		free(ctx->asn_client);
		ctx->asn_client = NULL;
	}

	if (ctx->asn_server)
	{
		free(ctx->asn_server);
		ctx->asn_server = NULL;
	}

	if (ctx->location_client)
	{
		free(ctx->location_client);
		ctx->location_client = NULL;
	}

	if (ctx->location_server)
	{
		free(ctx->location_server);
		ctx->location_server = NULL;
	}
	FREE(&ctx);
}

static int req_session_is_doh(const struct tfe_http_session *session, struct doh_ctx *ctx)
{
	// https://tools.ietf.org/html/rfc8484

	// https://tools.ietf.org/id/draft-ietf-doh-dns-over-https-10.html
	// Registration of application/dns-message Media Type

	// https://tools.ietf.org/id/draft-ietf-doh-dns-over-https-05.html
	// Registration of application/dns-udpwireformat Media Type

	// POST
	const char *cont_type_val = tfe_http_std_field_read(session->req, TFE_HTTP_CONT_TYPE);
	if (cont_type_val && strncasecmp(cont_type_val, "application/dns-message", strlen("application/dns-message")) == 0)
	{
		ctx->type = DOH_TYPE_MESSAGE;
		return 1;
	}
	if (cont_type_val && strncasecmp(cont_type_val, "application/dns-udpwireformat", strlen("application/dns-udpwireformat")) == 0)
	{
		ctx->type = DOH_TYPE_UDPWIREFORMAT;
		return 1;
	}

	// GET
	const char *accept_type_val = tfe_http_nonstd_field_read(session->req, "Accept");
	if (accept_type_val && strncasecmp(accept_type_val, "application/dns-message", strlen("application/dns-message")) == 0)
	{
		ctx->type = DOH_TYPE_MESSAGE;
		return 1;
	}
	if (accept_type_val && strncasecmp(accept_type_val, "application/dns-udpwireformat", strlen("application/dns-udpwireformat")) == 0)
	{
		ctx->type = DOH_TYPE_UDPWIREFORMAT;
		return 1;
	}

	return 0;
}

static void doh_process_req(const struct tfe_stream *stream, const struct tfe_http_session *session, void **pme)
{
	struct doh_ctx *ctx = *(struct doh_ctx **)pme;
	struct tfe_http_half *response = NULL;
	struct tfe_http_session *to_write = NULL;

	int rsp_len;
	int rsp_size;
	char *rsp_buff = NULL;
	char cont_len_str[20] = {0};
	char max_age_str[20] = {0};

	char *req_data = (char *)evbuffer_pullup(ctx->http_req_body, -1);
	size_t req_len = evbuffer_get_length(ctx->http_req_body);

	if (req_data && req_len)
	{
		ctx->doh_req = dns_new();
		assert(ctx->doh_req);

		if (dns_parser(ctx->doh_req, req_data, req_len) == -1)
		{
			// base64_len = (str_len / 3 + 1) * 4
			int temp_size = (req_len / 3 + 1) * 4;
			char *temp = (char *)ALLOC(char, temp_size);
			int len = base64_encode(temp, temp_size - 1, req_data, req_len);
			TFE_LOG_ERROR(g_doh_conf->local_logger, "%s Doh parser request failed, PASSTHROUGH, data:%s", stream->str_stream_info, len > 0 ? temp : "");
			free(temp);
			goto end;
		}
		TFE_LOG_DEBUG(g_doh_conf->local_logger, "%s qtype %d qname:%s",
					  stream->str_stream_info, ctx->doh_req->query_question.qtype, ctx->doh_req->query_question.qname);

		if (ctx->doh_req->query_question.qtype != DNS_TYPE_A && ctx->doh_req->query_question.qtype != DNS_TYPE_AAAA)
		{
			TFE_LOG_INFO(g_doh_conf->local_logger, "%s Doh qtype not A/AAAA, PASSTHROUGH", stream->str_stream_info);
			goto end;
		}
		if (strlen((char *)ctx->doh_req->query_question.qname) == 0)
		{
			TFE_LOG_INFO(g_doh_conf->local_logger, "%s Doh qname is empty, PASSTHROUGH", stream->str_stream_info);
			goto end;
		}

		doh_maat_scan(stream, session, ctx, (char *)ctx->doh_req->query_question.qname, ctx->doh_req->query_question.qtype);
		maat_state_free(ctx->scan_mid);
		ctx->scan_mid = NULL;
		if (!ctx->opts_num)
		{
			TFE_LOG_INFO(g_doh_conf->local_logger, "%s Doh no hit answer type, PASSTHROUGH", stream->str_stream_info);
			goto end;
		}

		rsp_size = sizeof(cheat_pkt_opt_t) * ctx->opts_num * 16;
		rsp_buff = (char *)ALLOC(char, rsp_size);
		rsp_len = dns_cheat_response(ctx->doh_req, ctx->opts, ctx->opts_num, rsp_buff, rsp_size - 1);
		if (rsp_len < 0)
		{
			TFE_LOG_ERROR(g_doh_conf->local_logger, "%s Doh cheat response failed: %d, PASSTHROUGH", stream->str_stream_info, rsp_len);
			goto end;
		}

		to_write = tfe_http_session_allow_write(session);
		if (to_write == NULL)
		{
			assert(0);
		}

		ctx->manipulate = 1;
		ATOMIC_INC(&(g_doh_conf->stat_val[STAT_ACTION_HIJACK]));
        increase_redirect_policy_hit_num();

        snprintf(cont_len_str, sizeof(cont_len_str), "%d", rsp_len);
		response = tfe_http_session_response_create(to_write, 200);
		tfe_http_std_field_write(response, TFE_HTTP_CONT_LENGTH, cont_len_str);
		if (ctx->type == DOH_TYPE_MESSAGE)
			tfe_http_std_field_write(response, TFE_HTTP_CONT_TYPE, "application/dns-message");
		else
			tfe_http_std_field_write(response, TFE_HTTP_CONT_TYPE, "application/dns-udpwireformat");

		/* The assigned freshness lifetime of a DoH HTTP response MUST be less
		 * than or equal to the smallest TTL in the Answer section of the DNS
		 * response. */
		snprintf(max_age_str, sizeof(max_age_str), "max-age=%d", ctx->min_ttl);
		tfe_http_std_field_write(response, TFE_HTTP_CACHE_CONTROL, max_age_str);
		tfe_http_half_append_body(response, rsp_buff, rsp_len, 0);
		tfe_http_half_append_body(response, NULL, 0, 0);

		tfe_http_session_response_set(to_write, response);
		tfe_http_session_detach(session);
	end:
		if (rsp_buff)
		{
			free(rsp_buff);
			rsp_buff = NULL;
		}
	}
}

int doh_on_init(struct tfe_proxy *proxy)
{
	const char *profile = "./conf/doh/doh.conf";

	g_doh_conf = ALLOC(struct doh_conf, 1);
	assert(g_doh_conf);

	MESA_load_profile_int_def(profile, "doh", "enable", &(g_doh_conf->enable), 1);

	if (!g_doh_conf->enable)
	{
		TFE_LOG_INFO(NULL, "Doh disabled.");
		goto success;
	}
	TFE_LOG_INFO(NULL, "Doh enabled.");

	g_doh_conf->thread_num = tfe_proxy_get_work_thread_count();
    g_doh_conf->local_logger = MESA_create_runtime_log_handle("doh", RLOG_LV_DEBUG);

	g_doh_conf->gc_evbase = tfe_proxy_get_gc_evbase();
	g_doh_conf->fs_handle = tfe_proxy_get_fs_handle();
	if (doh_field_init() != 0)
	{
		TFE_LOG_ERROR(NULL, "Doh init field stat failed.");
		goto error;
	}

	if (doh_kafka_init(profile, g_doh_conf) != 0)
	{
		TFE_LOG_ERROR(NULL, "Doh init kafka failed.");
		goto error;
	}

	if (doh_maat_init(profile, "maat") != 0)
	{
		TFE_LOG_ERROR(NULL, "Doh init maat failed.");
		goto error;
	}
	TFE_LOG_INFO(g_doh_conf->local_logger, "Doh init success.");

success:
	return 0;

error:
	TFE_LOG_ERROR(NULL, "Doh init failed.");
	return -1;
}

void doh_on_begin(const struct tfe_stream *stream, const struct tfe_http_session *session, unsigned int thread_id, void **pme)
{
	struct doh_ctx *ctx = NULL;

	if (!g_doh_conf->enable)
	{
		return;
	}

	ctx = *(struct doh_ctx **)pme;
	assert(ctx == NULL);
	ctx = doh_ctx_new(thread_id);

	*pme = ctx;
}

// return : NO_CALL_NEXT_PLUGIN
// return : CALL_NEXT_PLUGIN
int doh_on_data(const struct tfe_stream *stream, const struct tfe_http_session *session, enum tfe_http_event events, const unsigned char *body_frag, size_t frag_size, unsigned int thread_id, void **pme)
{
	if (!g_doh_conf->enable)
	{
		return CALL_NEXT_PLUGIN;
	}

	struct doh_ctx *ctx = *(struct doh_ctx **)pme;
	assert(ctx);

	if (!session->req || !req_session_is_doh(session, ctx))
	{
		return CALL_NEXT_PLUGIN;
	}

	if (!ctx->count)
	{
		ctx->count = 1;
		ATOMIC_INC(&(g_doh_conf->stat_val[STAT_SESSION]));
		TFE_LOG_DEBUG(g_doh_conf->local_logger, "%s method:%s content-type:%s accept:%s url:%s",
					  stream->str_stream_info,
					  http_std_method_to_string(session->req->req_spec.method),
					  tfe_http_std_field_read(session->req, TFE_HTTP_CONT_TYPE),
					  tfe_http_nonstd_field_read(session->req, "Accept"),
					  session->req->req_spec.url);
	}

	// request get
	if (REQ_METHOD_IS_GET(session->req->req_spec.method) && (events & EV_HTTP_REQ_HDR))
	{
		/* https://tools.ietf.org/html/draft-ietf-doh-dns-over-https-04
		 *
		 * When using the GET method the URI path MUST contain a query parameter
         * name-value pair [QUERYPARAMETER] with the name of "ct" and a value
         * indicating the media-format used for the dns parameter.  The value
         * may either be an explicit media type (e.g. ct=application/dns-
         * udpwireformat&dns=...) or it may be empty.  An empty value indicates
         * the default application/dns-udpwireformat type (e.g. ct&dns=...).
		 *
		 * NOTE: evhttp_parse_query_str()
		 *       support "ct=x&dns=xxx"
		 *       support "ct=&dns=xxx"
		 *   not support "ct&dns=xxx"
		 * So we parser "dns=" by self.
		 */
		int len;
		u_char *temp = NULL;
		u_char *dns_data = NULL;
		const char *uri_begin = session->req->req_spec.uri;
		const char *uri_end = session->req->req_spec.uri + strlen(session->req->req_spec.uri) - 1;
		char *dns_end = NULL;
		char *dns_begin = (char *)strstr(uri_begin, "dns=");
		if (dns_begin == NULL)
		{
			goto error;
		}

		if ((dns_begin == uri_begin || *(dns_begin - 1) == '&' || *(dns_begin - 1) == '?') && dns_begin + 4 < uri_end)
		{
			dns_begin = dns_begin + 4;
			dns_end = strstr(dns_begin, "&");
			if (dns_end == NULL)
			{
				dns_end = (char *)uri_end;
			}
			else
			{
				dns_end -= 1;
			}

			if (dns_end <= dns_begin)
			{
				goto error;
			}

			dns_data = (u_char *)calloc(1, dns_end - dns_begin + 1 + 1);
			memcpy(dns_data, dns_begin, dns_end - dns_begin + 1);
		}

		temp = (u_char *)ALLOC(u_char, strlen((char *)dns_data) * 2);
		len = tfe_decode_base64url(temp, dns_data);
		if (len == 0)
		{
			TFE_LOG_ERROR(g_doh_conf->local_logger, "%s Doh base64 decode uri failed:%s, PASSTHROUGH", stream->str_stream_info, session->req->req_spec.uri);
			goto error;
		}

		assert(ctx->http_req_body == NULL);
		ctx->http_req_body = evbuffer_new();
		evbuffer_add(ctx->http_req_body, temp, len);
		doh_process_req(stream, session, pme);

	error:
		if (dns_data)
		{
			free(dns_data);
			dns_data = NULL;
		}

		if (temp)
		{
			free(temp);
			temp = NULL;
		}
	}

	// request post
	if (REQ_METHOD_IS_POST(session->req->req_spec.method))
	{
		if (events & EV_HTTP_REQ_BODY_BEGIN)
		{
			assert(ctx->http_req_body == NULL);
			ctx->http_req_body = evbuffer_new();
		}

		if (events & EV_HTTP_REQ_BODY_CONT)
		{
			evbuffer_add(ctx->http_req_body, body_frag, frag_size);
		}

		if (events & EV_HTTP_REQ_BODY_END)
		{
			doh_process_req(stream, session, pme);
		}
	}

	return NO_CALL_NEXT_PLUGIN;
}

void doh_on_end(const struct tfe_stream *stream, const struct tfe_http_session *session, unsigned int thread_id, void **pme)
{
	if (!g_doh_conf->enable)
	{
		return;
	}

	struct doh_ctx *ctx = *(struct doh_ctx **)pme;
	if (ctx->manipulate)
	{
		int ret = doh_send_log(g_doh_conf, session, stream, ctx);
		if (ret > 0)
		{
			ATOMIC_ADD(&(g_doh_conf->stat_val[STAT_LOG_NUM]), ret);
		}
	}

	doh_ctx_free(ctx);
	*pme = NULL;
}

struct tfe_plugin doh_spec = {
	.symbol = NULL,
	.type = TFE_PLUGIN_TYPE_BUSINESS,
	.on_init = doh_on_init,
	.on_deinit = NULL,
	.on_open = NULL,
	.on_data = NULL,
	.on_close = NULL,
	.on_session_begin = doh_on_begin,
	.on_session_data = doh_on_data,
	.on_session_end = doh_on_end};
TFE_PLUGIN_REGISTER(DOH, doh_spec)