Update On Fri May 17 20:30:02 CEST 2024

echo/test/echo/echo.go

@@ -8,25 +8,28 @@ import (
 	"os"
 	"strconv"
 	"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 {
-			fmt.Println("read eof")
+			logger.Info("conn closed,read eof ", conn.RemoteAddr().String())
 			return
 		}
 		if err != nil {
-			fmt.Println(err.Error())
+			logger.Error(err.Error())
 			return
 		}
 		_, err = conn.Write(buf[:i])
 		if err != nil {
-			fmt.Println(err.Error())
+			logger.Error(err.Error())
 			return
 		}
 	}
@@ -114,6 +117,33 @@ func SendTcpMsg(msg []byte, address string) []byte {
 	return buf[:n]
 }
 
+func EchoTcpMsgLong(msg []byte, sleepTime time.Duration, address string) error {
+	logger := zap.S()
+	buf := make([]byte, len(msg))
+	conn, err := net.Dial("tcp", address)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	logger.Infof("conn start %s %s", conn.RemoteAddr().String(), conn.LocalAddr().String())
+	for i := 0; i < 10; i++ {
+		if _, err := conn.Write(msg); err != nil {
+			return err
+		}
+		n, err := conn.Read(buf)
+		if err != nil {
+			return err
+		}
+		if string(buf[:n]) != string(msg) {
+			return fmt.Errorf("msg not equal")
+		}
+		// to fake a long connection
+		time.Sleep(sleepTime)
+	}
+	logger.Infof("conn closed %s %s", conn.RemoteAddr().String(), conn.LocalAddr().String())
+	return nil
+}
+
 func SendUdpMsg(msg []byte, address string) []byte {
 	conn, err := net.Dial("udp", address)
 	if err != nil {

echo/test/relay_test.go

@@ -23,7 +23,8 @@ const (
 	ECHO_PORT   = 9002
 	ECHO_SERVER = "0.0.0.0:9002"
 
-	RAW_LISTEN = "0.0.0.0:1234"
+	RAW_LISTEN                     = "0.0.0.0:1234"
+	RAW_LISTEN_WITH_MAX_CONNECTION = "0.0.0.0:2234"
 
 	WS_LISTEN = "0.0.0.0:1235"
 	WS_REMOTE = "ws://0.0.0.0:2000"
@@ -61,6 +62,15 @@ func init() {
 				UDPRemotes:    []string{ECHO_SERVER},
 				TransportType: constant.RelayTypeRaw,
 			},
+			// raw cfg with max connection
+			{
+				Listen:        RAW_LISTEN_WITH_MAX_CONNECTION,
+				ListenType:    constant.RelayTypeRaw,
+				TCPRemotes:    []string{ECHO_SERVER},
+				UDPRemotes:    []string{ECHO_SERVER},
+				TransportType: constant.RelayTypeRaw,
+				MaxConnection: 1,
+			},
 
 			// ws
 			{
@@ -155,6 +165,24 @@ func TestRelayOverRaw(t *testing.T) {
 	// t.Log("test udp done!")
 }
 
+func TestRelayWithMaxConnectionCount(t *testing.T) {
+	msg := []byte("hello")
+
+	// first connection will be accepted
+	go func() {
+		err := echo.EchoTcpMsgLong(msg, time.Second, RAW_LISTEN_WITH_MAX_CONNECTION)
+		if err != nil {
+			t.Error(err)
+		}
+	}()
+
+	// second connection will be rejected
+	time.Sleep(time.Second) // wait for first connection
+	if err := echo.EchoTcpMsgLong(msg, time.Second, RAW_LISTEN_WITH_MAX_CONNECTION); err == nil {
+		t.Fatal("need error here")
+	}
+}
+
 func TestRelayWithDeadline(t *testing.T) {
 	logger, _ := zap.NewDevelopment()
 	msg := []byte("hello")