Add find_node query

5 files changed, 159 insertions, 22 deletions

diff --git a/src/proto.c b/src/proto.c
index fa8ddc5..8f4f737 100644
--- a/src/proto.c
+++ b/src/proto.c
@@ -107,7 +107,7 @@ uint8_t rand_byte() {
 	return val;
 }
 
-int create_ping(char* buff, size_t* buff_len, struct nodeid* self, struct nodeid* target, uint16_t tid) {
+int write_ping(char* buff, size_t* buff_len, struct nodeid* self, struct nodeid* target, uint16_t tid) {
 	char* buff_end = buff + *buff_len;
 
 	int rc = snprintf(buff, buff_end - buff, "d1:ad2:id20:");
@@ -163,9 +163,10 @@ int send_ping(struct dht* dht, struct nodeid* expected, time_t now, bool node_is
 	for(uint8_t *target_byte = (uint8_t*)&target; target_byte < ((uint8_t*)&target)+sizeof(target); target_byte++) {
 		*target_byte = rand_byte();
 	}
+	target = dht->self;
 
 	message->payload_len = 128;
-	int rc = create_ping(message->payload, &message->payload_len, &dht->self, &target, reqId);
+	int rc = write_ping(message->payload, &message->payload_len, &dht->self, &target, reqId);
 	if(rc != 0) {
 		return rc;
 	}
@@ -179,6 +180,7 @@ PROCESS_TIMEOUT(getclient_timeout) {
 	size_t reqId = (typeof(dht->requestdata[0])*)((void*)cont - offsetof(typeof(dht->requestdata[0]), cont)) - dht->requestdata;
 
 	if(cont->ping.attempt >= 2) {
+		dbg("Timing out request %d after %d attempts", reqId, cont->ping.attempt);
 		if(cont->ping.is_new)
 			return 0;
 
@@ -204,7 +206,7 @@ PROCESS_TIMEOUT(getclient_timeout) {
 	}
 
 	message->payload_len = 128;
-	int rc = create_ping(message->payload, &message->payload_len, &dht->self, &target, reqId);
+	int rc = write_ping(message->payload, &message->payload_len, &dht->self, &target, reqId);
 	if(rc != 0) {
 		fatal("Can't create ping");
 	}
@@ -393,18 +395,6 @@ void proto_begin(struct dht* dht, time_t now, struct message** output, const str
 	}
 
 	for(struct addrinfo* cur = res; cur != NULL; cur = cur->ai_next) {
-		char buff[128];
-		inet_ntop(cur->ai_family, &((struct sockaddr_in*)cur->ai_addr)->sin_addr, buff, 128);
-		dbg("IP %s", buff);
-
-		/* size_t i = 0; */
-		/* int rc = snprintf(buff+i, 128-i, "d1:ad2:id20:"); */
-		/* i += rc; */
-		/* memcpy(buff+i, &self, sizeof(struct nodeid)); */
-		/* i += sizeof(struct nodeid); */
-		/* rc = snprintf(buff+i, 128-i, "e1:q4:ping1:t2:ab1:y1:qe"); */
-		/* i += rc; */
-
 		send_ping(dht, NULL, now, true, cur->ai_addr, cur->ai_addrlen, &msgbuff);
 	}
 
diff --git a/src/proto.h b/src/proto.h
index 3ca7fc7..76718de 100644
--- a/src/proto.h
+++ b/src/proto.h
@@ -9,8 +9,8 @@
 #define MAX_DISC 32
 #define MAX_INFLIGHT 32
 
-#define PROTO_UNCTM 900
-#define PROTO_TMOUT 60
+#define PROTO_UNCTM 60
+#define PROTO_TMOUT 5
 
 struct ping {
 	struct nodeid remote_id;
diff --git a/src/query.c b/src/query.c
index 806c691..6012540 100644
--- a/src/query.c
+++ b/src/query.c
@@ -4,6 +4,7 @@
 
 #include <string.h>
 #include <errno.h>
+#include <assert.h>
 
 int handle_request(struct nodeid* self, const char* method, const char* packet, size_t packet_len, char** response, size_t response_len) {
 	if(strcmp(method, "ping") == 0) {
@@ -71,11 +72,107 @@ int handle_request(struct nodeid* self, const char* method, const char* packet,
 		*response += rc;
 		memcpy(*response, self, sizeof(struct nodeid));
 		*response += sizeof(struct nodeid);
-		rc = snprintf(*response, end-*response, "ee");
+		rc = snprintf(*response, end-*response, "e");
 		if(rc < 0)
 			return QUERY_EBADQ;
 		*response += rc;
-	}else if(strcmp(method, "get_peers") == 0) {
+	} else if(strcmp(method, "find_node") == 0) {
+		struct bcursor bcursor;
+		struct benc_node stream[256];
+		bcur_open(&bcursor, packet, packet+packet_len, stream, 256);
+
+		if(bcursor.readhead->type != BNT_DICT) {
+			err("Bad query: Packet is not a dict");
+			return QUERY_EBADQ;
+		}
+		if(bcur_next(&bcursor, 1) < 0) {
+			err("Bad query: No token after outer dict start");
+			return QUERY_EBADQ;
+		}
+
+		if(bcur_find_key(&bcursor, (const enum benc_nodetype[]){BNT_STRING}, (const char*[]){"a"}, (const size_t[]){1}, 1) != 0) {
+			err("Bad query: No arguments to request");
+			return QUERY_EBADQ;
+		}
+		bcur_next(&bcursor, 1);
+		if(bcursor.readhead->type != BNT_DICT) {
+			err("Bad query: Wrong value type for request");
+			return QUERY_EBADQ;
+		}
+		// Skip the dict element
+		bcur_next(&bcursor, 1);
+		bool target_set = false;
+		struct nodeid target;
+		while(bcursor.readhead->type != BNT_END) {
+			switch(bcur_find_key(&bcursor, (const enum benc_nodetype[]){BNT_STRING}, (const char*[]){"target"}, (const size_t[]){6}, 1)) {
+				case 0:
+					// Skip the key
+					bcur_next(&bcursor, 1);
+
+					if(bcursor.readhead->type != BNT_STRING) {
+						err("Bad query: Wrong value type for id");
+						return QUERY_EBADQ;
+					}
+
+					if(bcursor.readhead->size != 20) {
+						err("Bad query: Incorrect target length");
+						return QUERY_EBADQ;
+					}
+
+					target_set = true;
+					memcpy(&target, bcursor.readhead->loc, 20);
+
+					// Skip the value
+					bcur_next(&bcursor, 1);
+					break;
+			}
+		}
+
+		if(!target_set) {
+			err("Target argument not set");
+			return QUERY_EBADQ;
+		}
+
+		// Actually do the handling
+		struct entry* closest[8];
+		int found = routing_closest(&target, 8, closest);
+
+		char* end = (*response) + response_len;
+
+		int rc = snprintf(*response, end-*response, "d2:id20:");
+		if(rc < 0)
+			return QUERY_EBADQ;
+		*response += rc;
+		assert(*response < end);
+
+		memcpy(*response, self, sizeof(struct nodeid));
+		*response += sizeof(struct nodeid);
+		assert(*response < end);
+
+		rc = snprintf(*response, end-*response, "5:nodes%d:", found*26);
+		if(rc < 0)
+			return QUERY_EBADQ;
+		*response += rc;
+		assert(*response < end);
+
+		for(int i = 0; i < found; i++) {
+			memcpy(*response, &closest[i]->id, sizeof(struct nodeid));
+			*response += sizeof(struct nodeid); // 20
+			assert(*response < end);
+			memcpy(*response, &closest[i]->addr.ip, sizeof(uint32_t));
+			*response += sizeof(uint32_t); // 4
+			assert(*response < end);
+			memcpy(*response, &closest[i]->addr.port, sizeof(uint16_t));
+			*response += sizeof(uint16_t); // 2
+			assert(*response < end);
+		}
+
+		rc = snprintf(*response, end-*response, "e");
+		if(rc < 0)
+			return QUERY_EBADQ;
+		*response += rc;
+		assert(*response < end);
+	} else if(strcmp(method, "get_peers") == 0) {
 	} else {
 		return QUERY_EUNK;
 	}
diff --git a/test/proto.c b/test/proto.c
index ae42e36..277f68c 100644
--- a/test/proto.c
+++ b/test/proto.c
@@ -107,7 +107,7 @@ void test_ping() {
 	// We should have sent a response
 	TEST_ASSERT_EQUAL_PTR(message_cursor, outbuff+1);
 
-	TEST_ASSERT_EQUAL(48, outbuff[0].payload_len);
+	TEST_ASSERT_EQUAL(47, outbuff[0].payload_len);
 	TEST_ASSERT_EQUAL_CHAR_ARRAY("d1:t2:aa1:y1:r1:rd2:id20:BBBBBBBBBBBBBBBBBBBBee", outbuff[0].payload, 47);
 }
 
@@ -310,3 +310,53 @@ void test_ping_node_when_uncertain() {
 		TEST_ASSERT_GREATER_THAN(now, entry->expire);
 	}
 }
+
+void test_query_find_node() {
+	struct message outbuff[10] = {0};
+	time_t now = 0;
+
+	struct sockaddr_storage remote;
+	socklen_t remote_len;
+
+	struct dht dht;
+	dht.self = (struct nodeid){.inner={0x42424242, 0x42424242, 0x42424242, 0x42424242, 0x42424242}};
+	struct nodeid other = (struct nodeid){.inner={0x61616161, 0x61616161, 0x61616161, 0x61616161, 0x61616161}};
+
+	{
+		struct message* message_cursor = outbuff;
+		proto_begin(&dht, 0, &message_cursor, outbuff+2);
+		remote_len = outbuff[0].dest_len;
+		memcpy(&remote, &outbuff[0].dest, remote_len);
+	}
+	now += 10;
+
+	{
+		// The node responds
+		// We return no new nodes to stop any new pings from going out
+		char buff[] = "d1:y1:r1:t1:01:rd2:id20:aaaaaaaaaaaaaaaaaaaa5:nodes0:""ee";
+		struct message* message_cursor = outbuff;
+		proto_run(&dht, buff, sizeof(buff), (struct sockaddr_in*)&remote, remote_len, now, &message_cursor, outbuff+2);
+	}
+	struct entry* entry = routing_get(&other);
+	TEST_ASSERT_NOT_NULL(entry);
+
+	{
+		struct sockaddr_in other;
+		other.sin_family = AF_INET;
+		inet_pton(AF_INET, "255.255.255.255", &other.sin_addr.s_addr);
+		other.sin_port = htons(6881);
+
+		char buff[] = "d1:ad2:id20:abcdefghij01234567896:target20:aaaaaaaaaaaaaaaaaaaae1:q9:find_node1:t2:aa1:y1:qe";
+		struct message* message_cursor = outbuff;
+		proto_run(&dht, buff, sizeof(buff), &other, sizeof(other), 0, &message_cursor, outbuff+2);
+
+		// We should have sent a response
+		TEST_ASSERT_EQUAL_PTR(message_cursor, outbuff+1);
+
+		TEST_ASSERT_EQUAL(83, outbuff[0].payload_len);
+		TEST_ASSERT_EQUAL_CHAR_ARRAY("d1:t2:aa1:y1:r1:rd2:id20:BBBBBBBBBBBBBBBBBBBB5:nodes26:aaaaaaaaaaaaaaaaaaaa", outbuff[0].payload, 75);
+		TEST_ASSERT_EQUAL_MEMORY(&((struct sockaddr_in*)&remote)->sin_addr.s_addr, outbuff[0].payload+75, 4);
+		TEST_ASSERT_EQUAL_MEMORY(&((struct sockaddr_in*)&remote)->sin_port, outbuff[0].payload+79, 2);
+		TEST_ASSERT_EQUAL_CHAR_ARRAY("ee", outbuff[0].payload+81, 2);
+	}
+}
diff --git a/test/query.c b/test/query.c
index cccb639..5744fba 100644
--- a/test/query.c
+++ b/test/query.c
@@ -114,8 +114,8 @@ void test_ping() {
 	int rc = handle_request(&self, "ping", packet, packet_len, &response_cursor, response_end-response_cursor);
 
 	TEST_ASSERT_EQUAL(0, rc);
-	TEST_ASSERT_EQUAL(30, response_cursor - response);
-	TEST_ASSERT_EQUAL_CHAR_ARRAY("d2:id20:aaaaaaaaaaaaaaaaaaabee", response, 30);
+	TEST_ASSERT_EQUAL(29, response_cursor - response);
+	TEST_ASSERT_EQUAL_CHAR_ARRAY("d2:id20:aaaaaaaaaaaaaaaaaaabe", response, 30);
 }
 
 void test_bad_method() {