1 files changed, 40 insertions, 25 deletions

diff --git a/toxcore/DHT.c b/toxcore/DHT.c
index 9730771b..f5d93fc4 100644
--- a/toxcore/DHT.c
+++ b/toxcore/DHT.c
@@ -425,7 +425,9 @@ static void sort_list(Client_data *list, uint32_t length, uint8_t *comp_client_i
         list[i] = pairs[i].c2;
 }
 
-/* Replace the first good node that is further to the comp_client_id than that of the client_id in the list */
+/* Replace the first good node that is further to the comp_client_id than that of the client_id in the list
+ *
+ *  returns 0 when the item was stored, 1 otherwise */
 static int replace_good(   Client_data    *list,
                            uint32_t        length,
                            uint8_t        *client_id,
@@ -478,10 +480,12 @@ 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.
+ *
+ *  returns 1+ if the item is used in any list, 0 else
  */
-void addto_lists(DHT *dht, IP_Port ip_port, uint8_t *client_id)
+int addto_lists(DHT *dht, IP_Port ip_port, uint8_t *client_id)
 {
-    uint32_t i;
+    uint32_t i, used = 0;
 
     /* convert IPv4-in-IPv6 to IPv4 */
     if ((ip_port.ip.family == AF_INET6) && IN6_IS_ADDR_V4MAPPED(&ip_port.ip.ip6.in6_addr)) {
@@ -495,10 +499,13 @@ void addto_lists(DHT *dht, IP_Port ip_port, uint8_t *client_id)
     if (!client_or_ip_port_in_list(dht->close_clientlist, LCLIENT_LIST, client_id, ip_port)) {
         if (replace_bad(dht->close_clientlist, LCLIENT_LIST, client_id, ip_port)) {
             /* If we can't replace bad nodes we try replacing good ones. */
-            replace_good(dht->close_clientlist, LCLIENT_LIST, client_id, ip_port,
-                         dht->c->self_public_key);
-        }
-    }
+            if (!replace_good(dht->close_clientlist, LCLIENT_LIST, client_id, ip_port,
+                              dht->c->self_public_key))
+                used++;
+        } else
+            used++;
+    } else
+        used++;
 
     for (i = 0; i < dht->num_friends; ++i) {
         if (!client_or_ip_port_in_list(dht->friends_list[i].client_list,
@@ -507,17 +514,22 @@ void addto_lists(DHT *dht, IP_Port ip_port, uint8_t *client_id)
             if (replace_bad(dht->friends_list[i].client_list, MAX_FRIEND_CLIENTS,
                             client_id, ip_port)) {
                 /* If we can't replace bad nodes we try replacing good ones. */
-                replace_good(dht->friends_list[i].client_list, MAX_FRIEND_CLIENTS,
-                             client_id, ip_port, dht->friends_list[i].client_id);
-            }
-        }
+                if (!replace_good(dht->friends_list[i].client_list, MAX_FRIEND_CLIENTS,
+                                  client_id, ip_port, dht->friends_list[i].client_id))
+                    used++;
+            } else
+                used++;
+        } else
+            used++;
     }
+
+    return used;
 }
 
 /* If client_id is a friend or us, update ret_ip_port
  * nodeclient_id is the id of the node that sent us this info.
  */
-static void returnedip_ports(DHT *dht, IP_Port ip_port, uint8_t *client_id, uint8_t *nodeclient_id)
+static int returnedip_ports(DHT *dht, IP_Port ip_port, uint8_t *client_id, uint8_t *nodeclient_id)
 {
     uint32_t i, j;
     uint64_t temp_time = unix_time();
@@ -539,10 +551,9 @@ static void returnedip_ports(DHT *dht, IP_Port ip_port, uint8_t *client_id, uint
                     dht->close_clientlist[i].assoc6.ret_timestamp = temp_time;
                 }
 
-                return;
+                return 1;
             }
         }
-
     } else {
         for (i = 0; i < dht->num_friends; ++i) {
             if (id_equal(client_id, dht->friends_list[i].client_id)) {
@@ -556,13 +567,14 @@ static void returnedip_ports(DHT *dht, IP_Port ip_port, uint8_t *client_id, uint
                             dht->friends_list[i].client_list[j].assoc6.ret_timestamp = temp_time;
                         }
 
-                        return;
+                        return 1;
                     }
                 }
             }
         }
-
     }
+
+    return 0;
 }
 
 /* checks if ip/port or ping_id are already in the list to get nodes
@@ -841,14 +853,15 @@ static int handle_sendnodes_core(void *object, IP_Port source, uint8_t *packet,
         return 1;
 
     /* store the address the *request* was sent to */
-    addto_lists(dht, dht->send_nodes[send_nodes_index - 1].ip_port, packet + 1);
+    int used = addto_lists(dht, dht->send_nodes[send_nodes_index - 1].ip_port, packet + 1);
 
     if (dht->assoc) {
         IPPTs ippts;
 
         ippts.ip_port = dht->send_nodes[send_nodes_index - 1].ip_port;
         ippts.timestamp = dht->send_nodes[send_nodes_index - 1].timestamp;
-        Assoc_add_entry(dht->assoc, packet + 1, &ippts, &source);
+
+        Assoc_add_entry(dht->assoc, packet + 1, &ippts, &source, used ? 1 : 0);
     }
 
     *num_nodes_out = num_nodes;
@@ -879,13 +892,14 @@ static int handle_sendnodes(void *object, IP_Port source, uint8_t *packet, uint3
             ipp.port = nodes4_list[i].ip_port.port;
 
             send_ping_request(dht->ping, ipp, nodes4_list[i].client_id);
-            returnedip_ports(dht, ipp, nodes4_list[i].client_id, packet + 1);
+            int used = returnedip_ports(dht, ipp, nodes4_list[i].client_id, packet + 1);
 
             if (dht->assoc) {
                 IPPTs ippts;
                 ippts.ip_port = ipp;
                 ippts.timestamp = 0;
-                Assoc_add_entry(dht->assoc, nodes4_list[i].client_id, &ippts, NULL);
+
+                Assoc_add_entry(dht->assoc, nodes4_list[i].client_id, &ippts, NULL, used ? 1 : 0);
             }
         }
 
@@ -909,13 +923,14 @@ static int handle_sendnodes_ipv6(void *object, IP_Port source, uint8_t *packet,
     for (i = 0; i < num_nodes; i++)
         if (ipport_isset(&nodes_list[i].ip_port)) {
             send_ping_request(dht->ping, nodes_list[i].ip_port, nodes_list[i].client_id);
-            returnedip_ports(dht, nodes_list[i].ip_port, nodes_list[i].client_id, packet + 1);
+            int used = returnedip_ports(dht, nodes_list[i].ip_port, nodes_list[i].client_id, packet + 1);
 
             if (dht->assoc) {
                 IPPTs ippts;
                 ippts.ip_port = nodes_list[i].ip_port;
                 ippts.timestamp = 0;
-                Assoc_add_entry(dht->assoc, nodes_list[i].client_id, &ippts, NULL);
+
+                Assoc_add_entry(dht->assoc, nodes_list[i].client_id, &ippts, NULL, used ? 1 : 0);
             }
         }
 
@@ -1135,7 +1150,8 @@ void DHT_bootstrap(DHT *dht, IP_Port ip_port, uint8_t *public_key)
         IPPTs ippts;
         ippts.ip_port = ip_port;
         ippts.timestamp = 0;
-        Assoc_add_entry(dht->assoc, public_key, &ippts, NULL);
+
+        Assoc_add_entry(dht->assoc, public_key, &ippts, NULL, 0);
     }
 
     getnodes(dht, ip_port, public_key, dht->c->self_public_key);
@@ -1609,8 +1625,7 @@ DHT *new_DHT(Net_Crypto *c)
     init_cryptopackets(dht);
     cryptopacket_registerhandler(c, CRYPTO_PACKET_NAT_PING, &handle_NATping, dht);
 
-    /* dhtassoc is not mandatory for now */
-    dht->assoc = new_Assoc(dht);
+    dht->assoc = new_Assoc_default(dht->c->self_public_key);
 
     return dht;
 }