The non handshake initiator peer start the handshake after timeout

2026-04-22 02:06:39 +00:00 · 2025-01-20 15:28:14 +01:00
parent e9b3b6210d
commit bfa6df13c5
3 changed files with 284 additions and 28 deletions
--- a/client/internal/peer/endpoint_test.go
+++ b/client/internal/peer/endpoint_test.go
@@ -0,0 +1,178 @@
+package peer
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/mock"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/client/iface/configurer"
+	"github.com/netbirdio/netbird/client/iface/wgproxy"
+)
+
+type MockWgInterface struct {
+	mock.Mock
+
+	lastSetAddr *net.UDPAddr
+}
+
+func (m *MockWgInterface) GetStats(peerKey string) (configurer.WGStats, error) {
+	panic("implement me")
+}
+
+func (m *MockWgInterface) GetProxy() wgproxy.Proxy {
+	panic("implement me")
+}
+
+func (m *MockWgInterface) UpdatePeer(peerKey string, allowedIps string, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error {
+	args := m.Called(peerKey, allowedIps, keepAlive, endpoint, preSharedKey)
+	m.lastSetAddr = endpoint
+	return args.Error(0)
+}
+
+func (m *MockWgInterface) RemovePeer(publicKey string) error {
+	args := m.Called(publicKey)
+	return args.Error(0)
+}
+
+func Test_endpointUpdater_initiator(t *testing.T) {
+	mockWgInterface := &MockWgInterface{}
+	e := &endpointUpdater{
+		log: log.WithField("peer", "my-peer-key"),
+		wgConfig: WgConfig{
+			WgListenPort: 51820,
+			RemoteKey:    "secret-remote-key",
+			WgInterface:  mockWgInterface,
+			AllowedIps:   "172.16.254.1",
+		},
+		initiator: true,
+	}
+	addr := &net.UDPAddr{
+		IP:   net.ParseIP("127.0.0.1"),
+		Port: 1234,
+	}
+
+	mockWgInterface.On(
+		"UpdatePeer",
+		e.wgConfig.RemoteKey,
+		e.wgConfig.AllowedIps,
+		defaultWgKeepAlive,
+		addr,
+		(*wgtypes.Key)(nil),
+	).Return(nil)
+
+	if err := e.configureWGEndpoint(addr); err != nil {
+		t.Fatalf("updateWireGuardPeer() failed: %v", err)
+	}
+
+	mockWgInterface.AssertCalled(t, "UpdatePeer", e.wgConfig.RemoteKey, e.wgConfig.AllowedIps, defaultWgKeepAlive, addr, (*wgtypes.Key)(nil))
+}
+
+func Test_endpointUpdater_nonInitiator(t *testing.T) {
+	fallbackDelay = 1 * time.Second
+	mockWgInterface := &MockWgInterface{}
+	e := &endpointUpdater{
+		log: log.WithField("peer", "my-peer-key"),
+		wgConfig: WgConfig{
+			WgListenPort: 51820,
+			RemoteKey:    "secret-remote-key",
+			WgInterface:  mockWgInterface,
+			AllowedIps:   "172.16.254.1",
+		},
+		initiator: false,
+	}
+	addr := &net.UDPAddr{
+		IP:   net.ParseIP("127.0.0.1"),
+		Port: 1234,
+	}
+
+	mockWgInterface.On(
+		"UpdatePeer",
+		e.wgConfig.RemoteKey,
+		e.wgConfig.AllowedIps,
+		defaultWgKeepAlive,
+		(*net.UDPAddr)(nil),
+		(*wgtypes.Key)(nil),
+	).Return(nil)
+
+	mockWgInterface.On(
+		"UpdatePeer",
+		e.wgConfig.RemoteKey,
+		e.wgConfig.AllowedIps,
+		defaultWgKeepAlive,
+		addr,
+		(*wgtypes.Key)(nil),
+	).Return(nil)
+
+	err := e.configureWGEndpoint(addr)
+	if err != nil {
+		t.Fatalf("updateWireGuardPeer() failed: %v", err)
+	}
+	mockWgInterface.AssertCalled(t, "UpdatePeer", e.wgConfig.RemoteKey, e.wgConfig.AllowedIps, defaultWgKeepAlive, (*net.UDPAddr)(nil), (*wgtypes.Key)(nil))
+
+	time.Sleep(fallbackDelay + time.Second)
+
+	mockWgInterface.AssertCalled(t, "UpdatePeer", e.wgConfig.RemoteKey, e.wgConfig.AllowedIps, defaultWgKeepAlive, addr, (*wgtypes.Key)(nil))
+}
+
+func Test_endpointUpdater_overRule(t *testing.T) {
+	fallbackDelay = 1 * time.Second
+	mockWgInterface := &MockWgInterface{}
+	e := &endpointUpdater{
+		log: log.WithField("peer", "my-peer-key"),
+		wgConfig: WgConfig{
+			WgListenPort: 51820,
+			RemoteKey:    "secret-remote-key",
+			WgInterface:  mockWgInterface,
+			AllowedIps:   "172.16.254.1",
+		},
+		initiator: false,
+	}
+	addr1 := &net.UDPAddr{
+		IP:   net.ParseIP("127.0.0.1"),
+		Port: 1000,
+	}
+
+	addr2 := &net.UDPAddr{
+		IP:   net.ParseIP("127.0.0.1"),
+		Port: 1001,
+	}
+
+	mockWgInterface.On(
+		"UpdatePeer",
+		e.wgConfig.RemoteKey,
+		e.wgConfig.AllowedIps,
+		defaultWgKeepAlive,
+		(*net.UDPAddr)(nil),
+		(*wgtypes.Key)(nil),
+	).Return(nil)
+
+	mockWgInterface.On(
+		"UpdatePeer",
+		e.wgConfig.RemoteKey,
+		e.wgConfig.AllowedIps,
+		defaultWgKeepAlive,
+		addr2,
+		(*wgtypes.Key)(nil),
+	).Return(nil)
+
+	if err := e.configureWGEndpoint(addr1); err != nil {
+		t.Fatalf("updateWireGuardPeer() failed: %v", err)
+	}
+	mockWgInterface.AssertCalled(t, "UpdatePeer", e.wgConfig.RemoteKey, e.wgConfig.AllowedIps, defaultWgKeepAlive, (*net.UDPAddr)(nil), (*wgtypes.Key)(nil))
+
+	if err := e.configureWGEndpoint(addr2); err != nil {
+		t.Fatalf("updateWireGuardPeer() failed: %v", err)
+	}
+
+	time.Sleep(fallbackDelay + time.Second)
+
+	mockWgInterface.AssertCalled(t, "UpdatePeer", e.wgConfig.RemoteKey, e.wgConfig.AllowedIps, defaultWgKeepAlive, addr2, (*wgtypes.Key)(nil))
+
+	if mockWgInterface.lastSetAddr != addr2 {
+		t.Fatalf("lastSetAddr is not equal to addr2")
+	}
+}