5 files changed, 84 insertions, 80 deletions

diff --git a/core/DHT.c b/core/DHT.c
index 5eae147c..057ee9ce 100644
--- a/core/DHT.c
+++ b/core/DHT.c
@@ -932,49 +932,30 @@ static int send_NATping(uint8_t * public_key, uint64_t ping_id, uint8_t type)
 }
 
 /* Handle a received ping request for */
-static int handle_NATping(IP_Port source, uint8_t * packet, uint32_t length) 
+static int handle_NATping(IP_Port source, uint8_t * source_pubkey, uint8_t * packet, uint32_t length) 
 {
-    if (length < crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + ENCRYPTION_PADDING 
-            || length > MAX_DATA_SIZE + ENCRYPTION_PADDING)
-        return 1;
-
-    /* check if request is for us. */
-    if (id_equal(packet + 1, self_public_key)) {
-        uint8_t public_key[crypto_box_PUBLICKEYBYTES];
-        uint8_t data[MAX_DATA_SIZE];
-
-        int len = handle_request(public_key, data, packet, length);
-        if (len != sizeof(uint64_t) + 1)
-            return 1;
-
-        uint64_t ping_id;
-        memcpy(&ping_id, data + 1, sizeof(uint64_t));
+    uint64_t ping_id;
+    memcpy(&ping_id, packet + 1, sizeof(uint64_t));
 
-        int friendnumber = friend_number(public_key);
-        if (friendnumber == -1)
-            return 1;
+    int friendnumber = friend_number(source_pubkey);
+    if (friendnumber == -1)
+        return 1;
 
-        Friend * friend = &friends_list[friendnumber];
+    Friend * friend = &friends_list[friendnumber];
 
-        if (data[0] == 0) {
-            /* 1 is reply */
-            send_NATping(public_key, ping_id, 1);
-            friend->recvNATping_timestamp = unix_time();
+    if (packet[0] == 0) {
+        /* 1 is reply */
+        send_NATping(source_pubkey, ping_id, 1);
+        friend->recvNATping_timestamp = unix_time();
+        return 0;
+    } else if (packet[0] == 1) {
+        if (friend->NATping_id == ping_id) {
+            friend->NATping_id = ((uint64_t)random_int() << 32) + random_int();
+            friend->hole_punching = 1;
             return 0;
-        } else if (data[0] == 1) {
-            if (friend->NATping_id == ping_id) {
-                friend->NATping_id = ((uint64_t)random_int() << 32) + random_int();
-                friend->hole_punching = 1;
-                return 0;
-            }
         }
-        return 1;
     }
-
-    /* if request is not for us, try routing it. */
-    route_packet(packet + 1, packet, length);
-
-    return 0;
+    return 1;
 }
 
 /* Get the most common ip in the ip_portlist
@@ -1082,7 +1063,7 @@ void DHT_init(void)
     networking_registerhandler(1, &handle_ping_response);
     networking_registerhandler(2, &handle_getnodes);
     networking_registerhandler(3, &handle_sendnodes);
-    networking_registerhandler(254, &handle_NATping);
+    cryptopacket_registerhandler(254, &handle_NATping);
 }
 
 void doDHT(void)
diff --git a/core/friend_requests.c b/core/friend_requests.c
index 8276db29..ae19ebdd 100644
--- a/core/friend_requests.c
+++ b/core/friend_requests.c
@@ -123,39 +123,23 @@ static int request_received(uint8_t * client_id)
 }
 
 
-static int friendreq_handlepacket(IP_Port source, uint8_t * packet, uint32_t length)
+static int friendreq_handlepacket(IP_Port source, uint8_t * source_pubkey, uint8_t * packet, uint32_t length) 
 {
-    if (packet[0] == 32) {
-        if (length <= crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + 1 + ENCRYPTION_PADDING ||
-            length > MAX_DATA_SIZE + ENCRYPTION_PADDING)
-            return 1;
-        if (memcmp(packet + 1, self_public_key, crypto_box_PUBLICKEYBYTES) == 0) {// check if request is for us.
-            if (handle_friendrequest_isset == 0)
-                return 1;
-
-            uint8_t public_key[crypto_box_PUBLICKEYBYTES];
-            uint8_t data[MAX_DATA_SIZE];
-            int len = handle_request(public_key, data, packet, length);
-            if (len == -1)
-                return 1;
-            if (len <= sizeof(nospam))
-                return 1;
-            if (request_received(public_key))
-                return 1;
-            if (memcmp(data, &nospam, sizeof(nospam)) != 0)
-                return 1;
-
-            addto_receivedlist(public_key);
-            (*handle_friendrequest)(public_key, data + 4, len - 4, handle_friendrequest_userdata);
-        } else { /* if request is not for us, try routing it. */
-            if(route_packet(packet + 1, packet, length) == length)
-                return 0;
-        }
-    }
-    return 1;
+    if (handle_friendrequest_isset == 0)
+        return 1;
+    if (length <= sizeof(nospam))
+        return 1;
+    if (request_received(source_pubkey))
+        return 1;
+    if (memcmp(packet, &nospam, sizeof(nospam)) != 0)
+        return 1;
+
+    addto_receivedlist(source_pubkey);
+    (*handle_friendrequest)(source_pubkey, packet + 4, length - 4, handle_friendrequest_userdata);
+    return 0;
 }
 
 void friendreq_init(void)
 {
-    networking_registerhandler(32, &friendreq_handlepacket);
+    cryptopacket_registerhandler(32, &friendreq_handlepacket);
 }
diff --git a/core/net_crypto.c b/core/net_crypto.c
index ab18dd63..376708ab 100644
--- a/core/net_crypto.c
+++ b/core/net_crypto.c
@@ -221,15 +221,18 @@ int write_cryptpacket(int crypt_connection_id, uint8_t *data, uint32_t length)
    returns the length of the created packet on success */
 int create_request(uint8_t *packet, uint8_t *public_key, uint8_t *data, uint32_t length, uint8_t request_id)
 {
-    if (MAX_DATA_SIZE < length + 1 + crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + ENCRYPTION_PADDING)
+    if (MAX_DATA_SIZE < length + 1 + crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + 1 + ENCRYPTION_PADDING)
         return -1;
     uint8_t nonce[crypto_box_NONCEBYTES];
+    uint8_t temp[MAX_DATA_SIZE];
+    memcpy(temp + 1, data, length);
+    temp[0] = request_id;
     random_nonce(nonce);
-    int len = encrypt_data(public_key, self_secret_key, nonce, data, length,
+    int len = encrypt_data(public_key, self_secret_key, nonce, temp, length,
                            1 + crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + packet);
     if (len == -1)
         return -1;
-    packet[0] = request_id;
+    packet[0] = 32;
     memcpy(packet + 1, public_key, crypto_box_PUBLICKEYBYTES);
     memcpy(packet + 1 + crypto_box_PUBLICKEYBYTES, self_public_key, crypto_box_PUBLICKEYBYTES);
     memcpy(packet + 1 + crypto_box_PUBLICKEYBYTES * 2, nonce, crypto_box_NONCEBYTES);
@@ -241,7 +244,7 @@ int create_request(uint8_t *packet, uint8_t *public_key, uint8_t *data, uint32_t
    in data if a friend or ping request was sent to us and returns the length of the data.
    packet is the request packet and length is its length
    return -1 if not valid request. */
-int handle_request(uint8_t *public_key, uint8_t *data, uint8_t *packet, uint16_t length)
+static int handle_request(uint8_t *public_key, uint8_t *data, uint8_t *request_id, uint8_t *packet, uint16_t length)
 {
 
     if (length > crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + 1 + ENCRYPTION_PADDING &&
@@ -249,16 +252,51 @@ int handle_request(uint8_t *public_key, uint8_t *data, uint8_t *packet, uint16_t
         memcmp(packet + 1, self_public_key, crypto_box_PUBLICKEYBYTES) == 0) {
         memcpy(public_key, packet + 1 + crypto_box_PUBLICKEYBYTES, crypto_box_PUBLICKEYBYTES);
         uint8_t nonce[crypto_box_NONCEBYTES];
+        uint8_t temp[MAX_DATA_SIZE];
         memcpy(nonce, packet + 1 + crypto_box_PUBLICKEYBYTES * 2, crypto_box_NONCEBYTES);
         int len1 = decrypt_data(public_key, self_secret_key, nonce, packet + 1 + crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES,
-                                length - (crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + 1), data);
-        if(len1 == -1)
+                                length - (crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + 1), temp);
+        if(len1 == -1 || len1 == 0)
             return -1;
+        request_id[0] = temp[0];
+        --len1;
+        memcpy(data, temp + 1, len1);
         return len1;
     } else
         return -1;
 }
 
+static cryptopacket_handler_callback cryptopackethandlers[256] = {0};
+
+void cryptopacket_registerhandler(uint8_t byte, cryptopacket_handler_callback cb)
+{
+    cryptopackethandlers[byte] = cb;
+}
+
+static int cryptopacket_handle(IP_Port source, uint8_t * packet, uint32_t length)
+{
+    if (packet[0] == 32) {
+        if (length <= crypto_box_PUBLICKEYBYTES * 2 + crypto_box_NONCEBYTES + 1 + ENCRYPTION_PADDING ||
+            length > MAX_DATA_SIZE + ENCRYPTION_PADDING)
+            return 1;
+        if (memcmp(packet + 1, self_public_key, crypto_box_PUBLICKEYBYTES) == 0) {// check if request is for us.
+            uint8_t public_key[crypto_box_PUBLICKEYBYTES];
+            uint8_t data[MAX_DATA_SIZE];
+            uint8_t number;
+            int len = handle_request(public_key, data, &number, packet, length);
+            if (len == -1 || len == 0)
+                return 1;
+            if (!cryptopackethandlers[number]) return 1;
+            cryptopackethandlers[number](source, public_key, data, len - 1);
+            
+        } else { /* if request is not for us, try routing it. */
+            if(route_packet(packet + 1, packet, length) == length)
+                return 0;
+        }
+    }
+    return 1;
+}
+
 /* Send a crypto handshake packet containing an encrypted secret nonce and session public key
    to peer with connection_id and public_key
    the packet is encrypted with a random nonce which is sent in plain text with the packet */
@@ -579,6 +617,7 @@ void initNetCrypto(void)
 {
     memset(crypto_connections, 0 ,sizeof(crypto_connections));
     memset(incoming_connections, -1 ,sizeof(incoming_connections));
+    networking_registerhandler(32, &cryptopacket_handle);
     uint32_t i;
     for (i = 0; i < MAX_CRYPTO_CONNECTIONS; ++i)
         crypto_connections[i].number = ~0;
diff --git a/core/net_crypto.h b/core/net_crypto.h
index 135e099d..570b9373 100644
--- a/core/net_crypto.h
+++ b/core/net_crypto.h
@@ -25,6 +25,7 @@
 #define NET_CRYPTO_H
 
 #include "Lossless_UDP.h"
+#include "DHT.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -88,11 +89,10 @@ int write_cryptpacket(int crypt_connection_id, uint8_t *data, uint32_t length);
     returns the length of the created packet on success */
 int create_request(uint8_t *packet, uint8_t * public_key, uint8_t *data, uint32_t length, uint8_t request_id);
 
-/* puts the senders public key in the request in public_key, the data from the request
-    in data if a friend or ping request was sent to us and returns the length of the data.
-    packet is the request packet and length is its length
-    return -1 if not valid request. */
-int handle_request(uint8_t *public_key, uint8_t *data, uint8_t *packet, uint16_t length);
+
+typedef int (*cryptopacket_handler_callback)(IP_Port ip_port, uint8_t * source_pubkey, uint8_t *data, uint32_t len);
+/* Function to call when request beginning with byte is received */
+void cryptopacket_registerhandler(uint8_t byte, cryptopacket_handler_callback cb);
 
 /* Start a secure connection with other peer who has public_key and ip_port
     returns -1 if failure
diff --git a/testing/DHT_test.c b/testing/DHT_test.c
index 350093fd..c5b367cf 100644
--- a/testing/DHT_test.c
+++ b/testing/DHT_test.c
@@ -53,10 +53,10 @@ void print_clientlist()
     uint32_t i, j;
     IP_Port p_ip;
     printf("___________________CLOSE________________________________\n");
-    for(i = 0; i < 4; i++) {
+    for(i = 0; i < 32; i++) {
         printf("ClientID: ");
         for(j = 0; j < 32; j++) {
-            printf("%hhX", close_clientlist[i].client_id[j]);
+            printf("%02hhX", close_clientlist[i].client_id[j]);
         }
         p_ip = close_clientlist[i].ip_port;
         printf("\nIP: %u.%u.%u.%u Port: %u",p_ip.ip.c[0],p_ip.ip.c[1],p_ip.ip.c[2],p_ip.ip.c[3],ntohs(p_ip.port));