From ef58909fb73ae1dfd0dab7a6d5b7c7f5e0701959 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Thu, 21 May 2026 20:58:48 +0800 Subject: [PATCH] fix(dialer): correct inverted setMark on FreeBSD/OpenBSD Fix setMark early-return conditions that were inverted relative to Linux: on FreeBSD (SO_USER_COOKIE) and OpenBSD (SO_RTABLE), a non-zero mark was short-circuited instead of applied, disabling socket marking entirely. Also fix GetClientIP to trim leading whitespace from X-Forwarded-For entries (per RFC 7239), and fix Body.Read to subtract len(b) instead of n from recordSize when a single read exceeds the remaining quota. Add unit tests for internal/net/ (addr, dialer, http, ip, net, pipe, transport, resolve, proxyproto, udp). --- internal/net/addr_test.go | 255 ++++++++ internal/net/dialer/dialer_freebsd.go | 2 +- internal/net/dialer/dialer_linux_test.go | 40 ++ internal/net/dialer/dialer_openbsd.go | 2 +- internal/net/dialer/dialer_test.go | 137 +++++ internal/net/dialer/main_test.go | 31 + internal/net/http/http.go | 4 +- internal/net/http/http_test.go | 186 ++++++ internal/net/ip/ip_test.go | 36 ++ internal/net/main_test.go | 31 + internal/net/net_test.go | 224 +++++++ internal/net/pipe_test.go | 410 +++++++++++++ internal/net/proxyproto/proxyproto_test.go | 184 ++++++ internal/net/resovle_test.go | 140 +++++ internal/net/transport_test.go | 109 ++++ internal/net/udp/udp_test.go | 655 +++++++++++++++++++++ 16 files changed, 2442 insertions(+), 4 deletions(-) create mode 100644 internal/net/addr_test.go create mode 100644 internal/net/dialer/dialer_linux_test.go create mode 100644 internal/net/dialer/dialer_test.go create mode 100644 internal/net/dialer/main_test.go create mode 100644 internal/net/http/http_test.go create mode 100644 internal/net/ip/ip_test.go create mode 100644 internal/net/main_test.go create mode 100644 internal/net/net_test.go create mode 100644 internal/net/pipe_test.go create mode 100644 internal/net/proxyproto/proxyproto_test.go create mode 100644 internal/net/resovle_test.go create mode 100644 internal/net/transport_test.go create mode 100644 internal/net/udp/udp_test.go diff --git a/internal/net/addr_test.go b/internal/net/addr_test.go new file mode 100644 index 00000000..b40d2176 --- /dev/null +++ b/internal/net/addr_test.go @@ -0,0 +1,255 @@ +package net + +import ( + "net" + "net/netip" + "testing" +) + +func TestPortRange_Parse(t *testing.T) { + tests := []struct { + name string + input string + wantMin int + wantMax int + wantErr bool + }{ + {"single port", "8080", 8080, 8080, false}, + {"range", "8000-9000", 8000, 9000, false}, + {"zero", "0", 0, 0, false}, + {"max port", "65535", 65535, 65535, false}, + {"negative", "-1", 0, 0, true}, + {"too large", "65536", 0, 0, true}, + {"three parts", "1-2-3", 0, 0, true}, + {"invalid string", "abc", 0, 0, true}, + {"invalid min", "abc-123", 0, 0, true}, + {"invalid max", "123-abc", 0, 0, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var pr PortRange + err := pr.Parse(tt.input) + if (err != nil) != tt.wantErr { + t.Errorf("Parse() error = %v, wantErr %v", err, tt.wantErr) + return + } + if !tt.wantErr { + if pr.Min != tt.wantMin || pr.Max != tt.wantMax { + t.Errorf("Parse() got Min=%d Max=%d, want Min=%d Max=%d", pr.Min, pr.Max, tt.wantMin, tt.wantMax) + } + } + }) + } +} + +func TestPortRange_Contains(t *testing.T) { + pr := PortRange{Min: 8000, Max: 9000} + if !pr.Contains(8000) { + t.Error("should contain min") + } + if !pr.Contains(9000) { + t.Error("should contain max") + } + if !pr.Contains(8500) { + t.Error("should contain middle") + } + if pr.Contains(7999) { + t.Error("should not contain below min") + } + if pr.Contains(9001) { + t.Error("should not contain above max") + } +} + +func TestIPRange_Parse(t *testing.T) { + tests := []struct { + name string + input string + wantMin string + wantMax string + wantErr bool + }{ + {"single ipv4", "192.168.1.1", "192.168.1.1", "192.168.1.1", false}, + {"ipv4 range", "192.168.1.1-192.168.1.254", "192.168.1.1", "192.168.1.254", false}, + {"single ipv6", "::1", "::1", "::1", false}, + {"ipv6 range", "::1-::2", "::1", "::2", false}, + {"three parts", "1-2-3", "", "", true}, + {"invalid single", "not-an-ip", "", "", true}, + {"invalid min", "not-ip-192.168.1.2", "", "", true}, + {"invalid max", "192.168.1.1-not-ip", "", "", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var r IPRange + err := r.Parse(tt.input) + if (err != nil) != tt.wantErr { + t.Errorf("Parse() error = %v, wantErr %v", err, tt.wantErr) + return + } + if !tt.wantErr { + wantMin := netip.MustParseAddr(tt.wantMin) + wantMax := netip.MustParseAddr(tt.wantMax) + if r.Min != wantMin || r.Max != wantMax { + t.Errorf("Parse() got Min=%v Max=%v, want Min=%v Max=%v", r.Min, r.Max, wantMin, wantMax) + } + } + }) + } +} + +func TestIPRange_Contains(t *testing.T) { + r := IPRange{ + Min: netip.MustParseAddr("192.168.1.0"), + Max: netip.MustParseAddr("192.168.1.255"), + } + if !r.Contains(netip.MustParseAddr("192.168.1.0")) { + t.Error("should contain min") + } + if !r.Contains(netip.MustParseAddr("192.168.1.255")) { + t.Error("should contain max") + } + if !r.Contains(netip.MustParseAddr("192.168.1.100")) { + t.Error("should contain middle") + } + if r.Contains(netip.MustParseAddr("192.168.0.255")) { + t.Error("should not contain below min") + } + if r.Contains(netip.MustParseAddr("192.168.2.0")) { + t.Error("should not contain above max") + } +} + +func TestAddrPortRange_Addrs(t *testing.T) { + tests := []struct { + name string + apr AddrPortRange + want []string + }{ + {"with scheme", "http://192.168.1.1:8080", nil}, + {"tls scheme", "tls://192.168.1.1:8080", nil}, + {"invalid format", "not-a-valid-address", nil}, + {"single port", "192.168.1.1:8080", []string{"192.168.1.1:8080"}}, + {"port range", "192.168.1.1:8080-8082", []string{"192.168.1.1:8080", "192.168.1.1:8081", "192.168.1.1:8082"}}, + {"ipv6 single", "[::1]:8080", []string{"[::1]:8080"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := tt.apr.Addrs() + if len(got) != len(tt.want) { + t.Errorf("Addrs() = %v, want %v", got, tt.want) + return + } + for i := range got { + if got[i] != tt.want[i] { + t.Errorf("Addrs()[%d] = %s, want %s", i, got[i], tt.want[i]) + } + } + }) + } +} + +func Test_ipToAddr(t *testing.T) { + ip := net.ParseIP("192.168.1.1") + tests := []struct { + network string + want net.Addr + }{ + {"tcp", &net.TCPAddr{IP: ip, Port: 0}}, + {"tcp4", &net.TCPAddr{IP: ip, Port: 0}}, + {"tcp6", &net.TCPAddr{IP: ip, Port: 0}}, + {"udp", &net.UDPAddr{IP: ip, Port: 0}}, + {"udp4", &net.UDPAddr{IP: ip, Port: 0}}, + {"udp6", &net.UDPAddr{IP: ip, Port: 0}}, + {"ip", &net.IPAddr{IP: ip}}, + {"other", &net.IPAddr{IP: ip}}, + } + + for _, tt := range tests { + t.Run(tt.network, func(t *testing.T) { + got := ipToAddr(ip, tt.network) + if got.Network() != tt.want.Network() || got.String() != tt.want.String() { + t.Errorf("ipToAddr() = %#v, want %#v", got, tt.want) + } + }) + } +} + +func Test_findInterfaceByIP(t *testing.T) { + // Test with localhost IP + ip := net.ParseIP("127.0.0.1") + name, err := findInterfaceByIP(ip) + if err != nil { + t.Fatal(err) + } + // 127.0.0.1 should be on the loopback interface + if name == "" { + t.Log("127.0.0.1 not found on any interface (may happen in containers)") + } + + // Test with an IP that shouldn't be assigned + ip2 := net.ParseIP("203.0.113.1") + name2, err2 := findInterfaceByIP(ip2) + if err2 != nil { + t.Fatal(err2) + } + if name2 != "" { + t.Errorf("expected empty name for unassigned IP, got %s", name2) + } +} + +func TestParseInterfaceAddr(t *testing.T) { + // Empty interface name + ifce, addrs, err := ParseInterfaceAddr("", "tcp") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if ifce != "" { + t.Errorf("expected empty ifce, got %s", ifce) + } + if len(addrs) != 1 || addrs[0] != nil { + t.Errorf("expected [nil], got %v", addrs) + } + + // Non-existent interface + _, _, err = ParseInterfaceAddr("nonexistent_interface_xyz", "tcp") + if err == nil { + t.Error("expected error for non-existent interface") + } + + // IP string as interface name + ifce, addrs, err = ParseInterfaceAddr("127.0.0.1", "tcp") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if ifce == "" { + t.Log("127.0.0.1 not found on any interface (may happen in containers)") + } else if len(addrs) != 1 { + t.Errorf("expected 1 addr, got %d", len(addrs)) + } + + // IP that doesn't exist on any interface + ifce, addrs, err = ParseInterfaceAddr("203.0.113.1", "udp") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if ifce != "" { + t.Errorf("expected empty ifce for unassigned IP, got %s", ifce) + } + if len(addrs) != 1 { + t.Errorf("expected 1 addr, got %d", len(addrs)) + } + + // Test UDP network specific returns + ifce, addrs, err = ParseInterfaceAddr("127.0.0.1", "udp") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if ifce != "" { + if len(addrs) != 1 { + t.Errorf("expected 1 addr, got %d", len(addrs)) + } + } +} diff --git a/internal/net/dialer/dialer_freebsd.go b/internal/net/dialer/dialer_freebsd.go index 045b267e..7b9a1467 100644 --- a/internal/net/dialer/dialer_freebsd.go +++ b/internal/net/dialer/dialer_freebsd.go @@ -9,7 +9,7 @@ func bindDevice(network, address string, fd uintptr, ifceName string) error { } func setMark(fd uintptr, mark int) error { - if mark != 0 { + if mark == 0 { return nil } return unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_USER_COOKIE, mark) diff --git a/internal/net/dialer/dialer_linux_test.go b/internal/net/dialer/dialer_linux_test.go new file mode 100644 index 00000000..d66c948f --- /dev/null +++ b/internal/net/dialer/dialer_linux_test.go @@ -0,0 +1,40 @@ +package dialer + +import ( + "testing" +) + +func Test_bindDevice(t *testing.T) { + // Empty interface name - no-op + err := bindDevice("tcp", "127.0.0.1:80", 0, "") + if err != nil { + t.Errorf("expected nil for empty ifce, got %v", err) + } + + // Non-global unicast IP (loopback) with port - skips binding + err = bindDevice("tcp", "127.0.0.1:80", 0, "lo") + if err != nil { + t.Errorf("expected nil for non-global unicast, got %v", err) + } + + // ::1 (IPv6 loopback) is also non-global unicast + err = bindDevice("tcp", "[::1]:80", 0, "lo") + if err != nil { + t.Errorf("expected nil for IPv6 loopback, got %v", err) + } +} + +func Test_setMark(t *testing.T) { + // mark=0 is no-op + err := setMark(0, 0) + if err != nil { + t.Errorf("expected nil for mark=0, got %v", err) + } + + // Non-zero mark should succeed or fail based on fd validity + // We test with fd=0 which is stdin (should succeed with SO_MARK) + err = setMark(0, 1) + if err != nil { + t.Logf("setMark with fd=0: %v (expected on some systems)", err) + } +} diff --git a/internal/net/dialer/dialer_openbsd.go b/internal/net/dialer/dialer_openbsd.go index 83877d28..ada42152 100644 --- a/internal/net/dialer/dialer_openbsd.go +++ b/internal/net/dialer/dialer_openbsd.go @@ -9,7 +9,7 @@ func bindDevice(network, address string, fd uintptr, ifceName string) error { } func setMark(fd uintptr, mark int) error { - if mark != 0 { + if mark == 0 { return nil } return unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_RTABLE, mark) diff --git a/internal/net/dialer/dialer_test.go b/internal/net/dialer/dialer_test.go new file mode 100644 index 00000000..3b64e667 --- /dev/null +++ b/internal/net/dialer/dialer_test.go @@ -0,0 +1,137 @@ +package dialer + +import ( + "context" + "errors" + "net" + "testing" +) + +func TestDialer_Dial_DialFunc(t *testing.T) { + customErr := errors.New("custom dial") + d := &Dialer{ + DialFunc: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, customErr + }, + } + + _, err := d.Dial(context.Background(), "tcp", "127.0.0.1:0") + if err == nil { + t.Error("expected error from custom DialFunc") + } + if err != customErr { + t.Errorf("expected custom error, got %v", err) + } +} + +func TestDialer_Dial_DialFunc_Success(t *testing.T) { + expectedConn := &net.TCPConn{} + d := &Dialer{ + DialFunc: func(ctx context.Context, network, addr string) (net.Conn, error) { + return expectedConn, nil + }, + } + + conn, err := d.Dial(context.Background(), "tcp", "127.0.0.1:0") + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if conn != expectedConn { + t.Error("expected the connection from DialFunc") + } +} + +func TestDefaultTimeout(t *testing.T) { + if DefaultTimeout == 0 { + t.Error("DefaultTimeout should be non-zero") + } +} + +func TestDefaultNetDialer(t *testing.T) { + if DefaultNetDialer == nil { + t.Error("DefaultNetDialer should not be nil") + } +} + +func TestDialer_Dial_DialFunc_Unix(t *testing.T) { + d := &Dialer{ + DialFunc: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("unix error") + }, + } + + _, err := d.Dial(context.Background(), "unix", "/tmp/test.sock") + if err == nil { + t.Error("expected error from DialFunc") + } +} + +func TestDialer_Dial_NonStrictInterface(t *testing.T) { + // Non-existent interface, non-strict mode + d := &Dialer{ + Interface: "nonexistent_xyz", + DialFunc: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("dial failed") + }, + } + + _, err := d.Dial(context.Background(), "tcp", "127.0.0.1:0") + if err == nil { + t.Error("expected error from dial") + } +} + +func TestDialer_Dial_StrictInterfaceError(t *testing.T) { + // Strict mode: non-existent interface returns error immediately + d := &Dialer{ + Interface: "nonexistent_xyz!", + } + + _, err := d.Dial(context.Background(), "tcp", "127.0.0.1:0") + if err == nil { + t.Error("expected error from strict interface") + } +} + +func TestDialer_Dial_InterfaceWithIP(t *testing.T) { + // Use IP address as interface name in non-strict mode + d := &Dialer{ + Interface: "203.0.113.1", + DialFunc: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("dial failed") + }, + } + + _, err := d.Dial(context.Background(), "tcp", "127.0.0.1:0") + if err == nil { + t.Error("expected error from dial") + } +} + +func TestDialer_Dial_InterfaceLoopback(t *testing.T) { + // Use loopback interface + d := &Dialer{ + Interface: "lo", + DialFunc: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("dial failed") + }, + } + + _, err := d.Dial(context.Background(), "tcp", "127.0.0.1:0") + if err == nil { + t.Error("expected error from dial") + } +} + +func TestDialer_Dial_UDPInterface(t *testing.T) { + d := &Dialer{ + DialFunc: func(ctx context.Context, network, addr string) (net.Conn, error) { + return nil, errors.New("udp dial failed") + }, + } + + _, err := d.Dial(context.Background(), "udp", "") + if err == nil { + t.Error("expected error from dial") + } +} diff --git a/internal/net/dialer/main_test.go b/internal/net/dialer/main_test.go new file mode 100644 index 00000000..f8fd289a --- /dev/null +++ b/internal/net/dialer/main_test.go @@ -0,0 +1,31 @@ +package dialer + +import ( + "os" + "testing" + + "github.com/go-gost/core/logger" +) + +type nopLogger struct{} + +func (l *nopLogger) WithFields(map[string]any) logger.Logger { return l } +func (l *nopLogger) Trace(args ...any) {} +func (l *nopLogger) Tracef(format string, args ...any) {} +func (l *nopLogger) Debug(args ...any) {} +func (l *nopLogger) Debugf(format string, args ...any) {} +func (l *nopLogger) Info(args ...any) {} +func (l *nopLogger) Infof(format string, args ...any) {} +func (l *nopLogger) Warn(args ...any) {} +func (l *nopLogger) Warnf(format string, args ...any) {} +func (l *nopLogger) Error(args ...any) {} +func (l *nopLogger) Errorf(format string, args ...any) {} +func (l *nopLogger) Fatal(args ...any) {} +func (l *nopLogger) Fatalf(format string, args ...any) {} +func (l *nopLogger) GetLevel() logger.LogLevel { return logger.InfoLevel } +func (l *nopLogger) IsLevelEnabled(level logger.LogLevel) bool { return false } + +func TestMain(m *testing.M) { + logger.SetDefault(&nopLogger{}) + os.Exit(m.Run()) +} diff --git a/internal/net/http/http.go b/internal/net/http/http.go index 4cd48be6..39b96bfa 100644 --- a/internal/net/http/http.go +++ b/internal/net/http/http.go @@ -17,7 +17,7 @@ func GetClientIP(req *http.Request) net.IP { if sip == "" { ss := strings.Split(req.Header.Get("X-Forwarded-For"), ",") if len(ss) > 0 && ss[0] != "" { - sip = ss[0] + sip = strings.TrimSpace(ss[0]) } } if sip == "" { @@ -53,7 +53,7 @@ func (p *Body) Read(b []byte) (n int, err error) { b = b[:p.recordSize] } p.buf.Write(b) - p.recordSize -= n + p.recordSize -= len(b) } return diff --git a/internal/net/http/http_test.go b/internal/net/http/http_test.go new file mode 100644 index 00000000..76c9b3a2 --- /dev/null +++ b/internal/net/http/http_test.go @@ -0,0 +1,186 @@ +package http + +import ( + "bytes" + "io" + "net/http" + "strings" + "testing" +) + +func TestGetClientIP(t *testing.T) { + // nil request + if ip := GetClientIP(nil); ip != nil { + t.Errorf("expected nil for nil request, got %v", ip) + } + + tests := []struct { + name string + headers map[string]string + expected string + }{ + { + name: "CF-Connecting-IP takes priority", + headers: map[string]string{"CF-Connecting-IP": "1.2.3.4", "X-Forwarded-For": "5.6.7.8", "X-Real-Ip": "9.10.11.12"}, + expected: "1.2.3.4", + }, + { + name: "X-Forwarded-For fallback single", + headers: map[string]string{"X-Forwarded-For": "10.0.0.1"}, + expected: "10.0.0.1", + }, + { + name: "X-Forwarded-For fallback with multiple", + headers: map[string]string{"X-Forwarded-For": "10.0.0.1, 10.0.0.2, 10.0.0.3"}, + expected: "10.0.0.1", + }, + { + name: "X-Real-Ip fallback", + headers: map[string]string{"X-Real-Ip": "172.16.0.1"}, + expected: "172.16.0.1", + }, + { + name: "no headers", + headers: map[string]string{}, + expected: "", + }, + { + name: "empty X-Forwarded-For", + headers: map[string]string{"X-Forwarded-For": ""}, + expected: "", + }, + { + name: "X-Forwarded-For with leading space", + headers: map[string]string{"X-Forwarded-For": " 10.0.0.1, 10.0.0.2"}, + expected: "10.0.0.1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req, err := http.NewRequest("GET", "http://example.com", nil) + if err != nil { + t.Fatal(err) + } + for k, v := range tt.headers { + req.Header.Set(k, v) + } + + got := GetClientIP(req) + if tt.expected == "" { + if got != nil { + t.Errorf("expected nil, got %v", got) + } + } else { + if got == nil || got.String() != tt.expected { + t.Errorf("expected %s, got %v", tt.expected, got) + } + } + }) + } +} + +func TestBody_Read(t *testing.T) { + data := "hello world, this is a test message" + reader := io.NopCloser(strings.NewReader(data)) + + body := NewBody(reader, 10) + buf := make([]byte, 5) + + // First read - 5 bytes + n, err := body.Read(buf) + if err != nil { + t.Fatal(err) + } + if n != 5 { + t.Errorf("expected 5 bytes, got %d", n) + } + + // Second read - another 5 bytes (fills record size) + buf2 := make([]byte, 5) + n, err = body.Read(buf2) + if err != nil { + t.Fatal(err) + } + if n != 5 { + t.Errorf("expected 5 bytes, got %d", n) + } + + // Read remaining - record size exceeded, content not recorded + buf3 := make([]byte, 100) + n, err = body.Read(buf3) + if err != nil && err != io.EOF { + t.Fatal(err) + } + + content := body.Content() + if len(content) != 10 { + t.Errorf("content capped at 10, got %d bytes", len(content)) + } + + if body.Length() != int64(len(data)) { + t.Errorf("expected length %d, got %d", len(data), body.Length()) + } + + if err := body.Close(); err != nil { + t.Errorf("unexpected error on close: %v", err) + } +} + +func TestBody_ReadLargerThanRecordSize(t *testing.T) { + // Test the branch where a single read returns more bytes than remaining recordSize + data := "hello world!" + reader := io.NopCloser(strings.NewReader(data)) + + // recordSize=3, but we read in 8-byte chunks → n(8) > recordSize(3) + body := NewBody(reader, 3) + buf := make([]byte, 8) + n, err := body.Read(buf) + if err != nil && err != io.EOF { + t.Fatal(err) + } + if n != 8 { + t.Errorf("expected 8 bytes, got %d", n) + } + // Content should be capped at 3 bytes (recordSize) + content := body.Content() + if len(content) != 3 { + t.Errorf("expected 3 bytes recorded, got %d: %q", len(content), string(content)) + } +} + +func TestBody_NoRecord(t *testing.T) { + data := "hello world" + reader := io.NopCloser(strings.NewReader(data)) + + body := NewBody(reader, 0) + buf := make([]byte, 100) + n, err := body.Read(buf) + if err != nil && err != io.EOF { + t.Fatal(err) + } + + if body.Length() != int64(n) { + t.Errorf("expected length %d, got %d", n, body.Length()) + } + + content := body.Content() + if len(content) != 0 { + t.Errorf("expected empty content when recordSize=0, got %v", content) + } +} + +func TestBody_Content(t *testing.T) { + body := &Body{buf: *bytes.NewBufferString("recorded")} + content := body.Content() + if string(content) != "recorded" { + t.Errorf("expected 'recorded', got '%s'", string(content)) + } +} + +func TestBody_Length(t *testing.T) { + body := &Body{length: 42} + if body.Length() != 42 { + t.Errorf("expected 42, got %d", body.Length()) + } +} diff --git a/internal/net/ip/ip_test.go b/internal/net/ip/ip_test.go new file mode 100644 index 00000000..9ddd2d96 --- /dev/null +++ b/internal/net/ip/ip_test.go @@ -0,0 +1,36 @@ +package ip + +import ( + "testing" + + "github.com/songgao/water/waterutil" +) + +func TestProtocol(t *testing.T) { + tests := []struct { + name string + p waterutil.IPProtocol + want string + }{ + {"HOPOPT", waterutil.HOPOPT, "HOPOPT"}, + {"ICMP", waterutil.ICMP, "ICMP"}, + {"IGMP", waterutil.IGMP, "IGMP"}, + {"GGP", waterutil.GGP, "GGP"}, + {"TCP", waterutil.TCP, "TCP"}, + {"UDP", waterutil.UDP, "UDP"}, + {"IPv6-Route", waterutil.IPv6_Route, "IPv6-Route"}, + {"IPv6-Frag", waterutil.IPv6_Frag, "IPv6-Frag"}, + {"IPv6-ICMP", waterutil.IPv6_ICMP, "IPv6-ICMP"}, + {"unknown protocol", waterutil.IPProtocol(255), "unknown(255)"}, + {"unknown value", waterutil.IPProtocol(100), "unknown(100)"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := Protocol(tt.p) + if got != tt.want { + t.Errorf("Protocol() = %q, want %q", got, tt.want) + } + }) + } +} diff --git a/internal/net/main_test.go b/internal/net/main_test.go new file mode 100644 index 00000000..02c15314 --- /dev/null +++ b/internal/net/main_test.go @@ -0,0 +1,31 @@ +package net + +import ( + "os" + "testing" + + "github.com/go-gost/core/logger" +) + +type nopLogger struct{} + +func (l *nopLogger) WithFields(map[string]any) logger.Logger { return l } +func (l *nopLogger) Trace(args ...any) {} +func (l *nopLogger) Tracef(format string, args ...any) {} +func (l *nopLogger) Debug(args ...any) {} +func (l *nopLogger) Debugf(format string, args ...any) {} +func (l *nopLogger) Info(args ...any) {} +func (l *nopLogger) Infof(format string, args ...any) {} +func (l *nopLogger) Warn(args ...any) {} +func (l *nopLogger) Warnf(format string, args ...any) {} +func (l *nopLogger) Error(args ...any) {} +func (l *nopLogger) Errorf(format string, args ...any) {} +func (l *nopLogger) Fatal(args ...any) {} +func (l *nopLogger) Fatalf(format string, args ...any) {} +func (l *nopLogger) GetLevel() logger.LogLevel { return logger.InfoLevel } +func (l *nopLogger) IsLevelEnabled(level logger.LogLevel) bool { return false } + +func TestMain(m *testing.M) { + logger.SetDefault(&nopLogger{}) + os.Exit(m.Run()) +} diff --git a/internal/net/net_test.go b/internal/net/net_test.go new file mode 100644 index 00000000..061feeed --- /dev/null +++ b/internal/net/net_test.go @@ -0,0 +1,224 @@ +package net + +import ( + "errors" + "io" + "net" + "testing" + + xio "github.com/go-gost/x/internal/io" +) + +func TestIsIPv4(t *testing.T) { + tests := []struct { + address string + want bool + }{ + {"192.168.1.1", true}, + {"10.0.0.1", true}, + {"127.0.0.1:8080", true}, + {"::1", false}, + {"[::1]:8080", false}, + {"", false}, + {":8080", false}, + } + + for _, tt := range tests { + t.Run(tt.address, func(t *testing.T) { + if got := IsIPv4(tt.address); got != tt.want { + t.Errorf("IsIPv4(%q) = %v, want %v", tt.address, got, tt.want) + } + }) + } +} + +// mockConn implements net.Conn for testing +type mockConn struct { + net.Conn + readErr error + writeErr error +} + +func (m *mockConn) Read(p []byte) (int, error) { + if m.readErr != nil { + return 0, m.readErr + } + return 0, nil +} + +func (m *mockConn) Write(p []byte) (int, error) { + if m.writeErr != nil { + return 0, m.writeErr + } + return len(p), nil +} + +func (m *mockConn) Close() error { return nil } + +type mockAddr struct{} + +func (m mockAddr) Network() string { return "tcp" } +func (m mockAddr) String() string { return "127.0.0.1:0" } + +// mockReadWriteCloser implements io.ReadWriteCloser + CloseRead + CloseWrite +type mockReadWriteCloser struct { + data []byte + readPos int + writeBuf []byte + closed bool + readClosed bool + writeClosed bool + readErr error + writeErr error + writeReadErr error + deadlineFunc func() error +} + +func (m *mockReadWriteCloser) Read(p []byte) (int, error) { + if m.readClosed { + return 0, net.ErrClosed + } + if m.readErr != nil { + return 0, m.readErr + } + if m.readPos >= len(m.data) { + return 0, io.EOF + } + n := copy(p, m.data[m.readPos:]) + m.readPos += n + return n, nil +} + +func (m *mockReadWriteCloser) Write(p []byte) (int, error) { + if m.writeClosed { + return 0, net.ErrClosed + } + if m.writeErr != nil { + return 0, m.writeErr + } + m.writeBuf = append(m.writeBuf, p...) + return len(p), nil +} + +func (m *mockReadWriteCloser) Close() error { + m.closed = true + return nil +} + +func (m *mockReadWriteCloser) CloseRead() error { + m.readClosed = true + return nil +} + +func (m *mockReadWriteCloser) CloseWrite() error { + m.writeClosed = true + return xio.ErrUnsupported +} + +func (m *mockReadWriteCloser) SetReadDeadline(t interface{}) error { + if m.deadlineFunc != nil { + return m.deadlineFunc() + } + return nil +} + +func TestNewReadWriteConn(t *testing.T) { + r := &mockReadWriteCloser{data: []byte("read data")} + w := &mockReadWriteCloser{} + c := &mockConn{} + + rwc := NewReadWriteConn(r, w, c) + + // Test Read + buf := make([]byte, 9) + n, err := rwc.Read(buf) + if err != nil { + t.Error(err) + } + if n != 9 || string(buf[:n]) != "read data" { + t.Errorf("expected 'read data', got %q", string(buf[:n])) + } + + // Test Write + n, err = rwc.Write([]byte("write data")) + if err != nil { + t.Error(err) + } + if n != 10 { + t.Errorf("expected 10 bytes written, got %d", n) + } +} + +func Test_readWriteConn_CloseRead_CloseWrite(t *testing.T) { + r := &mockReadWriteCloser{} + w := &mockReadWriteCloser{} + c := &mockConn{} + + conn := NewReadWriteConn(r, w, c) + + if cr, ok := conn.(xio.CloseRead); ok { + err := cr.CloseRead() + if err == nil { + t.Error("expected error when CloseRead on conn without CloseRead") + } + } else { + t.Error("should implement CloseRead") + } + + if cw, ok := conn.(xio.CloseWrite); ok { + err := cw.CloseWrite() + if err == nil { + t.Error("expected error when CloseWrite on conn without CloseWrite") + } + } else { + t.Error("should implement CloseWrite") + } +} + +// closeReadWriteConn implements both CloseRead and CloseWrite +type closeReadWriteConn struct { + net.Conn +} + +func (c *closeReadWriteConn) CloseRead() error { return nil } +func (c *closeReadWriteConn) CloseWrite() error { return nil } +func (c *closeReadWriteConn) Close() error { return nil } + +func Test_readWriteConn_CloseRead_CloseWrite_supported(t *testing.T) { + r := &mockReadWriteCloser{} + w := &mockReadWriteCloser{} + c := &closeReadWriteConn{} + + conn := NewReadWriteConn(r, w, c) + + if cr, ok := conn.(xio.CloseRead); ok { + err := cr.CloseRead() + if err != nil { + t.Errorf("unexpected error: %v", err) + } + } else { + t.Error("should implement CloseRead") + } + + if cw, ok := conn.(xio.CloseWrite); ok { + err := cw.CloseWrite() + if err != nil { + t.Errorf("unexpected error: %v", err) + } + } else { + t.Error("should implement CloseWrite") + } +} + +// errReadWriteCloser returns errors on read/write, for error path testing +type errReadWriteCloser struct { + *mockReadWriteCloser +} + +func (m *errReadWriteCloser) Read(p []byte) (int, error) { + return 0, errors.New("read error") +} + +func (m *errReadWriteCloser) Write(p []byte) (int, error) { + return 0, errors.New("write error") +} diff --git a/internal/net/pipe_test.go b/internal/net/pipe_test.go new file mode 100644 index 00000000..e9abcb60 --- /dev/null +++ b/internal/net/pipe_test.go @@ -0,0 +1,410 @@ +package net + +import ( + "context" + "errors" + "io" + "net" + "testing" + "time" +) + +// pipeConn is a fully controllable in-memory connection for testing Pipe. +// It supports proper half-close semantics: CloseRead only stops reads, +// CloseWrite only stops writes, and Close stops both. +type pipeConn struct { + readBuf []byte + readPos int + writeBuf []byte + readClosed bool + writeClosed bool + closedCh chan struct{} + readDead time.Time + writeDead time.Time + readErr error // error to return on next Read + eofAfter int // return io.EOF after this many bytes total read (-1 = never) + bytesRead int +} + +func newPipeConn(data []byte) *pipeConn { + return &pipeConn{ + readBuf: data, + closedCh: make(chan struct{}), + } +} + +func (pc *pipeConn) Read(p []byte) (int, error) { + select { + case <-pc.closedCh: + return 0, net.ErrClosed + default: + } + if pc.readClosed { + return 0, net.ErrClosed + } + if pc.readErr != nil { + err := pc.readErr + pc.readErr = nil + return 0, err + } + if pc.readPos >= len(pc.readBuf) { + // eofAfter=-1 means sleep briefly then return (0, nil) to simulate + // a connection that's still open but has no data yet. + // This avoids busy-looping while still preventing EOF. + if pc.eofAfter < 0 { + time.Sleep(10 * time.Millisecond) + return 0, nil // no data, no error, no EOF + } + return 0, io.EOF + } + n := copy(p, pc.readBuf[pc.readPos:]) + pc.readPos += n + pc.bytesRead += n + return n, nil +} + +func (pc *pipeConn) Write(p []byte) (int, error) { + select { + case <-pc.closedCh: + return 0, net.ErrClosed + default: + } + if pc.writeClosed { + return 0, net.ErrClosed + } + pc.writeBuf = append(pc.writeBuf, p...) + return len(p), nil +} + +func (pc *pipeConn) Close() error { + select { + case <-pc.closedCh: + default: + close(pc.closedCh) + } + return nil +} + +func (pc *pipeConn) CloseRead() error { + pc.readClosed = true + return nil +} + +func (pc *pipeConn) CloseWrite() error { + pc.writeClosed = true + return nil +} + +func (pc *pipeConn) SetReadDeadline(t time.Time) error { + pc.readDead = t + return nil +} + +func (pc *pipeConn) SetWriteDeadline(t time.Time) error { + pc.writeDead = t + return nil +} + +func TestPipe_NormalTransfer(t *testing.T) { + ctx := context.Background() + c1 := newPipeConn([]byte("hello from c1")) + c2 := newPipeConn([]byte("hello from c2")) + + err := Pipe(ctx, c1, c2) + if err != nil { + t.Errorf("expected nil error, got %v", err) + } + + // c2's writeBuf has data from c1 + if string(c2.writeBuf) != "hello from c1" { + t.Errorf("c2.writeBuf = %q, want %q", string(c2.writeBuf), "hello from c1") + } + // c1's writeBuf has data from c2 + if string(c1.writeBuf) != "hello from c2" { + t.Errorf("c1.writeBuf = %q, want %q", string(c1.writeBuf), "hello from c2") + } +} + +func TestPipe_ImmediateEOF(t *testing.T) { + ctx := context.Background() + // Empty read buffers produce immediate EOF, simulating closed connections + c1 := newPipeConn(nil) + c2 := newPipeConn(nil) + + err := Pipe(ctx, c1, c2) + if err != nil { + t.Errorf("expected nil error (EOF is not an error from Pipe), got %v", err) + } +} + +func TestPipe_ContextCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + + // These never produce EOF and never error, so Pipe will run until canceled + c1 := newPipeConn(nil) + c2 := newPipeConn(nil) + c1.eofAfter = -1 // disable EOF + c2.eofAfter = -1 + + // Start cancellation in a goroutine + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + err := Pipe(ctx, c1, c2) + if err == nil { + t.Error("expected context error, got nil") + } +} + +func TestPipe_ContextAlreadyCanceled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + c1 := newPipeConn(nil) + c1.eofAfter = -1 + c2 := newPipeConn(nil) + c2.eofAfter = -1 + + err := Pipe(ctx, c1, c2) + if err == nil { + t.Error("expected context error, got nil") + } +} + +func TestPipe_ReadError(t *testing.T) { + ctx := context.Background() + readErr := errors.New("read failure") + c1 := newPipeConn(nil) + c1.readErr = readErr + // c2 also has a read error so both goroutines terminate + otherErr := errors.New("other error") + c2 := newPipeConn(nil) + c2.readErr = otherErr + + err := Pipe(ctx, c1, c2) + if err == nil { + t.Error("expected error, got nil") + } +} + +func TestPipe_WriteError(t *testing.T) { + ctx := context.Background() + c1 := newPipeConn([]byte("data")) + // c2 has its write side pre-closed + c2 := newPipeConn(nil) + c2.writeClosed = true + + err := Pipe(ctx, c1, c2) + if err == nil { + t.Error("expected write error, got nil") + } +} + +func Test_readDeadliner(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + + r := &mockReadWriteCloser{data: []byte("testdata")} + rd := &readDeadliner{Reader: r, ctx: ctx} + + // Normal read + buf := make([]byte, 8) + n, err := rd.Read(buf) + if err != nil && err != io.EOF { + t.Fatal(err) + } + if n != 8 || string(buf[:n]) != "testdata" { + t.Errorf("expected 'testdata', got %q", string(buf[:n])) + } + + // Context canceled + cancel() + _, err = rd.Read(buf) + if err == nil { + t.Error("expected context error, got nil") + } +} + +func Test_halfClose(t *testing.T) { + // pipeConn supports CloseRead and CloseWrite (success). + // CloseWrite returns nil (successful half-close), which sets a read deadline on dst. + src := newPipeConn(nil) + dst := newPipeConn(nil) + halfClose(src, dst) + if !src.readClosed { + t.Error("src should have CloseRead called") + } + if !dst.writeClosed { + t.Error("dst should have CloseWrite called (success path)") + } + if dst.readDead.IsZero() { + t.Error("dst should have read deadline set after successful CloseWrite") + } +} + +func Test_halfClose_MockReadWriteCloser(t *testing.T) { + // mockReadWriteCloser has CloseRead, but CloseWrite returns ErrUnsupported -> dst.Close() + src2 := &mockReadWriteCloser{} + dst2 := &mockReadWriteCloser{} + halfClose(src2, dst2) + if !src2.readClosed { + t.Error("src should have CloseRead called") + } + if !dst2.closed { + t.Error("dst should be fully closed when CloseWrite is ErrUnsupported") + } +} + +func Test_forceClose(t *testing.T) { + c1 := newPipeConn(nil) + c2 := newPipeConn(nil) + forceClose(c1, c2, nil) + // Check that closedCh is closed + select { + case <-c1.closedCh: + default: + t.Error("c1 should be closed") + } + select { + case <-c2.closedCh: + default: + t.Error("c2 should be closed") + } +} + +func Test_pipeHalf(t *testing.T) { + ctx := context.Background() + src := newPipeConn([]byte("hello world")) + dst := newPipeConn(nil) + + err := pipeHalf(ctx, src, dst) + if err != nil { + t.Fatal(err) + } + if string(dst.writeBuf) != "hello world" { + t.Errorf("dst.writeBuf = %q, want %q", string(dst.writeBuf), "hello world") + } +} + +func Test_pipeHalf_ReadError(t *testing.T) { + ctx := context.Background() + readErr := errors.New("read failed") + src := newPipeConn(nil) + src.readErr = readErr + dst := newPipeConn(nil) + + err := pipeHalf(ctx, src, dst) + if err == nil { + t.Error("expected error, got nil") + } +} + +func Test_pipeHalf_WriteError(t *testing.T) { + ctx := context.Background() + src := newPipeConn([]byte("data")) + dst := newPipeConn(nil) + dst.writeClosed = true + + err := pipeHalf(ctx, src, dst) + if err == nil { + t.Error("expected error, got nil") + } +} + +func Test_pipeHalf_ContextCanceled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + src := newPipeConn([]byte("data")) + src.eofAfter = -1 + dst := newPipeConn(nil) + + err := pipeHalf(ctx, src, dst) + if err == nil { + t.Error("expected context error, got nil") + } +} + +func TestPipe_CtxDoneWithError(t *testing.T) { + // Test the ctx.Done() path where firstErr is non-nil + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + c1 := newPipeConn(nil) + c1.readErr = errors.New("error before cancel") + c1.eofAfter = -1 + c2 := newPipeConn(nil) + c2.eofAfter = -1 + + done := make(chan error, 1) + go func() { + done <- Pipe(ctx, c1, c2) + }() + + time.Sleep(20 * time.Millisecond) + cancel() + + err := <-done + if err == nil { + t.Error("expected error, got nil") + } +} + +func TestPipe_CtxDoneWithoutError(t *testing.T) { + // Standard ctx.Done path with no prior error → returns ctx.Err() + ctx, cancel := context.WithCancel(context.Background()) + + c1 := newPipeConn(nil) + c1.eofAfter = -1 + c2 := newPipeConn(nil) + c2.eofAfter = -1 + + go func() { + time.Sleep(30 * time.Millisecond) + cancel() + }() + + err := Pipe(ctx, c1, c2) + if err == nil { + t.Error("expected context error, got nil") + } +} + +func TestTransport_ReturnsError(t *testing.T) { + // Both directions must return errors since Transport reads one from errc + readErr := errors.New("not EOF error") + rw1 := &mockReadWriter{readErr: readErr} + rw2 := &mockReadWriter{readErr: errors.New("other error")} + + err := Transport(rw1, rw2) + if err == nil { + t.Error("expected error, got nil") + } +} + +func Test_halfClose_NoCloseWriteInterface(t *testing.T) { + // Test the path where dst doesn't implement CloseWrite + type noCloseWrite struct { + *mockReadWriteCloser + } + + src := &mockReadWriteCloser{} + // Create a type that doesn't have CloseWrite (use a type without the method) + // Actually, we need a type that implements io.ReadWriteCloser but NOT CloseWrite + dst := &noCloseWriteCloser{} + halfClose(src, dst) + if !dst.closed { + // halfClose checks cw, ok := dst.(xio.CloseWrite) → false + // falls through to dst.Close() + t.Error("dst should be closed when CloseWrite is not available") + } +} + +type noCloseWriteCloser struct { + closed bool +} + +func (c *noCloseWriteCloser) Read(p []byte) (int, error) { return 0, nil } +func (c *noCloseWriteCloser) Write(p []byte) (int, error) { return len(p), nil } +func (c *noCloseWriteCloser) Close() error { c.closed = true; return nil } diff --git a/internal/net/proxyproto/proxyproto_test.go b/internal/net/proxyproto/proxyproto_test.go new file mode 100644 index 00000000..c92ceac8 --- /dev/null +++ b/internal/net/proxyproto/proxyproto_test.go @@ -0,0 +1,184 @@ +package proxyproto + +import ( + "net" + "testing" + "time" + + xio "github.com/go-gost/x/internal/io" + proxyproto "github.com/pires/go-proxyproto" +) + +type testConn struct { + net.Conn +} + +func (c *testConn) Close() error { return nil } +func (c *testConn) RemoteAddr() net.Addr { return &net.TCPAddr{IP: net.IPv4(1, 2, 3, 4), Port: 1234} } +func (c *testConn) LocalAddr() net.Addr { return &net.TCPAddr{IP: net.IPv4(5, 6, 7, 8), Port: 5678} } +func (c *testConn) CloseRead() error { return xio.ErrUnsupported } +func (c *testConn) CloseWrite() error { return xio.ErrUnsupported } + +type testListener struct { + net.Listener +} + +func (l *testListener) Accept() (net.Conn, error) { + return &testConn{}, nil +} + +func (l *testListener) Close() error { return nil } +func (l *testListener) Addr() net.Addr { return &net.TCPAddr{} } + +func TestWrapListener_NoProxyProto(t *testing.T) { + ln := &testListener{} + result := WrapListener(0, ln, 0) + if result != ln { + t.Error("expected original listener when ppv=0") + } +} + +func TestWrapListener_WithProxyProto(t *testing.T) { + ln := &testListener{} + result := WrapListener(1, ln, 5*time.Second) + if result == ln { + t.Error("expected wrapped listener when ppv>0") + } + _, ok := result.(*listener) + if !ok { + t.Errorf("expected *listener, got %T", result) + } +} + +func Test_serverConn_Context(t *testing.T) { + conn := &testConn{} + sc := &serverConn{Conn: conn} + if sc.Context() != nil { + // When no context set, it should be nil + _ = sc.Context() + } +} + +func TestWrapClientConn_NoProxyProto(t *testing.T) { + c := &testConn{} + result := WrapClientConn(0, nil, nil, c) + if result != c { + t.Error("expected original conn when ppv=0") + } +} + +func TestWrapClientConn_NilConn(t *testing.T) { + result := WrapClientConn(1, nil, nil, nil) + if result != nil { + t.Error("expected nil for nil conn") + } +} + +func TestWrapClientConn_NilSrc(t *testing.T) { + c := &testConn{} + result := WrapClientConn(1, nil, &net.TCPAddr{IP: net.IPv4(1, 1, 1, 1), Port: 80}, c) + if result != c { + t.Error("expected original conn when src is nil") + } +} + +func TestWrapClientConn_NilDst(t *testing.T) { + c := &testConn{} + result := WrapClientConn(1, &net.TCPAddr{IP: net.IPv4(1, 1, 1, 1), Port: 80}, nil, c) + if result != c { + t.Error("expected original conn when dst is nil") + } +} + +func Test_convertAddr(t *testing.T) { + // nil addr + if addr := convertAddr(nil); addr != nil { + t.Errorf("expected nil, got %v", addr) + } + + // TCP addr with valid IP + tcpAddr := &net.TCPAddr{IP: net.IPv4(192, 168, 1, 1), Port: 8080} + result := convertAddr(tcpAddr) + tcpResult, ok := result.(*net.TCPAddr) + if !ok { + t.Errorf("expected *net.TCPAddr, got %T", result) + } else { + if !tcpResult.IP.Equal(net.IPv4(192, 168, 1, 1)) || tcpResult.Port != 8080 { + t.Errorf("got %v, want 192.168.1.1:8080", result) + } + } + + // UDP addr + udpAddr := &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 53} + result = convertAddr(udpAddr) + udpResult, ok := result.(*net.UDPAddr) + if !ok { + t.Errorf("expected *net.UDPAddr, got %T", result) + } else { + if !udpResult.IP.Equal(net.IPv4(10, 0, 0, 1)) || udpResult.Port != 53 { + t.Errorf("got %v, want 10.0.0.1:53", result) + } + } + + // addr with IPv6 zero (unspecified) + ip6zero := &net.TCPAddr{IP: net.IPv6zero, Port: 9999} + result = convertAddr(ip6zero) + tcpResult2, ok := result.(*net.TCPAddr) + if !ok { + t.Errorf("expected *net.TCPAddr, got %T", result) + } else if !tcpResult2.IP.Equal(net.IPv4zero) { + t.Errorf("expected IPv4zero, got %v", tcpResult2.IP) + } + + // addr with nil IP + nilIP := &net.TCPAddr{IP: nil, Port: 1234} + result = convertAddr(nilIP) + tcpResult3, ok := result.(*net.TCPAddr) + if !ok { + t.Errorf("expected *net.TCPAddr, got %T", result) + } else if !tcpResult3.IP.Equal(net.IPv4zero) { + t.Errorf("expected IPv4zero for nil IP, got %v", tcpResult3.IP) + } +} + +func Test_serverConn_RemoteAddr_LocalAddr(t *testing.T) { + conn := &testConn{} + sc := &serverConn{Conn: conn} + if sc.RemoteAddr().String() != "1.2.3.4:1234" { + t.Errorf("unexpected RemoteAddr: %s", sc.RemoteAddr()) + } + if sc.LocalAddr().String() != "5.6.7.8:5678" { + t.Errorf("unexpected LocalAddr: %s", sc.LocalAddr()) + } +} + +func Test_serverConn_RemoteAddr_LocalAddr_WithProxyproto(t *testing.T) { + // When the inner conn is a proxyproto.Conn + rawConn := &testConn{} + ppConn := proxyproto.NewConn(rawConn) + sc := &serverConn{Conn: ppConn} + if sc.RemoteAddr().String() != "1.2.3.4:1234" { + t.Errorf("unexpected RemoteAddr: %s", sc.RemoteAddr()) + } + if sc.LocalAddr().String() != "5.6.7.8:5678" { + t.Errorf("unexpected LocalAddr: %s", sc.LocalAddr()) + } +} + +func Test_serverConn_CloseRead(t *testing.T) { + conn := &testConn{} + sc := &serverConn{Conn: conn} + err := sc.CloseRead() + if err != xio.ErrUnsupported { + t.Errorf("expected ErrUnsupported, got %v", err) + } +} + +func Test_serverConn_CloseWrite(t *testing.T) { + conn := &testConn{} + sc := &serverConn{Conn: conn} + err := sc.CloseWrite() + if err != xio.ErrUnsupported { + t.Errorf("expected ErrUnsupported, got %v", err) + } +} diff --git a/internal/net/resovle_test.go b/internal/net/resovle_test.go new file mode 100644 index 00000000..6726555e --- /dev/null +++ b/internal/net/resovle_test.go @@ -0,0 +1,140 @@ +package net + +import ( + "context" + "net" + "testing" + + "github.com/go-gost/core/hosts" + "github.com/go-gost/core/resolver" +) + +type mockResolver struct { + ips []net.IP + err error +} + +func (m *mockResolver) Resolve(ctx context.Context, network, host string, opts ...resolver.Option) ([]net.IP, error) { + if m.err != nil { + return nil, m.err + } + return m.ips, nil +} + +type mockHostMapper struct { + ips []net.IP +} + +func (m *mockHostMapper) Lookup(ctx context.Context, network, host string, opts ...hosts.Option) ([]net.IP, bool) { + if len(m.ips) == 0 { + return nil, false + } + return m.ips, true +} + +func TestResolve_EmptyAddr(t *testing.T) { + ctx := context.Background() + addr, err := Resolve(ctx, "tcp", "", nil, nil, nil) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if addr != "" { + t.Errorf("expected empty, got %s", addr) + } +} + +func TestResolve_NoHost(t *testing.T) { + ctx := context.Background() + addr, err := Resolve(ctx, "tcp", ":8080", nil, nil, nil) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if addr != ":8080" { + t.Errorf("expected :8080, got %s", addr) + } +} + +func TestResolve_HostMapperMatch(t *testing.T) { + ctx := context.Background() + mapper := &mockHostMapper{ips: []net.IP{net.ParseIP("10.0.0.1")}} + r := &mockResolver{ips: []net.IP{net.ParseIP("192.168.1.1")}} + + addr, err := Resolve(ctx, "tcp", "example.com:8080", r, mapper, nil) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if addr != "10.0.0.1:8080" { + t.Errorf("expected 10.0.0.1:8080, got %s", addr) + } +} + +func TestResolve_Resolver(t *testing.T) { + ctx := context.Background() + r := &mockResolver{ips: []net.IP{net.ParseIP("192.168.1.1")}} + + addr, err := Resolve(ctx, "tcp", "example.com:8080", r, nil, nil) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if addr != "192.168.1.1:8080" { + t.Errorf("expected 192.168.1.1:8080, got %s", addr) + } +} + +func TestResolve_ResolverInvalid(t *testing.T) { + ctx := context.Background() + r := &mockResolver{err: resolver.ErrInvalid} + + addr, err := Resolve(ctx, "tcp", "example.com:8080", r, nil, nil) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if addr != "example.com:8080" { + t.Errorf("expected example.com:8080, got %s", addr) + } +} + +func TestResolve_ResolverOtherError(t *testing.T) { + ctx := context.Background() + r := &mockResolver{err: net.UnknownNetworkError("test-error")} + + _, err := Resolve(ctx, "tcp", "example.com:8080", r, nil, nil) + if err == nil { + t.Error("expected error, got nil") + } +} + +func TestResolve_ResolverEmptyResult(t *testing.T) { + ctx := context.Background() + r := &mockResolver{ips: nil} + + _, err := Resolve(ctx, "tcp", "example.com:8080", r, nil, nil) + if err == nil { + t.Error("expected error, got nil") + } +} + +func TestResolve_NoResolverNoHosts(t *testing.T) { + ctx := context.Background() + addr, err := Resolve(ctx, "tcp", "example.com:8080", nil, nil, nil) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if addr != "example.com:8080" { + t.Errorf("expected example.com:8080, got %s", addr) + } +} + +func TestResolve_HostsNoMatchUsesResolver(t *testing.T) { + ctx := context.Background() + mapper := &mockHostMapper{} + r := &mockResolver{ips: []net.IP{net.ParseIP("10.0.0.1")}} + + addr, err := Resolve(ctx, "tcp", "example.com:8080", r, mapper, nil) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if addr != "10.0.0.1:8080" { + t.Errorf("expected 10.0.0.1:8080, got %s", addr) + } +} diff --git a/internal/net/transport_test.go b/internal/net/transport_test.go new file mode 100644 index 00000000..54d9d3a5 --- /dev/null +++ b/internal/net/transport_test.go @@ -0,0 +1,109 @@ +package net + +import ( + "errors" + "io" + "testing" +) + +// mockReadWriter implements io.ReadWriter +type mockReadWriter struct { + data []byte + readPos int + readErr error + + writeBuf []byte +} + +func (m *mockReadWriter) Read(p []byte) (int, error) { + if m.readErr != nil { + return 0, m.readErr + } + if m.readPos >= len(m.data) { + return 0, io.EOF + } + n := copy(p, m.data[m.readPos:]) + m.readPos += n + return n, nil +} + +func (m *mockReadWriter) Write(p []byte) (int, error) { + m.writeBuf = append(m.writeBuf, p...) + return len(p), nil +} + +func TestTransport_Normal(t *testing.T) { + rw1 := &mockReadWriter{data: []byte("from rw1")} + rw2 := &mockReadWriter{data: []byte("from rw2")} + + done := make(chan struct{}, 2) + + go func() { + _ = CopyBuffer(rw1, rw2, 1024) + done <- struct{}{} + }() + go func() { + _ = CopyBuffer(rw2, rw1, 1024) + done <- struct{}{} + }() + + // Wait for both goroutines to complete + <-done + <-done + + // rw2 receives data from rw1 + if string(rw2.writeBuf) != "from rw1" { + t.Errorf("rw2.writeBuf = %q, want %q", string(rw2.writeBuf), "from rw1") + } + // rw1 receives data from rw2 + if string(rw1.writeBuf) != "from rw2" { + t.Errorf("rw1.writeBuf = %q, want %q", string(rw1.writeBuf), "from rw2") + } +} + +func TestTransport_ReadError(t *testing.T) { + readErr := errors.New("read failure") + rw1 := &mockReadWriter{readErr: readErr} + rw2 := &mockReadWriter{data: []byte("from rw2")} + + err := Transport(rw1, rw2) + if err == nil { + t.Error("expected error, got nil") + } +} + +func TestCopyBuffer(t *testing.T) { + src := &mockReadWriter{data: []byte("test data")} + dst := &mockReadWriter{} + + err := CopyBuffer(dst, src, 1024) + if err != nil { + t.Fatal(err) + } + if string(dst.writeBuf) != "test data" { + t.Errorf("dst.writeBuf = %q, want %q", string(dst.writeBuf), "test data") + } +} + +func TestCopyBuffer_ReadError(t *testing.T) { + readErr := errors.New("read failure") + src := &mockReadWriter{readErr: readErr} + dst := &mockReadWriter{} + + err := CopyBuffer(dst, src, 1024) + if err == nil { + t.Error("expected error, got nil") + } +} + +func TestTransport_EOFIgnored(t *testing.T) { + // Transport treats io.EOF as non-error. + // Both goroutines return nil/EOF → Transport should return nil. + rw1 := &mockReadWriter{data: []byte("data")} + rw2 := &mockReadWriter{data: []byte("")} // empty data → immediate EOF + + err := Transport(rw1, rw2) + if err != nil { + t.Errorf("expected nil (EOF is not an error), got %v", err) + } +} diff --git a/internal/net/udp/udp_test.go b/internal/net/udp/udp_test.go new file mode 100644 index 00000000..8f8d1b67 --- /dev/null +++ b/internal/net/udp/udp_test.go @@ -0,0 +1,655 @@ +package udp + +import ( + "context" + "errors" + "net" + "sync" + "testing" + "time" + + "github.com/go-gost/core/bypass" + "github.com/go-gost/core/logger" +) + +// nopLogger is a no-op logger for testing +type nopLogger struct{} + +func (l *nopLogger) WithFields(map[string]any) logger.Logger { return l } +func (l *nopLogger) Trace(args ...any) {} +func (l *nopLogger) Tracef(format string, args ...any) {} +func (l *nopLogger) Debug(args ...any) {} +func (l *nopLogger) Debugf(format string, args ...any) {} +func (l *nopLogger) Info(args ...any) {} +func (l *nopLogger) Infof(format string, args ...any) {} +func (l *nopLogger) Warn(args ...any) {} +func (l *nopLogger) Warnf(format string, args ...any) {} +func (l *nopLogger) Error(args ...any) {} +func (l *nopLogger) Errorf(format string, args ...any) {} +func (l *nopLogger) Fatal(args ...any) {} +func (l *nopLogger) Fatalf(format string, args ...any) {} +func (l *nopLogger) GetLevel() logger.LogLevel { return logger.InfoLevel } +func (l *nopLogger) IsLevelEnabled(level logger.LogLevel) bool { return false } + +// mockPacketConn implements net.PacketConn for unit testing +type mockPacketConn struct { + readFn func(b []byte) (int, net.Addr, error) + writeFn func(b []byte, addr net.Addr) (int, error) + closeErr error + localAddr net.Addr + closed bool + mu sync.Mutex +} + +func newMockPacketConn() *mockPacketConn { + return &mockPacketConn{ + localAddr: &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 10000}, + } +} + +func (m *mockPacketConn) ReadFrom(b []byte) (int, net.Addr, error) { + if m.readFn != nil { + return m.readFn(b) + } + // Default: return on first call, then error + return 0, &net.UDPAddr{IP: net.IPv4(1, 2, 3, 4), Port: 1234}, errors.New("no readFn set") +} + +func (m *mockPacketConn) WriteTo(b []byte, addr net.Addr) (int, error) { + if m.writeFn != nil { + return m.writeFn(b, addr) + } + return len(b), nil +} + +func (m *mockPacketConn) Close() error { + m.mu.Lock() + defer m.mu.Unlock() + m.closed = true + return m.closeErr +} + +func (m *mockPacketConn) LocalAddr() net.Addr { return m.localAddr } +func (m *mockPacketConn) SetDeadline(t time.Time) error { return nil } +func (m *mockPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (m *mockPacketConn) SetWriteDeadline(t time.Time) error { return nil } +func (m *mockPacketConn) SetReadBuffer(n int) error { return nil } +func (m *mockPacketConn) SetWriteBuffer(n int) error { return nil } +func (m *mockPacketConn) SyscallConn() (interface{}, error) { + return nil, nil +} +func (m *mockPacketConn) RemoteAddr() net.Addr { + return &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 20000} +} + +func Test_connPool_New(t *testing.T) { + p := newConnPool(time.Second) + if p == nil { + t.Fatal("expected non-nil pool") + } + p.Close() +} + +func Test_connPool_NilReceiver(t *testing.T) { + var p *connPool + if c, ok := p.Get("key"); c != nil || ok { + t.Error("nil receiver should return nil, false") + } + p.Set("key", &conn{}) + p.Delete("key") + p.Close() +} + +func Test_connPool_SetGetDelete(t *testing.T) { + p := newConnPool(time.Hour) + defer p.Close() + + c := newConn(nil, nil, nil, 1, false) + + got, ok := p.Get("key") + if ok || got != nil { + t.Error("expected no value for non-existent key") + } + + p.Set("key", c) + got, ok = p.Get("key") + if !ok || got != c { + t.Error("expected to get the set value") + } + + p.Delete("key") + got, ok = p.Get("key") + if ok || got != nil { + t.Error("expected no value after delete") + } +} + +func Test_connPool_CloseCleansUp(t *testing.T) { + p := newConnPool(time.Hour) + c := newConn(nil, nil, nil, 1, false) + p.Set("key", c) + p.Close() + // Double close should be safe + p.Close() +} + +func Test_connPool_idleCheck(t *testing.T) { + p := newConnPool(50 * time.Millisecond) + p.WithLogger(&nopLogger{}) + defer p.Close() + + c := newConn(nil, nil, nil, 1, false) + c.SetIdle(true) + p.Set("key", c) + + time.Sleep(150 * time.Millisecond) + + got, ok := p.Get("key") + if ok || got != nil { + t.Error("idle connection should have been cleaned up") + } +} + +func Test_conn_Close(t *testing.T) { + c := newConn(newMockPacketConn(), &net.UDPAddr{Port: 1000}, &net.UDPAddr{Port: 2000}, 10, false) + if c.isClosed() { + t.Error("should not be closed initially") + } + err := c.Close() + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if !c.isClosed() { + t.Error("should be closed after Close()") + } + err = c.Close() + if err != nil { + t.Error("double close should not error") + } +} + +func Test_conn_IsIdle_SetIdle(t *testing.T) { + c := newConn(nil, nil, nil, 1, false) + if c.IsIdle() { + t.Error("should not be idle initially") + } + c.SetIdle(true) + if !c.IsIdle() { + t.Error("should be idle after SetIdle(true)") + } + c.SetIdle(false) + if c.IsIdle() { + t.Error("should not be idle after SetIdle(false)") + } +} + +func Test_conn_ReadFrom(t *testing.T) { + c := newConn(newMockPacketConn(), &net.UDPAddr{Port: 1000}, &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 2000}, 10, false) + + testData := []byte("test data") + err := c.WriteQueue(testData) + if err != nil { + t.Fatal(err) + } + + buf := make([]byte, 1024) + n, addr, err := c.ReadFrom(buf) + if err != nil { + t.Fatal(err) + } + if n != len(testData) { + t.Errorf("expected %d bytes, got %d", len(testData), n) + } + if string(buf[:n]) != "test data" { + t.Errorf("expected 'test data', got %q", string(buf[:n])) + } + if addr.String() != "10.0.0.1:2000" { + t.Errorf("expected remote addr, got %v", addr) + } +} + +func Test_conn_ReadFrom_QueueFull(t *testing.T) { + c := newConn(nil, nil, nil, 0, false) + err := c.WriteQueue([]byte("data")) + if err == nil { + t.Error("expected error when queue is full (size 0)") + } +} + +func Test_conn_ReadFrom_Closed(t *testing.T) { + c := newConn(nil, nil, nil, 0, false) + c.Close() + + err := c.WriteQueue([]byte("data")) + if err == nil { + t.Error("expected error (queue full or closed), got nil") + } + + _, _, err = c.ReadFrom(make([]byte, 10)) + if err != net.ErrClosed { + t.Errorf("expected ErrClosed, got %v", err) + } +} + +func Test_conn_Read(t *testing.T) { + c := newConn(newMockPacketConn(), nil, &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 2000}, 10, false) + + testData := []byte("hello") + c.WriteQueue(testData) + + buf := make([]byte, 1024) + n, err := c.Read(buf) + if err != nil { + t.Fatal(err) + } + if n != len(testData) { + t.Errorf("expected %d bytes, got %d", len(testData), n) + } +} + +func Test_conn_Write(t *testing.T) { + pc := newMockPacketConn() + var written []byte + pc.writeFn = func(b []byte, addr net.Addr) (int, error) { + written = append(written, b...) + return len(b), nil + } + + c := newConn(pc, nil, &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 2000}, 10, true) + + n, err := c.Write([]byte("data")) + if err != nil { + t.Fatal(err) + } + if n != 4 { + t.Errorf("expected 4 bytes, got %d", n) + } + if string(written) != "data" { + t.Errorf("expected 'data', got %q", string(written)) + } + + if c.isClosed() { + t.Error("should not close after Write when keepalive=true") + } +} + +func Test_conn_WriteTo_NoKeepalive(t *testing.T) { + pc := newMockPacketConn() + c := newConn(pc, nil, &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 2000}, 10, false) + + _, err := c.WriteTo([]byte("data"), &net.UDPAddr{IP: net.IPv4(1, 2, 3, 4), Port: 9999}) + if err != nil { + t.Fatal(err) + } + + if !c.isClosed() { + t.Error("expected conn to close after WriteTo when keepalive=false") + } +} + +func Test_conn_WriteTo_Keepalive(t *testing.T) { + pc := newMockPacketConn() + c := newConn(pc, nil, &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 2000}, 10, true) + + _, err := c.WriteTo([]byte("data"), &net.UDPAddr{IP: net.IPv4(1, 2, 3, 4), Port: 9999}) + if err != nil { + t.Fatal(err) + } + + if c.isClosed() { + t.Error("should not close after WriteTo when keepalive=true") + } +} + +func Test_conn_LocalAddr_RemoteAddr(t *testing.T) { + laddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8000} + raddr := &net.UDPAddr{IP: net.IPv4(192, 168, 1, 1), Port: 9000} + c := newConn(nil, laddr, raddr, 1, false) + + if c.LocalAddr().String() != laddr.String() { + t.Errorf("expected %v, got %v", laddr, c.LocalAddr()) + } + if c.RemoteAddr().String() != raddr.String() { + t.Errorf("expected %v, got %v", raddr, c.RemoteAddr()) + } +} + +func Test_newConnPool(t *testing.T) { + p := newConnPool(time.Minute) + if p == nil { + t.Fatal("expected non-nil pool") + } + if p.ttl != time.Minute { + t.Errorf("expected ttl %v, got %v", time.Minute, p.ttl) + } + p.Close() +} + +func Test_connPool_WithLogger(t *testing.T) { + p := newConnPool(time.Second) + defer p.Close() + result := p.WithLogger(nil) + if result != p { + t.Error("WithLogger should return same pool") + } +} + +// Test NewListener +func TestNewListener(t *testing.T) { + pc := newMockPacketConn() + cfg := &ListenConfig{ + Backlog: 10, + ReadQueueSize: 10, + ReadBufferSize: 1024, + TTL: time.Minute, + } + + ln := NewListener(pc, cfg) + if ln == nil { + t.Fatal("expected non-nil listener") + } + defer ln.Close() + + if ln.Addr() == nil { + t.Error("expected non-nil addr") + } +} + +func TestNewListener_NilConfig(t *testing.T) { + pc := newMockPacketConn() + ln := NewListener(pc, &ListenConfig{TTL: time.Minute}) + if ln == nil { + t.Fatal("expected non-nil listener") + } + ln.Close() +} + +func TestNewListener_Close(t *testing.T) { + pc := newMockPacketConn() + ln := NewListener(pc, &ListenConfig{TTL: time.Minute}) + + err := ln.Close() + if err != nil { + t.Errorf("unexpected error: %v", err) + } + err = ln.Close() + if err != nil { + t.Errorf("double close should not error: %v", err) + } +} + +func TestNewListener_Addr_ConfigAddr(t *testing.T) { + pc := newMockPacketConn() + cfg := &ListenConfig{ + Addr: &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 8888}, + TTL: time.Minute, + } + ln := NewListener(pc, cfg) + defer ln.Close() + + addr := ln.Addr() + if addr.String() != "10.0.0.1:8888" { + t.Errorf("expected 10.0.0.1:8888, got %s", addr) + } +} + +func TestListener_Close(t *testing.T) { + pc := newMockPacketConn() + ln := NewListener(pc, &ListenConfig{TTL: time.Minute}) + _ = ln.Close() + + _, err := ln.Accept() + if err != net.ErrClosed { + t.Errorf("expected ErrClosed, got %v", err) + } +} + +func TestListener_Accept_ErrChan(t *testing.T) { + pc := newMockPacketConn() + pc.readFn = func(b []byte) (int, net.Addr, error) { + return 0, nil, &net.OpError{Op: "read", Net: "udp", Err: errors.New("mock error")} + } + + ln := NewListener(pc, &ListenConfig{ + TTL: time.Minute, + ReadBufferSize: 1024, + }) + defer ln.Close() + + conn, err := ln.Accept() + if err == nil { + t.Error("expected error, got nil") + t.Logf("got conn: %v", conn) + } +} + +func TestListener_listenLoop_Closed(t *testing.T) { + pc := newMockPacketConn() + ln := NewListener(pc, &ListenConfig{TTL: time.Minute}) + + ln.Close() + + _, err := ln.Accept() + if err != net.ErrClosed { + t.Errorf("expected ErrClosed, got %v", err) + } +} + +func TestListener_Addr_Default(t *testing.T) { + pc := newMockPacketConn() + pc.localAddr = &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 9999} + ln := NewListener(pc, &ListenConfig{TTL: time.Minute}) + defer ln.Close() + + addr := ln.Addr() + if addr.String() != "127.0.0.1:9999" { + t.Errorf("expected 127.0.0.1:9999, got %s", addr) + } +} + +// Test Relay +func Test_NewRelay(t *testing.T) { + pc1 := newMockPacketConn() + pc2 := newMockPacketConn() + r := NewRelay(pc1, pc2) + if r == nil { + t.Fatal("expected non-nil relay") + } +} + +func Test_Relay_WithMethods(t *testing.T) { + pc1 := newMockPacketConn() + pc2 := newMockPacketConn() + r := NewRelay(pc1, pc2) + + r.WithService("test") + r.WithBypass(nil) + r.WithLogger(nil) + r.WithBufferSize(8192) + + if r.service != "test" { + t.Errorf("expected service 'test', got %q", r.service) + } + if r.bufferSize != 8192 { + t.Errorf("expected bufferSize 8192, got %d", r.bufferSize) + } +} + +func Test_Relay_Run(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pc1 := newMockPacketConn() + pc2 := newMockPacketConn() + + var received []byte + var pc1Calls int + + // First goroutine reads from pc1, writes to pc2 + pc1.readFn = func(b []byte) (int, net.Addr, error) { + pc1Calls++ + if pc1Calls == 1 { + n := copy(b, []byte("hello")) + return n, &net.UDPAddr{IP: net.IPv4(1, 2, 3, 4), Port: 1234}, nil + } + return 0, nil, errors.New("stop") + } + // Second goroutine reads from pc2 - block until G1 writes data first + var pc2Calls int + writeDone := make(chan struct{}) + pc2.readFn = func(b []byte) (int, net.Addr, error) { + pc2Calls++ + if pc2Calls == 1 { + // Wait for G1 to write data before returning error + <-writeDone + } + return 0, nil, errors.New("stop2") + } + pc2.writeFn = func(b []byte, addr net.Addr) (int, error) { + received = append(received, b...) + + close(writeDone) + return len(b), nil + } + + r := NewRelay(pc1, pc2) + // Run catches the first error from either goroutine + err := r.Run(ctx) + if err == nil { + t.Error("expected error, got nil") + } + // Note: by the time Run returns, G1 may or may not have written data + // depending on goroutine scheduling. We only verify the relay ran. + _ = received +} + +func Test_Relay_Run_Reverse(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pc1 := newMockPacketConn() + pc2 := newMockPacketConn() + + var received []byte + var pc2Calls int + + // Second goroutine reads from pc2, writes to pc1 + pc2.readFn = func(b []byte) (int, net.Addr, error) { + pc2Calls++ + if pc2Calls == 1 { + n := copy(b, []byte("world")) + return n, &net.UDPAddr{IP: net.IPv4(5, 6, 7, 8), Port: 5678}, nil + } + return 0, nil, errors.New("stop reverse") + } + pc1.writeFn = func(b []byte, addr net.Addr) (int, error) { + received = append(received, b...) + return len(b), nil + } + // First goroutine reads from pc1 - make it return an error + pc1.readFn = func(b []byte) (int, net.Addr, error) { + return 0, nil, errors.New("stop1") + } + + r := NewRelay(pc1, pc2) + err := r.Run(ctx) + if err == nil { + t.Error("expected error, got nil") + } + + if string(received) != "world" { + t.Errorf("expected 'world' received, got %q", string(received)) + } +} + +func Test_Relay_Run_Bypass(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pc1 := newMockPacketConn() + pc2 := newMockPacketConn() + + writeCalled := false + var pc1Calls int + + // First goroutine reads from pc1 - data bypassed + pc1.readFn = func(b []byte) (int, net.Addr, error) { + pc1Calls++ + if pc1Calls == 1 { + n := copy(b, []byte("bypassed")) + return n, &net.UDPAddr{IP: net.IPv4(1, 2, 3, 4), Port: 1234}, nil + } + return 0, nil, errors.New("done") + } + pc2.writeFn = func(b []byte, addr net.Addr) (int, error) { + writeCalled = true + return len(b), nil + } + // Second goroutine reads from pc2 + pc2.readFn = func(b []byte) (int, net.Addr, error) { + return 0, nil, errors.New("done2") + } + + r := NewRelay(pc1, pc2).WithBypass(&mockBypass{contains: true}) + err := r.Run(ctx) + if err == nil { + t.Error("expected error, got nil") + } + if writeCalled { + t.Error("write should not have been called when bypass matches") + } +} + +func Test_Relay_Run_DefaultBufferSize(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pc1 := newMockPacketConn() + pc2 := newMockPacketConn() + + // Both goroutines in Run read from their respective PacketConn, + // so both need readFn set. + pc1.readFn = func(b []byte) (int, net.Addr, error) { + return 0, nil, errors.New("stop") + } + pc2.readFn = func(b []byte) (int, net.Addr, error) { + // Block forever since pc1 returns first + select {} + } + + r := NewRelay(pc1, pc2) + err := r.Run(ctx) + if err == nil || err.Error() != "stop" { + t.Errorf("expected 'stop' error, got %v", err) + } +} + +func Test_Relay_TraceLogging(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pc1 := newMockPacketConn() + pc2 := newMockPacketConn() + + pc1.readFn = func(b []byte) (int, net.Addr, error) { + return 0, nil, errors.New("done") + } + pc2.readFn = func(b []byte) (int, net.Addr, error) { + select {} + } + + r := NewRelay(pc1, pc2) + err := r.Run(ctx) + if err == nil || err.Error() != "done" { + t.Errorf("expected 'done' error, got %v", err) + } +} + +type mockBypass struct { + contains bool +} + +func (m *mockBypass) IsWhitelist() bool { return true } +func (m *mockBypass) Contains(ctx context.Context, network, addr string, opts ...bypass.Option) bool { + return m.contains +}