peer: store total number of peers instead of iterating

This is faster, since it means adding a new peer is O(1) instead of
O(n). It's also safe to do because we're holding the device_update_lock
on both the ++ and the --.

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

3 files changed, 4 insertions, 15 deletions

diff --git a/src/device.h b/src/device.h
index 4932477..a74f58b 100644
--- a/src/device.h
+++ b/src/device.h
@@ -49,7 +49,7 @@ struct wireguard_device {
 	struct routing_table peer_routing_table;
 	struct mutex device_update_lock, socket_update_lock;
 	struct list_head device_list, peer_list;
-	unsigned int device_update_gen;
+	unsigned int num_peers, device_update_gen;
 	u32 fwmark;
 	u16 incoming_port;
 };
diff --git a/src/peer.c b/src/peer.c
index e42de44..55efde3 100644
--- a/src/peer.c
+++ b/src/peer.c
@@ -20,8 +20,9 @@ struct wireguard_peer *peer_create(struct wireguard_device *wg, const u8 public_
 
 	lockdep_assert_held(&wg->device_update_lock);
 
-	if (peer_total_count(wg) >= MAX_PEERS_PER_DEVICE)
+	if (wg->num_peers >= MAX_PEERS_PER_DEVICE)
 		return NULL;
+	++wg->num_peers;
 
 	peer = kzalloc(sizeof(struct wireguard_peer), GFP_KERNEL);
 	if (!peer)
@@ -89,6 +90,7 @@ void peer_remove(struct wireguard_peer *peer)
 	flush_workqueue(peer->device->packet_crypt_wq); /* The first flush is for encrypt/decrypt step. */
 	flush_workqueue(peer->device->packet_crypt_wq); /* The second flush is for send/receive step. */
 	flush_workqueue(peer->device->handshake_send_wq);
+	--peer->device->num_peers;
 	peer_put(peer);
 }
 
@@ -127,14 +129,3 @@ void peer_remove_all(struct wireguard_device *wg)
 	list_for_each_entry_safe (peer, temp, &wg->peer_list, peer_list)
 		peer_remove(peer);
 }
-
-unsigned int peer_total_count(struct wireguard_device *wg)
-{
-	unsigned int i = 0;
-	struct wireguard_peer *peer;
-
-	lockdep_assert_held(&wg->device_update_lock);
-	list_for_each_entry (peer, &wg->peer_list, peer_list)
-		++i;
-	return i;
-}
diff --git a/src/peer.h b/src/peer.h
index 994c8b6..49784d1 100644
--- a/src/peer.h
+++ b/src/peer.h
@@ -67,6 +67,4 @@ void peer_remove_all(struct wireguard_device *wg);
 
 struct wireguard_peer *peer_lookup_by_index(struct wireguard_device *wg, u32 index);
 
-unsigned int peer_total_count(struct wireguard_device *wg);
-
 #endif /* _WG_PEER_H */