ratelimiter: prevent init/uninit race

Fixes a classic ABA problem that isn't actually reachable because of
rtnl_lock, but it's good to be correct anyway.

Reported-by: Jann Horn <jann@thejh.net>
Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

2 files changed, 8 insertions, 4 deletions

diff --git a/src/ratelimiter.c b/src/ratelimiter.c
index 140d99e..638da2b 100644
--- a/src/ratelimiter.c
+++ b/src/ratelimiter.c
@@ -144,10 +144,10 @@ err_oom:
 
 int ratelimiter_init(void)
 {
+	mutex_lock(&init_lock);
 	if (atomic64_inc_return(&refcnt) != 1)
-		return 0;
+		goto out;
 
-	mutex_lock(&init_lock);
 	entry_cache = KMEM_CACHE(ratelimiter_entry, 0);
 	if (!entry_cache)
 		goto err;
@@ -174,6 +174,7 @@ int ratelimiter_init(void)
 
 	queue_delayed_work(system_power_efficient_wq, &gc_work, HZ);
 	get_random_bytes(&key, sizeof(key));
+out:
 	mutex_unlock(&init_lock);
 	return 0;
 
@@ -187,10 +188,10 @@ err:
 
 void ratelimiter_uninit(void)
 {
+	mutex_lock(&init_lock);
 	if (atomic64_dec_if_positive(&refcnt))
-		return;
+		goto out;
 
-	mutex_lock(&init_lock);
 	cancel_delayed_work_sync(&gc_work);
 	gc_entries(NULL);
 	rcu_barrier();
@@ -199,6 +200,7 @@ void ratelimiter_uninit(void)
 	kvfree(table_v6);
 #endif
 	kmem_cache_destroy(entry_cache);
+out:
 	mutex_unlock(&init_lock);
 }
 
diff --git a/src/timers.c b/src/timers.c
index d351fde..f1549a7 100644
--- a/src/timers.c
+++ b/src/timers.c
@@ -9,6 +9,8 @@
 #include "queueing.h"
 #include "socket.h"
 
+#include <linux/ktime.h>
+
 /*
  * Timer for retransmitting the handshake if we don't hear back after `REKEY_TIMEOUT + jitter` ms
  * Timer for sending empty packet if we have received a packet but after have not sent one for `KEEPALIVE_TIMEOUT` ms