4 files changed, 63 insertions, 59 deletions

diff --git a/core/DHT.c b/core/DHT.c
index be3ac186..b5224b8f 100644
--- a/core/DHT.c
+++ b/core/DHT.c
@@ -352,7 +352,7 @@ static int replace_good(   Client_data *   list,
 /* Attempt to add client with ip_port and client_id to the friends client list 
  * and close_clientlist 
  */
-static void addto_lists(IP_Port ip_port, uint8_t * client_id)
+void addto_lists(IP_Port ip_port, uint8_t * client_id)
 {
     uint32_t i;
 
@@ -554,62 +554,6 @@ static int sendnodes(IP_Port ip_port, uint8_t * public_key, uint8_t * client_id,
     return sendpacket(ip_port, data, 1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES + len);
 }
 
-/* Packet handling functions, one to handle each types of packets we receive
- * Returns 0 if handled correctly, 1 if packet is bad.
- */
-static int handle_pingreq(uint8_t * packet, uint32_t length, IP_Port source)
-{
-    uint64_t ping_id;
-    if(length != 1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES + sizeof(ping_id) + ENCRYPTION_PADDING)
-        return 1;
-
-    /* check if packet is from ourself. */
-    if(id_equal(packet + 1, self_public_key))
-        return 1;
-
-    int len = decrypt_data( packet + 1, 
-                            self_secret_key, 
-                            packet + 1 + CLIENT_ID_SIZE,
-                            packet + 1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES,
-                            sizeof(ping_id) + ENCRYPTION_PADDING, 
-                            (uint8_t *)&ping_id );
-
-    if(len != sizeof(ping_id))
-        return 1;
-
-    send_ping_response(source, (clientid_t*) (packet + 1), ping_id);
-    send_ping_request(source, (clientid_t*) (packet + 1)); /* TODO: make this smarter? */
-
-    return 0;
-}
-
-static int handle_pingres(uint8_t * packet, uint32_t length, IP_Port source)
-{
-    uint64_t ping_id;
-    if(length != 1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES + sizeof(ping_id) + ENCRYPTION_PADDING)
-        return 1;
-
-    /* check if packet is from ourself. */
-    if(id_equal(packet + 1, self_public_key))
-        return 1;
-
-    int len = decrypt_data( packet + 1, 
-                            self_secret_key, 
-                            packet + 1 + CLIENT_ID_SIZE,
-                            packet + 1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES,
-                            sizeof(ping_id) + ENCRYPTION_PADDING, 
-                            (uint8_t *)&ping_id );
-
-    if(len != sizeof(ping_id))
-        return 1;
-
-    if(is_pinging(source, ping_id)) {
-        addto_lists(source, packet + 1);
-        return 0;
-    }
-    return 1;
-}
-
 static int handle_getnodes(uint8_t * packet, uint32_t length, IP_Port source)
 {
     uint64_t ping_id;
@@ -1134,10 +1078,10 @@ int DHT_handlepacket(uint8_t * packet, uint32_t length, IP_Port source)
 {
     switch (packet[0]) {
     case 0:
-        return handle_pingreq(packet, length, source);
+        return handle_ping_request(packet, length, source);
 
     case 1:
-        return handle_pingres(packet, length, source);
+        return handle_ping_response(packet, length, source);
 
     case 2:
         return handle_getnodes(packet, length, source);
diff --git a/core/DHT.h b/core/DHT.h
index 0edaebf3..2308abd8 100644
--- a/core/DHT.h
+++ b/core/DHT.h
@@ -104,6 +104,8 @@ int DHT_load(uint8_t *data, uint32_t size);
     returns 1 if we are */
 int DHT_isconnected();
 
+void addto_lists(IP_Port ip_port, uint8_t * client_id);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/core/ping.c b/core/ping.c
index a687f2fb..6a1fbb7e 100644
--- a/core/ping.c
+++ b/core/ping.c
@@ -162,3 +162,59 @@ int send_ping_response(IP_Port ipp, clientid_t* client_id, uint64_t ping_id)
 
     return sendpacket(ipp, (uint8_t*) &pk, sizeof(pk));
 }
+
+int handle_ping_request(uint8_t* packet, uint32_t length, IP_Port source)
+{
+    pingreq_t* p = (pingreq_t*) packet;
+    int        rc;
+    uint64_t   ping_id;
+
+    if (length != sizeof(pingreq_t) || id_eq(&p->client_id, self_id))
+        return 1;
+
+    // Decrypt ping_id
+    rc = decrypt_data((uint8_t*) &p->client_id,
+                      self_secret_key,
+                      (uint8_t*) &p->nonce,
+                      (uint8_t*) &p->ping_id,
+                      sizeof(ping_id) + ENCRYPTION_PADDING,
+                      (uint8_t*) &ping_id);
+
+    if (rc != sizeof(ping_id))
+        return 1;
+
+    // Send response
+    send_ping_response(source, &p->client_id, ping_id);
+    send_ping_request(source, &p->client_id); // Make this smarter?
+
+    return 0;
+}
+
+int handle_ping_response(uint8_t* packet, uint32_t length, IP_Port source)
+{
+    pingres_t* p = (pingres_t*) packet;
+    int       rc;
+    uint64_t  ping_id;
+
+    if (length != sizeof(pingres_t) || id_eq(&p->client_id, self_id))
+        return 1;
+
+    // Decrypt ping_id
+    rc = decrypt_data((uint8_t*) &p->client_id,
+                      self_secret_key,
+                      (uint8_t*) &p->nonce,
+                      (uint8_t*) &p->ping_id,
+                      sizeof(ping_id) + ENCRYPTION_PADDING,
+                      (uint8_t*) &ping_id);
+
+    if (rc != sizeof(ping_id))
+        return 1;
+
+    // Make sure ping_id is correct
+    if(!is_pinging(source, ping_id))
+        return 1;
+
+    // Associate source ip with client_id
+    addto_lists(source, (uint8_t*) &p->client_id);
+    return 0;
+}
diff --git a/core/ping.h b/core/ping.h
index f2770a00..2cab7d59 100644
--- a/core/ping.h
+++ b/core/ping.h
@@ -12,3 +12,5 @@ uint64_t add_ping(IP_Port ipp);
 bool is_pinging(IP_Port ipp, uint64_t ping_id);
 int send_ping_request(IP_Port ipp, clientid_t* client_id);
 int send_ping_response(IP_Port ipp, clientid_t* client_id, uint64_t ping_id);
+int handle_ping_request(uint8_t* packet, uint32_t length, IP_Port source);
+int handle_ping_response(uint8_t* packet, uint32_t length, IP_Port source);