Save DHT nodes in a smarter way.

2 files changed, 45 insertions, 110 deletions

diff --git a/toxcore/DHT.c b/toxcore/DHT.c
index 8cd8852c..dd7c0acd 100644
--- a/toxcore/DHT.c
+++ b/toxcore/DHT.c
@@ -2282,8 +2282,7 @@ void kill_DHT(DHT *dht)
     ping_array_free_all(&dht->dht_harden_ping_array);
     kill_ping(dht->ping);
     free(dht->friends_list);
-    free(dht->loaded_friends_list);
-    free(dht->loaded_clients_list);
+    free(dht->loaded_nodes_list);
     free(dht);
 }
 
@@ -2292,23 +2291,20 @@ void kill_DHT(DHT *dht)
 #define DHT_STATE_COOKIE_GLOBAL 0x159000d
 
 #define DHT_STATE_COOKIE_TYPE      0x11ce
-#define DHT_STATE_TYPE_FRIENDS_ASSOC46  3
-#define DHT_STATE_TYPE_CLIENTS_ASSOC46  4
+#define DHT_STATE_TYPE_NODES       4
 
 /* Get the size of the DHT (for saving). */
 uint32_t DHT_size(const DHT *dht)
 {
     uint32_t num = 0, i;
 
-    for (i = 0; i < LCLIENT_LIST; ++i)
-        if ((dht->close_clientlist[i].assoc4.timestamp != 0) ||
-                (dht->close_clientlist[i].assoc6.timestamp != 0))
-            num++;
+    for (i = 0; i < LCLIENT_LIST; ++i) {
+        num += (dht->close_clientlist[i].assoc4.timestamp != 0) + (dht->close_clientlist[i].assoc6.timestamp != 0);
+    }
 
     uint32_t size32 = sizeof(uint32_t), sizesubhead = size32 * 2;
     return size32
-           + sizesubhead + sizeof(DHT_Friend) * dht->num_friends
-           + sizesubhead + sizeof(Client_data) * num;
+           + sizesubhead + sizeof(Node_format) * num;
 }
 
 static uint8_t *z_state_save_subheader(uint8_t *data, uint32_t len, uint16_t type)
@@ -2328,71 +2324,32 @@ void DHT_save(DHT *dht, uint8_t *data)
     *(uint32_t *)data = DHT_STATE_COOKIE_GLOBAL;
     data += sizeof(uint32_t);
 
-    len = sizeof(DHT_Friend) * dht->num_friends;
-    type = DHT_STATE_TYPE_FRIENDS_ASSOC46;
-    data = z_state_save_subheader(data, len, type);
-    memcpy(data, dht->friends_list, len);
-    data += len;
-
     uint32_t num = 0, i;
 
-    for (i = 0; i < LCLIENT_LIST; ++i)
-        if ((dht->close_clientlist[i].assoc4.timestamp != 0) ||
-                (dht->close_clientlist[i].assoc6.timestamp != 0))
-            num++;
+    //TODO: don't only save close nodes.
+    for (i = 0; i < LCLIENT_LIST; ++i) {
+        num += (dht->close_clientlist[i].assoc4.timestamp != 0) + (dht->close_clientlist[i].assoc6.timestamp != 0);
+    }
 
-    len = num * sizeof(Client_data);
-    type = DHT_STATE_TYPE_CLIENTS_ASSOC46;
+    len = num * sizeof(Node_format);
+    type = DHT_STATE_TYPE_NODES;
     data = z_state_save_subheader(data, len, type);
 
     if (num) {
-        Client_data *clients = (Client_data *)data;
-
-        for (num = 0, i = 0; i < LCLIENT_LIST; ++i)
-            if ((dht->close_clientlist[i].assoc4.timestamp != 0) ||
-                    (dht->close_clientlist[i].assoc6.timestamp != 0))
-                memcpy(&clients[num++], &dht->close_clientlist[i], sizeof(Client_data));
-    }
-}
-
-static void DHT_bootstrap_loaded_clients(DHT *dht)
-{
-    if (!dht->loaded_clients_list)
-        return;
-
-    uint32_t i;
-
-    Client_data *client_list = dht->loaded_clients_list;
-    uint32_t client_count = dht->loaded_num_clients;
-
-    for (i = 0; i < client_count; ++i) {
-        if (client_list[i].assoc4.timestamp != 0)
-            DHT_bootstrap(dht, client_list[i].assoc4.ip_port, client_list[i].client_id);
-
-        if (client_list[i].assoc6.timestamp != 0)
-            DHT_bootstrap(dht, client_list[i].assoc6.ip_port, client_list[i].client_id);
-    }
-}
-
-static void getnodes_of_loaded_friend_clients(DHT *dht)
-{
-    if (!dht->loaded_friends_list)
-        return;
-
-    uint32_t i, j;
-
-    DHT_Friend *friend_list = dht->loaded_friends_list;
-    uint32_t friend_count = dht->loaded_num_friends;
+        Node_format *clients = (Node_format *)data;
 
-    for (i = 0; i < friend_count; ++i) {
-        for (j = 0; j < MAX_FRIEND_CLIENTS; ++j) {
-            Client_data *client = &friend_list[i].client_list[j];
-
-            if (client->assoc4.timestamp != 0)
-                getnodes(dht, client->assoc4.ip_port, client->client_id, friend_list[i].client_id, NULL);
+        for (num = 0, i = 0; i < LCLIENT_LIST; ++i) {
+            if (dht->close_clientlist[i].assoc4.timestamp != 0) {
+                memcpy(&clients[num].public_key, &dht->close_clientlist[i].client_id, crypto_box_PUBLICKEYBYTES);
+                clients[num].ip_port = dht->close_clientlist[i].assoc4.ip_port;
+                ++num;
+            }
 
-            if (client->assoc6.timestamp != 0)
-                getnodes(dht, client->assoc6.ip_port, client->client_id, friend_list[i].client_id, NULL);
+            if (dht->close_clientlist[i].assoc6.timestamp != 0) {
+                memcpy(&clients[num].public_key, &dht->close_clientlist[i].client_id, crypto_box_PUBLICKEYBYTES);
+                clients[num].ip_port = dht->close_clientlist[i].assoc6.ip_port;
+                ++num;
+            }
         }
     }
 }
@@ -2403,17 +2360,19 @@ int DHT_connect_after_load(DHT *dht)
     if (dht == NULL)
         return -1;
 
-    getnodes_of_loaded_friend_clients(dht);
-    DHT_bootstrap_loaded_clients(dht);
+    if (!dht->loaded_nodes_list)
+        return -1;
+
+    unsigned int i;
 
-    // Loaded lists were allocd, free them
-    free(dht->loaded_friends_list);
-    dht->loaded_friends_list = NULL;
-    dht->loaded_num_friends = 0;
+    for (i = 0; i < dht->loaded_num_nodes; ++i) {
+        DHT_bootstrap(dht, dht->loaded_nodes_list[i].ip_port, dht->loaded_nodes_list[i].public_key);
+    }
 
-    free(dht->loaded_clients_list);
-    dht->loaded_clients_list = NULL;
-    dht->loaded_num_clients = 0;
+    // Loaded lists were allocd, free them
+    free(dht->loaded_nodes_list);
+    dht->loaded_nodes_list = NULL;
+    dht->loaded_num_nodes = 0;
 
     return 0;
 }
@@ -2424,44 +2383,22 @@ static int dht_load_state_callback(void *outer, const uint8_t *data, uint32_t le
     uint32_t num, i;
 
     switch (type) {
-        case DHT_STATE_TYPE_FRIENDS_ASSOC46:
-            if (length % sizeof(DHT_Friend) != 0)
-                break;
-
-            { /* localize declarations */
-                DHT_Friend *friend_list = (DHT_Friend *)data;
-                num = length / sizeof(DHT_Friend);
-
-                free(dht->loaded_friends_list);
-                // Copy to loaded_friends_list
-                dht->loaded_friends_list = calloc(num, sizeof(DHT_Friend));
-
-                for (i = 0; i < num; i++)
-                    memcpy(&(dht->loaded_friends_list[i]), &(friend_list[i]), sizeof(DHT_Friend));
-
-                dht->loaded_num_friends = num;
-
-                dht->has_loaded_friends_clients = 1;
-            } /* localize declarations */
-
-            break;
-
-        case DHT_STATE_TYPE_CLIENTS_ASSOC46:
-            if ((length % sizeof(Client_data)) != 0)
+        case DHT_STATE_TYPE_NODES:
+            if ((length % sizeof(Node_format)) != 0)
                 break;
 
             { /* localize declarations */
-                num = length / sizeof(Client_data);
-                Client_data *client_list = (Client_data *)data;
+                num = length / sizeof(Node_format);
+                Node_format *client_list = (Node_format *)data;
 
-                free(dht->loaded_clients_list);
+                free(dht->loaded_nodes_list);
                 // Copy to loaded_clients_list
-                dht->loaded_clients_list = calloc(num, sizeof(Client_data));
+                dht->loaded_nodes_list = calloc(num, sizeof(Node_format));
 
                 for (i = 0; i < num; i++)
-                    memcpy(&(dht->loaded_clients_list[i]), &(client_list[i]), sizeof(Client_data));
+                    memcpy(&(dht->loaded_nodes_list[i]), &(client_list[i]), sizeof(Node_format));
 
-                dht->loaded_num_clients = num;
+                dht->loaded_num_nodes = num;
 
                 dht->has_loaded_friends_clients = 1;
             } /* localize declarations */
diff --git a/toxcore/DHT.h b/toxcore/DHT.h
index 318a6002..981728a1 100644
--- a/toxcore/DHT.h
+++ b/toxcore/DHT.h
@@ -214,10 +214,8 @@ typedef struct {
     // Used after loading of file (tox_load), but no longer needed after connect (tox_connect)
     // Unsure if friends_list and num_friends could just be used instead?
     int has_loaded_friends_clients; // Whether or not we have loaded on the first do_DHT
-    DHT_Friend    *loaded_friends_list;
-    uint32_t       loaded_num_friends;
-    Client_data   *loaded_clients_list;
-    uint32_t       loaded_num_clients;
+    Node_format   *loaded_nodes_list;
+    uint32_t       loaded_num_nodes;
 
     Shared_Keys shared_keys_recv;
     Shared_Keys shared_keys_sent;