base-golang/pkg/proxy/tcp.go
2024-07-23 10:23:43 +08:00

95 lines
1.8 KiB
Go

package proxy
import (
"net"
"strconv"
)
var _ TCP = (*tcp)(nil)
type TCPArgs struct {
Local string //监听地址
Parent string //被代理地址
Timeout int //拨号超时(毫秒)
OutCallback func() bool //回调 //是否允许代理
}
type TCP interface {
Start() (err error)
Close()
callback(inConn net.Conn)
outToTCP(inConn net.Conn) (err error)
}
type tcp struct {
inConn net.Conn
outConn net.Conn
listen Listener
cfg TCPArgs
}
func (s *tcp) Start() (err error) {
host, port, _ := net.SplitHostPort(s.cfg.Local)
p, _ := strconv.Atoi(port)
s.listen = NewListener(host, p)
err = s.listen.ListenTCP(s.callback)
if err != nil {
return
}
return
}
func (s *tcp) Close() {
if s.inConn != nil {
CloseConn(s.inConn)
}
if s.outConn != nil {
CloseConn(s.outConn)
}
_ = s.listen.CloseListen()
}
func (s *tcp) callback(inConn net.Conn) {
defer func() {
if err := recover(); err != nil {
logger.Sugar().Infof("conn handler crashed with err : %s", err)
}
}()
if s.cfg.OutCallback != nil {
if !s.cfg.OutCallback() {
CloseConn(inConn)
return
}
}
err := s.outToTCP(inConn)
if err != nil {
CloseConn(inConn)
}
s.inConn = inConn
}
func (s *tcp) outToTCP(inConn net.Conn) error {
outConn, err := ConnectHost(s.cfg.Parent, s.cfg.Timeout)
if err != nil {
return err
}
inAddr := inConn.RemoteAddr().String()
inLocalAddr := inConn.LocalAddr().String()
outAddr := outConn.RemoteAddr().String()
outLocalAddr := outConn.LocalAddr().String()
IoBind(inConn, outConn, func(isSrcErr bool, err error) {
CloseConn(inConn)
CloseConn(outConn)
logger.Sugar().Infof("conn %s - %s - %s -%s released", inAddr, inLocalAddr, outLocalAddr, outAddr)
}, func(n int, d bool) {}, 0)
logger.Sugar().Infof("conn %s - %s - %s -%s connected", inAddr, inLocalAddr, outLocalAddr, outAddr)
return nil
}