From b6fe49e59a39cb12334cbcad82b07a6247eb22d5 Mon Sep 17 00:00:00 2001
From: irungentoo <irungentoo@gmail.com>
Date: Tue, 4 Mar 2014 15:36:29 -0500
Subject: [PATCH] CPU usage improvements.

Shared keys for recieved DHT packets are now stored in a special array
so that they don't need to be recomputed for nodes that send us lots of
packets.
---
 toxcore/DHT.c  | 47 ++++++++++++++++++++++++++++++++++++++++++++++-
 toxcore/DHT.h  | 20 ++++++++++++++++++++
 toxcore/ping.c |  2 +-
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/toxcore/DHT.c b/toxcore/DHT.c
index d7e626c4..be68116c 100644
--- a/toxcore/DHT.c
+++ b/toxcore/DHT.c
@@ -114,6 +114,51 @@ static int client_id_cmp(ClientPair p1, ClientPair p2)
     return c;
 }
 
+/* Copy shared_key to decrypt DHT packet from client_id into shared_key
+ */
+void DHT_get_shared_key(DHT *dht, uint8_t *shared_key, uint8_t *client_id)
+{
+    uint32_t i, num = ~0, curr = 0;
+
+    for (i = 0; i < MAX_KEYS_PER_SLOT; ++i) {
+        int index = client_id[30] * MAX_KEYS_PER_SLOT + i;
+
+        if (dht->shared_keys.keys[index].stored) {
+            if (memcmp(client_id, dht->shared_keys.keys[index].client_id, CLIENT_ID_SIZE) == 0) {
+                memcpy(shared_key, dht->shared_keys.keys[index].shared_key, crypto_box_BEFORENMBYTES);
+                ++dht->shared_keys.keys[index].times_requested;
+                dht->shared_keys.keys[index].time_last_requested = unix_time();
+                return;
+            }
+
+            if (num != 0) {
+                if (is_timeout(dht->shared_keys.keys[index].time_last_requested, KEYS_TIMEOUT)) {
+                    num = 0;
+                    curr = index;
+                } else if (num > dht->shared_keys.keys[index].times_requested) {
+                    num = dht->shared_keys.keys[index].times_requested;
+                    curr = index;
+                }
+            }
+        } else {
+            if (num != 0) {
+                num = 0;
+                curr = index;
+            }
+        }
+    }
+
+    encrypt_precompute(client_id, dht->self_secret_key, shared_key);
+
+    if (num != (uint32_t)~0) {
+        dht->shared_keys.keys[curr].stored = 1;
+        dht->shared_keys.keys[curr].times_requested = 1;
+        memcpy(dht->shared_keys.keys[curr].client_id, client_id, CLIENT_ID_SIZE);
+        memcpy(dht->shared_keys.keys[curr].shared_key, shared_key, crypto_box_BEFORENMBYTES);
+        dht->shared_keys.keys[curr].time_last_requested = unix_time();
+    }
+}
+
 /* Check if client with client_id is already in list of length length.
  * If it is then set its corresponding timestamp to current time.
  * If the id is already in the list with a different ip_port, update it.
@@ -991,7 +1036,7 @@ static int handle_getnodes(void *object, IP_Port source, uint8_t *packet, uint32
     uint8_t plain[CLIENT_ID_SIZE + NODES_ENCRYPTED_MESSAGE_LENGTH];
     uint8_t shared_key[crypto_box_BEFORENMBYTES];
 
-    encrypt_precompute(packet + 1, dht->self_secret_key, shared_key);
+    DHT_get_shared_key(dht, shared_key, packet + 1);
     int len = decrypt_data_fast( shared_key,
                                  packet + 1 + CLIENT_ID_SIZE,
                                  packet + 1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES,
diff --git a/toxcore/DHT.h b/toxcore/DHT.h
index 490d2306..32abdc64 100644
--- a/toxcore/DHT.h
+++ b/toxcore/DHT.h
@@ -143,7 +143,20 @@ typedef struct {
 } Node_format;
 
 /*----------------------------------------------------------------------------------*/
+/* struct to store some shared keys so we don't have to regenerate them for each request. */
+#define MAX_KEYS_PER_SLOT 4
+#define KEYS_TIMEOUT 600
+struct SHARED_KEYS {
+    struct {
+        uint8_t client_id[CLIENT_ID_SIZE];
+        uint8_t shared_key[crypto_box_BEFORENMBYTES];
+        uint32_t times_requested;
+        uint8_t  stored; /* 0 if not, 1 if is */
+        uint64_t time_last_requested;
+    } keys[256 * MAX_KEYS_PER_SLOT];
+};
 
+/*----------------------------------------------------------------------------------*/
 
 typedef struct {
     Net_Crypto  *c;
@@ -162,6 +175,8 @@ typedef struct {
     DHT_Friend    *friends_list;
     uint16_t       num_friends;
 
+    struct SHARED_KEYS shared_keys;
+
     struct PING   *ping;
 #ifdef ENABLE_ASSOC_DHT
     struct Assoc  *assoc;
@@ -170,6 +185,11 @@ typedef struct {
 } DHT;
 /*----------------------------------------------------------------------------------*/
 
+/* Copy shared_key to decrypt DHT packet from client_id into shared_key
+ */
+void DHT_get_shared_key(DHT *dht, uint8_t *shared_key, uint8_t *client_id);
+
+
 void DHT_getnodes(DHT *dht, IP_Port *from_ipp, uint8_t *from_id, uint8_t *which_id);
 
 /* Add a new friend to the friends list.
diff --git a/toxcore/ping.c b/toxcore/ping.c
index 2bc1f07d..ff410f39 100644
--- a/toxcore/ping.c
+++ b/toxcore/ping.c
@@ -218,7 +218,7 @@ static int handle_ping_request(void *_dht, IP_Port source, uint8_t *packet, uint
     uint8_t shared_key[crypto_box_BEFORENMBYTES];
 
     // Decrypt ping_id
-    encrypt_precompute(packet + 1, ping->dht->self_secret_key, shared_key);
+    DHT_get_shared_key(dht, shared_key, packet + 1);
     rc = decrypt_data_fast(shared_key,
                            packet + 1 + CLIENT_ID_SIZE,
                            packet + 1 + CLIENT_ID_SIZE + crypto_box_NONCEBYTES,