2 files changed, 52 insertions, 33 deletions

diff --git a/toxcore/DHT.c b/toxcore/DHT.c
index c286567f..12f38f18 100644
--- a/toxcore/DHT.c
+++ b/toxcore/DHT.c
@@ -120,6 +120,7 @@ static int client_in_list(Client_data *list, uint32_t length, uint8_t *client_id
     uint32_t i;
 
     for (i = 0; i < length; i++)
+
         /* Dead nodes are considered dead (not in the list)*/
         if (!is_timeout(list[i].assoc4.timestamp, KILL_NODE_TIMEOUT) ||
                 !is_timeout(list[i].assoc6.timestamp, KILL_NODE_TIMEOUT))
@@ -278,6 +279,7 @@ static void get_close_nodes_inner(DHT *dht, uint8_t *client_id, Node_format *nod
             continue;
 
         IPPTsPng *ipptp = NULL;
+
         if (sa_family == AF_INET)
             ipptp = &client->assoc4;
         else
@@ -570,25 +572,26 @@ static void returnedip_ports(DHT *dht, IP_Port ip_port, uint8_t *client_id, uint
     }
 }
 
+/* checks if ip/port or ping_id are already in the list to get nodes
+ * if both are set, both must match, otherwise the set must match
+ *
+ *  returns 0 if neither is set or no match was found
+ *  returns the (index + 1) of the match if one was found
+ */
 static int is_gettingnodes(DHT *dht, IP_Port ip_port, uint64_t ping_id)
 {
-    uint32_t i;
-    uint8_t pinging;
+    uint8_t ip_valid = ip_isset(&ip_port.ip);
 
-    for (i = 0; i < LSEND_NODES_ARRAY; ++i ) {
-        if (!is_timeout(dht->send_nodes[i].timestamp, PING_TIMEOUT)) {
-            pinging = 0;
-
-            if (ping_id != 0 && dht->send_nodes[i].id == ping_id)
-                ++pinging;
+    if (!ip_valid && !ping_id)
+        return 0;
 
-            if (ip_isset(&ip_port.ip) && ipport_equal(&dht->send_nodes[i].ip_port, &ip_port))
-                ++pinging;
+    uint32_t i;
 
-            if (pinging == (ping_id != 0) + ip_isset(&ip_port.ip))
-                return 1;
-        }
-    }
+    for (i = 0; i < LSEND_NODES_ARRAY; i++)
+        if (!is_timeout(dht->send_nodes[i].timestamp, PING_TIMEOUT))
+            if (!ping_id || (dht->send_nodes[i].id == ping_id))
+                if (!ip_valid || ipport_equal(&dht->send_nodes[i].ip_port, &ip_port))
+                    return i + 1;
 
     return 0;
 }
@@ -835,7 +838,9 @@ static int handle_sendnodes(void *object, IP_Port source, uint8_t *packet, uint3
 
     memcpy(&ping_id, plain, sizeof(ping_id));
 
-    if (!is_gettingnodes(dht, source, ping_id))
+    int send_nodes_index = is_gettingnodes(dht, source, ping_id);
+
+    if (!send_nodes_index)
         return 1;
 
     Node4_format *nodes4_list = (Node4_format *)(plain + sizeof(ping_id));
@@ -858,7 +863,8 @@ static int handle_sendnodes(void *object, IP_Port source, uint8_t *packet, uint3
         num_nodes = num_nodes_ok;
     }
 
-    addto_lists(dht, source, packet + 1);
+    /* store the address the *request* was sent to */
+    addto_lists(dht, dht->send_nodes[send_nodes_index - 1].ip_port, packet + 1);
 
     for (i = 0; i < num_nodes; ++i)  {
         send_ping_request(dht->ping, nodes_list[i].ip_port, nodes_list[i].client_id);
@@ -1193,12 +1199,15 @@ static int friend_iplist(DHT *dht, IP_Port *ip_portlist, uint16_t friend_num)
 
 #ifdef FRIEND_IPLIST_PAD
     memcpy(ip_portlist, ipv6s, num_ipv6s * sizeof(IP_Port));
+
     if (num_ipv6s == MAX_FRIEND_CLIENTS)
         return MAX_FRIEND_CLIENTS;
 
     int num_ipv4s_used = MAX_FRIEND_CLIENTS - num_ipv6s;
+
     if (num_ipv4s_used > num_ipv4s)
         num_ipv4s_used = num_ipv4s;
+
     memcpy(&ip_portlist[num_ipv6s], ipv4s, num_ipv4s_used * sizeof(IP_Port));
     return num_ipv6s + num_ipv4s_used;
 
@@ -1808,6 +1817,7 @@ static int dht_load_state_callback(void *outer, uint8_t *data, uint32_t length,
             break;
 
 #ifdef DEBUG
+
         default:
             fprintf(stderr, "Load state (DHT): contains unrecognized part (len %u, type %u)\n",
                     length, type);
diff --git a/toxcore/ping.c b/toxcore/ping.c
index c7b829b4..bd754c53 100644
--- a/toxcore/ping.c
+++ b/toxcore/ping.c
@@ -3,7 +3,7 @@
  *
  * This file is donated to the Tox Project.
  * Copyright 2013  plutooo
- * 
+ *
  *  Copyright (C) 2013 Tox project All Rights Reserved.
  *
  *  This file is part of Tox.
@@ -20,7 +20,7 @@
  *
  *  You should have received a copy of the GNU General Public License
  *  along with Tox.  If not, see <http://www.gnu.org/licenses/>.
- * 
+ *
  */
 
 #ifdef HAVE_CONFIG_H
@@ -108,28 +108,35 @@ static uint64_t add_ping(PING *ping, IP_Port ipp)  // O(n)
     return ping->pings[p].id;
 }
 
-static bool is_pinging(PING *ping, IP_Port ipp, uint64_t ping_id)    // O(n) TODO: Replace this with something else.
+/* checks if ip/port or ping_id are already in the list to ping
+ * if both are set, both must match, otherwise the set must match
+ *
+ *  returns 0 if neither is set or no match was found
+ *  returns the (index + 1) of the match if one was found
+ */
+static int is_pinging(PING *ping, IP_Port ipp, uint64_t ping_id)
 {
+    // O(n) TODO: Replace this with something else.
 
-    /* shouldn't that be an OR ? */
-    if (!ip_isset(&ipp.ip) && ping_id == 0)
-        return false;
+    /* at least one MUST be set */
+    uint8_t ip_valid = ip_isset(&ipp.ip);
 
-    size_t i, id;
+    if (!ip_valid && !ping_id)
+        return 0;
+
+    size_t i;
 
     remove_timeouts(ping);
 
     for (i = 0; i < ping->num_pings; i++) {
-        id = (ping->pos_pings + i) % PING_NUM_MAX;
+        size_t id = (ping->pos_pings + i) % PING_NUM_MAX;
 
-        /* ping_id = 0 means match any id. */
-        if ((!ip_isset(&ipp.ip) || ipport_equal(&ping->pings[id].ip_port, &ipp)) &&
-                (ping->pings[id].id == ping_id || ping_id == 0)) {
-            return true;
-        }
+        if (!ping_id || (ping->pings[id].id == ping_id))
+            if (!ip_valid || ipport_equal(&ping->pings[id].ip_port, &ipp))
+                return id + 1;
     }
 
-    return false;
+    return 0;
 }
 
 #define DHT_PING_SIZE (1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES + sizeof(uint64_t) + ENCRYPTION_PADDING)
@@ -246,11 +253,13 @@ static int handle_ping_response(void *_dht, IP_Port source, uint8_t *packet, uin
         return 1;
 
     /* Make sure ping_id is correct. */
-    if (!is_pinging(ping, source, ping_id))
+    int ping_index = is_pinging(ping, source, ping_id);
+
+    if (!ping_index)
         return 1;
 
-    // Associate source ip with client_id
-    addto_lists(dht, source, packet + 1);
+    /* Associate client_id with the ip the request was sent to */
+    addto_lists(dht, ping->pings[ping_index - 1].ip_port, packet + 1);
     return 0;
 }