timers: take reference like a lookup table

5 files changed, 42 insertions, 46 deletions

diff --git a/src/config.c b/src/config.c
index 4ca33ce..7f217ff 100644
--- a/src/config.c
+++ b/src/config.c
@@ -70,16 +70,9 @@ static int set_peer(struct wireguard_device *wg, void __user *user_peer, size_t
 	if (!peer) { /* Peer doesn't exist yet. Add a new one. */
 		if (in_peer.remove_me)
 			return -ENODEV; /* Tried to remove a non existing peer. */
-		peer = peer_create(wg, in_peer.public_key);
+		peer = peer_rcu_get(peer_create(wg, in_peer.public_key));
 		if (!peer)
 			return -ENOMEM;
-		rcu_read_lock();
-		peer = peer_get(peer);
-		rcu_read_unlock();
-		if (!peer) {
-			pr_err("Peer disappeared while creating\n");
-			return -EAGAIN;
-		}
 		if (netdev_pub(wg)->flags & IFF_UP)
 			timers_init_peer(peer);
 	} else
diff --git a/src/peer.c b/src/peer.c
index aa630a4..e1aaf6b 100644
--- a/src/peer.c
+++ b/src/peer.c
@@ -49,15 +49,26 @@ struct wireguard_peer *peer_get(struct wireguard_peer *peer)
 	return peer;
 }
 
+struct wireguard_peer *peer_rcu_get(struct wireguard_peer *peer)
+{
+	rcu_read_lock();
+	peer = peer_get(peer);
+	rcu_read_unlock();
+	return peer;
+}
+
+/* We have a separate "remove" function to get rid of the final reference because
+ * peer_list, clearing handshakes, and flushing all require mutexes which requires
+ * sleeping, which must only be done from certain contexts. */
 void peer_remove(struct wireguard_peer *peer)
 {
 	if (unlikely(!peer))
 		return;
 	lockdep_assert_held(&peer->device->device_update_lock);
-
-	list_del(&peer->peer_list);
 	noise_handshake_clear(&peer->handshake);
 	noise_keypairs_clear(&peer->keypairs);
+	list_del(&peer->peer_list);
+	timers_uninit_peer(peer);
 	routing_table_remove_by_peer(&peer->device->peer_routing_table, peer);
 	pubkey_hashtable_remove(&peer->device->peer_hashtable, peer);
 	if (peer->device->workqueue)
@@ -70,12 +81,10 @@ static void rcu_release(struct rcu_head *rcu)
 {
 	struct wireguard_peer *peer = container_of(rcu, struct wireguard_peer, rcu);
 	pr_debug("Peer %Lu (%pISpfsc) destroyed\n", peer->internal_id, &peer->endpoint_addr);
-	timers_uninit_peer(peer);
 	skb_queue_purge(&peer->tx_packet_queue);
 	if (peer->endpoint_dst)
 		dst_release(peer->endpoint_dst);
-	memzero_explicit(peer, sizeof(struct wireguard_peer));
-	kfree(peer);
+	kzfree(peer);
 }
 
 static void kref_release(struct kref *refcount)
@@ -98,9 +107,7 @@ int peer_for_each_unlocked(struct wireguard_device *wg, int (*fn)(struct wiregua
 
 	lockdep_assert_held(&wg->device_update_lock);
 	list_for_each_entry_safe(peer, temp, &wg->peer_list, peer_list) {
-		rcu_read_lock();
-		peer = peer_get(peer);
-		rcu_read_unlock();
+		peer = peer_rcu_get(peer);
 		if (unlikely(!peer))
 			continue;
 		ret = fn(peer, data);
diff --git a/src/peer.h b/src/peer.h
index 306f9b3..994b523 100644
--- a/src/peer.h
+++ b/src/peer.h
@@ -46,6 +46,8 @@ struct wireguard_peer {
 struct wireguard_peer *peer_create(struct wireguard_device *wg, const u8 public_key[static NOISE_PUBLIC_KEY_LEN]);
 
 struct wireguard_peer *peer_get(struct wireguard_peer *peer);
+struct wireguard_peer *peer_rcu_get(struct wireguard_peer *peer);
+
 void peer_put(struct wireguard_peer *peer);
 void peer_remove(struct wireguard_peer *peer);
 void peer_remove_all(struct wireguard_device *wg);
diff --git a/src/send.c b/src/send.c
index a7d505d..cbe8d95 100644
--- a/src/send.c
+++ b/src/send.c
@@ -33,8 +33,11 @@ void packet_send_handshake_initiation(struct wireguard_peer *peer)
 
 void packet_send_handshake_initiation_ratelimited(struct wireguard_peer *peer)
 {
-	if (time_is_before_jiffies64(peer->last_sent_handshake + REKEY_TIMEOUT))
-		packet_queue_send_handshake_initiation(peer);
+	if (time_is_before_jiffies64(peer->last_sent_handshake + REKEY_TIMEOUT)) {
+		peer = peer_rcu_get(peer);
+		if (likely(peer))
+			packet_queue_send_handshake_initiation(peer);
+	}
 }
 
 void packet_send_handshake_response(struct wireguard_peer *peer)
@@ -62,13 +65,9 @@ void packet_send_queued_handshakes(struct work_struct *work)
 	peer_put(peer);
 }
 
+/* Consumes peer reference. */
 void packet_queue_send_handshake_initiation(struct wireguard_peer *peer)
 {
-	rcu_read_lock();
-	peer = peer_get(peer);
-	rcu_read_unlock();
-	if (!peer)
-		return;
 	/* Queues up calling packet_send_queued_handshakes(peer), where we do a peer_put(peer) after: */
 	if (!queue_work(peer->device->workqueue, &peer->transmit_handshake_work))
 		peer_put(peer); /* If the work was already queued, we want to drop the extra reference */
diff --git a/src/timers.c b/src/timers.c
index 300becb..3f176e3 100644
--- a/src/timers.c
+++ b/src/timers.c
@@ -19,10 +19,14 @@ static inline unsigned long slack_time(unsigned long time)
 	return time & ~(roundup_pow_of_two(HZ / 4) - 1);
 }
 
+#define peer_get_from_ptr(ptr) \
+	struct wireguard_peer *peer = peer_rcu_get((struct wireguard_peer *)ptr); \
+	if (unlikely(!peer)) \
+		return;
+
 static void expired_retransmit_handshake(unsigned long ptr)
 {
-	struct wireguard_peer *peer = (struct wireguard_peer *)ptr;
-
+	peer_get_from_ptr(ptr);
 	pr_debug("Handshake for peer %Lu (%pISpfsc) did not complete after %d seconds, retrying\n", peer->internal_id, &peer->endpoint_addr, REKEY_TIMEOUT / HZ);
 	if (peer->timer_handshake_attempts > MAX_TIMER_HANDSHAKES) {
 		del_timer(&peer->timer_send_keepalive);
@@ -33,48 +37,40 @@ static void expired_retransmit_handshake(unsigned long ptr)
 		 * of a partial exchange. */
 		if (likely(peer->timer_kill_ephemerals.data))
 			mod_timer(&peer->timer_kill_ephemerals, jiffies + (REJECT_AFTER_TIME * 3));
+		peer_put(peer);
 		return;
 	}
-	packet_queue_send_handshake_initiation(peer);
+	packet_queue_send_handshake_initiation(peer); /* Takes our reference. */
 	++peer->timer_handshake_attempts;
 }
 
 static void expired_send_keepalive(unsigned long ptr)
 {
-	struct wireguard_peer *peer = (struct wireguard_peer *)ptr;
-
+	peer_get_from_ptr(ptr);
 	packet_send_keepalive(peer);
 	if (peer->timer_need_another_keepalive) {
 		peer->timer_need_another_keepalive = false;
 		mod_timer(&peer->timer_send_keepalive, jiffies + KEEPALIVE_TIMEOUT);
 	}
+	peer_put(peer);
 }
 
 static void expired_new_handshake(unsigned long ptr)
 {
-	struct wireguard_peer *peer = (struct wireguard_peer *)ptr;
-
+	peer_get_from_ptr(ptr);
 	pr_debug("Retrying handshake with peer %Lu (%pISpfsc) because we stopped hearing back after %d seconds\n", peer->internal_id, &peer->endpoint_addr, (KEEPALIVE_TIMEOUT + REKEY_TIMEOUT) / HZ);
-	packet_queue_send_handshake_initiation(peer);
+	packet_queue_send_handshake_initiation(peer); /* Takes our reference. */
 }
 
 static void expired_kill_ephemerals(unsigned long ptr)
 {
-	struct wireguard_peer *peer = (struct wireguard_peer *)ptr;
-
-	rcu_read_lock();
-	peer = peer_get(peer);
-	rcu_read_unlock();
-	if (!peer)
-		return;
-
-	if (!queue_work(peer->device->workqueue, &peer->clear_peer_work))
+	peer_get_from_ptr(ptr);
+	if (!queue_work(peer->device->workqueue, &peer->clear_peer_work)) /* Takes our reference. */
 		peer_put(peer); /* If the work was already on the queue, we want to drop the extra reference */
 }
 static void queued_expired_kill_ephemerals(struct work_struct *work)
 {
 	struct wireguard_peer *peer = container_of(work, struct wireguard_peer, clear_peer_work);
-
 	pr_debug("Zeroing out all keys for peer %Lu (%pISpfsc), since we haven't received a new one in %d seconds\n", peer->internal_id, &peer->endpoint_addr, (REJECT_AFTER_TIME * 3) / HZ);
 	noise_handshake_clear(&peer->handshake);
 	noise_keypairs_clear(&peer->keypairs);
@@ -83,11 +79,10 @@ static void queued_expired_kill_ephemerals(struct work_struct *work)
 
 static void expired_send_persistent_keepalive(unsigned long ptr)
 {
-	struct wireguard_peer *peer = (struct wireguard_peer *)ptr;
-
-	if (unlikely(!peer->persistent_keepalive_interval))
-		return;
-	packet_send_keepalive(peer);
+	peer_get_from_ptr(ptr);
+	if (likely(peer->persistent_keepalive_interval))
+		packet_send_keepalive(peer);
+	peer_put(peer);
 }
 
 /* Should be called after an authenticated data packet is sent. */