1 files changed, 41 insertions, 14 deletions

diff --git a/toxcore/net_crypto.c b/toxcore/net_crypto.c
index 48bf5164..63012ce6 100644
--- a/toxcore/net_crypto.c
+++ b/toxcore/net_crypto.c
@@ -1238,7 +1238,11 @@ static int handle_packet_connection(Net_Crypto *c, int crypt_connection_id, cons
 
                 conn->status = CRYPTO_CONN_NOT_CONFIRMED;
                 /* Status needs to be CRYPTO_CONN_NOT_CONFIRMED for this to work. */
-                set_connection_dht_public_key(c, crypt_connection_id, dht_public_key, current_time_monotonic());
+                set_connection_dht_public_key(c, crypt_connection_id, dht_public_key);
+
+                if (conn->dht_pk_callback)
+                    conn->dht_pk_callback(conn->dht_pk_callback_object, conn->dht_pk_callback_number, dht_public_key);
+
             } else {
                 return -1;
             }
@@ -1473,7 +1477,11 @@ static int handle_new_connection_handshake(Net_Crypto *c, IP_Port source, const
             if (create_send_handshake(c, crypt_connection_id, n_c.cookie, n_c.dht_public_key) == 0) {
                 conn->status = CRYPTO_CONN_NOT_CONFIRMED;
                 /* Status needs to be CRYPTO_CONN_NOT_CONFIRMED for this to work. */
-                set_connection_dht_public_key(c, crypt_connection_id, n_c.dht_public_key, current_time_monotonic());
+                set_connection_dht_public_key(c, crypt_connection_id, n_c.dht_public_key);
+
+                if (conn->dht_pk_callback)
+                    conn->dht_pk_callback(conn->dht_pk_callback_object, conn->dht_pk_callback_number, n_c.dht_public_key);
+
                 ret = 0;
             }
         }
@@ -1522,7 +1530,7 @@ int accept_crypto_connection(Net_Crypto *c, New_Connection *n_c)
 
     conn->status = CRYPTO_CONN_NOT_CONFIRMED;
     /* Status needs to be CRYPTO_CONN_NOT_CONFIRMED for this to work. */
-    set_connection_dht_public_key(c, crypt_connection_id, n_c->dht_public_key, current_time_monotonic());
+    set_connection_dht_public_key(c, crypt_connection_id, n_c->dht_public_key);
     conn->packet_send_rate = CRYPTO_PACKET_MIN_RATE;
     conn->packets_left = CRYPTO_MIN_QUEUE_LENGTH;
     crypto_connection_add_source(c, crypt_connection_id, n_c->source);
@@ -1618,9 +1626,9 @@ static int connect_peer_tcp(Net_Crypto *c, int crypt_connection_id)
 /* Copy friends DHT public key into dht_key.
  *
  * return 0 on failure (no key copied).
- * return timestamp on success (key copied).
+ * return 1 on success (key copied).
  */
-uint64_t get_connection_dht_key(const Net_Crypto *c, int crypt_connection_id, uint8_t *dht_public_key)
+unsigned int get_connection_dht_key(const Net_Crypto *c, int crypt_connection_id, uint8_t *dht_public_key)
 {
     Crypto_Connection *conn = get_crypto_connection(c, crypt_connection_id);
 
@@ -1631,28 +1639,22 @@ uint64_t get_connection_dht_key(const Net_Crypto *c, int crypt_connection_id, ui
         return 0;
 
     memcpy(dht_public_key, conn->dht_public_key, crypto_box_PUBLICKEYBYTES);
-    return conn->dht_public_key_timestamp;
+    return 1;
 }
 
 
 /* Set the DHT public key of the crypto connection.
- * timestamp is the time (current_time_monotonic()) at which the key was last confirmed belonging to
- * the other peer.
  *
  * return -1 on failure.
  * return 0 on success.
  */
-int set_connection_dht_public_key(Net_Crypto *c, int crypt_connection_id, const uint8_t *dht_public_key,
-                                  uint64_t timestamp)
+int set_connection_dht_public_key(Net_Crypto *c, int crypt_connection_id, const uint8_t *dht_public_key)
 {
     Crypto_Connection *conn = get_crypto_connection(c, crypt_connection_id);
 
     if (conn == 0)
         return -1;
 
-    if (timestamp <= conn->dht_public_key_timestamp)
-        return -1;
-
     if (conn->dht_public_key_set == 1 && memcmp(conn->dht_public_key, dht_public_key, crypto_box_PUBLICKEYBYTES) == 0)
         return -1;
 
@@ -1662,7 +1664,6 @@ int set_connection_dht_public_key(Net_Crypto *c, int crypt_connection_id, const
 
     memcpy(conn->dht_public_key, dht_public_key, crypto_box_PUBLICKEYBYTES);
     conn->dht_public_key_set = 1;
-    conn->dht_public_key_timestamp = timestamp;
 
     if (conn->status == CRYPTO_CONN_COOKIE_REQUESTING) {
         conn->cookie_request_number = random_64b();
@@ -1692,6 +1693,9 @@ int set_direct_ip_port(Net_Crypto *c, int crypt_connection_id, IP_Port ip_port)
     if (conn == 0)
         return -1;
 
+    if (ip_port.ip.family != AF_INET && ip_port.ip.family != AF_INET6)
+        return -1;
+
     if (!ipport_equal(&ip_port, &conn->ip_port)) {
         if (bs_list_add(&c->ip_port_list, &ip_port, crypt_connection_id)) {
             bs_list_remove(&c->ip_port_list, &conn->ip_port, crypt_connection_id);
@@ -2243,6 +2247,29 @@ int connection_lossy_data_handler(Net_Crypto *c, int crypt_connection_id,
     return 0;
 }
 
+
+/* Set the function for this friend that will be callbacked with object and number
+ * when that friend gives us his DHT temporary public key.
+ *
+ * object and number will be passed as argument to this function.
+ *
+ * return -1 on failure.
+ * return 0 on success.
+ */
+int nc_dht_pk_callback(Net_Crypto *c, int crypt_connection_id, void (*function)(void *data, int32_t number,
+                       const uint8_t *dht_public_key), void *object, uint32_t number)
+{
+    Crypto_Connection *conn = get_crypto_connection(c, crypt_connection_id);
+
+    if (conn == 0)
+        return -1;
+
+    conn->dht_pk_callback = function;
+    conn->dht_pk_callback_object = object;
+    conn->dht_pk_callback_number = number;
+    return 0;
+}
+
 /* Get the crypto connection id from the ip_port.
  *
  * return -1 on failure.