Separate IP_Port packing from pack_nodes() and unpack_nodes()

Allows us to pack IP_Port structs that are part of arbitrarily structured data.

1 files changed, 145 insertions, 92 deletions

diff --git a/toxcore/DHT.c b/toxcore/DHT.c
index 2e475d01..1d9ad54f 100644
--- a/toxcore/DHT.c
+++ b/toxcore/DHT.c
@@ -27,9 +27,7 @@
 #include "config.h"
 #endif
 
-#ifdef TOX_DEBUG
 #include <assert.h>
-#endif
 
 #include "logger.h"
 
@@ -233,6 +231,129 @@ int packed_node_size(uint8_t ip_family)
 }
 
 
+/* Packs an IP_Port structure into data of max size length.
+ *
+ * Returns size of packed IP_Port data on success
+ * Return -1 on failure.
+ */
+static int pack_ip_port(uint8_t *data, uint16_t length, const IP_Port *ip_port)
+{
+    if (data == NULL) {
+        return -1;
+    }
+
+    int ipv6 = -1;
+    uint8_t net_family;
+
+    if (ip_port->ip.family == AF_INET) {
+        // FIXME use functions to convert endianness
+        ipv6 = 0;
+        net_family = TOX_AF_INET;
+    } else if (ip_port->ip.family == TCP_INET) {
+        ipv6 = 0;
+        net_family = TOX_TCP_INET;
+    } else if (ip_port->ip.family == AF_INET6) {
+        ipv6 = 1;
+        net_family = TOX_AF_INET6;
+    } else if (ip_port->ip.family == TCP_INET6) {
+        ipv6 = 1;
+        net_family = TOX_TCP_INET6;
+    } else {
+        return -1;
+    }
+
+    if (ipv6 == 0) {
+        uint32_t size = 1 + SIZE_IP4 + sizeof(uint16_t);
+
+        if (size > length) {
+            return -1;
+        }
+
+        data[0] = net_family;
+        memcpy(data + 1, &ip_port->ip.ip4, SIZE_IP4);
+        memcpy(data + 1 + SIZE_IP4, &ip_port->port, sizeof(uint16_t));
+        return size;
+    } else if (ipv6 == 1) {
+        uint32_t size = 1 + SIZE_IP6 + sizeof(uint16_t);
+
+        if (size > length) {
+            return -1;
+        }
+
+        data[0] = net_family;
+        memcpy(data + 1, &ip_port->ip.ip6, SIZE_IP6);
+        memcpy(data + 1 + SIZE_IP6, &ip_port->port, sizeof(uint16_t));
+        return size;
+    } else {
+        return -1;
+    }
+}
+
+/* Unpack IP_Port structure from data of max size length into ip_port.
+ *
+ * Return size of unpacked ip_port on success.
+ * Return -1 on failure.
+ */
+static int unpack_ip_port(IP_Port *ip_port, const uint8_t *data, uint16_t length, uint8_t tcp_enabled)
+{
+    if (data == NULL) {
+        return -1;
+    }
+
+    int ipv6 = -1;
+    uint8_t host_family;
+
+    if (data[0] == TOX_AF_INET) {
+        ipv6 = 0;
+        host_family = AF_INET;
+    } else if (data[0] == TOX_TCP_INET) {
+        if (!tcp_enabled) {
+            return -1;
+        }
+
+        ipv6 = 0;
+        host_family = TCP_INET;
+    } else if (data[0] == TOX_AF_INET6) {
+        ipv6 = 1;
+        host_family = AF_INET6;
+    } else if (data[0] == TOX_TCP_INET6) {
+        if (!tcp_enabled) {
+            return -1;
+        }
+
+        ipv6 = 1;
+        host_family = TCP_INET6;
+    } else {
+        return -1;
+    }
+
+    if (ipv6 == 0) {
+        uint32_t size = 1 + SIZE_IP4 + sizeof(uint16_t);
+
+        if (size > length) {
+            return -1;
+        }
+
+        ip_port->ip.family = host_family;
+        memcpy(&ip_port->ip.ip4, data + 1, SIZE_IP4);
+        memcpy(&ip_port->port, data + 1 + SIZE_IP4, sizeof(uint16_t));
+        return size;
+    } else if (ipv6 == 1) {
+        uint32_t size = 1 + SIZE_IP6 + sizeof(uint16_t);
+
+        if (size > length) {
+            return -1;
+        }
+
+        ip_port->ip.family = host_family;
+        memcpy(&ip_port->ip.ip6, data + 1, SIZE_IP6);
+        memcpy(&ip_port->port, data + 1 + SIZE_IP6, sizeof(uint16_t));
+        return size;
+    } else {
+        return -1;
+    }
+}
+
 /* Pack number of nodes into data of maxlength length.
  *
  * return length of packed nodes on success.
@@ -242,54 +363,24 @@ int pack_nodes(uint8_t *data, uint16_t length, const Node_format *nodes, uint16_
 {
     uint32_t i, packed_length = 0;
 
-    for (i = 0; i < number; ++i) {
-        int ipv6 = -1;
-        uint8_t net_family;
+    for (i = 0; i < number && packed_length < length; ++i) {
+        int ipp_size = pack_ip_port(data + packed_length, length - packed_length, &nodes[i].ip_port);
 
-        // FIXME use functions to convert endianness
-        if (nodes[i].ip_port.ip.family == AF_INET) {
-            ipv6 = 0;
-            net_family = TOX_AF_INET;
-        } else if (nodes[i].ip_port.ip.family == TCP_INET) {
-            ipv6 = 0;
-            net_family = TOX_TCP_INET;
-        } else if (nodes[i].ip_port.ip.family == AF_INET6) {
-            ipv6 = 1;
-            net_family = TOX_AF_INET6;
-        } else if (nodes[i].ip_port.ip.family == TCP_INET6) {
-            ipv6 = 1;
-            net_family = TOX_TCP_INET6;
-        } else {
+        if (ipp_size == -1) {
             return -1;
         }
 
-        if (ipv6 == 0) {
-            uint32_t size = PACKED_NODE_SIZE_IP4;
-
-            if (packed_length + size > length) {
-                return -1;
-            }
+        packed_length += ipp_size;
 
-            data[packed_length] = net_family;
-            memcpy(data + packed_length + 1, &nodes[i].ip_port.ip.ip4, SIZE_IP4);
-            memcpy(data + packed_length + 1 + SIZE_IP4, &nodes[i].ip_port.port, sizeof(uint16_t));
-            memcpy(data + packed_length + 1 + SIZE_IP4 + sizeof(uint16_t), nodes[i].public_key, crypto_box_PUBLICKEYBYTES);
-            packed_length += size;
-        } else if (ipv6 == 1) {
-            uint32_t size = PACKED_NODE_SIZE_IP6;
-
-            if (packed_length + size > length) {
-                return -1;
-            }
-
-            data[packed_length] = net_family;
-            memcpy(data + packed_length + 1, &nodes[i].ip_port.ip.ip6, SIZE_IP6);
-            memcpy(data + packed_length + 1 + SIZE_IP6, &nodes[i].ip_port.port, sizeof(uint16_t));
-            memcpy(data + packed_length + 1 + SIZE_IP6 + sizeof(uint16_t), nodes[i].public_key, crypto_box_PUBLICKEYBYTES);
-            packed_length += size;
-        } else {
+        if (packed_length + crypto_box_PUBLICKEYBYTES > length) {
             return -1;
         }
+
+        memcpy(data + packed_length, nodes[i].public_key, crypto_box_PUBLICKEYBYTES);
+        packed_length += crypto_box_PUBLICKEYBYTES;
+
+        uint32_t increment = ipp_size + crypto_box_PUBLICKEYBYTES;
+        assert(increment == PACKED_NODE_SIZE_IP4 || increment == PACKED_NODE_SIZE_IP6);
     }
 
     return packed_length;
@@ -308,62 +399,24 @@ int unpack_nodes(Node_format *nodes, uint16_t max_num_nodes, uint16_t *processed
     uint32_t num = 0, len_processed = 0;
 
     while (num < max_num_nodes && len_processed < length) {
-        int ipv6 = -1;
-        uint8_t host_family;
-
-        if (data[len_processed] == TOX_AF_INET) {
-            ipv6 = 0;
-            host_family = AF_INET;
-        } else if (data[len_processed] == TOX_TCP_INET) {
-            if (!tcp_enabled) {
-                return -1;
-            }
+        int ipp_size = unpack_ip_port(&nodes[num].ip_port, data + len_processed, length - len_processed, tcp_enabled);
 
-            ipv6 = 0;
-            host_family = TCP_INET;
-        } else if (data[len_processed] == TOX_AF_INET6) {
-            ipv6 = 1;
-            host_family = AF_INET6;
-        } else if (data[len_processed] == TOX_TCP_INET6) {
-            if (!tcp_enabled) {
-                return -1;
-            }
-
-            ipv6 = 1;
-            host_family = TCP_INET6;
-        } else {
+        if (ipp_size == -1) {
             return -1;
         }
 
-        if (ipv6 == 0) {
-            uint32_t size = PACKED_NODE_SIZE_IP4;
+        len_processed += ipp_size;
 
-            if (len_processed + size > length) {
-                return -1;
-            }
-
-            nodes[num].ip_port.ip.family = host_family;
-            memcpy(&nodes[num].ip_port.ip.ip4, data + len_processed + 1, SIZE_IP4);
-            memcpy(&nodes[num].ip_port.port, data + len_processed + 1 + SIZE_IP4, sizeof(uint16_t));
-            memcpy(nodes[num].public_key, data + len_processed + 1 + SIZE_IP4 + sizeof(uint16_t), crypto_box_PUBLICKEYBYTES);
-            len_processed += size;
-            ++num;
-        } else if (ipv6 == 1) {
-            uint32_t size = PACKED_NODE_SIZE_IP6;
-
-            if (len_processed + size > length) {
-                return -1;
-            }
-
-            nodes[num].ip_port.ip.family = host_family;
-            memcpy(&nodes[num].ip_port.ip.ip6, data + len_processed + 1, SIZE_IP6);
-            memcpy(&nodes[num].ip_port.port, data + len_processed + 1 + SIZE_IP6, sizeof(uint16_t));
-            memcpy(nodes[num].public_key, data + len_processed + 1 + SIZE_IP6 + sizeof(uint16_t), crypto_box_PUBLICKEYBYTES);
-            len_processed += size;
-            ++num;
-        } else {
+        if (len_processed + crypto_box_PUBLICKEYBYTES > length) {
             return -1;
         }
+
+        memcpy(nodes[num].public_key, data + len_processed, crypto_box_PUBLICKEYBYTES);
+        len_processed += crypto_box_PUBLICKEYBYTES;
+        ++num;
+
+        uint32_t increment = ipp_size + crypto_box_PUBLICKEYBYTES;
+        assert(increment == PACKED_NODE_SIZE_IP4 || increment == PACKED_NODE_SIZE_IP6);
     }
 
     if (processed_data_len) {