Added udp semaphore control

yosebyte · Dec 2, 2024 · 4e6409d · 4e6409d
1 parent 9026e6c
commit 4e6409d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/internal/forward/tcp.go b/internal/forward/tcp.go
@@ -28,7 +28,7 @@ func HandleTCP(parsedURL *url.URL, whiteList *sync.Map) error {
 		return err
 	}
 	defer linkListen.Close()
-	tempSlot := make(chan struct{}, 1024)
+	semaphore := make(chan struct{}, 1024)
 	for {
 		linkConn, err := linkListen.AcceptTCP()
 		if err != nil {
@@ -37,9 +37,9 @@ func HandleTCP(parsedURL *url.URL, whiteList *sync.Map) error {
 			continue
 		}
 		linkConn.SetNoDelay(true)
-		tempSlot <- struct{}{}
+		semaphore <- struct{}{}
 		go func(linkConn net.Conn) {
-			defer func() { <-tempSlot }()
+			defer func() { <-semaphore }()
 			clientAddr := linkConn.RemoteAddr().String()
 			log.Info("Client connection established: [%v]", clientAddr)
 			if parsedURL.Fragment != "" {

diff --git a/internal/forward/udp.go b/internal/forward/udp.go
@@ -27,6 +27,7 @@ func HandleUDP(parsedURL *url.URL, whiteList *sync.Map) error {
 		return err
 	}
 	defer linkConn.Close()
+	semaphore := make(chan struct{}, 1024)
 	for {
 		buffer := make([]byte, 8192)
 		n, clientAddr, err := linkConn.ReadFromUDP(buffer)
@@ -41,7 +42,9 @@ func HandleUDP(parsedURL *url.URL, whiteList *sync.Map) error {
 				continue
 			}
 		}
-		go func() {
+		semaphore <- struct{}{}
+		go func(buffer []byte, n int, clientAddr *net.UDPAddr) {
+			defer func() { <-semaphore }()
 			targetConn, err := net.DialUDP("udp", nil, targetAddr)
 			if err != nil {
 				log.Error("Unable to dial target address: [%v] %v", targetAddr, err)
@@ -71,6 +74,6 @@ func HandleUDP(parsedURL *url.URL, whiteList *sync.Map) error {
 				return
 			}
 			log.Info("Transfer completed successfully")
-		}()
+		}(buffer, n, clientAddr)
 	}
 }