127 lines
2.3 KiB
Go
127 lines
2.3 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"git.makemake.in/kzkzzzz/mycommon/myconf"
|
|
"git.makemake.in/kzkzzzz/mycommon/mylog"
|
|
"io"
|
|
"math/rand"
|
|
"net"
|
|
"os"
|
|
"os/signal"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
type ProxyItem struct {
|
|
RemoteAddr []string
|
|
LocalAddr string
|
|
}
|
|
|
|
var (
|
|
mainWg = &sync.WaitGroup{}
|
|
mainCtx, mainCancel = context.WithCancel(context.Background())
|
|
|
|
proxyInfo map[string]*ProxyItem
|
|
)
|
|
|
|
func Run() {
|
|
err := myconf.Conf().UnmarshalKey("proxy", &proxyInfo)
|
|
if err != nil {
|
|
mylog.Error(err)
|
|
return
|
|
}
|
|
|
|
names := make([]string, 0, len(proxyInfo))
|
|
for name := range proxyInfo {
|
|
names = append(names, name)
|
|
}
|
|
|
|
for _, name := range names {
|
|
go listenTcp(name, proxyInfo[name])
|
|
}
|
|
|
|
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("tcp4", item.LocalAddr)
|
|
if err != nil {
|
|
mylog.Errorf("[%s] listen err: %s", name, err)
|
|
return
|
|
}
|
|
|
|
mylog.Infof("[%s] forward %s -> %+v", name, item.LocalAddr, item.RemoteAddr)
|
|
|
|
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.Debugf("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.Debugf("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.Debug(err)
|
|
}
|
|
|
|
}
|