Add tests for dynamic-destination flag.

kzemek · Mar 24, 2024 · 9746172 · 9746172
1 parent 6188bdf
commit 9746172
Show file tree

Hide file tree

Showing 2 changed files with 112 additions and 4 deletions.
diff --git 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
@@ -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)
+	}
+}