diff --git a/tests/tcp_test.go b/tests/tcp_test.go index 1a5d3ce..fea0ac6 100644 --- a/tests/tcp_test.go +++ b/tests/tcp_test.go @@ -75,7 +75,7 @@ func TestListen(t *testing.T) { receivedData4 := make(chan listenResult, 1) go runServer(t, "127.0.0.1:54321", receivedData4) - time.Sleep(1 * time.Second) + time.Sleep(100 * time.Millisecond) conn, err := net.Dial("tcp", "127.0.0.1:12345") if err != nil { @@ -123,7 +123,7 @@ func TestListen_unknown(t *testing.T) { receivedData4 := make(chan listenResult, 1) go runServer(t, "127.0.0.1:54322", receivedData4) - time.Sleep(1 * time.Second) + time.Sleep(100 * time.Millisecond) conn, err := net.Dial("tcp", "127.0.0.1:12346") if err != nil { @@ -171,7 +171,7 @@ func TestListen_proxyV2(t *testing.T) { receivedData4 := make(chan listenResult, 1) go runServer(t, "127.0.0.1:54323", receivedData4) - time.Sleep(1 * time.Second) + time.Sleep(100 * time.Millisecond) conn, err := net.Dial("tcp", "127.0.0.1:12347") if err != nil { @@ -200,3 +200,52 @@ func TestListen_proxyV2(t *testing.T) { t.Errorf("Unexpected source address: %v", result.saddr) } } + +func TestTCPListen_DynamicDestination(t *testing.T) { + opts := utils.Options{ + Protocol: utils.TCP, + ListenAddr: netip.MustParseAddrPort("0.0.0.0:12350"), + TargetAddr4: netip.MustParseAddrPort("127.0.0.1:443"), + TargetAddr6: netip.MustParseAddrPort("[::1]:443"), + DynamicDestination: true, + Mark: 0, + AllowedSubnets: nil, + Verbose: 2, + } + + lvl := slog.LevelInfo + if opts.Verbose > 0 { + lvl = slog.LevelDebug + } + + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: lvl})) + + listenConfig := net.ListenConfig{} + errors := make(chan error, 1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go tcp.Listen(ctx, &listenConfig, &opts, logger, errors) + + receivedData4 := make(chan listenResult, 1) + go runServer(t, "127.0.0.1:56324", receivedData4) + + time.Sleep(100 * time.Millisecond) + + conn, err := net.Dial("tcp", "127.0.0.1:12350") + if err != nil { + t.Fatalf("Failed to connect to server: %v", err) + } + defer conn.Close() + + conn.Write([]byte("PROXY TCP4 192.168.0.1 127.0.0.1 56324 56324\r\nmoredata")) + result := <-receivedData4 + + if !reflect.DeepEqual(result.data, []byte("moredata")) { + t.Errorf("Unexpected data: %v", result.data) + } + + if result.saddr.String() != "192.168.0.1:56324" { + t.Errorf("Unexpected source address: %v", result.saddr) + } +} diff --git a/tests/udp_test.go b/tests/udp_test.go index 0ece86b..cbca211 100644 --- a/tests/udp_test.go +++ b/tests/udp_test.go @@ -65,7 +65,7 @@ func TestListenUDP(t *testing.T) { receivedData4 := make(chan listenResult, 1) go runUDPServer(t, "127.0.0.1:54323", receivedData4) - time.Sleep(1 * time.Second) + time.Sleep(100 * time.Millisecond) conn, err := net.Dial("udp", "127.0.0.1:12347") if err != nil { @@ -94,3 +94,62 @@ func TestListenUDP(t *testing.T) { t.Errorf("Unexpected source address: %v", result.saddr) } } + +func TestListenUDP_DynamicDestination(t *testing.T) { + opts := utils.Options{ + Protocol: utils.UDP, + ListenAddr: netip.MustParseAddrPort("0.0.0.0:12348"), + TargetAddr4: netip.MustParseAddrPort("127.0.0.1:443"), + TargetAddr6: netip.MustParseAddrPort("[::1]:443"), + DynamicDestination: true, + Mark: 0, + AllowedSubnets: nil, + Verbose: 2, + } + + lvl := slog.LevelInfo + if opts.Verbose > 0 { + lvl = slog.LevelDebug + } + + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: lvl})) + + listenConfig := net.ListenConfig{} + errors := make(chan error, 1) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go udp.Listen(ctx, &listenConfig, &opts, logger, errors) + + receivedData4 := make(chan listenResult, 1) + go runUDPServer(t, "127.0.0.1:56324", receivedData4) + + time.Sleep(100 * time.Millisecond) + + conn, err := net.Dial("udp", "127.0.0.1:12348") + if err != nil { + t.Fatalf("Failed to connect to server: %v", err) + } + defer conn.Close() + + buf := []byte{0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A} + buf = append(buf, 0x21) // PROXY + buf = append(buf, 0x12) // UDP4 + buf = append(buf, 0x00, 0x0C) // 12 bytes + buf = append(buf, 192, 168, 0, 1) // saddr + buf = append(buf, 127, 0, 0, 1) // daddr + buf = append(buf, 0xDC, 0x04) // sport 56324 + buf = append(buf, 0xDC, 0x04) // sport 56324 + buf = append(buf, []byte("moredata")...) + + conn.Write(buf) + result := <-receivedData4 + + if !reflect.DeepEqual(result.data, []byte("moredata")) { + t.Errorf("Unexpected data: %v", result.data) + } + + if result.saddr.String() != "192.168.0.1:56324" { + t.Errorf("Unexpected source address: %v", result.saddr) + } +}