Add unit test

This commit is contained in:
Zoltán Papp 2025-02-17 21:26:37 +01:00
parent 775b4feb7e
commit 360c7134f7
2 changed files with 132 additions and 130 deletions

View File

@ -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)
}
})
}
}

View File

@ -1,117 +1,173 @@
//go:build linux //go:build linux && !android
package wgproxy package wgproxy
import ( import (
"context" "context"
"io"
"net" "net"
"os"
"runtime"
"testing" "testing"
"time"
"github.com/netbirdio/netbird/client/iface/wgproxy/ebpf" "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" "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() 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 { tests := []struct {
name string name string
proxy Proxy 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)
if err := ebpfProxy.Listen(); err != nil {
t.Fatalf("failed to initialize ebpf proxy: %s", err)
}
defer func() { defer func() {
if err := ebpfProxy.Free(); err != nil { _ = relayedConn.Close()
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 { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
relayedConn := newMockConn()
err := tt.proxy.AddTurnConn(ctx, nil, relayedConn) err := tt.proxy.AddTurnConn(ctx, nil, relayedConn)
if err != nil { if err != nil {
t.Errorf("error: %v", err) t.Errorf("error: %v", err)