diff --git a/pkg/wgtunnel/tunnels.go b/pkg/wgtunnel/tunnels.go index 1e5aff1b4..570bb016f 100644 --- a/pkg/wgtunnel/tunnels.go +++ b/pkg/wgtunnel/tunnels.go @@ -3,6 +3,7 @@ package wgtunnel import ( "errors" "fmt" + "maps" "net" "os" "os/exec" @@ -120,8 +121,8 @@ func GetUsableIP(network *net.IPNet) (net.IP, error) { // If the peer already exists, it returns the existing IP address. // Otherwise, it allocates a new IP address for the peer and stores the peer configuration. func (m *InterfaceManager) IpForPeer(peerName string, publicKey string) string { - m.peersMutex.RLock() - defer m.peersMutex.RUnlock() + m.peersMutex.Lock() + defer m.peersMutex.Unlock() log.Debug().Msgf("Allocating IP for peer: PeerName=%s, PublicKey=%s\n", peerName, publicKey) if _, ok := m.peers[peerName]; !ok { // Peer not found. Store the peer and return the IP. @@ -160,7 +161,9 @@ func (m *InterfaceManager) RemovePeer(peerName string) error { func (m *InterfaceManager) GetPeers() map[string]PeerConfig { m.peersMutex.RLock() defer m.peersMutex.RUnlock() - return m.peers + peers := make(map[string]PeerConfig, len(m.peers)) + maps.Copy(peers, m.peers) + return peers } func (m *InterfaceManager) PublicKey() (string, error) { diff --git a/pkg/wgtunnel/tunnels_test.go b/pkg/wgtunnel/tunnels_test.go new file mode 100644 index 000000000..d1a203eac --- /dev/null +++ b/pkg/wgtunnel/tunnels_test.go @@ -0,0 +1,96 @@ +package wgtunnel + +import ( + "fmt" + "net" + "sync" + "testing" +) + +func newTestInterfaceManager(t *testing.T) *InterfaceManager { + t.Helper() + + _, network, err := net.ParseCIDR("10.89.0.0/16") + if err != nil { + t.Fatalf("failed to parse test network: %v", err) + } + + allocator, err := NewIPAllocator(network.String()) + if err != nil { + t.Fatalf("failed to create allocator: %v", err) + } + + serverIP := net.IPAddr{IP: net.ParseIP("10.89.0.1")} + if err := allocator.Reserve(serverIP); err != nil { + t.Fatalf("failed to reserve server IP: %v", err) + } + + return &InterfaceManager{ + interfaceName: "wg0", + network: *network, + ipAddress: serverIP, + peers: make(map[string]PeerConfig), + peersMutex: sync.RWMutex{}, + ipManager: allocator, + } +} + +func TestIpForPeerConcurrentAllocations(t *testing.T) { + manager := newTestInterfaceManager(t) + + const peerCount = 512 + var ready sync.WaitGroup + var start sync.WaitGroup + var done sync.WaitGroup + ready.Add(peerCount) + start.Add(1) + done.Add(peerCount) + + allocated := make(chan string, peerCount) + for i := range peerCount { + go func() { + defer done.Done() + ready.Done() + start.Wait() + + peerName := fmt.Sprintf("10.1.%d.%d", i/256, i%256) + allocated <- manager.IpForPeer(peerName, fmt.Sprintf("key-%d", i)) + }() + } + + ready.Wait() + start.Done() + done.Wait() + close(allocated) + + seen := make(map[string]bool, peerCount) + for ip := range allocated { + if ip == "" { + t.Fatal("expected allocated IP, got empty string") + } + if seen[ip] { + t.Fatalf("IP %s allocated more than once", ip) + } + seen[ip] = true + } + if len(seen) != peerCount { + t.Fatalf("expected %d allocations, got %d", peerCount, len(seen)) + } +} + +func TestGetPeersReturnsCopy(t *testing.T) { + manager := newTestInterfaceManager(t) + peerIP := manager.IpForPeer("10.1.0.1", "key-1") + if peerIP == "" { + t.Fatal("expected allocated peer IP") + } + + peers := manager.GetPeers() + peers["10.1.0.2"] = PeerConfig{PublicKey: "key-2"} + + manager.peersMutex.RLock() + defer manager.peersMutex.RUnlock() + if _, found := manager.peers["10.1.0.2"]; found { + t.Fatal("GetPeers returned mutable internal peers map") + } +}