// git version 5effe725

#include <stdio.h>
#include <arpa/inet.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>

#include "dns.h"

#define DNS_ERROR_LOG(msg) printf("%s, %s:%d", msg, __FILE__, __LINE__)

/////////////////////////////////////////////////////////////////////////////////
// compress/decompress API
/////////////////////////////////////////////////////////////////////////////////

static int dns_compress_rr_str(u_char *domain, int domain_len, u_char *result)
{
	int section_len = 0;
	int result_pos = 1;
	int domain_pos = 0;

	if (domain_len < 0 || domain_len > DNS_MAX_NAME + 1 || '.' == domain[0] || '.' == domain[domain_len - 1])
	{
		return -1;
	}

	while ((domain[domain_pos] != '\n') || (domain[domain_pos] != '\0'))
	{
		section_len = 0;
		while ((domain[domain_pos] != '.') && (domain[domain_pos] != '\n') && (domain[domain_pos] != '\0'))
		{
			result[result_pos] = domain[domain_pos];
			result_pos++;
			domain_pos++;
			section_len++;
		}
		result[result_pos - section_len - 1] = section_len;

		if ((domain[domain_pos] == '\n') || (domain[domain_pos] == '\0'))
			break;
		result_pos++;
		domain_pos++;
	}

	result[result_pos] = '\0';

	if (result_pos >= domain_len)
	{
		return result_pos + 1;
	}
	else
	{
		return result_pos;
	}

	return 0;
}

static int dns_name_compress(u_char *name, int name_len, char *payload)
{
	int used_len = 0;
	u_int16_t compress_len = 0;
	u_char compress_name[DNS_MAX_NAME + 1] = {0};

	if (name_len)
	{
		compress_len = dns_compress_rr_str(name, name_len, compress_name);
		if (compress_len <= 0)
		{
			DNS_ERROR_LOG("dns_compress_rr_str()");
			return -1;
		}
		NS_SETLEN(compress_name, compress_len, payload, used_len);
	}
	else
	{
		NS_SET8(0, payload, used_len);
	}

	return used_len;
}

static int dns_name_decompress(char *msg, u_char **ptr, u_char *buf, int buflen, char *end)
{
	u_char *p = NULL;
	int index = 0, len = 0;
	int np = 0, tot_len = 0;

	p = *ptr;
	*ptr = NULL;
	index = 0;
	np = 0;
	while (1)
	{
		if (0 == p[0])
		{
			break;
		}

		if (0x0c0 == (p[0] & 0x0c0))
		{
			if (p + 2 > (u_char *)end)
				return -1;

			len = ((p[0] & 0x03f) << 8) + p[1];

			if (NULL == *ptr)
			{
				tot_len += 2;
				*ptr = p + 2;
			}

			p = (u_char *)msg + len;
			if (p > (u_char *)end)
				return -1;

			/* too many pointers. */
			if (np++ > 16)
				return -1;

			continue;
		}

		len = p[0];
		p++;
		tot_len++;

		if (p + len > (u_char *)end)
			return -1;

		if (index + len >= buflen - 1)
			return -1;

		memcpy(buf + index, p, len);
		index += len;
		buf[index++] = '.';
		p += len;
		tot_len += len;
	}

	if (NULL == *ptr)
	{
		*ptr = p + 1;
		tot_len++;
	}
	/*
	 * omit last '.'
	 */
	if (index > 0)
	{
		buf[index - 1] = '\0';
	}
	else
	{
		buf[0] = '\0';
	}

	return tot_len;
}

/////////////////////////////////////////////////////////////////////////////////
// get API
/////////////////////////////////////////////////////////////////////////////////

static void get_dns_hdr_info(dns_hdr_t *dns_hdr, char *payload)
{
	dns_hdr_t *tmp = ((dns_hdr_t *)payload);
	memset(dns_hdr, 0, sizeof(dns_hdr_t));

	dns_hdr->qr = tmp->qr;
	dns_hdr->opcode = tmp->opcode;
	dns_hdr->aa = tmp->aa;
	dns_hdr->tc = tmp->tc;
	dns_hdr->rd = tmp->rd;
	dns_hdr->ra = tmp->ra;
	dns_hdr->z = tmp->z;
	dns_hdr->rcode = tmp->rcode;

	dns_hdr->id = ntohs(tmp->id);
	dns_hdr->qdcount = ntohs(tmp->qdcount);
	dns_hdr->ancount = ntohs(tmp->ancount);
	dns_hdr->aucount = ntohs(tmp->aucount);
	dns_hdr->adcount = ntohs(tmp->adcount);
}

static void get_rr_type_nsec3(char **ptr, nsec3_t *nsec3, char *end)
{
	NS_GET8(nsec3->hash_algo, *ptr);
	NS_GET8(nsec3->flags, *ptr);
	NS_GET16(nsec3->iteration, *ptr);
	NS_GET8(nsec3->salt_len, *ptr);
	nsec3->salt_value = *(u_char **)ptr;
	*ptr += nsec3->salt_len; /* jump salt_value */

	NS_GET8(nsec3->hash_len, *ptr);
	nsec3->next_hash_owner = *(u_char **)ptr;
	*ptr += nsec3->hash_len; /* jump next_hash_owner */
}

// unused
__attribute__((unused))static int get_rr_signer(u_char **ptr, u_char *buf, int buflen, char *end)
{
	u_char *p = NULL;
	int len = 0, i = 0;

	p = *ptr;
	if (0 == p[0])
	{
		len = strlen("Root");
		memcpy(buf, "Root", len);
		buf[len] = '\0';
		*ptr += 1;
		return 1;
	}

	p = *ptr;
	*ptr = NULL;
	len = 0;

	while (1)
	{
		len++;
		if (0x03 == p[0])
		{
			p += 1;
			continue;
		}

		buf[i++] = p[0];
		p += 1;

		if (0 == p[0])
		{
			len++;
			p += 1;
			break;
		}
	}

	*ptr = p;
	buf[i] = '\0';

	return len;
}

static void get_rr_type_rrsig(char **ptr, rrsig_t *rrsig, char *end)
{
	NS_GET16(rrsig->type_covered, *ptr);
	NS_GET8(rrsig->algo, *ptr);
	NS_GET8(rrsig->labels, *ptr);
	NS_GET32(rrsig->original_ttl, *ptr);
	NS_GET32(rrsig->sig_expiration, *ptr);
	NS_GET32(rrsig->sig_inception, *ptr);
	NS_GET16(rrsig->key_tag, *ptr);
}

static int get_rr_type_wks(char **ptr, wks_t *wks, char *end)
{
	if (*ptr + 4 > end)
		return -1;
	NS_GET32(wks->addr, *ptr);
	if (*ptr + 1 > end)
		return -1;
	NS_GET8(wks->protocol, *ptr);
	wks->bitmap = *(u_char **)ptr;

	return 0;
}

static void get_rr_type_info(char **ptr, hinfo_t *hinfo, char *end)
{
	int len = 0;

	NS_GET8(hinfo->cpu_len, *ptr);
	len = MIN(hinfo->cpu_len, sizeof(DNS_HINFO_MAX_CPU - 1));
	memcpy((char *)hinfo->cpu, *ptr, len);
	hinfo->cpu[len] = '\0';
	*ptr += hinfo->cpu_len;
	hinfo->cpu_len = len;

	NS_GET8(hinfo->os_len, *ptr);
	len = MIN(hinfo->os_len, sizeof(DNS_HINFO_MAX_OS - 1));
	memcpy((char *)hinfo->os, *ptr, len);
	hinfo->os[len] = '\0';
	*ptr += hinfo->os_len;
	hinfo->os_len = len;
}

static int get_rr_type_soa(char *msg, char **ptr, soa_t *soa, char *end)
{
	if (0 >= dns_name_decompress(msg, (u_char **)ptr, soa->mname, sizeof(soa->mname), end))
		return -1;
	if (0 >= dns_name_decompress(msg, (u_char **)ptr, soa->rname, sizeof(soa->rname), end))
		return -1;

	if (*ptr + 4 > end)
		return -1;
	NS_GET32(soa->serial, *ptr);
	if (*ptr + 4 > end)
		return -1;
	NS_GET32(soa->refresh, *ptr);
	if (*ptr + 4 > end)
		return -1;
	NS_GET32(soa->retry, *ptr);
	if (*ptr + 4 > end)
		return -1;
	NS_GET32(soa->expire, *ptr);
	if (*ptr + 4 > end)
		return -1;
	NS_GET32(soa->minimum, *ptr);

	return 0;
}

static int get_rr_common_field(char *msg, char **ptr, dns_rr_t *rr, char *end)
{
	char *p = NULL;
	if (*ptr == NULL)
	{
		return -1;
	}

	dns_name_decompress(msg, (u_char **)ptr, rr->name, DNS_MAX_NAME + 1, end);
#if 0
	if(0 >= dns_name_decompress(msg, (u_char**)ptr, rr->name, DNS_MAX_NAME+1, end))
	{
		return -1;
	}
#endif
	if (*ptr == NULL || *ptr + 2 > end)
	{
		return -1;
	}
	NS_GET16(rr->type, *ptr);

	if (*ptr == NULL || *ptr + 2 > end)
	{
		return -1;
	}

	NS_GET16(rr->rr_class, *ptr);

	if (DNS_CLASS_UNKNOWN == rr->rr_class)
	{
		return -1;
	}

	if (*ptr == NULL || *ptr + 4 > end)
	{
		return -1;
	}
	NS_GET32(rr->ttl, *ptr);
	if (*ptr == NULL || *ptr + 2 > end)
	{
		return -1;
	}
	NS_GET16(rr->rdlength, *ptr);
	p = *ptr + rr->rdlength;
	if (*ptr == NULL || p > end)
	{
		return -1;
	}

	return 0;
}

static int get_dns_query_question(char *msg, char **ptr, dns_query_question_t *q, char *end)
{
	int used_len = 0;
	if (0 >= dns_name_decompress(msg, (u_char **)ptr, q->qname, DNS_MAX_NAME + 1, end))
	{
		return -1;
	}

	used_len = strlen((char *)q->qname);
	if (q->qname[used_len - 1] == '.')
		return -2;

	if (*ptr + 2 > end)
	{
		return -1;
	}

	NS_GET16(q->qtype, *ptr);

	if (*ptr + 2 > end)
	{
		return -1;
	}

	NS_GET16(q->qclass, *ptr);

	return 0;
}

static int get_one_resource_record(char *msg, char **ptr, dns_rr_t *rr, char *end)
{
	u_int32_t len = 0, byte = 0;
	u_char *original_ptr = NULL;

	switch (rr->type)
	{
	case DNS_TYPE_CNAME:
		original_ptr = (u_char *)*ptr;
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.cname, DNS_MAX_NAME + 1, end))
		{
			return 0;
		}

		break;
	case DNS_TYPE_HINFO:
		get_rr_type_info(ptr, &(rr->rdata.hinfo), end);
		break;
	case DNS_TYPE_MB:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.mb, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_MD:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.md, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_MF:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.mf, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_MG:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.mg, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_MINFO:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.minfo.rmailbx, DNS_MAX_NAME + 1, end))
			return 0;
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.minfo.emailbx, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_MR:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.mr, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_MX:
		if (*ptr + 2 > end)
			return 0;
		NS_GET16(rr->rdata.mx.preference, *ptr);
		if (rr->rdlength - 2 < ((u_char *)*ptr)[0])
		{
			if (rr->rdlength < 2)
			{
				*ptr += rr->rdlength;
				break;
			}
			len = MIN(DNS_MAX_NAME - 1, rr->rdlength - 2); /* size=1byte */
			memcpy(rr->rdata.mx.exchange, *ptr, len);	   /* error labels */
			rr->rdata.mx.exchange[len] = '\0';
			*ptr += rr->rdlength - 2;
		}
		else
		{
			if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.mx.exchange, DNS_MAX_NAME + 1, end))
				return 0;
		}
		break;
	case DNS_TYPE_NS:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.ns, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_PTR:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.ptr, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_SOA:
		original_ptr = (u_char *)*ptr;
		if (0 != get_rr_type_soa(msg, ptr, &(rr->rdata.soa), end))
			return 0;
		if ((char *)original_ptr + rr->rdlength != *ptr)
		{
			*ptr = (char *)original_ptr + rr->rdlength;
		}
		break;
	case DNS_TYPE_A:
		if (*ptr + 4 > end)
			return 0;
		memcpy(rr->rdata.a, *ptr, NS_INT32SZ);
		(*ptr) += NS_INT32SZ;
		break;
	case DNS_TYPE_AAAA:
		if (*ptr + 16 > end)
			return -1;
		memcpy(rr->rdata.aaaa, *ptr, 16);
		(*ptr) += 16;
		break;
	case DNS_TYPE_DNAME:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.dname, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_ISDN:
		memcpy(rr->rdata.isdn, *ptr, sizeof(u_char));
		(*ptr) += 1;
		break;
	case DNS_TYPE_TXT:
		len = MIN(DNS_MAX_NAME - 1, rr->rdlength - 1); /* size=1byte */
		memcpy(rr->rdata.txt.txt, *ptr + 1, len);
		rr->rdata.txt.size = len;
		*ptr += rr->rdlength;
		break;
	case DNS_TYPE_RP:
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.rp.mailbox, DNS_MAX_NAME + 1, end))
			return 0;
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.rp.txt_rr, DNS_MAX_NAME + 1, end))
			return 0;
		break;
	case DNS_TYPE_NULL:
		len = MIN(DNS_MAX_NAME - 1, rr->rdlength - 1); /* size=1byte */
		memcpy(rr->rdata.null.null, *ptr + 1, len);
		rr->rdata.null.size = len;
		*ptr += rr->rdlength;
		break;
	case DNS_TYPE_WKS:
		if (0 != get_rr_type_wks(ptr, &(rr->rdata.wks), end))
			return 0;
		rr->rdata.wks.size = rr->rdlength - 5;
		*ptr += rr->rdlength - 5;
	case DNS_TYPE_SRV:
		NS_GET16(rr->rdata.srv.priority, *ptr);
		NS_GET16(rr->rdata.srv.weight, *ptr);
		NS_GET16(rr->rdata.srv.port, *ptr);
		if (0 >= dns_name_decompress(msg, (u_char **)ptr, rr->rdata.srv.target, DNS_MAX_TARGET, end))
			return 0;
		break;
	case DNS_TYPE_DS:
	case DNS_TYPE_DLV:
		if (*ptr + 4 > end)
			return 0;
		NS_GET16(rr->rdata.ds.key_tag, *ptr);
		NS_GET8(rr->rdata.ds.algo, *ptr);
		NS_GET8(rr->rdata.ds.digest_type, *ptr);
		rr->rdata.ds.digest = *(u_char **)ptr;
		rr->rdata.ds.digest_len = rr->rdlength - 4;
		*ptr += rr->rdlength - 4;
		break;
	case DNS_TYPE_RRSIG:
		if (*ptr + 18 > end)
			return 0;
		get_rr_type_rrsig(ptr, &(rr->rdata.rrsig), end);
		len = dns_name_decompress(msg, (u_char **)ptr, rr->rdata.rrsig.signer_name, DNS_MAX_SIGNER_NAME, end);
		if (len <= 0)
		{
			return -1;
		}
		rr->rdata.rrsig.signature = *(u_char **)ptr;
		rr->rdata.rrsig.signature_len = rr->rdlength - 18 - len;
		*ptr += rr->rdlength - 18 - len;
		break;
	case DNS_TYPE_NSEC:
		original_ptr = (u_char *)*ptr;
		len = dns_name_decompress(msg, (u_char **)ptr, rr->rdata.nsec.next_domain, DNS_MAX_OWNER, end);
		if (len <= 0)
		{
			return -1;
		}

		if ((original_ptr + rr->rdlength != (u_char *)*ptr) && (*ptr != NULL))
		{
			NS_GET16(len, *ptr);
			byte = MIN(DNS_MAX_MAPS - 1, len);
			memcpy(rr->rdata.nsec.type_bit_maps, *ptr, byte);
			rr->rdata.nsec.type_bit_maps[byte] = '\0';
			*ptr += len;
			rr->rdata.nsec.maps_temp_len = len;
			rr->rdata.nsec.maps_len = len;
			len = byte;

			byte = ((u_char *)ptr)[0];
			if ((byte & 0xFF) == 0xFF || byte == 128)
			{
				*ptr += 1; /* jump 0xFF */
				byte = ((u_char *)ptr)[0];
				*ptr += 1; /* jump 1 byte of len */
				len = MIN(DNS_MAX_MAPS - 1 - len, byte);
				memcpy(rr->rdata.nsec.type_bit_maps + rr->rdata.nsec.maps_len, *ptr, len);
				(rr->rdata.nsec.type_bit_maps + rr->rdata.nsec.maps_len)[len] = '\0';
				*ptr += byte; /* jump  byte */
				rr->rdata.nsec.maps_len += len;
			}
		}
		break;
	case DNS_TYPE_DNSKEY:
		if (*ptr + 4 > end)
			return 0;
		NS_GET16(rr->rdata.dnskey.flags, *ptr);
		NS_GET8(rr->rdata.dnskey.protocol, *ptr);
		NS_GET8(rr->rdata.dnskey.algo, *ptr);
		rr->rdata.dnskey.public_key = *(u_char **)ptr;
		rr->rdata.dnskey.public_key_len = rr->rdlength - 4; /* sizeof(flags)+sizeof(protocol)+sizeof(algo) */
		*ptr += rr->rdlength - 4;							/* todo add log */
		break;
	case DNS_TYPE_NSEC3:
		if (*ptr + 5 > end)
			return 0;
		original_ptr = (u_char *)*ptr;
		get_rr_type_nsec3(ptr, &(rr->rdata.nsec3), end);

		if ((original_ptr + rr->rdlength != (u_char *)*ptr) && (*ptr != NULL))
		{
			NS_GET16(len, *ptr);
			byte = MIN(DNS_MAX_MAPS - 1, len);
			memcpy(rr->rdata.nsec3.type_bit_maps, *ptr, byte);
			rr->rdata.nsec3.type_bit_maps[byte] = '\0';
			*ptr += len;
			rr->rdata.nsec3.maps_temp_len = byte;
			rr->rdata.nsec3.maps_len = byte;
			len = byte;

			byte = ((u_char *)*ptr)[0];
			if ((byte & 0xFF) == 0xFF || byte == 128)
			{
				*ptr += 1; /* jump 0xFF */
				byte = ((u_char *)*ptr)[0];
				*ptr += 1; /* jump 1 byte of len */
				len = MIN(DNS_MAX_MAPS - 1 - len, byte);
				memcpy(rr->rdata.nsec3.type_bit_maps + rr->rdata.nsec3.maps_len, *ptr, len);
				(rr->rdata.nsec3.type_bit_maps + rr->rdata.nsec3.maps_len)[len] = '\0';
				*ptr += byte; /* jump  byte */
				rr->rdata.nsec3.maps_len += len;
			}
		}
		break;
	case DNS_TYPE_NSEC3PARAM:
		NS_GET8(rr->rdata.nsec3param.hash_algo, *ptr);
		NS_GET8(rr->rdata.nsec3param.flags, *ptr);
		NS_GET16(rr->rdata.nsec3param.iteration, *ptr);
		rr->rdata.nsec3param.salt_len = rr->rdlength - 4 - 1;
		*ptr += 1;
		rr->rdata.nsec3param.salt_value = *(u_char **)ptr;
		*ptr += rr->rdlength - 5;
		break;
	case DNS_TYPE_OPT:
		/* fail through */
	case DNS_TYPE_UNKNOWN:
		/* fail through */
	default:
		memcpy(rr->rdata.unknown_data, *ptr, rr->rdlength);
		rr->rdata.unknown_data[rr->rdlength] = '\0';
		(*ptr) += rr->rdlength;
		// printf("No support respone type, type: %d", rr->type);
		break;
	}

	return 1;
}

static int parse_resource_record(dns_info_t *dns_info, char *payload, int payload_len, char **cur_pos)
{
	int i = 0;

	dns_info->rr_count = dns_info->hdr_info.ancount + dns_info->hdr_info.adcount + dns_info->hdr_info.aucount;
	if (dns_info->rr_count == 0)
	{
		return 0;
	}

	dns_info->rr_count = MIN(dns_info->rr_count, MAX_RR_NUM);
	for (i = 0; i < dns_info->rr_count; i++)
	{
		if ((u_char *)*cur_pos >= (u_char *)payload + payload_len || (u_char *)*cur_pos < (u_char *)payload)
		{
			DNS_ERROR_LOG("parse_resource_record()");
			return -1;
		}

		if (0 != get_rr_common_field(payload, cur_pos, &dns_info->rr[i], payload + payload_len))
		{
			if (dns_info->rr[i].rr_class == DNS_CLASS_UNKNOWN)
			{
				DNS_ERROR_LOG("get_rr_common_field()");
				return -1; /* error packet */
			}
			else
			{
				DNS_ERROR_LOG("get_rr_common_field()");
				dns_info->rr_count -= 1;
				i -= 1;
				continue;
			}
		}

		if (dns_info->rr[i].rdlength == 0)
		{
			continue;
		}

		if (get_one_resource_record(payload, cur_pos, &dns_info->rr[i], payload + payload_len) != 1)
		{
			DNS_ERROR_LOG("get_one_resource_record()");
			dns_info->rr_count -= 1;
			i -= 1;
		}
	}

	assert(i == dns_info->rr_count);
	return 0;
}

/////////////////////////////////////////////////////////////////////////////////
// set API
/////////////////////////////////////////////////////////////////////////////////

static void set_dns_hdr_info(dns_hdr_t *src, char *payload)
{
	dns_hdr_t *dst = ((dns_hdr_t *)payload);
	memset(dst, 0, sizeof(dns_hdr_t));

	dst->qr = src->qr;
	dst->opcode = src->opcode;
	dst->aa = src->aa;
	dst->tc = src->tc;
	dst->rd = src->rd;
	dst->ra = src->ra;
	dst->z = src->z;
	dst->rcode = src->rcode;

	dst->id = htons(src->id);
	dst->qdcount = htons(src->qdcount); // 16bits: QDCOUNT: number of questions
	dst->ancount = htons(src->ancount); // 16bits: ANCOUNT: number of answer resource records
	dst->aucount = htons(src->aucount); // 16bits: NSCOUNT: number of authority resource records
	dst->adcount = htons(src->adcount); // 16bits: ARCOUNT: number of additional resource records
}

static int set_dns_pkt_rr_header(char *payload, int payload_len, dns_rr_t *rr, int rdlength, char **rr_header_length_ptr)
{
	int ret;
	int used_len = 0;

	if (rr->name)
	{
		ret = dns_name_compress(rr->name, strlen((char *)rr->name), payload);
		if (ret == -1)
		{
			return -1;
		}
		used_len += ret;
	}
	else
	{
		NS_SET16(0xc00c, payload, used_len);
	}

	NS_SET16(rr->type, payload, used_len);
	NS_SET16(rr->rr_class, payload, used_len);
	NS_SET32(rr->ttl, payload, used_len);
	*rr_header_length_ptr = payload + used_len;
	NS_SET16(rr->rdlength, payload, used_len);

	return used_len;
}

static void reset_rr_header_length(char **rr_length_ptr, int rdlength)
{
	u_int16_t seg_16 = 0;
	seg_16 = htons(rdlength);
	memset(*rr_length_ptr, 0, sizeof(seg_16));
	memcpy(*rr_length_ptr, &seg_16, sizeof(seg_16));
}

/////////////////////////////////////////////////////////////////////////////////
// cheat API
/////////////////////////////////////////////////////////////////////////////////

static void set_cheat_pkt_header(dns_hdr_t *dns_hdr)
{
	dns_hdr->qr = 1;	 // 1bit: Response
	dns_hdr->opcode = 0; // 4bits: Query
	dns_hdr->aa = 0;	 // 1bit: authoritative answer
	dns_hdr->tc = 0;	 // 1bit: Not truncated
	dns_hdr->rd = 1;	 // 1bit: Recursion Desired
	dns_hdr->ra = 1;	 // 1bit: Recursion Available
	dns_hdr->z = 0;		 // 3bits: Reserved for future use: Must be zero in all queries and responses
	dns_hdr->rcode = 0;	 // 4bits: 0: No error condition

	dns_hdr->id = htons(dns_hdr->id);
	dns_hdr->qdcount = htons(dns_hdr->qdcount); // 16bits: QDCOUNT: number of questions
	dns_hdr->ancount = htons(dns_hdr->ancount); // 16bits: ANCOUNT: number of answer resource records
	dns_hdr->aucount = htons(dns_hdr->aucount); // 16bits: NSCOUNT: number of authority resource records
	dns_hdr->adcount = htons(dns_hdr->adcount); // 16bits: ARCOUNT: number of additional resource records
}

static int set_cheat_pkt_question(char *payload, dns_query_question_t *query, int query_num)
{
	int ret;
	int used_len = 0;

	/* 只处理一个请求 */
	ret = dns_name_compress(query->qname, strlen((char *)query->qname), payload);
	if (ret == -1)
	{
		return -1;
	}
	used_len += ret;

	NS_SET16(query->qtype, payload, used_len);
	NS_SET16(query->qclass, payload, used_len);

	return used_len;
}

// NOTE 返回值 TODO
static int set_cheat_pkt_rr_pdata_type_str(char *payload, cheat_pkt_opt_t *cheat_opt)
{
	int used_len = 0;
	u_int16_t q_class = 1;
	u_int16_t compress_len = 0;
	u_char compress_name[DNS_MAX_NAME + 1] = {0};

	NS_SET16(0xc00c, payload, used_len);
	NS_SET16(cheat_opt->res_type, payload, used_len);
	NS_SET16(q_class, payload, used_len);
	NS_SET32(cheat_opt->ttl, payload, used_len);

	compress_len = dns_compress_rr_str(cheat_opt->res_info, cheat_opt->res_len, compress_name);
	if (compress_len <= 0)
	{
		DNS_ERROR_LOG("dns_compress_rr_str()");
	}

	NS_SET16(compress_len, payload, used_len);
	NS_SETLEN(compress_name, compress_len, payload, used_len);

	return used_len;
}

static int set_cheat_pkt_rr_pdata_type_ip(char *payload, cheat_pkt_opt_t *cheat_opt)
{
	int used_len = 0;
	u_int16_t q_class = 1;

	NS_SET16(0xc00c, payload, used_len);
	NS_SET16(cheat_opt->res_type, payload, used_len);
	NS_SET16(q_class, payload, used_len);
	NS_SET32(cheat_opt->ttl, payload, used_len);
	NS_SET16(cheat_opt->res_len, payload, used_len);
	NS_SETLEN((char *)cheat_opt->res_info, cheat_opt->res_len, payload, used_len);

	return used_len;
}

/////////////////////////////////////////////////////////////////////////////////
// public API
/////////////////////////////////////////////////////////////////////////////////

dns_info_t *dns_new(void)
{
	dns_info_t *dns_info = (dns_info_t *)calloc(sizeof(dns_info_t), 1);
	return dns_info;
}

void dns_free(dns_info_t *dns_info)
{
	if (dns_info)
	{
		free(dns_info);
		dns_info = NULL;
	}
}

int dns_parser(dns_info_t *dns_info, char *in_buff, int buff_len)
{
	char *cur_pos = NULL;

	get_dns_hdr_info(&dns_info->hdr_info, in_buff);
	cur_pos = in_buff + sizeof(dns_hdr_t);

	if (0 == dns_info->hdr_info.qdcount || dns_info->hdr_info.qdcount > 1)
	{
		return -1;
	}

	if (0 != get_dns_query_question(in_buff, &cur_pos, &(dns_info->query_question), in_buff + buff_len))
	{
		DNS_ERROR_LOG("get_dns_query_question()");
		return -1;
	}

	if (parse_resource_record(dns_info, in_buff, buff_len, &cur_pos) != 0)
	{
		DNS_ERROR_LOG("parse_resource_record()");
		return -1;
	}

	return 0;
}

int dns_package(dns_info_t *dns_info, char *out_buff, int buff_size)
{
	char *rr_header_length_ptr = NULL;
	int i = 0;
	int rr_count = 0;
	dns_hdr_t *dns_hdr = &dns_info->hdr_info;
	int ret = 0;
	int used_len = 0;
	dns_rr_t *rr;

	set_dns_hdr_info(dns_hdr, out_buff);
	used_len += sizeof(dns_hdr_t);
	ret = set_cheat_pkt_question(out_buff + used_len, &dns_info->query_question, dns_hdr->qdcount);
	if (ret < 0)
	{
		return -1;
	}
	used_len += ret;

	rr_count = dns_hdr->adcount + dns_hdr->ancount + dns_hdr->aucount;
	for (i = 0; i < rr_count; i++)
	{
		rr = &dns_info->rr[i];
		ret = set_dns_pkt_rr_header(out_buff + used_len, buff_size, rr, 0, &rr_header_length_ptr);
		if (ret == -1)
		{
			return -1;
		}
		used_len += ret;

		if (rr->rdlength == 0)
		{
			continue;
		}

		int rr_header_used_len = used_len;
		switch (rr->type)
		{
		case DNS_TYPE_A:
			NS_SETLEN(rr->rdata.a, strlen((char *)rr->rdata.a), out_buff, used_len);
			break;
		case DNS_TYPE_AAAA:
			NS_SETLEN(rr->rdata.aaaa, 16, out_buff, used_len);
			break;
		case DNS_TYPE_CNAME:
			ret = dns_name_compress(rr->rdata.cname, strlen((char *)rr->rdata.cname), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_MB:
			ret = dns_name_compress(rr->rdata.mb, strlen((char *)rr->rdata.mb), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_MD:
			ret = dns_name_compress(rr->rdata.md, strlen((char *)rr->rdata.md), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_MF:
			ret = dns_name_compress(rr->rdata.mf, strlen((char *)rr->rdata.mf), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_MG:
			ret = dns_name_compress(rr->rdata.mg, strlen((char *)rr->rdata.mg), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_MINFO:
			ret = dns_name_compress(rr->rdata.minfo.rmailbx, strlen((char *)rr->rdata.minfo.rmailbx), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;

			ret = dns_name_compress(rr->rdata.minfo.emailbx, strlen((char *)rr->rdata.minfo.emailbx), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_MR:
			ret = dns_name_compress(rr->rdata.mr, strlen((char *)rr->rdata.mr), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_MX:
			NS_SET16(rr->rdata.mx.preference, out_buff, used_len);

			ret = dns_name_compress(rr->rdata.mx.exchange, strlen((char *)rr->rdata.mx.exchange), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_NS:
			ret = dns_name_compress(rr->rdata.ns, strlen((char *)rr->rdata.ns), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_PTR:
			ret = dns_name_compress(rr->rdata.ptr, strlen((char *)rr->rdata.ptr), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_SOA:
			ret = dns_name_compress(rr->rdata.soa.mname, strlen((char *)rr->rdata.soa.mname), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;

			ret = dns_name_compress(rr->rdata.soa.rname, strlen((char *)rr->rdata.soa.rname), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;

			NS_SET32(rr->rdata.soa.serial, out_buff, used_len);
			NS_SET32(rr->rdata.soa.refresh, out_buff, used_len);
			NS_SET32(rr->rdata.soa.retry, out_buff, used_len);
			NS_SET32(rr->rdata.soa.expire, out_buff, used_len);
			NS_SET32(rr->rdata.soa.minimum, out_buff, used_len);
			break;
		case DNS_TYPE_DNAME:
			ret = dns_name_compress(rr->rdata.dname, strlen((char *)rr->rdata.dname), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_ISDN:
			NS_SETLEN(rr->rdata.isdn, strlen((char *)rr->rdata.isdn), out_buff, used_len);
			break;
		case DNS_TYPE_RP:
			ret = dns_name_compress(rr->rdata.rp.mailbox, strlen((char *)rr->rdata.rp.mailbox), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;

			ret = dns_name_compress(rr->rdata.rp.txt_rr, strlen((char *)rr->rdata.rp.txt_rr), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_SRV:
			NS_SET16(rr->rdata.srv.priority, out_buff, used_len);
			NS_SET16(rr->rdata.srv.weight, out_buff, used_len);
			NS_SET16(rr->rdata.srv.port, out_buff, used_len);

			ret = dns_name_compress(rr->rdata.srv.target, strlen((char *)rr->rdata.srv.target), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			break;
		case DNS_TYPE_RRSIG:
			NS_SET16(rr->rdata.rrsig.type_covered, out_buff, used_len);
			NS_SET8(rr->rdata.rrsig.algo, out_buff, used_len);
			NS_SET8(rr->rdata.rrsig.labels, out_buff, used_len);
			NS_SET32(rr->rdata.rrsig.original_ttl, out_buff, used_len);
			NS_SET32(rr->rdata.rrsig.sig_expiration, out_buff, used_len);
			NS_SET32(rr->rdata.rrsig.sig_inception, out_buff, used_len);
			NS_SET16(rr->rdata.rrsig.key_tag, out_buff, used_len);

			ret = dns_name_compress(rr->rdata.rrsig.signer_name, strlen((char *)rr->rdata.rrsig.signer_name), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;
			NS_SETLEN(rr->rdata.rrsig.signature, rr->rdata.rrsig.signature_len, out_buff, used_len);
			break;
		case DNS_TYPE_DNSKEY:
			NS_SET16(rr->rdata.dnskey.flags, out_buff, used_len);
			NS_SET8(rr->rdata.dnskey.protocol, out_buff, used_len);
			NS_SET8(rr->rdata.dnskey.algo, out_buff, used_len);
			NS_SETLEN(rr->rdata.dnskey.public_key, rr->rdata.dnskey.public_key_len, out_buff, used_len);
			break;
		case DNS_TYPE_NSEC3PARAM:
			NS_SET8(rr->rdata.nsec3param.hash_algo, out_buff, used_len);
			NS_SET8(rr->rdata.nsec3param.flags, out_buff, used_len);
			NS_SET16(rr->rdata.nsec3param.iteration, out_buff, used_len);
			NS_SET8(rr->rdata.nsec3param.salt_len, out_buff, used_len);
			NS_SETLEN(rr->rdata.nsec3param.salt_value, rr->rdata.nsec3param.salt_len, out_buff, used_len);
			break;
		case DNS_TYPE_DS:
			/* fail through */
		case DNS_TYPE_DLV:
			NS_SET16(rr->rdata.ds.key_tag, out_buff, used_len);
			NS_SET8(rr->rdata.ds.algo, out_buff, used_len);
			NS_SET8(rr->rdata.ds.digest_type, out_buff, used_len);
			NS_SETLEN(rr->rdata.ds.digest, rr->rdata.ds.digest_len, out_buff, used_len);
			break;
		case DNS_TYPE_TXT:
			NS_SET8(rr->rdata.txt.size, out_buff, used_len);
			NS_SETLEN(rr->rdata.txt.txt, rr->rdata.txt.size, out_buff, used_len);
			break;
		case DNS_TYPE_NULL:
			NS_SET8(rr->rdata.null.size, out_buff, used_len);
			NS_SETLEN(rr->rdata.null.null, rr->rdata.null.size, out_buff, used_len);
			break;
		case DNS_TYPE_WKS:
			NS_SET32(rr->rdata.wks.addr, out_buff, used_len);
			NS_SET8(rr->rdata.wks.protocol, out_buff, used_len);
			NS_SETLEN(rr->rdata.wks.bitmap, rr->rdata.wks.size, out_buff, used_len);
			break;
		case DNS_TYPE_HINFO:
			NS_SET8(rr->rdata.hinfo.cpu_len, out_buff, used_len);
			NS_SETLEN(rr->rdata.hinfo.cpu, rr->rdata.hinfo.cpu_len, out_buff, used_len);
			NS_SET8(rr->rdata.hinfo.os_len, out_buff, used_len);
			NS_SETLEN(rr->rdata.hinfo.os, rr->rdata.hinfo.os_len, out_buff, used_len);
			break;
		case DNS_TYPE_NSEC:
			ret = dns_name_compress(rr->rdata.nsec.next_domain, strlen((char *)rr->rdata.nsec.next_domain), out_buff + used_len);
			if (ret == -1)
			{
				return -1;
			}
			used_len += ret;

			if (rr->rdata.nsec.maps_len)
			{
				NS_SET16(rr->rdata.nsec.maps_temp_len, out_buff, used_len);
				NS_SETLEN(rr->rdata.nsec.type_bit_maps, rr->rdata.nsec.maps_temp_len, out_buff, used_len);

				if (rr->rdata.nsec.maps_len > rr->rdata.nsec.maps_temp_len)
				{
					NS_SET8(0xFF, out_buff, used_len);
					NS_SET8(rr->rdata.nsec.maps_len - rr->rdata.nsec.maps_temp_len, out_buff, used_len);
					NS_SETLEN(rr->rdata.nsec.type_bit_maps + rr->rdata.nsec.maps_temp_len, rr->rdata.nsec.maps_len - rr->rdata.nsec.maps_temp_len, out_buff, used_len);
				}
			}
			break;
		case DNS_TYPE_NSEC3:
			NS_SET8(rr->rdata.nsec3.hash_algo, out_buff, used_len);
			NS_SET8(rr->rdata.nsec3.flags, out_buff, used_len);
			NS_SET16(rr->rdata.nsec3.iteration, out_buff, used_len);
			NS_SET8(rr->rdata.nsec3.salt_len, out_buff, used_len);
			NS_SETLEN(rr->rdata.nsec3.salt_value, rr->rdata.nsec3.salt_len, out_buff, used_len);
			NS_SET8(rr->rdata.nsec3.hash_len, out_buff, used_len);
			NS_SETLEN(rr->rdata.nsec3.next_hash_owner, rr->rdata.nsec3.hash_len, out_buff, used_len);

			if (rr->rdata.nsec3.maps_len)
			{
				NS_SET16(rr->rdata.nsec3.maps_temp_len, out_buff, used_len);
				NS_SETLEN(rr->rdata.nsec3.type_bit_maps, rr->rdata.nsec3.maps_temp_len, out_buff, used_len);

				if (rr->rdata.nsec3.maps_len > rr->rdata.nsec3.maps_temp_len)
				{
					NS_SET8(0xFF, out_buff, used_len);
					NS_SET8(rr->rdata.nsec3.maps_len - rr->rdata.nsec3.maps_temp_len, out_buff, used_len);
					NS_SETLEN(rr->rdata.nsec3.type_bit_maps + rr->rdata.nsec3.maps_temp_len, rr->rdata.nsec3.maps_len - rr->rdata.nsec3.maps_temp_len, out_buff, used_len);
				}
			}
			break;
		case DNS_TYPE_OPT:
			/* fail through */
		case DNS_TYPE_UNKNOWN:
			/* fail through */
		default:
			NS_SETLEN(rr->rdata.unknown_data, rr->rdlength, out_buff, used_len);
			break;
		}

		if (used_len - rr_header_used_len != rr->rdlength)
		{
			reset_rr_header_length(&rr_header_length_ptr, used_len - rr_header_used_len);
		}
		if (used_len > buff_size)
		{
			// printf("dns_package buff to short, %s, %d\n", __FILE__, __LINE__);
			return -1;
		}
	}

	return used_len;
}

int dns_cheat_response(dns_info_t *dns_info, cheat_pkt_opt_t *cheat_opt, int cheat_opt_num, char *out_buff, int buff_size)
{
	int i = 0;
	dns_hdr_t *dns_hdr = NULL;
	int ret = 0, used_len = 0;

	memset(out_buff, 0, buff_size);
	dns_hdr = (dns_hdr_t *)out_buff;
	dns_hdr->id = dns_info->hdr_info.id;
	dns_hdr->qdcount = 1;
	dns_hdr->ancount = cheat_opt_num;
	used_len += sizeof(dns_hdr_t);

	ret = set_cheat_pkt_question(out_buff + used_len, &dns_info->query_question, dns_hdr->qdcount);
	if (ret < 0)
	{
		return -1;
	}
	used_len += ret;

	for (i = 0; i < cheat_opt_num; i++)
	{
		switch (cheat_opt[i].res_type)
		{
		case DNS_TYPE_A:
		case DNS_TYPE_AAAA:
			used_len += set_cheat_pkt_rr_pdata_type_ip(out_buff + used_len, &cheat_opt[i]);
			break;
		case DNS_TYPE_CNAME:
		case DNS_TYPE_NS:
			used_len += set_cheat_pkt_rr_pdata_type_str(out_buff + used_len, &cheat_opt[i]);
			break;
		}

		if (used_len > buff_size)
		{
			return -2;
		}
	}

	set_cheat_pkt_header(dns_hdr);

	return used_len;
}