Try to fix sync bug

yosebyte · Dec 6, 2024 · 8d0bfb5 · 8d0bfb5
1 parent ca63180
commit 8d0bfb5
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/internal/tunnel/server.go b/internal/tunnel/server.go
@@ -38,13 +38,12 @@ func Server(parsedURL *url.URL, whiteList *sync.Map) error {
 	}
 	defer linkConn.Close()
 	log.Info("Tunnel connection established from: [%v]", linkConn.RemoteAddr().String())
-	mu := &sync.Mutex{}
 	errChan := make(chan error, 2)
 	go func() {
-		errChan <- ServeTCP(parsedURL, whiteList, linkAddr, targetTCPAddr, linkListen, linkConn, mu)
+		errChan <- ServeTCP(parsedURL, whiteList, linkAddr, targetTCPAddr, linkListen, linkConn)
 	}()
 	go func() {
-		errChan <- ServeUDP(parsedURL, whiteList, linkAddr, targetUDPAddr, linkListen, linkConn, mu)
+		errChan <- ServeUDP(parsedURL, whiteList, linkAddr, targetUDPAddr, linkListen, linkConn)
 	}()
 	return <-errChan
 }
diff --git a/internal/tunnel/tcp.go b/internal/tunnel/tcp.go
@@ -9,13 +9,14 @@ import (
 	"github.com/yosebyte/passport/pkg/log"
 )
 
-func ServeTCP(parsedURL *url.URL, whiteList *sync.Map, linkAddr, targetAddr *net.TCPAddr, linkListen *net.TCPListener, linkConn *net.TCPConn, mu *sync.Mutex) error {
+func ServeTCP(parsedURL *url.URL, whiteList *sync.Map, linkAddr, targetAddr *net.TCPAddr, linkListen *net.TCPListener, linkConn *net.TCPConn) error {
 	targetListen, err := net.ListenTCP("tcp", targetAddr)
 	if err != nil {
 		log.Error("Unable to listen target address: [%v]", targetAddr)
 		return err
 	}
 	defer targetListen.Close()
+	var mu sync.Mutex
 	semaphore := make(chan struct{}, 1024)
 	for {
 		targetConn, err := targetListen.AcceptTCP()

diff --git a/internal/tunnel/udp.go b/internal/tunnel/udp.go
@@ -8,20 +8,21 @@ import (
 	"github.com/yosebyte/passport/pkg/log"
 )
 
-func ServeUDP(parsedURL *url.URL, whiteList *sync.Map, linkAddr *net.TCPAddr, targetAddr *net.UDPAddr, linkListen *net.TCPListener, linkConn *net.TCPConn, mu *sync.Mutex) error {
+func ServeUDP(parsedURL *url.URL, whiteList *sync.Map, linkAddr *net.TCPAddr, targetAddr *net.UDPAddr, linkListen *net.TCPListener, linkConn *net.TCPConn) error {
 	targetConn, err := net.ListenUDP("udp", targetAddr)
 	if err != nil {
 		log.Error("Unable to listen target address: [%v]", targetAddr)
 		return err
 	}
 	defer targetConn.Close()
+	var mu sync.Mutex
 	semaphore := make(chan struct{}, 1024)
 	for {
 		buffer := make([]byte, 8192)
 		n, clientAddr, err := targetConn.ReadFromUDP(buffer)
 		if err != nil {
 			log.Error("Unable to read from client address: [%v] %v", clientAddr, err)
-			continue
+			break
 		}
 		if parsedURL.Fragment != "" {
 			clientIP := clientAddr.IP.String()
@@ -65,4 +66,5 @@ func ServeUDP(parsedURL *url.URL, whiteList *sync.Map, linkAddr *net.TCPAddr, ta
 			log.Info("Transfer completed successfully")
 		}(buffer, n, clientAddr)
 	}
+	return nil
 }