update

app/proxy.go

@@ -1,14 +1,23 @@
 package app
 
 import (
+	"context"
 	"git.makemake.in/test/mycommon/mylog"
 	"io"
 	"math/rand"
 	"net"
+	"os"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
 )
 
+var (
+	mainWg              = &sync.WaitGroup{}
+	mainCtx, mainCancel = context.WithCancel(context.Background())
+)
+
 func Run() {
 	rand.Seed(time.Now().UnixNano())
 
@@ -16,7 +25,16 @@ func Run() {
 		go listenTcp(name, item)
 	}
 
-	select {}
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
+
+	v := <-ch
+	mylog.Infof("捕获退出信号: %s", v.String())
+	mainCancel()
+
+	mainWg.Wait()
+	mylog.Info("stop")
+
 }
 
 func listenTcp(name string, item ProxyItem) {
@@ -36,11 +54,14 @@ func listenTcp(name string, item ProxyItem) {
 			return
 		}
 
+		mainWg.Add(1)
 		go proxyTcp(localConn, item)
 	}
 }
 
 func proxyTcp(localConn net.Conn, item ProxyItem) {
+	defer mainWg.Done()
+
 	randRemoteAddr := item.RemoteAddr[rand.Intn(len(item.RemoteAddr))]
 
 	remoteConn, err := net.DialTimeout("tcp", randRemoteAddr, time.Second*3)
@@ -52,11 +73,23 @@ func proxyTcp(localConn net.Conn, item ProxyItem) {
 
 	mylog.Infof("start proxy %s -> %s", localConn.RemoteAddr(), randRemoteAddr)
 
+	waitCh := make(chan struct{})
+	go func() {
 		wg0 := &sync.WaitGroup{}
 		wg0.Add(2)
 		go copyConn(wg0, localConn, remoteConn)
 		go copyConn(wg0, remoteConn, localConn)
 		wg0.Wait()
+		close(waitCh)
+	}()
+
+	select {
+	case <-waitCh:
+
+	case <-mainCtx.Done():
+		localConn.Close()
+		remoteConn.Close()
+	}
 
 	mylog.Infof("stop proxy %s -> %s", localConn.RemoteAddr(), randRemoteAddr)
 

main.go

@@ -14,6 +14,6 @@ func main() {
 	flag.Parse()
 
 	myconf.Init(configPath, &app.ProxyMap)
-	mylog.Init("debug", mylog.DefaultConfig)
+	defer mylog.Flush()
 	app.Run()
 }