[PATCH] Don't delete peers if not needed on `setconf`

[Grzegorz Nosek]

This prevents dropping sessions during `wg setconf` without manually
removing dead peers.

Signed-off-by: Grzegorz Nosek <root at localdomain.pl>
---
 src/netlink.c |  6 +++++-
 src/peer.c    | 31 +++++++++++++++++++++++++++++++
 src/peer.h    |  3 +++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/netlink.c b/src/netlink.c
index 190e405..a249380 100644
--- a/src/netlink.c
+++ b/src/netlink.c
@@ -434,6 +434,8 @@ static int set_peer(struct wg_device *wg, struct nlattr **attrs)
 		goto out;
 	}
 
+	peer->remove_me = false;
+
 	if (preshared_key) {
 		down_write(&peer->handshake.lock);
 		memcpy(&peer->handshake.preshared_key, preshared_key,
@@ -543,7 +545,7 @@ static int wg_set_device(struct sk_buff *skb, struct genl_info *info)
 	}
 
 	if (flags & WGDEVICE_F_REPLACE_PEERS)
-		wg_peer_remove_all(wg);
+		wg_peer_mark(wg);
 
 	if (info->attrs[WGDEVICE_A_PRIVATE_KEY] &&
 	    nla_len(info->attrs[WGDEVICE_A_PRIVATE_KEY]) ==
@@ -599,6 +601,8 @@ skip_set_private_key:
 	}
 	ret = 0;
 
+	if (flags & WGDEVICE_F_REPLACE_PEERS)
+		wg_peer_sweep(wg);
 out:
 	mutex_unlock(&wg->device_update_lock);
 	rtnl_unlock();
diff --git a/src/peer.c b/src/peer.c
index 071eedf..26e6df2 100644
--- a/src/peer.c
+++ b/src/peer.c
@@ -195,6 +195,37 @@ void wg_peer_remove_all(struct wg_device *wg)
 		peer_remove_after_dead(peer);
 }
 
+void wg_peer_mark(struct wg_device *wg)
+{
+	struct wg_peer *peer, *temp;
+
+	lockdep_assert_held(&wg->device_update_lock);
+	list_for_each_entry_safe(peer, temp, &wg->peer_list, peer_list) {
+		peer->remove_me = true;
+	}
+}
+
+void wg_peer_sweep(struct wg_device *wg)
+{
+	struct wg_peer *peer, *temp;
+	LIST_HEAD(dead_peers);
+
+	lockdep_assert_held(&wg->device_update_lock);
+
+	list_for_each_entry_safe(peer, temp, &wg->peer_list, peer_list) {
+		if (peer->remove_me) {
+			peer_make_dead(peer);
+			list_add_tail(&peer->peer_list, &dead_peers);
+		}
+	}
+	if (list_empty(&dead_peers))
+		return;
+
+	synchronize_rcu();
+	list_for_each_entry_safe(peer, temp, &dead_peers, peer_list)
+		peer_remove_after_dead(peer);
+}
+
 static void rcu_release(struct rcu_head *rcu)
 {
 	struct wg_peer *peer = container_of(rcu, struct wg_peer, rcu);
diff --git a/src/peer.h b/src/peer.h
index 23af409..b67dd73 100644
--- a/src/peer.h
+++ b/src/peer.h
@@ -64,6 +64,7 @@ struct wg_peer {
 	u64 internal_id;
 	struct napi_struct napi;
 	bool is_dead;
+	bool remove_me;
 };
 
 struct wg_peer *wg_peer_create(struct wg_device *wg,
@@ -79,5 +80,7 @@ static inline struct wg_peer *wg_peer_get(struct wg_peer *peer)
 void wg_peer_put(struct wg_peer *peer);
 void wg_peer_remove(struct wg_peer *peer);
 void wg_peer_remove_all(struct wg_device *wg);
+void wg_peer_mark(struct wg_device *wg);
+void wg_peer_sweep(struct wg_device *wg);
 
 #endif /* _WG_PEER_H */
-- 
2.17.1