proxyport/app/proxy.go

111 lines
2.0 KiB
Go

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())
for name, item := range ProxyMap {
go listenTcp(name, item)
}
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) {
listen, err := net.Listen("tcp", item.LocalAddr)
if err != nil {
mylog.Errorf("[%s] listen err: %s", name, err)
return
}
mylog.Infof("[%s] listen %s", name, item.LocalAddr)
for {
localConn, err := listen.Accept()
if err != nil {
mylog.Errorf("conn accept err: %s", item.LocalAddr, err)
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)
if err != nil {
localConn.Close()
mylog.Errorf("connect remote err: %s", err)
return
}
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)
}
func copyConn(wg0 *sync.WaitGroup, localConn, remoteConn net.Conn) {
defer func() {
wg0.Done()
localConn.Close()
remoteConn.Close()
}()
_, err := io.Copy(remoteConn, localConn)
if err != nil {
mylog.Error(err)
}
}