From a2eb33549e848edde2bb0e001a37f493e2a8a826 Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Mon, 28 Jul 2025 14:49:44 +0200
Subject: [PATCH] Instead of nil address remove and add back the peer

---
 client/iface/configurer/kernel_unix.go | 30 +++++++++++++----
 client/iface/configurer/usp.go         | 45 ++++++++++++++++++++++++--
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/client/iface/configurer/kernel_unix.go b/client/iface/configurer/kernel_unix.go
index 0816206d4..96b286175 100644
--- a/client/iface/configurer/kernel_unix.go
+++ b/client/iface/configurer/kernel_unix.go
@@ -78,18 +78,36 @@ func (c *KernelConfigurer) RemoveEndpointAddress(peerKey string) error {
 	if err != nil {
 		return err
 	}
-	peer := wgtypes.PeerConfig{
+
+	// Get the existing peer to preserve its allowed IPs
+	existingPeer, err := c.getPeer(c.deviceName, peerKey)
+	if err != nil {
+		return fmt.Errorf("get peer: %w", err)
+	}
+
+	removePeerCfg := wgtypes.PeerConfig{
 		PublicKey: peerKeyParsed,
-		Endpoint:  nil,
+		Remove:    true,
 	}
 
-	config := wgtypes.Config{
-		Peers: []wgtypes.PeerConfig{peer},
+	if err := c.configure(wgtypes.Config{Peers: []wgtypes.PeerConfig{removePeerCfg}}); err != nil {
+		return fmt.Errorf(`error removing peer %s from interface %s: %w`, peerKey, c.deviceName, err)
 	}
 
-	if err := c.configure(config); err != nil {
-		return fmt.Errorf("failed to remove endpoint address: %w", err)
+	//Re-add the peer without the endpoint but same AllowedIPs
+	reAddPeerCfg := wgtypes.PeerConfig{
+		PublicKey:         peerKeyParsed,
+		AllowedIPs:        existingPeer.AllowedIPs,
+		ReplaceAllowedIPs: true,
 	}
+
+	if err := c.configure(wgtypes.Config{Peers: []wgtypes.PeerConfig{reAddPeerCfg}}); err != nil {
+		return fmt.Errorf(
+			`error re-adding peer %s to interface %s with allowed IPs %v: %w`,
+			peerKey, c.deviceName, existingPeer.AllowedIPs, err,
+		)
+	}
+
 	return nil
 }
 
diff --git a/client/iface/configurer/usp.go b/client/iface/configurer/usp.go
index c898da9a2..33011cd54 100644
--- a/client/iface/configurer/usp.go
+++ b/client/iface/configurer/usp.go
@@ -109,20 +109,59 @@ func (c *WGUSPConfigurer) UpdatePeer(peerKey string, allowedIps []netip.Prefix,
 func (c *WGUSPConfigurer) RemoveEndpointAddress(peerKey string) error {
 	peerKeyParsed, err := wgtypes.ParseKey(peerKey)
 	if err != nil {
-		return err
+		return fmt.Errorf("parse peer key: %w", err)
 	}
 
+	ipcStr, err := c.device.IpcGet()
+	if err != nil {
+		return fmt.Errorf("get IPC config: %w", err)
+	}
+
+	// Parse current status to get allowed IPs for the peer
+	stats, err := parseStatus(c.deviceName, ipcStr)
+	if err != nil {
+		return fmt.Errorf("parse IPC config: %w", err)
+	}
+
+	var allowedIPs []net.IPNet
+	found := false
+	for _, peer := range stats.Peers {
+		if peer.PublicKey == peerKey {
+			allowedIPs = peer.AllowedIPs
+			found = true
+			break
+		}
+	}
+	if !found {
+		return fmt.Errorf("peer %s not found", peerKey)
+	}
+
+	// remove the peer from the WireGuard configuration
 	peer := wgtypes.PeerConfig{
 		PublicKey: peerKeyParsed,
-		Endpoint:  nil,
+		Remove:    true,
 	}
 
 	config := wgtypes.Config{
 		Peers: []wgtypes.PeerConfig{peer},
 	}
+	if ipcErr := c.device.IpcSet(toWgUserspaceString(config)); ipcErr != nil {
+		return fmt.Errorf("failed to remove peer: %s", ipcErr)
+	}
+
+	// Build the peer config
+	peer = wgtypes.PeerConfig{
+		PublicKey:         peerKeyParsed,
+		ReplaceAllowedIPs: true,
+		AllowedIPs:        allowedIPs,
+	}
+
+	config = wgtypes.Config{
+		Peers: []wgtypes.PeerConfig{peer},
+	}
 
 	if err := c.device.IpcSet(toWgUserspaceString(config)); err != nil {
-		return fmt.Errorf("failed to remove endpoint address: %w", err)
+		return fmt.Errorf("remove endpoint address: %w", err)
 	}
 
 	return nil