diff --git a/app/proxy.go b/app/proxy.go index 2947b56..d9284fb 100644 --- a/app/proxy.go +++ b/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) - wg0 := &sync.WaitGroup{} - wg0.Add(2) - go copyConn(wg0, localConn, remoteConn) - go copyConn(wg0, remoteConn, localConn) - wg0.Wait() + 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) diff --git a/main.go b/main.go index 30a7d8b..7882d09 100644 --- a/main.go +++ b/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() }