5 files changed, 335 insertions, 23 deletions

diff --git a/toxcore/DHT.c b/toxcore/DHT.c
index db9bfe3c..bb746fb0 100644
--- a/toxcore/DHT.c
+++ b/toxcore/DHT.c
@@ -1457,13 +1457,13 @@ void kill_DHT(DHT *dht)
 }
 
 /* Get the size of the DHT (for saving). */
-uint32_t DHT_size(DHT *dht)
+uint32_t DHT_size_old(DHT *dht)
 {
     return sizeof(dht->close_clientlist) + sizeof(DHT_Friend) * dht->num_friends;
 }
 
 /* Save the DHT in data where data is an array of size DHT_size(). */
-void DHT_save(DHT *dht, uint8_t *data)
+void DHT_save_old(DHT *dht, uint8_t *data)
 {
     memcpy(data, dht->close_clientlist, sizeof(dht->close_clientlist));
     memcpy(data + sizeof(dht->close_clientlist), dht->friends_list, sizeof(DHT_Friend) * dht->num_friends);
@@ -1474,7 +1474,7 @@ void DHT_save(DHT *dht, uint8_t *data)
  *  return -1 if failure.
  *  return 0 if success.
  */
-int DHT_load(DHT *dht, uint8_t *data, uint32_t size)
+int DHT_load_old(DHT *dht, uint8_t *data, uint32_t size)
 {
     if (size < sizeof(dht->close_clientlist)) {
         fprintf(stderr, "DHT_load: Expected at least %u bytes, got %u.\n", sizeof(dht->close_clientlist), size);
@@ -1518,6 +1518,129 @@ int DHT_load(DHT *dht, uint8_t *data, uint32_t size)
     return 0;
 }
 
+
+/* new DHT format for load/save, more robust and forward compatible */
+
+#define DHT_STATE_COOKIE_GLOBAL 0x159000d
+
+#define DHT_STATE_COOKIE_TYPE      0x11ce
+#define DHT_STATE_TYPE_FRIENDS     1
+#define DHT_STATE_TYPE_CLIENTS     2
+
+/* Get the size of the DHT (for saving). */
+uint32_t DHT_size(DHT *dht)
+{
+    uint32_t num = 0, i;
+    for (i = 0; i < LCLIENT_LIST; ++i)
+        if (dht->close_clientlist[i].timestamp != 0)
+            num++;
+
+    uint32_t size32 = sizeof(uint32_t), sizesubhead = size32 * 2;
+    return size32
+         + sizesubhead + sizeof(DHT_Friend) * dht->num_friends
+         + sizesubhead + sizeof(Client_data) * num;
+}
+
+static uint8_t *z_state_save_subheader(uint8_t *data, uint32_t len, uint16_t type)
+{
+    uint32_t *data32 = (uint32_t *)data;
+    data32[0] = len;
+    data32[1] = (DHT_STATE_COOKIE_TYPE << 16) | type;
+    data += sizeof(uint32_t) * 2;
+    return data;
+}
+
+/* Save the DHT in data where data is an array of size DHT_size(). */
+void DHT_save(DHT *dht, uint8_t *data)
+{
+    uint32_t len;
+    uint16_t type;
+    *(uint32_t *)data = DHT_STATE_COOKIE_GLOBAL;
+    data += sizeof(uint32_t);
+
+    len = sizeof(DHT_Friend) * dht->num_friends;
+    type = DHT_STATE_TYPE_FRIENDS;
+    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].timestamp != 0)
+            num++;
+
+    if (!num)
+        return;
+
+    len = num * sizeof(Client_data);
+    type = DHT_STATE_TYPE_CLIENTS;
+    data = z_state_save_subheader(data, len, type);
+    Client_data *clients = (Client_data *)data;
+    for (num = 0, i = 0; i < LCLIENT_LIST; ++i)
+        if (dht->close_clientlist[i].timestamp != 0)
+            memcpy(&clients[num++], &dht->close_clientlist[i], sizeof(Client_data));
+    data += len;
+}
+
+static int dht_load_state_callback(void *outer, uint8_t *data, uint32_t length, uint16_t type)
+{
+    DHT *dht = outer;
+    uint32_t num, i, j;
+    switch(type) {
+        case DHT_STATE_TYPE_FRIENDS:
+            if (length % sizeof(DHT_Friend) != 0)
+                break;
+
+            DHT_Friend *friend_list = (DHT_Friend *)data;
+            num = length / sizeof(DHT_Friend);
+            for (i = 0; i < num; ++i) {
+                DHT_addfriend(dht, friend_list[i].client_id);
+                for (j = 0; j < MAX_FRIEND_CLIENTS; ++j) {
+                    Client_data *client = &friend_list[i].client_list[j];
+                    if (client->timestamp != 0)
+                        getnodes(dht, client->ip_port, client->client_id, friend_list[i].client_id);
+                }
+            }
+
+            break;
+
+        case DHT_STATE_TYPE_CLIENTS:
+            if ((length % sizeof(Client_data)) != 0)
+                break;
+
+            num = length / sizeof(Client_data);
+            Client_data *client_list = (Client_data *)data;
+            for (i = 0; i < num; ++i)
+                if (client_list[i].timestamp != 0)
+                    DHT_bootstrap(dht, client_list[i].ip_port, client_list[i].client_id);
+
+            break;
+
+        default:
+            fprintf(stderr, "Load state (DHT): contains unrecognized part (len %u, type %u)\n",
+                            length, type);
+    }
+
+    return 0;
+}
+
+/* Load the DHT from data of size size.
+ *
+ *  return -1 if failure.
+ *  return 0 if success.
+ */
+int DHT_load_new(DHT *dht, uint8_t *data, uint32_t length)
+{
+    uint32_t cookie_len = sizeof(uint32_t);
+    if (length > cookie_len) {
+        uint32_t *data32 = (uint32_t *)data;
+        if (data32[0] == DHT_STATE_COOKIE_GLOBAL)
+            return load_state(dht_load_state_callback, dht, data + cookie_len,
+                                    length - cookie_len, DHT_STATE_COOKIE_TYPE);
+    }
+
+    return DHT_load_old(dht, data, length);
+}
 /*  return 0 if we are not connected to the DHT.
  *  return 1 if we are.
  */
diff --git a/toxcore/DHT.h b/toxcore/DHT.h
index ff20c892..e6f227f7 100644
--- a/toxcore/DHT.h
+++ b/toxcore/DHT.h
@@ -226,11 +226,13 @@ DHT *new_DHT(Net_Crypto *c);
 void kill_DHT(DHT *dht);
 
 /* Load the DHT from data of size size.
+ *  old/new: version of config file
  *
  *  return -1 if failure.
  *  return 0 if success.
  */
-int DHT_load(DHT *dht, uint8_t *data, uint32_t size);
+int DHT_load_old(DHT *dht, uint8_t *data, uint32_t size);
+int DHT_load_new(DHT *dht, uint8_t *data, uint32_t size);
 
 /*  return 0 if we are not connected to the DHT.
  *  return 1 if we are.
diff --git a/toxcore/Messenger.c b/toxcore/Messenger.c
index c46e3938..f2be6008 100644
--- a/toxcore/Messenger.c
+++ b/toxcore/Messenger.c
@@ -1331,7 +1331,7 @@ void doMessenger(Messenger *m)
 }
 
 /*  return size of the messenger data (for saving) */
-uint32_t Messenger_size(Messenger *m)
+uint32_t Messenger_size_old(Messenger *m)
 {
     return crypto_box_PUBLICKEYBYTES + crypto_box_SECRETKEYBYTES
            + sizeof(uint32_t)                  // nospam.
@@ -1344,8 +1344,8 @@ uint32_t Messenger_size(Messenger *m)
            ;
 }
 
-/* Save the messenger in data of size Messenger_size(). */
-void Messenger_save(Messenger *m, uint8_t *data)
+/* Save the messenger in data of size Messenger_size(). Old version without cookies. */
+static void Messenger_save_old(Messenger *m, uint8_t *data)
 {
     save_keys(m->net_crypto, data);
     data += crypto_box_PUBLICKEYBYTES + crypto_box_SECRETKEYBYTES;
@@ -1372,8 +1372,8 @@ void Messenger_save(Messenger *m, uint8_t *data)
     memcpy(data, m->name, small_size);
 }
 
-/* Load the messenger from data of size length. */
-int Messenger_load(Messenger *m, uint8_t *data, uint32_t length)
+/* Load the messenger from data of size length. Old version without cookies. */
+static int Messenger_load_old(Messenger *m, uint8_t *data, uint32_t length)
 {
     if (length == ~((uint32_t)0))
         return -1;
@@ -1393,6 +1393,9 @@ int Messenger_load(Messenger *m, uint8_t *data, uint32_t length)
     length -= sizeof(nospam);
 
     uint32_t size;
+    if (length < sizeof(size))
+        return -1;
+
     memcpy(&size, data, sizeof(size));
     data += sizeof(size);
     length -= sizeof(size);
@@ -1400,10 +1403,9 @@ int Messenger_load(Messenger *m, uint8_t *data, uint32_t length)
     if (length < size)
         return -1;
 
-    if (DHT_load(m->dht, data, size) == -1)
+    if (DHT_load_old(m->dht, data, size) == -1)
         fprintf(stderr, "Data file: Something wicked happened to the stored connections...\n");
-
-    /* go on, friends still might be intact */
+     /* DO go on, friends/name still might be intact */
 
     data += size;
     length -= size;
@@ -1420,24 +1422,22 @@ int Messenger_load(Messenger *m, uint8_t *data, uint32_t length)
 
     if (!(size % sizeof(Friend))) {
         uint16_t num = size / sizeof(Friend);
-        Friend temp[num];
-        memcpy(temp, data, size);
+        Friend *friend_list = (Friend *)data;
 
         uint32_t i;
-
         for (i = 0; i < num; ++i) {
-            if (temp[i].status >= 3) {
-                int fnum = m_addfriend_norequest(m, temp[i].client_id);
-                setfriendname(m, fnum, temp[i].name, temp[i].name_length);
+            if (friend_list[i].status >= 3) {
+                int fnum = m_addfriend_norequest(m, friend_list[i].client_id);
+                setfriendname(m, fnum, friend_list[i].name, friend_list[i].name_length);
                 /* set_friend_statusmessage(fnum, temp[i].statusmessage, temp[i].statusmessage_length); */
-            } else if (temp[i].status != 0) {
+            } else if (friend_list[i].status != 0) {
                 /* TODO: This is not a good way to do this. */
                 uint8_t address[FRIEND_ADDRESS_SIZE];
-                memcpy(address, temp[i].client_id, crypto_box_PUBLICKEYBYTES);
-                memcpy(address + crypto_box_PUBLICKEYBYTES, &(temp[i].friendrequest_nospam), sizeof(uint32_t));
+                memcpy(address, friend_list[i].client_id, crypto_box_PUBLICKEYBYTES);
+                memcpy(address + crypto_box_PUBLICKEYBYTES, &(friend_list[i].friendrequest_nospam), sizeof(uint32_t));
                 uint16_t checksum = address_checksum(address, FRIEND_ADDRESS_SIZE - sizeof(checksum));
                 memcpy(address + crypto_box_PUBLICKEYBYTES + sizeof(uint32_t), &checksum, sizeof(checksum));
-                m_addfriend(m, address, temp[i].info, temp[i].info_size);
+                m_addfriend(m, address, friend_list[i].info, friend_list[i].info_size);
             }
         }
     }
@@ -1454,7 +1454,7 @@ int Messenger_load(Messenger *m, uint8_t *data, uint32_t length)
     data += sizeof(small_size);
     length -= sizeof(small_size);
 
-    if (length != small_size)
+    if (length < small_size)
         return -1;
 
     setname(m, data, small_size);
@@ -1462,6 +1462,148 @@ int Messenger_load(Messenger *m, uint8_t *data, uint32_t length)
     return 0;
 }
 
+
+/* new messenger format for load/save, more robust and forward compatible */
+
+#define MESSENGER_STATE_COOKIE_GLOBAL 0x15ed1b1e
+
+#define MESSENGER_STATE_COOKIE_TYPE      0x01ce
+#define MESSENGER_STATE_TYPE_NOSPAMKEYS  1
+#define MESSENGER_STATE_TYPE_DHT         2
+#define MESSENGER_STATE_TYPE_FRIENDS     3
+#define MESSENGER_STATE_TYPE_NAME        4
+
+/*  return size of the messenger data (for saving) */
+uint32_t Messenger_size(Messenger *m)
+{
+    uint32_t size32 = sizeof(uint32_t), sizesubhead = size32 * 2;
+    return   size32 * 2                                      // global cookie
+           + sizesubhead + sizeof(uint32_t) +  crypto_box_PUBLICKEYBYTES + crypto_box_SECRETKEYBYTES
+           + sizesubhead + DHT_size(m->dht)                  // DHT
+           + sizesubhead + sizeof(Friend) * m->numfriends    // Friendlist itself.
+           + sizesubhead + m->name_length                    // Own nickname.
+           ;
+}
+
+static uint8_t *z_state_save_subheader(uint8_t *data, uint32_t len, uint16_t type)
+{
+    uint32_t *data32 = (uint32_t *)data;
+    data32[0] = len;
+    data32[1] = (MESSENGER_STATE_COOKIE_TYPE << 16) | type;
+    data += sizeof(uint32_t) * 2;
+    return data;
+}
+
+/* Save the messenger in data of size Messenger_size(). */
+void Messenger_save(Messenger *m, uint8_t *data)
+{
+    uint32_t len;
+    uint16_t type;
+    uint32_t *data32, size32 = sizeof(uint32_t);
+
+    data32 = (uint32_t *)data;
+    data32[0] = 0;
+    data32[1] = MESSENGER_STATE_COOKIE_GLOBAL;
+    data += size32 * 2;
+
+#ifdef DEBUG
+    assert(sizeof(get_nospam(&(m->fr))) == sizeof(uint32_t));
+#endif
+    len = size32 + crypto_box_PUBLICKEYBYTES + crypto_box_SECRETKEYBYTES;
+    type = MESSENGER_STATE_TYPE_NOSPAMKEYS;
+    data = z_state_save_subheader(data, len, type);
+    *(uint32_t *)data = get_nospam(&(m->fr));
+    save_keys(m->net_crypto, data + size32);
+    data += len;
+
+    len = DHT_size(m->dht);
+    type = MESSENGER_STATE_TYPE_DHT;
+    data = z_state_save_subheader(data, len, type);
+    DHT_save(m->dht, data);
+    data += len;
+
+    len = sizeof(Friend) * m->numfriends;
+    type = MESSENGER_STATE_TYPE_FRIENDS;
+    data = z_state_save_subheader(data, len, type);
+    memcpy(data, m->friendlist, len);
+    data += len;
+
+    len = m->name_length;
+    type = MESSENGER_STATE_TYPE_NAME;
+    data = z_state_save_subheader(data, len, type);
+    memcpy(data, m->name, len);
+    data += len;
+}
+
+static int messenger_load_state_callback(void *outer, uint8_t *data, uint32_t length, uint16_t type)
+{
+    Messenger *m = outer;
+    switch(type) {
+        case MESSENGER_STATE_TYPE_NOSPAMKEYS:
+            if (length == crypto_box_PUBLICKEYBYTES + crypto_box_SECRETKEYBYTES + sizeof(uint32_t)) {
+                set_nospam(&(m->fr), *(uint32_t *)data);
+                load_keys(m->net_crypto, &data[sizeof(uint32_t)]);
+            }
+            else
+                return -1;    /* critical */
+            break;
+
+        case MESSENGER_STATE_TYPE_DHT:
+            DHT_load_new(m->dht, data, length);
+            break;
+
+        case MESSENGER_STATE_TYPE_FRIENDS:
+            if (!(length % sizeof(Friend))) {
+                uint16_t num = length / sizeof(Friend);
+                Friend *friends = (Friend *)data;
+                uint32_t i;
+                for (i = 0; i < num; ++i) {
+                    if (friends[i].status >= 3) {
+                        int fnum = m_addfriend_norequest(m, friends[i].client_id);
+                        setfriendname(m, fnum, friends[i].name, friends[i].name_length);
+                        /* set_friend_statusmessage(fnum, temp[i].statusmessage, temp[i].statusmessage_length); */
+                    } else if (friends[i].status != 0) {
+                        /* TODO: This is not a good way to do this. */
+                        uint8_t address[FRIEND_ADDRESS_SIZE];
+                        memcpy(address, friends[i].client_id, crypto_box_PUBLICKEYBYTES);
+                        memcpy(address + crypto_box_PUBLICKEYBYTES, &(friends[i].friendrequest_nospam), sizeof(uint32_t));
+                        uint16_t checksum = address_checksum(address, FRIEND_ADDRESS_SIZE - sizeof(checksum));
+                        memcpy(address + crypto_box_PUBLICKEYBYTES + sizeof(uint32_t), &checksum, sizeof(checksum));
+                        m_addfriend(m, address, friends[i].info, friends[i].info_size);
+                    }
+                }
+            }
+            break;
+
+        case MESSENGER_STATE_TYPE_NAME:
+            if ((length > 0) && (length < MAX_NAME_LENGTH)) {
+                setname(m, data, length);
+            }
+            break;
+
+        default:
+            fprintf(stderr, "Load state: contains unrecognized part (len %u, type %u)\n",
+                            length, type);
+    }
+
+    return 0;
+}
+
+/* Load the messenger from data of size length. */
+int Messenger_load(Messenger *m, uint8_t *data, uint32_t length)
+{
+    uint32_t cookie_len = 2 * sizeof(uint32_t);
+    if (length < cookie_len)
+        return -1;
+
+    uint32_t *data32 = (uint32_t *)data;
+    if (!data32[0] && (data32[1] == MESSENGER_STATE_COOKIE_GLOBAL))
+        return load_state(messenger_load_state_callback, m, data + cookie_len,
+                          length - cookie_len, MESSENGER_STATE_COOKIE_TYPE);
+    else       /* old state file */
+        return Messenger_load_old(m, data, length);
+}
+
 /* Allocate and return a list of valid friend id's. List must be freed by the
  * caller.
  *
diff --git a/toxcore/util.c b/toxcore/util.c
index e1c46ee5..28f7d26e 100644
--- a/toxcore/util.c
+++ b/toxcore/util.c
@@ -43,6 +43,47 @@ void id_cpy(uint8_t *dest, uint8_t *src)
     memcpy(dest, src, CLIENT_ID_SIZE);
 }
 
+int load_state(load_state_callback_func load_state_callback, void *outer,
+        uint8_t *data, uint32_t length, uint16_t cookie_inner)
+{
+    if (!load_state_callback || !data) {
+        fprintf(stderr, "load_state() called with invalid args.\n");
+        return -1;
+    }
+
+
+    uint16_t type;
+    uint32_t length_sub, cookie_type;
+    uint32_t size32 = sizeof(uint32_t), size_head = size32 * 2;
+    while (length > size_head) {
+        length_sub = *(uint32_t *)data;
+        cookie_type = *(uint32_t *)(data + size32);
+        data += size_head;
+        length -= size_head;
+
+        if (length < length_sub) {
+            /* file truncated */
+            fprintf(stderr, "state file too short: %u < %u\n", length, length_sub);
+            return -1;
+        }
+
+        if ((cookie_type >> 16) != cookie_inner) {
+            /* something is not matching up in a bad way, give up */
+            fprintf(stderr, "state file garbeled: %04hx != %04hx\n", (cookie_type >> 16), cookie_inner);
+            return -1;
+        }
+
+        type = cookie_type & 0xFFFF;
+        if (-1 == load_state_callback(outer, data, length_sub, type))
+            return -1;
+
+        data += length_sub;
+        length -= length_sub;
+    }
+
+    return length == 0 ? 0 : - 1;
+};
+
 #ifdef LOGGING
 time_t starttime = 0;
 char logbuffer[512];
diff --git a/toxcore/util.h b/toxcore/util.h
index 71be84ca..6937f7d4 100644
--- a/toxcore/util.h
+++ b/toxcore/util.h
@@ -16,6 +16,10 @@ uint64_t random_64b();
 bool id_eq(uint8_t *dest, uint8_t *src);
 void id_cpy(uint8_t *dest, uint8_t *src);
 
+typedef int (*load_state_callback_func)(void *outer, uint8_t *data, uint32_t len, uint16_t type);
+int load_state(load_state_callback_func load_state_callback, void *outer,
+                        uint8_t *data, uint32_t length, uint16_t cookie_inner);
+
 #undef LOGGING
 /* #define LOGGING */
 #ifdef LOGGING