Add unit test

client/iface/wgproxy/proxy_linux_test.go

@@ -1,54 +0,0 @@
-//go:build linux && !android
-
-package wgproxy
-
-import (
-	"context"
-	"os"
-	"testing"
-
-	"github.com/netbirdio/netbird/client/iface/wgproxy/ebpf"
-)
-
-func TestProxyCloseByRemoteConnEBPF(t *testing.T) {
-	if os.Getenv("GITHUB_ACTIONS") != "true" {
-		t.Skip("Skipping test as it requires root privileges")
-	}
-	ctx := context.Background()
-
-	ebpfProxy := ebpf.NewWGEBPFProxy(51831)
-	if err := ebpfProxy.Listen(); err != nil {
-		t.Fatalf("failed to initialize ebpf proxy: %s", err)
-	}
-
-	defer func() {
-		if err := ebpfProxy.Free(); err != nil {
-			t.Errorf("failed to free ebpf proxy: %s", err)
-		}
-	}()
-
-	tests := []struct {
-		name  string
-		proxy Proxy
-	}{
-		{
-			name:  "ebpf proxy",
-			proxy: ebpf.NewProxyWrapper(ebpfProxy),
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			relayedConn := newMockConn()
-			err := tt.proxy.AddTurnConn(ctx, nil, relayedConn)
-			if err != nil {
-				t.Errorf("error: %v", err)
-			}
-
-			_ = relayedConn.Close()
-			if err := tt.proxy.CloseConn(); err != nil {
-				t.Errorf("error: %v", err)
-			}
-		})
-	}
-}

client/iface/wgproxy/proxy_test.go

@@ -1,117 +1,173 @@
-//go:build linux
+//go:build linux && !android
 
 package wgproxy
 
 import (
 	"context"
-	"io"
 	"net"
-	"os"
-	"runtime"
 	"testing"
-	"time"
 
 	"github.com/netbirdio/netbird/client/iface/wgproxy/ebpf"
-	udpProxy "github.com/netbirdio/netbird/client/iface/wgproxy/udp"
+	"github.com/netbirdio/netbird/client/iface/wgproxy/udp"
 	"github.com/netbirdio/netbird/util"
 )
 
-func TestMain(m *testing.M) {
+func init() {
-	_ = util.InitLog("trace", "console")
+	_ = util.InitLog("debug", "console")
-	code := m.Run()
-	os.Exit(code)
 }
+func TestProxyRedirect(t *testing.T) {
+	ebpfProxy := ebpf.NewWGEBPFProxy(51831)
+	if err := ebpfProxy.Listen(); err != nil {
+		t.Fatalf("failed to initialize ebpf proxy: %s", err)
+	}
 
-type mocConn struct {
+	defer func() {
-	closeChan chan struct{}
+		if err := ebpfProxy.Free(); err != nil {
-	closed    bool
+			t.Errorf("failed to free ebpf proxy: %s", err)
-}
+		}
+	}()
 
-func newMockConn() *mocConn {
+	tests := []struct {
-	return &mocConn{
+		name   string
-		closeChan: make(chan struct{}),
+		proxy  Proxy
+		wgPort int
+	}{
+		{
+			name:   "ebpf kernel proxy",
+			proxy:  ebpf.NewProxyWrapper(ebpfProxy),
+			wgPort: 51831,
+		},
+		{
+			name:   "udp kernel proxy",
+			proxy:  udp.NewWGUDPProxy(51832),
+			wgPort: 51832,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			redirectTraffic(t, tt.proxy, tt.wgPort)
+		})
 	}
 }
 
-func (m *mocConn) Read(b []byte) (n int, err error) {
+func redirectTraffic(t *testing.T, proxy Proxy, wgPort int) {
-	<-m.closeChan
+	t.Helper()
-	return 0, io.EOF
-}
 
-func (m *mocConn) Write(b []byte) (n int, err error) {
+	msgHelloFromRelay := []byte("hello from relay")
-	<-m.closeChan
+	msgRedirected := [][]byte{
-	return 0, io.EOF
+		[]byte("hello 1. to p2p"),
-}
+		[]byte("hello 2. to p2p"),
-
+		[]byte("hello 3. to p2p"),
-func (m *mocConn) Close() error {
-	if m.closed == true {
-		return nil
 	}
 
-	m.closed = true
+	dummyWgListener, err := net.ListenUDP("udp", &net.UDPAddr{
-	close(m.closeChan)
+		IP:   net.IPv4(127, 0, 0, 1),
-	return nil
+		Port: wgPort})
-}
+	if err != nil {
+		t.Fatalf("failed to listen on udp port: %s", err)
+	}
 
-func (m *mocConn) LocalAddr() net.Addr {
+	relayedServer, err := net.ListenUDP("udp",
-	panic("implement me")
+		&net.UDPAddr{
-}
+			IP:   net.IPv4(127, 0, 0, 1),
+			Port: 1234,
+		},
+	)
 
-func (m *mocConn) RemoteAddr() net.Addr {
+	relayedConn, err := net.Dial("udp", "127.0.0.1:1234")
-	return &net.UDPAddr{
+
-		IP: net.ParseIP("172.16.254.1"),
+	defer func() {
+		_ = dummyWgListener.Close()
+		_ = relayedConn.Close()
+		_ = relayedServer.Close()
+	}()
+
+	if err := proxy.AddTurnConn(context.Background(), nil, relayedConn); err != nil {
+		t.Errorf("error: %v", err)
+	}
+	defer func() {
+		if err := proxy.CloseConn(); err != nil {
+			t.Errorf("error: %v", err)
+		}
+	}()
+
+	proxy.Work()
+
+	if _, err := relayedServer.WriteTo(msgHelloFromRelay, relayedConn.LocalAddr()); err != nil {
+		t.Errorf("error relayedServer.Write(msgHelloFromRelay): %v", err)
+	}
+
+	n, err := dummyWgListener.Read(make([]byte, 1024))
+	if err != nil {
+		t.Errorf("error: %v", err)
+	}
+
+	if n != len(msgHelloFromRelay) {
+		t.Errorf("expected %d bytes, got %d", len(msgHelloFromRelay), n)
+	}
+
+	p2pEndpointAddr := &net.UDPAddr{
+		IP:   net.IPv4(192, 168, 0, 56),
+		Port: 1234,
+	}
+	proxy.RedirectAs(p2pEndpointAddr)
+
+	for _, msg := range msgRedirected {
+		if _, err := relayedServer.WriteTo(msg, relayedConn.LocalAddr()); err != nil {
+			t.Errorf("error: %v", err)
+		}
+	}
+
+	for i := 0; i < len(msgRedirected); i++ {
+		buf := make([]byte, 1024)
+		n, rAddr, err := dummyWgListener.ReadFrom(buf)
+		if err != nil {
+			t.Errorf("error: %v", err)
+		}
+
+		if rAddr.String() != p2pEndpointAddr.String() {
+			t.Errorf("expected %s, got %s", p2pEndpointAddr.String(), rAddr.String())
+		}
+		if string(buf[:n]) != string(msgRedirected[i]) {
+			t.Errorf("expected %s, got %s", string(msgRedirected[i]), string(buf[:n]))
+		}
 	}
 }
 
-func (m *mocConn) SetDeadline(t time.Time) error {
+func TestProxyCloseByRemoteConnEBPF(t *testing.T) {
-	panic("implement me")
-}
-
-func (m *mocConn) SetReadDeadline(t time.Time) error {
-	panic("implement me")
-}
-
-func (m *mocConn) SetWriteDeadline(t time.Time) error {
-	panic("implement me")
-}
-
-func TestProxyCloseByRemoteConn(t *testing.T) {
 	ctx := context.Background()
 
+	ebpfProxy := ebpf.NewWGEBPFProxy(51831)
+	if err := ebpfProxy.Listen(); err != nil {
+		t.Fatalf("failed to initialize ebpf proxy: %s", err)
+	}
+
+	defer func() {
+		if err := ebpfProxy.Free(); err != nil {
+			t.Errorf("failed to free ebpf proxy: %s", err)
+		}
+	}()
+
 	tests := []struct {
 		name  string
 		proxy Proxy
 	}{
 		{
-			name:  "userspace proxy",
+			name:  "ebpf proxy",
-			proxy: udpProxy.NewWGUDPProxy(51830),
+			proxy: ebpf.NewProxyWrapper(ebpfProxy),
+		},
+		{
+			name:  "udp proxy",
+			proxy: udp.NewWGUDPProxy(51832),
 		},
 	}
 
-	if runtime.GOOS == "linux" && os.Getenv("GITHUB_ACTIONS") != "true" {
+	relayedConn, _ := net.Dial("udp", "127.0.0.1:1234")
-		ebpfProxy := ebpf.NewWGEBPFProxy(51831)
+	defer func() {
-		if err := ebpfProxy.Listen(); err != nil {
+		_ = relayedConn.Close()
-			t.Fatalf("failed to initialize ebpf proxy: %s", err)
+	}()
-		}
-		defer func() {
-			if err := ebpfProxy.Free(); err != nil {
-				t.Errorf("failed to free ebpf proxy: %s", err)
-			}
-		}()
-		proxyWrapper := ebpf.NewProxyWrapper(ebpfProxy)
-
-		tests = append(tests, struct {
-			name  string
-			proxy Proxy
-		}{
-			name:  "ebpf proxy",
-			proxy: proxyWrapper,
-		})
-	}
-
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			relayedConn := newMockConn()
 			err := tt.proxy.AddTurnConn(ctx, nil, relayedConn)
 			if err != nil {
 				t.Errorf("error: %v", err)