From ebcb0f87a2f1a81eca3255c56ba6fcee9a2eadd6 Mon Sep 17 00:00:00 2001
From: lizifeng <lzfengyyy@gmail.com>
Date: Tue, 18 Apr 2023 11:07:36 +0800
Subject: [PATCH] update

---
 app/proxy.go | 64 ++++++++++++++++++++--------------------------------
 1 file changed, 25 insertions(+), 39 deletions(-)

diff --git a/app/proxy.go b/app/proxy.go
index d20756c..2947b56 100644
--- a/app/proxy.go
+++ b/app/proxy.go
@@ -9,23 +9,17 @@ import (
 	"time"
 )
 
-var (
-	wg = &sync.WaitGroup{}
-)
-
 func Run() {
 	rand.Seed(time.Now().UnixNano())
 
 	for name, item := range ProxyMap {
-		wg.Add(1)
 		go listenTcp(name, item)
 	}
 
-	wg.Wait()
+	select {}
 }
 
 func listenTcp(name string, item ProxyItem) {
-	defer wg.Done()
 
 	listen, err := net.Listen("tcp", item.LocalAddr)
 	if err != nil {
@@ -36,56 +30,48 @@ func listenTcp(name string, item ProxyItem) {
 	mylog.Infof("[%s] listen %s", name, item.LocalAddr)
 
 	for {
-		conn, err := listen.Accept()
+		localConn, err := listen.Accept()
 		if err != nil {
 			mylog.Errorf("conn accept err: %s", item.LocalAddr, err)
 			return
 		}
 
-		go proxyTcp(conn, item)
+		go proxyTcp(localConn, item)
 	}
 }
 
-func proxyTcp(conn net.Conn, item ProxyItem) {
-
+func proxyTcp(localConn net.Conn, item ProxyItem) {
 	randRemoteAddr := item.RemoteAddr[rand.Intn(len(item.RemoteAddr))]
 
-	remoteConn, err := net.Dial("tcp", randRemoteAddr)
+	remoteConn, err := net.DialTimeout("tcp", randRemoteAddr, time.Second*3)
 	if err != nil {
-		mylog.Errorf("connect remote err: %s\n", err)
+		localConn.Close()
+		mylog.Errorf("connect remote err: %s", err)
 		return
 	}
 
-	go copyConn(randRemoteAddr, conn, remoteConn)
+	mylog.Infof("start proxy %s -> %s", localConn.RemoteAddr(), randRemoteAddr)
+
+	wg0 := &sync.WaitGroup{}
+	wg0.Add(2)
+	go copyConn(wg0, localConn, remoteConn)
+	go copyConn(wg0, remoteConn, localConn)
+	wg0.Wait()
+
+	mylog.Infof("stop proxy %s -> %s", localConn.RemoteAddr(), randRemoteAddr)
 
 }
 
-func copyConn(addr string, conn, remoteConn net.Conn) {
-	mylog.Infof("start proxy: %s", addr)
-
-	wg2 := &sync.WaitGroup{}
-	wg2.Add(2)
-
-	go func() {
-		_, err := io.Copy(conn, remoteConn)
-		if err != nil {
-			mylog.Error(err)
-			return
-		}
-		wg2.Done()
+func copyConn(wg0 *sync.WaitGroup, localConn, remoteConn net.Conn) {
+	defer func() {
+		wg0.Done()
+		localConn.Close()
+		remoteConn.Close()
 	}()
 
-	go func() {
-		_, err := io.Copy(remoteConn, conn)
-		if err != nil {
-			mylog.Error(err)
-			return
-		}
-		wg2.Done()
-	}()
-
-	wg2.Wait()
-
-	mylog.Infof("stop proxy: %s", addr)
+	_, err := io.Copy(remoteConn, localConn)
+	if err != nil {
+		mylog.Error(err)
+	}
 
 }