diff --git a/app/config.go b/app/config.go index 3997348..5892061 100644 --- a/app/config.go +++ b/app/config.go @@ -5,7 +5,7 @@ type ( ProxyItem struct { RemoteAddr []string - LocalPort int + LocalAddr string } ) diff --git a/app/proxy.go b/app/proxy.go index 92b87e3..d20756c 100644 --- a/app/proxy.go +++ b/app/proxy.go @@ -1,34 +1,44 @@ package app import ( - "fmt" "git.makemake.in/test/mycommon/mylog" "io" + "math/rand" "net" + "sync" + "time" +) + +var ( + wg = &sync.WaitGroup{} ) func Run() { - for _, item := range ProxyMap { - go listenTcp(item) + rand.Seed(time.Now().UnixNano()) + + for name, item := range ProxyMap { + wg.Add(1) + go listenTcp(name, item) } - select {} + wg.Wait() } -func listenTcp(item ProxyItem) { - localAddr := fmt.Sprintf(":%d", item.LocalPort) - listen, err := net.Listen("tcp", localAddr) +func listenTcp(name string, item ProxyItem) { + defer wg.Done() + + listen, err := net.Listen("tcp", item.LocalAddr) if err != nil { - mylog.Errorf("listen err: %s", localAddr, err) + mylog.Errorf("[%s] listen err: %s", name, err) return } - mylog.Infof("listen %s", localAddr) + mylog.Infof("[%s] listen %s", name, item.LocalAddr) for { conn, err := listen.Accept() if err != nil { - mylog.Errorf("conn accept err: %s", localAddr, err) + mylog.Errorf("conn accept err: %s", item.LocalAddr, err) return } @@ -37,36 +47,45 @@ func listenTcp(item ProxyItem) { } func proxyTcp(conn net.Conn, item ProxyItem) { - //defer conn.Close() - for _, addr := range item.RemoteAddr { - remoteConn, err := net.Dial("tcp", addr) - if err != nil { - mylog.Errorf("connect remote err: %s\n", err) - return - } + randRemoteAddr := item.RemoteAddr[rand.Intn(len(item.RemoteAddr))] - mylog.Infof("connect remote %s", addr) - - go copyConn(conn, remoteConn) - go copyConn(remoteConn, conn) - } - -} - -func copyConn(src, dst net.Conn) { - defer func() { - src.Close() - dst.Close() - - mylog.Warn("close conn") - }() - - _, err := io.Copy(src, dst) + remoteConn, err := net.Dial("tcp", randRemoteAddr) if err != nil { - mylog.Error(err) + mylog.Errorf("connect remote err: %s\n", err) return } - //fmt.Println(written) + go copyConn(randRemoteAddr, conn, remoteConn) + +} + +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() + }() + + go func() { + _, err := io.Copy(remoteConn, conn) + if err != nil { + mylog.Error(err) + return + } + wg2.Done() + }() + + wg2.Wait() + + mylog.Infof("stop proxy: %s", addr) + } diff --git a/go.mod b/go.mod index 1e0678a..2ac1231 100644 --- a/go.mod +++ b/go.mod @@ -2,7 +2,7 @@ module proxyport go 1.18 -require git.makemake.in/test/mycommon v0.0.0-20230114162108-0fdfb59809db +require git.makemake.in/test/mycommon v0.0.0-20230119033226-3679bd502caf require ( github.com/fsnotify/fsnotify v1.6.0 // indirect diff --git a/main.go b/main.go index c3cfcf7..30a7d8b 100644 --- a/main.go +++ b/main.go @@ -1,13 +1,19 @@ package main import ( + "flag" "git.makemake.in/test/mycommon/myconf" "git.makemake.in/test/mycommon/mylog" "proxyport/app" ) +var configPath string + func main() { - myconf.Init("config.toml", &app.ProxyMap) + flag.StringVar(&configPath, "c", "config.toml", "配置文件路径") + flag.Parse() + + myconf.Init(configPath, &app.ProxyMap) mylog.Init("debug", mylog.DefaultConfig) app.Run() }