memory leak fixes, now answers

removed forking to improve performance
fixed memory leaks
added header construct function
now return MXDOMAIN to all valid DNS queries

3 files changed, 138 insertions, 48 deletions

diff --git a/src/dns.c b/src/dns.c
index 6cdae43..b8e4d22 100644
--- a/src/dns.c
+++ b/src/dns.c
@@ -5,8 +5,52 @@
 
 #include "dns.h"
 
+int dns_construct_header ( struct dns_header* _header, char* _buffer, int _bufflen )
+{
+	if ( !_buffer || !_header || _bufflen < 12 )
+		return -1;
+
+	*((uint16_t*)_buffer) = _header->id; //Since only copied, no flipping necessary
+	_buffer[2] = 
+		((_header->QR & 0x01) << 7) |
+		((_header->OPCODE & 0x0F) << 3) |
+		((_header->AA & 0x01) << 2) |
+		((_header->TC & 0x01) << 1) |
+		( _header->RD & 0x01);
+	_buffer[3] =
+		((_header->RA & 0x01) << 7) |
+		((_header->Z  & 0x07) << 4) |
+		( _header->RCODE & 0x0F);
+	*((uint16_t*)(_buffer + 4 )) = FLIP_BYTES(_header->question_count);
+	*((uint16_t*)(_buffer + 6 )) = FLIP_BYTES(_header->answer_count);
+	*((uint16_t*)(_buffer + 8 )) = FLIP_BYTES(_header->authorative_count);
+	*((uint16_t*)(_buffer + 10)) = FLIP_BYTES(_header->additional_count);
+
+	return 12;
+}
+
+int dns_destroy_struct ( struct dns_message* _msg )
+{
+	if ( !_msg )
+		return -1;
+
+	if ( _msg->question_count > 0 && _msg->question) {
+		free ( _msg->question );
+		_msg->question = NULL;
+	}
+
+	if ( _msg->answer_count > 0 && _msg->answer) {
+		free ( _msg->answer );
+		_msg->answer = NULL;
+	}
+
+	return 0;
+}
+
 int dns_parse_packet ( char* _buffer, int _bufflen, struct dns_message* _msg )
 {
+	//TODO refactor
+
 	if ( !_buffer || !_bufflen || !_msg )
 		return 1; //Invalid input
 
@@ -23,16 +67,25 @@ int dns_parse_packet ( char* _buffer, int _bufflen, struct dns_message* _msg )
 	_msg->header.RA = (0x80 & *( (uint8_t*) (_buffer + 3))) >> 7;
 	_msg->header.Z  = (0x70 & *( (uint8_t*) (_buffer + 3))) >> 4;
 	_msg->header.RCODE = (0x0F & *( (uint8_t*) (_buffer + 3)));
-	_msg->header.question_count	= (*((uint8_t*) (_buffer + 4 )) << 8) | *((uint8_t*) (_buffer + 5 ));
-	_msg->header.answer_count	= (*((uint8_t*) (_buffer + 6 )) << 8) | *((uint8_t*) (_buffer + 7 ));
+	_msg->question_count = _msg->header.question_count = (*((uint8_t*) (_buffer + 4 )) << 8) | *((uint8_t*) (_buffer + 5 ));
+	_msg->answer_count   = _msg->header.answer_count   = (*((uint8_t*) (_buffer + 6 )) << 8) | *((uint8_t*) (_buffer + 7 ));
 	_msg->header.authorative_count	= (*((uint8_t*) (_buffer + 8 )) << 8) | *((uint8_t*) (_buffer + 9 ));
 	_msg->header.additional_count	= (*((uint8_t*) (_buffer + 10)) << 8) | *((uint8_t*) (_buffer + 11));
+
 	//TODO remove
 	printf("ANSWER %i\n", _msg->header.answer_count);
 	printf("QUESTI %i\n", _msg->header.question_count);
 	printf("AUTHOR %i\n", _msg->header.authorative_count);
 	printf("ADDITI %i\n", _msg->header.additional_count);
 
+	//Check for sensible QD, AN, NS and ARCOUNTS before massive memory allocation
+	if(	_msg->header.question_count > 4 ||
+		_msg->header.answer_count > 32 ||
+		_msg->header.authorative_count > 32 ||
+		_msg->header.additional_count > 32 ) {
+		return 1;
+	}
+
 	//Allocate question array
 	//TODO Only implements question section.
 	size_t qsize = sizeof(typeof(*(_msg->question))) * _msg->question_count;
@@ -40,6 +93,9 @@ int dns_parse_packet ( char* _buffer, int _bufflen, struct dns_message* _msg )
 	_msg->question = malloc ( qsize );
 	memset( _msg->question, 0, qsize );
 
+	if (!_msg->question) //malloc failed
+		return 1;
+
 	int ptr = 12; //byte counter
 	for (int i = 0; i < _msg->question_count; i++) {
 		int qname_len = qname_check ( (_buffer + ptr), _bufflen - ptr);
@@ -49,12 +105,15 @@ int dns_parse_packet ( char* _buffer, int _bufflen, struct dns_message* _msg )
 
 		_msg->question[i].qname = _buffer + ptr;
 		ptr += qname_len;
-		_msg->question[i].qtype = (uint16_t*) (_buffer + ptr);
-		ptr += 2;
-		_msg->question[i].qclass = (uint16_t*) (_buffer + ptr);
 
-		if( ptr >= _bufflen ) //Out of bounds check
+		if( ptr >= (_bufflen - 4) ) //Out of bounds check
 			return 1;
+
+		_msg->question[i].qtype =  ((uint8_t)*(_buffer + ptr) << 8) | ((uint8_t)*(_buffer + ptr + 1));
+		ptr += 2;
+		_msg->question[i].qclass = ((uint8_t)*(_buffer + ptr) << 8) | ((uint8_t)*(_buffer + ptr + 1));
+		ptr += 2;
+
 	}
 
 	return 0;
@@ -92,7 +151,6 @@ int fqdn_to_qname( char* _source, int _sourcelen, char* _sink ,int _sinklen )
 	return i+2;
 }
 
-
 int qname_to_fqdn( char* _source, int _sourcelen, char* _sink, int _sinklen )
 {
 	if ( !_sourcelen || !_sinklen ) {
diff --git a/src/dns.h b/src/dns.h
index f93480e..fe871e0 100644
--- a/src/dns.h
+++ b/src/dns.h
@@ -12,18 +12,49 @@
 //TODO remove
 #include <stdio.h>
 
+//Resource Records
+#define RR_A	1
+#define RR_NS	2
+#define RR_CNAME 5
+#define RR_SOA	6
+#define RR_MX	15
+#define RR_TXT	16
+#define RR_AAAA	28
+#define RR_SRV	33
+
+//Record Classes
+#define CL_IN	1 //Internet
+#define CL_CS	2 //CSNET (Onsolete)
+#define CL_CH	3 //CHAOS
+#define CL_HS	4 //Hesiod
+
+//OPCODES
+#define OP_Q	0 //Query
+#define OP_IQ	1 //Inverse Query
+#define OP_STAT	2 //Status request
+
+//Responsecode
+#define RCODE_NOERR	0
+#define RCODE_FORMAT	1
+#define RCODE_SERVFAIL	2
+#define RCODE_NAMEERR	3
+#define RCODE_NI	4 //Not implemented
+#define RCODE_REFUSED	5
+
+#define FLIP_BYTES(u)	(((0x00FF & u) << 8) | ((0xFF00 & u) >> 8))
+
 /**
  * Data is COPIED
  * */
 struct dns_header;
 
 /**
- * Data is REFERENCED
+ * QNAME is REFERENCED
  * */
 struct dns_question;
 
 /**
- * Data is REFERENCED
+ * NAME is REFERENCED
  * */
 struct dns_answer;
 
@@ -55,16 +86,16 @@ struct dns_header {
 
 struct dns_question {
 	const char* qname;
-	const uint16_t* qtype;
-	uint16_t* qclass;
+	uint16_t qtype;
+	uint16_t qclass;
 };
 
 struct dns_answer {
 	char* name; //in qname format
-	uint16_t* type;
-	uint16_t* class;
-	uint32_t* ttl;
-	uint16_t* rdlength;
+	uint16_t type;
+	uint16_t class;
+	uint32_t ttl;
+	uint16_t rdlength;
 	char* rdata;
 };
 
@@ -78,10 +109,24 @@ struct dns_message {
 	struct dns_answer* answer;
 };
 
+int dns_construct_header (
+		struct	dns_header* _header,
+		char*	_buffer,
+		int	_bufflen
+		);
+
+/**
+ * Frees all malloced memory
+ * */
+int dns_destroy_struct ( struct dns_message* _msg );
+
 /**
  * Parse the packet in _buffer and populate the dns_message struct
  * Struct may still be written to on failure but contents are invalid
  * returns: 0 on success, !=0 on failure
+ *
+ * ONLY WRITES QUESTION SECTION. ALL OTHER ARE IGNORED
+ *
  * */
 int dns_parse_packet ( char* _buffer, int _bufflen, struct dns_message* _msg );
 
diff --git a/src/main.c b/src/main.c
index 1e96f33..75574a0 100644
--- a/src/main.c
+++ b/src/main.c
@@ -158,17 +158,11 @@ int main(	int argc,
 
 		LOGPRINTF(_LOG_DEBUG, "UDP Packet size %i", ret);
 
-		if ( (handle_connection (sock_server,
-					&sock_client_addr,
-					sock_client_addr_len,
-					recv_buffer,
-					ret ) ) ) {
-			LOGPRINTF(_LOG_ERROR, "handle_connection() failed");
-			return errno;
-		}
-		else {
-			LOGPRINTF ( _LOG_DEBUG, "forked " );
-		}
+		handle_connection ( sock_server,
+				&sock_client_addr,
+				sock_client_addr_len,
+				recv_buffer,
+				ret );
 	}
 
 	close( sock_server );
@@ -182,35 +176,28 @@ int handle_connection (	int _socket,
 			char* _buffer,
 			int _bufflen )
 {
-	pid_t pid = fork();
-
-	if ( pid > 0)
-		return 0;
-	else if ( pid < 0 )
-		return errno;
-
-	//Change signal handler. signal_term shuts down entire socket on TERM
-	signal ( SIGTERM, signal_term_child );
-
-	//echo request
-	sendto (_socket, _buffer, _bufflen, 0, (struct sockaddr*) sockaddr_client, sockaddr_client_len);
-
-	//Testing recieved question
-	//TODO remove
-
 	struct dns_message msg;
 
 	if (dns_parse_packet (_buffer, _bufflen, &msg) ) {
 		LOGPRINTF (_LOG_DEBUG, "Malformed packet recieved. parsing failed");
-		close ( _socket );
-		exit( 1 );
+		return 1;
 	}
 
-	char out[128];
-	qname_to_fqdn( msg.question[0].qname, 100, out, 128);
-	printf("%s\n", out);
+	if(msg.question_count > 0) {
+		char out[128];
+		qname_to_fqdn( msg.question[0].qname, 100, out, 128);
+		printf("%s %i\n", out, msg.question[0].qtype);
+	}
+	
+	dns_destroy_struct ( &msg );
+
+	//Always return NXDOMAIN
+	struct dns_header head = {msg.header.id,1,OP_Q,0,0,0,0,0,RCODE_NAMEERR,0,0,0,0};
+	char ret[20];
+	int retlen = dns_construct_header ( &head, ret, 20 );
+	sendto (_socket, ret, retlen, 0, (struct sockaddr*) sockaddr_client, sockaddr_client_len);
 	
-	exit ( 0 );
+	return 0;
 }
 
 void signal_term ( ) {