/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015 by Sergey Fetisov <fsenok@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

/*
 * version: 1.0 demo (7.02.2015)
 * brief:   tiny dns ipv4 server using lwip (pcb)
 */

#include "dnserver.h"

#define DNS_MAX_HOST_NAME_LEN 128

static struct udp_pcb *pcb = NULL;
dns_query_proc_t query_proc = NULL;

#pragma pack(push, 1)
typedef struct
{
#if BYTE_ORDER == LITTLE_ENDIAN
	uint8_t rd: 1,     /* Recursion Desired */
	        tc: 1,     /* Truncation Flag */
	        aa: 1,     /* Authoritative Answer Flag */
	        opcode: 4, /* Operation code */
	        qr: 1;     /* Query/Response Flag */
	uint8_t rcode: 4,  /* Response Code */
	        z: 3,      /* Zero */
	        ra: 1;     /* Recursion Available */
#else
	uint8_t qr: 1,     /* Query/Response Flag */
	        opcode: 4, /* Operation code */
	        aa: 1,     /* Authoritative Answer Flag */
	        tc: 1,     /* Truncation Flag */
	        rd: 1;     /* Recursion Desired */
	uint8_t ra: 1,     /* Recursion Available */
	        z: 3,      /* Zero */
	        rcode: 4;  /* Response Code */
#endif
} dns_header_flags_t;

typedef struct
{
	uint16_t id;
	dns_header_flags_t flags;
	uint16_t n_record[4];
} dns_header_t;

typedef struct dns_answer
{
	uint16_t name;
	uint16_t type;
	uint16_t Class;
	uint32_t ttl;
	uint16_t len;
	uint32_t addr;
} dns_answer_t;
#pragma pack(pop)

typedef struct dns_query
{
	char name[DNS_MAX_HOST_NAME_LEN];
	uint16_t type;
	uint16_t Class;
} dns_query_t;

static uint16_t get_uint16(const uint8_t *pnt)
{
  uint16_t result;
  memcpy(&result, pnt, sizeof(result));
  return result;
}

static int parse_next_query(void *data, int size, dns_query_t *query)
{
	int len;
	int labels;
	uint8_t *ptr;

	len = 0;
	labels = 0;
	ptr = (uint8_t *)data;

	while (true)
	{
		uint8_t lable_len;
		if (size <= 0) return -1;
		lable_len = *ptr++;
		size--;
		if (lable_len == 0) break;
		if (labels > 0)
		{
			if (len == DNS_MAX_HOST_NAME_LEN) return -2;
			query->name[len++] = '.';
		}
		if (lable_len > size) return -1;
		if (len + lable_len >= DNS_MAX_HOST_NAME_LEN) return -2;
		memcpy(&query->name[len], ptr, lable_len);
		len += lable_len;
		ptr += lable_len;
		size -= lable_len;
		labels++;
	}

	if (size < 4) return -1;
	query->name[len] = 0;
	query->type = get_uint16(ptr);
	ptr += 2;
	query->Class = get_uint16(ptr);
	ptr += 2;
	return ptr - (uint8_t *)data;
}

static void udp_recv_proc(void *arg, struct udp_pcb *upcb, struct pbuf *p, const ip_addr_t *addr, u16_t port)
{
	int len;
	dns_header_t *header;
	static dns_query_t query;
	struct pbuf *out;
	ip4_addr_t host_addr;
	dns_answer_t *answer;

	(void)arg;

	if (p->len <= sizeof(dns_header_t)) goto error;
	header = (dns_header_t *)p->payload;
	if (header->flags.qr != 0) goto error;
	if (ntohs(header->n_record[0]) != 1) goto error;

	len = parse_next_query(header + 1, p->len - sizeof(dns_header_t), &query);
	if (len < 0) goto error;
	if (!query_proc(query.name, &host_addr)) goto error;

	len += sizeof(dns_header_t);
	out = pbuf_alloc(PBUF_TRANSPORT, len + 16, PBUF_POOL);
	if (out == NULL) goto error;

	memcpy(out->payload, p->payload, len);
	header = (dns_header_t *)out->payload;
	header->flags.qr = 1;
	header->n_record[1] = htons(1);
	answer = (struct dns_answer *)((uint8_t *)out->payload + len);
	answer->name = htons(0xC00C);
	answer->type = htons(1);
	answer->Class = htons(1);
	answer->ttl = htonl(32);
	answer->len = htons(4);
	answer->addr = host_addr.addr;

	udp_sendto(upcb, out, addr, port);
	pbuf_free(out);

error:
	pbuf_free(p);
}

err_t dnserv_init(const ip_addr_t *bind, uint16_t port, dns_query_proc_t qp)
{
	err_t err;
	udp_init();
	dnserv_free();
	pcb = udp_new();
	if (pcb == NULL)
		return ERR_MEM;
	err = udp_bind(pcb, bind, port);
	if (err != ERR_OK)
	{
		dnserv_free();
		return err;
	}
	udp_recv(pcb, udp_recv_proc, NULL);
	query_proc = qp;
	return ERR_OK;
}

void dnserv_free(void)
{
	if (pcb == NULL) return;
	udp_remove(pcb);
	pcb = NULL;
}