mirror of
https://github.com/netbirdio/netbird.git
synced 2026-04-20 17:26:40 +00:00
[management,proxy,client] Add L4 capabilities (TLS/TCP/UDP) (#5530)
This commit is contained in:
Viktor Liu
493
proxy/internal/udp/relay_test.go
Normal file
493
proxy/internal/udp/relay_test.go
Normal file
@@ -0,0 +1,493 @@
|
||||
package udp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/types"
|
||||
)
|
||||
|
||||
func TestRelay_BasicPacketExchange(t *testing.T) {
|
||||
// Set up a UDP backend that echoes packets.
|
||||
backend, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backend.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, err := backend.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backend.WriteTo(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
// Set up the relay's public-facing listener.
|
||||
listener, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer listener.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
backendAddr := backend.LocalAddr().String()
|
||||
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
relay := New(ctx, RelayConfig{Logger: logger, Listener: listener, Target: backendAddr, DialFunc: dialFunc})
|
||||
go relay.Serve()
|
||||
defer relay.Close()
|
||||
|
||||
// Create a client and send a packet to the relay.
|
||||
client, err := net.Dial("udp", listener.LocalAddr().String())
|
||||
require.NoError(t, err)
|
||||
defer client.Close()
|
||||
|
||||
testData := []byte("hello UDP relay")
|
||||
_, err = client.Write(testData)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the echoed response.
|
||||
if err := client.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
buf := make([]byte, 1024)
|
||||
n, err := client.Read(buf)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, testData, buf[:n], "should receive echoed packet")
|
||||
}
|
||||
|
||||
func TestRelay_MultipleClients(t *testing.T) {
|
||||
backend, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backend.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, err := backend.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backend.WriteTo(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
listener, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer listener.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
relay := New(ctx, RelayConfig{Logger: logger, Listener: listener, Target: backend.LocalAddr().String(), DialFunc: dialFunc})
|
||||
go relay.Serve()
|
||||
defer relay.Close()
|
||||
|
||||
// Two clients, each should get their own session.
|
||||
for i, msg := range []string{"client-1", "client-2"} {
|
||||
client, err := net.Dial("udp", listener.LocalAddr().String())
|
||||
require.NoError(t, err, "client %d", i)
|
||||
defer client.Close()
|
||||
|
||||
_, err = client.Write([]byte(msg))
|
||||
require.NoError(t, err)
|
||||
|
||||
if err := client.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
buf := make([]byte, 1024)
|
||||
n, err := client.Read(buf)
|
||||
require.NoError(t, err, "client %d read", i)
|
||||
assert.Equal(t, msg, string(buf[:n]), "client %d should get own echo", i)
|
||||
}
|
||||
|
||||
// Verify two sessions were created.
|
||||
relay.mu.RLock()
|
||||
sessionCount := len(relay.sessions)
|
||||
relay.mu.RUnlock()
|
||||
assert.Equal(t, 2, sessionCount, "should have two sessions")
|
||||
}
|
||||
|
||||
func TestRelay_Close(t *testing.T) {
|
||||
listener, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
relay := New(ctx, RelayConfig{Logger: logger, Listener: listener, Target: "127.0.0.1:9999", DialFunc: dialFunc})
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
relay.Serve()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
relay.Close()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Serve did not return after Close")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRelay_SessionCleanup(t *testing.T) {
|
||||
backend, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backend.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, err := backend.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backend.WriteTo(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
listener, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer listener.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
relay := New(ctx, RelayConfig{Logger: logger, Listener: listener, Target: backend.LocalAddr().String(), DialFunc: dialFunc})
|
||||
go relay.Serve()
|
||||
defer relay.Close()
|
||||
|
||||
// Create a session.
|
||||
client, err := net.Dial("udp", listener.LocalAddr().String())
|
||||
require.NoError(t, err)
|
||||
_, err = client.Write([]byte("hello"))
|
||||
require.NoError(t, err)
|
||||
|
||||
if err := client.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
buf := make([]byte, 1024)
|
||||
_, err = client.Read(buf)
|
||||
require.NoError(t, err)
|
||||
client.Close()
|
||||
|
||||
// Verify session exists.
|
||||
relay.mu.RLock()
|
||||
assert.Equal(t, 1, len(relay.sessions))
|
||||
relay.mu.RUnlock()
|
||||
|
||||
// Make session appear idle by setting lastSeen to the past.
|
||||
relay.mu.Lock()
|
||||
for _, sess := range relay.sessions {
|
||||
sess.lastSeen.Store(time.Now().Add(-2 * DefaultSessionTTL).UnixNano())
|
||||
}
|
||||
relay.mu.Unlock()
|
||||
|
||||
// Trigger cleanup manually.
|
||||
relay.cleanupIdleSessions()
|
||||
|
||||
relay.mu.RLock()
|
||||
assert.Equal(t, 0, len(relay.sessions), "idle sessions should be cleaned up")
|
||||
relay.mu.RUnlock()
|
||||
}
|
||||
|
||||
// TestRelay_CloseAndRecreate verifies that closing a relay and creating a new
|
||||
// one on the same port works cleanly (simulates port mapping modify cycle).
|
||||
func TestRelay_CloseAndRecreate(t *testing.T) {
|
||||
backend, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backend.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, err := backend.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backend.WriteTo(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
// First relay.
|
||||
ln1, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
|
||||
relay1 := New(ctx, RelayConfig{Logger: logger, Listener: ln1, Target: backend.LocalAddr().String(), DialFunc: dialFunc})
|
||||
go relay1.Serve()
|
||||
|
||||
client1, err := net.Dial("udp", ln1.LocalAddr().String())
|
||||
require.NoError(t, err)
|
||||
_, err = client1.Write([]byte("relay1"))
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, client1.SetReadDeadline(time.Now().Add(2*time.Second)))
|
||||
buf := make([]byte, 1024)
|
||||
n, err := client1.Read(buf)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "relay1", string(buf[:n]))
|
||||
client1.Close()
|
||||
|
||||
// Close first relay.
|
||||
relay1.Close()
|
||||
|
||||
// Second relay on same port.
|
||||
port := ln1.LocalAddr().(*net.UDPAddr).Port
|
||||
ln2, err := net.ListenPacket("udp", fmt.Sprintf("127.0.0.1:%d", port))
|
||||
require.NoError(t, err)
|
||||
|
||||
relay2 := New(ctx, RelayConfig{Logger: logger, Listener: ln2, Target: backend.LocalAddr().String(), DialFunc: dialFunc})
|
||||
go relay2.Serve()
|
||||
defer relay2.Close()
|
||||
|
||||
client2, err := net.Dial("udp", ln2.LocalAddr().String())
|
||||
require.NoError(t, err)
|
||||
defer client2.Close()
|
||||
_, err = client2.Write([]byte("relay2"))
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, client2.SetReadDeadline(time.Now().Add(2*time.Second)))
|
||||
n, err = client2.Read(buf)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "relay2", string(buf[:n]), "second relay should work on same port")
|
||||
}
|
||||
|
||||
func TestRelay_SessionLimit(t *testing.T) {
|
||||
backend, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backend.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, err := backend.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backend.WriteTo(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
listener, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer listener.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
// Create a relay with a max of 2 sessions.
|
||||
relay := New(ctx, RelayConfig{Logger: logger, Listener: listener, Target: backend.LocalAddr().String(), DialFunc: dialFunc, MaxSessions: 2})
|
||||
go relay.Serve()
|
||||
defer relay.Close()
|
||||
|
||||
// Create 2 clients to fill up the session limit.
|
||||
for i := range 2 {
|
||||
client, err := net.Dial("udp", listener.LocalAddr().String())
|
||||
require.NoError(t, err, "client %d", i)
|
||||
defer client.Close()
|
||||
|
||||
_, err = client.Write([]byte("hello"))
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, client.SetReadDeadline(time.Now().Add(2*time.Second)))
|
||||
buf := make([]byte, 1024)
|
||||
_, err = client.Read(buf)
|
||||
require.NoError(t, err, "client %d should get response", i)
|
||||
}
|
||||
|
||||
relay.mu.RLock()
|
||||
assert.Equal(t, 2, len(relay.sessions), "should have exactly 2 sessions")
|
||||
relay.mu.RUnlock()
|
||||
|
||||
// Third client should get its packet dropped (session creation fails).
|
||||
client3, err := net.Dial("udp", listener.LocalAddr().String())
|
||||
require.NoError(t, err)
|
||||
defer client3.Close()
|
||||
|
||||
_, err = client3.Write([]byte("should be dropped"))
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, client3.SetReadDeadline(time.Now().Add(500*time.Millisecond)))
|
||||
buf := make([]byte, 1024)
|
||||
_, err = client3.Read(buf)
|
||||
assert.Error(t, err, "third client should time out because session was rejected")
|
||||
|
||||
relay.mu.RLock()
|
||||
assert.Equal(t, 2, len(relay.sessions), "session count should not exceed limit")
|
||||
relay.mu.RUnlock()
|
||||
}
|
||||
|
||||
// testObserver records UDP session lifecycle events for test assertions.
|
||||
type testObserver struct {
|
||||
mu sync.Mutex
|
||||
started int
|
||||
ended int
|
||||
rejected int
|
||||
dialErr int
|
||||
packets int
|
||||
bytes int
|
||||
}
|
||||
|
||||
func (o *testObserver) UDPSessionStarted(types.AccountID) { o.mu.Lock(); o.started++; o.mu.Unlock() }
|
||||
func (o *testObserver) UDPSessionEnded(types.AccountID) { o.mu.Lock(); o.ended++; o.mu.Unlock() }
|
||||
func (o *testObserver) UDPSessionDialError(types.AccountID) { o.mu.Lock(); o.dialErr++; o.mu.Unlock() }
|
||||
func (o *testObserver) UDPSessionRejected(types.AccountID) { o.mu.Lock(); o.rejected++; o.mu.Unlock() }
|
||||
func (o *testObserver) UDPPacketRelayed(_ types.RelayDirection, b int) {
|
||||
o.mu.Lock()
|
||||
o.packets++
|
||||
o.bytes += b
|
||||
o.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestRelay_CloseFiresObserverEnded(t *testing.T) {
|
||||
backend, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backend.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, err := backend.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backend.WriteTo(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
listener, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer listener.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
obs := &testObserver{}
|
||||
relay := New(ctx, RelayConfig{Logger: logger, Listener: listener, Target: backend.LocalAddr().String(), AccountID: "test-acct", DialFunc: dialFunc})
|
||||
relay.SetObserver(obs)
|
||||
go relay.Serve()
|
||||
|
||||
// Create two sessions.
|
||||
for i := range 2 {
|
||||
client, err := net.Dial("udp", listener.LocalAddr().String())
|
||||
require.NoError(t, err, "client %d", i)
|
||||
|
||||
_, err = client.Write([]byte("hello"))
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, client.SetReadDeadline(time.Now().Add(2*time.Second)))
|
||||
buf := make([]byte, 1024)
|
||||
_, err = client.Read(buf)
|
||||
require.NoError(t, err)
|
||||
client.Close()
|
||||
}
|
||||
|
||||
obs.mu.Lock()
|
||||
assert.Equal(t, 2, obs.started, "should have 2 started events")
|
||||
obs.mu.Unlock()
|
||||
|
||||
// Close should fire UDPSessionEnded for all remaining sessions.
|
||||
relay.Close()
|
||||
|
||||
obs.mu.Lock()
|
||||
assert.Equal(t, 2, obs.ended, "Close should fire UDPSessionEnded for each session")
|
||||
obs.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestRelay_SessionRateLimit(t *testing.T) {
|
||||
backend, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backend.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, err := backend.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backend.WriteTo(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
listener, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer listener.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
logger := log.NewEntry(log.StandardLogger())
|
||||
dialFunc := func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial(network, address)
|
||||
}
|
||||
|
||||
obs := &testObserver{}
|
||||
// High max sessions (1000) but the relay uses a rate limiter internally
|
||||
// (default: 50/s burst 100). We exhaust the burst by creating sessions
|
||||
// rapidly, then verify that subsequent creates are rejected.
|
||||
relay := New(ctx, RelayConfig{Logger: logger, Listener: listener, Target: backend.LocalAddr().String(), AccountID: "test-acct", DialFunc: dialFunc, MaxSessions: 1000})
|
||||
relay.SetObserver(obs)
|
||||
go relay.Serve()
|
||||
defer relay.Close()
|
||||
|
||||
// Exhaust the burst by calling getOrCreateSession directly with
|
||||
// synthetic addresses. This is faster than real UDP round-trips.
|
||||
for i := range sessionCreateBurst + 20 {
|
||||
addr := &net.UDPAddr{IP: net.IPv4(10, 0, byte(i/256), byte(i%256)), Port: 10000 + i}
|
||||
_, _ = relay.getOrCreateSession(addr)
|
||||
}
|
||||
|
||||
obs.mu.Lock()
|
||||
rejected := obs.rejected
|
||||
obs.mu.Unlock()
|
||||
|
||||
assert.Greater(t, rejected, 0, "some sessions should be rate-limited")
|
||||
}
|
||||
Reference in New Issue
Block a user