Add DHT_create_packet

1 files changed, 36 insertions, 34 deletions

diff --git a/toxcore/DHT.c b/toxcore/DHT.c
index 3753c387..9a6996b2 100644
--- a/toxcore/DHT.c
+++ b/toxcore/DHT.c
@@ -365,6 +365,29 @@ static int pack_ip_port(uint8_t *data, uint16_t length, const IP_Port *ip_port)
     return -1;
 }
 
+static int DHT_create_packet(const uint8_t public_key[crypto_box_PUBLICKEYBYTES],
+                             const uint8_t *shared_key, const uint8_t type, uint8_t *plain, size_t plain_length, uint8_t *packet)
+{
+    uint8_t encrypted[plain_length + crypto_box_MACBYTES];
+    uint8_t nonce[crypto_box_NONCEBYTES];
+
+    random_nonce(nonce);
+
+    int encrypted_length = encrypt_data_symmetric(shared_key, nonce,
+                           plain, plain_length, encrypted);
+
+    if (encrypted_length == -1) {
+        return -1;
+    }
+
+    packet[0] = type;
+    memcpy(packet + 1, public_key, crypto_box_PUBLICKEYBYTES);
+    memcpy(packet + 1 + crypto_box_PUBLICKEYBYTES, nonce, crypto_box_NONCEBYTES);
+    memcpy(packet + 1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES, encrypted, encrypted_length);
+
+    return 1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES + encrypted_length;
+}
+
 /* Unpack IP_Port structure from data of max size length into ip_port.
  *
  * Return size of unpacked ip_port on success.
@@ -1190,8 +1213,7 @@ static int getnodes(DHT *dht, IP_Port ip_port, const uint8_t *public_key, const
     }
 
     uint8_t plain[crypto_box_PUBLICKEYBYTES + sizeof(ping_id)];
-    uint8_t encrypt[sizeof(plain) + crypto_box_MACBYTES];
-    uint8_t data[1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES + sizeof(encrypt)];
+    uint8_t data[1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES + sizeof(plain) + crypto_box_MACBYTES];
 
     memcpy(plain, client_id, crypto_box_PUBLICKEYBYTES);
     memcpy(plain + crypto_box_PUBLICKEYBYTES, &ping_id, sizeof(ping_id));
@@ -1199,25 +1221,14 @@ static int getnodes(DHT *dht, IP_Port ip_port, const uint8_t *public_key, const
     uint8_t shared_key[crypto_box_BEFORENMBYTES];
     DHT_get_shared_key_sent(dht, shared_key, public_key);
 
-    uint8_t nonce[crypto_box_NONCEBYTES];
-    random_nonce(nonce);
+    int len = DHT_create_packet(dht->self_public_key, shared_key, NET_PACKET_GET_NODES,
+                                plain, sizeof(plain), data);
 
-    int len = encrypt_data_symmetric(shared_key,
-                                     nonce,
-                                     plain,
-                                     sizeof(plain),
-                                     encrypt);
-
-    if (len != sizeof(encrypt)) {
+    if (len != sizeof(data)) {
         return -1;
     }
 
-    data[0] = NET_PACKET_GET_NODES;
-    memcpy(data + 1, dht->self_public_key, crypto_box_PUBLICKEYBYTES);
-    memcpy(data + 1 + crypto_box_PUBLICKEYBYTES, nonce, crypto_box_NONCEBYTES);
-    memcpy(data + 1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES, encrypt, len);
-
-    return sendpacket(dht->net, ip_port, data, sizeof(data));
+    return sendpacket(dht->net, ip_port, data, len);
 }
 
 /* Send a send nodes response: message for IPv6 nodes */
@@ -1234,16 +1245,11 @@ static int sendnodes_ipv6(const DHT *dht, IP_Port ip_port, const uint8_t *public
     }
 
     size_t Node_format_size = sizeof(Node_format);
-    uint8_t data[1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES
-                 + Node_format_size * MAX_SENT_NODES + length + crypto_box_MACBYTES];
 
     Node_format nodes_list[MAX_SENT_NODES];
     uint32_t num_nodes = get_close_nodes(dht, client_id, nodes_list, 0, LAN_ip(ip_port.ip) == 0, 1);
 
     uint8_t plain[1 + Node_format_size * MAX_SENT_NODES + length];
-    uint8_t encrypt[sizeof(plain) + crypto_box_MACBYTES];
-    uint8_t nonce[crypto_box_NONCEBYTES];
-    random_nonce(nonce);
 
     int nodes_length = 0;
 
@@ -1257,22 +1263,18 @@ static int sendnodes_ipv6(const DHT *dht, IP_Port ip_port, const uint8_t *public
 
     plain[0] = num_nodes;
     memcpy(plain + 1 + nodes_length, sendback_data, length);
-    int len = encrypt_data_symmetric(shared_encryption_key,
-                                     nonce,
-                                     plain,
-                                     1 + nodes_length + length,
-                                     encrypt);
 
-    if (len != 1 + nodes_length + length + crypto_box_MACBYTES) {
+    uint8_t data[1 + nodes_length + length + 1 + crypto_box_PUBLICKEYBYTES
+                 + crypto_box_NONCEBYTES + crypto_box_MACBYTES];
+
+    int len = DHT_create_packet(dht->self_public_key, shared_encryption_key, NET_PACKET_SEND_NODES_IPV6,
+                                plain, 1 + nodes_length + length, data);
+
+    if (len != sizeof(data)) {
         return -1;
     }
 
-    data[0] = NET_PACKET_SEND_NODES_IPV6;
-    memcpy(data + 1, dht->self_public_key, crypto_box_PUBLICKEYBYTES);
-    memcpy(data + 1 + crypto_box_PUBLICKEYBYTES, nonce, crypto_box_NONCEBYTES);
-    memcpy(data + 1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES, encrypt, len);
-
-    return sendpacket(dht->net, ip_port, data, 1 + crypto_box_PUBLICKEYBYTES + crypto_box_NONCEBYTES + len);
+    return sendpacket(dht->net, ip_port, data, len);
 }
 
 static int handle_getnodes(void *object, IP_Port source, const uint8_t *packet, uint16_t length, void *userdata)