122 lines
2.8 KiB
Go
122 lines
2.8 KiB
Go
package handler
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
|
|
"git.kingecg.top/kingecg/gologger"
|
|
)
|
|
|
|
// TcpProxyHandler TCP代理处理器
|
|
type TcpProxyHandler struct {
|
|
Name string
|
|
Listen string
|
|
Target string
|
|
listener net.Listener
|
|
wg sync.WaitGroup
|
|
stopCh chan struct{}
|
|
}
|
|
|
|
func NewTcpProxyHandler(name, listen, target string) *TcpProxyHandler {
|
|
return &TcpProxyHandler{
|
|
Name: name,
|
|
Listen: listen,
|
|
Target: target,
|
|
stopCh: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
// Start 开始监听并处理TCP连接
|
|
func (t *TcpProxyHandler) Start() error {
|
|
l := gologger.GetLogger("TcpProxy")
|
|
l.Info(fmt.Sprintf("TCP proxy [%s] starting, listen: %s, target: %s", t.Name, t.Listen, t.Target))
|
|
|
|
var err error
|
|
t.listener, err = net.Listen("tcp", t.Listen)
|
|
if err != nil {
|
|
l.Error(fmt.Sprintf("TCP proxy [%s] failed to listen on %s: %v", t.Name, t.Listen, err))
|
|
return err
|
|
}
|
|
|
|
go t.acceptLoop()
|
|
return nil
|
|
}
|
|
|
|
// acceptLoop 接受连接并处理
|
|
func (t *TcpProxyHandler) acceptLoop() {
|
|
l := gologger.GetLogger("TcpProxy")
|
|
for {
|
|
conn, err := t.listener.Accept()
|
|
if err != nil {
|
|
select {
|
|
case <-t.stopCh:
|
|
l.Info(fmt.Sprintf("TCP proxy [%s] stopped", t.Name))
|
|
return
|
|
default:
|
|
l.Error(fmt.Sprintf("TCP proxy [%s] accept error: %v", t.Name, err))
|
|
continue
|
|
}
|
|
}
|
|
t.wg.Add(1)
|
|
go t.handleConn(conn)
|
|
}
|
|
}
|
|
|
|
// handleConn 处理单个TCP连接
|
|
func (t *TcpProxyHandler) handleConn(clientConn net.Conn) {
|
|
defer t.wg.Done()
|
|
l := gologger.GetLogger("TcpProxy")
|
|
|
|
l.Debug(fmt.Sprintf("TCP proxy [%s] new connection from %s", t.Name, clientConn.RemoteAddr()))
|
|
|
|
// 连接到目标服务器
|
|
targetConn, err := net.DialTimeout("tcp", t.Target, 10*1000000000) // 10秒超时
|
|
if err != nil {
|
|
l.Error(fmt.Sprintf("TCP proxy [%s] failed to connect to target %s: %v", t.Name, t.Target, err))
|
|
clientConn.Close()
|
|
return
|
|
}
|
|
|
|
// 使用管道模式转发数据
|
|
go t.pipe(clientConn, targetConn, "client->target")
|
|
go t.pipe(targetConn, clientConn, "target->client")
|
|
|
|
l.Debug(fmt.Sprintf("TCP proxy [%s] connection established", t.Name))
|
|
}
|
|
|
|
// pipe 在两个连接之间转发数据
|
|
func (t *TcpProxyHandler) pipe(dst, src net.Conn, direction string) {
|
|
defer dst.Close()
|
|
l := gologger.GetLogger("TcpProxy")
|
|
buf := make([]byte, 8192)
|
|
for {
|
|
n, err := src.Read(buf)
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
l.Debug(fmt.Sprintf("TCP proxy [%s] %s read error: %v", t.Name, direction, err))
|
|
}
|
|
return
|
|
}
|
|
_, err = dst.Write(buf[:n])
|
|
if err != nil {
|
|
l.Debug(fmt.Sprintf("TCP proxy [%s] %s write error: %v", t.Name, direction, err))
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// Stop 停止TCP代理服务器
|
|
func (t *TcpProxyHandler) Stop() error {
|
|
l := gologger.GetLogger("TcpProxy")
|
|
l.Info(fmt.Sprintf("TCP proxy [%s] stopping...", t.Name))
|
|
close(t.stopCh)
|
|
if t.listener != nil {
|
|
t.listener.Close()
|
|
}
|
|
t.wg.Wait()
|
|
l.Info(fmt.Sprintf("TCP proxy [%s] stopped", t.Name))
|
|
return nil
|
|
}
|