test: opt many tests

internal/conn/relay_conn_test.go

@@ -67,7 +67,7 @@ func TestInnerConn_ReadWrite(t *testing.T) {
 	assert.Equal(t, int64(len(testData)), innerC.stats.Down)
 }
 
-func TestCopyConn(t *testing.T) {
+func TestCopyTCPConn(t *testing.T) {
 	// 设置监听端口，模拟外部服务器
 	echoServer, err := net.Listen("tcp", "127.0.0.1:0") // 0 表示自动选择端口
 	assert.NoError(t, err)
@@ -121,3 +121,64 @@ func TestCopyConn(t *testing.T) {
 	// wait for the copyConn to finish
 	<-done
 }
+
+func TestCopyUDPConn(t *testing.T) {
+	// 设置监听地址，模拟外部UDP服务器
+	serverAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
+	assert.NoError(t, err)
+
+	echoServer, err := net.ListenUDP("udp", serverAddr)
+	assert.NoError(t, err)
+	defer echoServer.Close()
+
+	msg := "Hello, UDP!"
+
+	go func() {
+		buffer := make([]byte, 1024)
+		for {
+			n, remoteAddr, err := echoServer.ReadFromUDP(buffer)
+			if err != nil {
+				return
+			}
+			_, err = echoServer.WriteToUDP(buffer[:n], remoteAddr)
+			if err != nil {
+				return
+			}
+		}
+	}()
+
+	clientConn, err := net.DialUDP("udp", nil, echoServer.LocalAddr().(*net.UDPAddr))
+	assert.NoError(t, err)
+	defer clientConn.Close()
+
+	remoteConn, err := net.DialUDP("udp", nil, echoServer.LocalAddr().(*net.UDPAddr))
+	assert.NoError(t, err)
+	defer remoteConn.Close()
+
+	c1 := &innerConn{Conn: clientConn, remoteLabel: "client", stats: &Stats{}}
+	c2 := &innerConn{Conn: remoteConn, remoteLabel: "server", stats: &Stats{}}
+
+	done := make(chan struct{})
+	go func() {
+		if err := copyConn(c1, c2); err != nil {
+			t.Log(err)
+		}
+		done <- struct{}{}
+		close(done)
+	}()
+
+	_, err = clientConn.Write([]byte(msg))
+	assert.NoError(t, err)
+
+	buffer := make([]byte, len(msg))
+	n, _, err := clientConn.ReadFromUDP(buffer)
+	assert.NoError(t, err)
+	assert.Equal(t, msg, string(buffer[:n]))
+
+	// 关闭连接
+	_ = clientConn.Close()
+	_ = remoteConn.Close()
+
+	// 等待 copyConn 完成
+	<-done
+}

internal/relay/conf/cfg.go

@@ -89,10 +89,6 @@ func (r *Config) Validate() error {
 		return errors.New("both tcp and udp remotes are empty")
 	}
 
-	if len(r.UDPRemotes) > 0 {
-		zap.S().Warn("UDP RELAY WAS DISABLED FOR NOW, THIS FEATURE WILL BE AVAILABLE IN THE FUTURE")
-	}
-
 	for _, protocol := range r.BlockedProtocols {
 		if protocol != ProtocolHTTP && protocol != ProtocolTLS {
 			return fmt.Errorf("invalid blocked protocol:%s", protocol)

test/cmd/echo/main.go

@@ -8,5 +8,6 @@ import (
 
 func main() {
 	log.Println("start tcp.udp echo server at: 0.0.0.0:2333")
-	echo.RunEchoServer("0.0.0.0", 2333)
+	es := echo.NewEchoServer("0.0.0.0", 2333)
+	_ = es.Run()
 }

test/echo/echo.go

@@ -1,101 +1,145 @@
+//nolint:errcheck
 package echo
 
 import (
 	"fmt"
 	"io"
 	"log"
 	"net"
-	"os"
 	"strconv"
+	"sync"
 	"time"
 
 	"go.uber.org/zap"
 )
 
-func echo(conn net.Conn) {
-	logger := zap.S().Named(("echo-test-server"))
-	defer conn.Close()
-	defer fmt.Println("conn closed", conn.RemoteAddr().String())
-	buf := make([]byte, 10)
-	for {
-		i, err := conn.Read(buf)
-		if err == io.EOF {
-			logger.Info("conn closed,read eof ", conn.RemoteAddr().String())
-			return
-		}
-		if err != nil {
-			logger.Error(err.Error())
-			return
-		}
-		println("echo server receive", string(buf[:i]))
-		_, err = conn.Write(buf[:i])
-		if err != nil {
-			logger.Error(err.Error())
-			return
-		}
+type EchoServer struct {
+	host        string
+	port        int
+	tcpListener net.Listener
+	udpConn     *net.UDPConn
+	stopChan    chan struct{}
+	wg          sync.WaitGroup
+	logger      *zap.SugaredLogger
+}
+
+func NewEchoServer(host string, port int) *EchoServer {
+	return &EchoServer{
+		host:     host,
+		port:     port,
+		stopChan: make(chan struct{}),
+		logger:   zap.S().Named("echo-test-server"),
 	}
 }
 
-func ServeTcp(l net.Listener) {
-	for {
-		conn, err := l.Accept()
-		if err != nil {
-			fmt.Println("accept err", err.Error())
-			continue
-		}
-		go echo(conn)
+func (s *EchoServer) Run() error {
+	addr := s.host + ":" + strconv.Itoa(s.port)
+	var err error
+
+	// Start TCP server
+	s.tcpListener, err = net.Listen("tcp", addr)
+	if err != nil {
+		return fmt.Errorf("failed to start TCP server: %w", err)
+	}
+
+	// Start UDP server
+	udpAddr := net.UDPAddr{IP: net.ParseIP(s.host), Port: s.port}
+	s.udpConn, err = net.ListenUDP("udp", &udpAddr)
+	if err != nil {
+		return fmt.Errorf("failed to start UDP server: %w", err)
+	}
+
+	s.logger.Infof("Echo server started at: %s", addr)
+
+	s.wg.Add(2)
+	go s.serveTCP()
+	go s.serveUDP()
+
+	return nil
+}
+
+func (s *EchoServer) Stop() {
+	close(s.stopChan)
+	if s.tcpListener != nil {
+		s.tcpListener.Close()
 	}
+	if s.udpConn != nil {
+		s.udpConn.Close()
+	}
+	s.wg.Wait()
+	s.logger.Info("Echo server stopped")
 }
 
-func ServeUdp(conn *net.UDPConn) {
-	buf := make([]byte, 1500)
+func (s *EchoServer) serveTCP() {
+	defer s.wg.Done()
 	for {
-		number, remote, err := conn.ReadFromUDP(buf)
-		if err != nil {
-			fmt.Printf("net.ReadFromUDP() error: %s\n", err)
-		}
-		_, writeErr := conn.WriteTo(buf[0:number], remote)
-		if writeErr != nil {
-			fmt.Printf("net.WriteTo() error: %s\n", writeErr)
+		select {
+		case <-s.stopChan:
+			return
+		default:
+			conn, err := s.tcpListener.Accept()
+			if err != nil {
+				select {
+				case <-s.stopChan:
+					return
+				default:
+					s.logger.Errorf("Failed to accept TCP connection: %v", err)
+				}
+				continue
+			}
+			go s.handleTCPConn(conn)
 		}
 	}
 }
 
-func RunEchoServer(host string, port int) {
-	var err error
-	tcpAddr := host + ":" + strconv.Itoa(port)
-	l, err := net.Listen("tcp", tcpAddr)
-	defer func() {
-		err = l.Close()
+func (s *EchoServer) handleTCPConn(conn net.Conn) {
+	defer conn.Close()
+	s.logger.Infof("New TCP connection from: %s", conn.RemoteAddr())
+
+	buf := make([]byte, 1024)
+	for {
+		n, err := conn.Read(buf)
+		if err == io.EOF {
+			s.logger.Infof("Connection closed by client: %s", conn.RemoteAddr())
+			return
+		}
 		if err != nil {
-			fmt.Println(err.Error())
+			s.logger.Errorf("Error reading from connection: %v", err)
+			return
 		}
-	}()
 
-	if err != nil {
-		fmt.Println("ERROR", err)
-		os.Exit(1)
-	}
+		s.logger.Infof("Received from %s: %s", conn.RemoteAddr(), string(buf[:n]))
 
-	udpAddr := net.UDPAddr{Port: port, IP: net.ParseIP(host)}
-	udpConn, err := net.ListenUDP("udp", &udpAddr)
-	defer func() {
-		err = udpConn.Close()
+		_, err = conn.Write(buf[:n])
 		if err != nil {
-			fmt.Println(err.Error())
+			s.logger.Errorf("Error writing to connection: %v", err)
+			return
 		}
-	}()
-
-	if err != nil {
-		fmt.Println("ERROR", err)
-		os.Exit(1)
 	}
+}
 
-	fmt.Println("start echo server at:", tcpAddr)
-	stop := make(chan error)
-	go ServeTcp(l)
-	go ServeUdp(udpConn)
-	<-stop
+func (s *EchoServer) serveUDP() {
+	defer s.wg.Done()
+	buf := make([]byte, 1024)
+	for {
+		select {
+		case <-s.stopChan:
+			return
+		default:
+			n, remoteAddr, err := s.udpConn.ReadFromUDP(buf)
+			if err != nil {
+				s.logger.Errorf("Error reading UDP: %v", err)
+				continue
+			}
+
+			s.logger.Infof("Received UDP from %s: %s", remoteAddr, string(buf[:n]))
+
+			_, err = s.udpConn.WriteToUDP(buf[:n], remoteAddr)
+			if err != nil {
+				s.logger.Errorf("Error writing UDP: %v", err)
+			}
+		}
+	}
 }
 
 func SendTcpMsg(msg []byte, address string) []byte {