base-golang/pkg/proxy/tcp.go

package proxy

import (
	"net"
	"runtime/debug"
	"strconv"
)

var _ TCP = (*tcp)(nil)

type TCPArgs struct {
	Local               string      //监听地址
	Parent              string      //被代理地址
	Timeout             int         //拨号超时(毫秒)
	PoolSize            int         //代理池数
	CheckParentInterval int         //检查被代理连接间隔(秒)
	OutCallback         func() bool //回调 //是否允许代理
}

type TCP interface {
	Start() (err error)
	Close()

	cleanOutPool()
	callback(inConn net.Conn)
	outToTCP(inConn *net.Conn) (err error)
	initOutConnPool()
}

type tcp struct {
	inConn  *net.Conn
	outPool OutPool
	listen  ServerChannel
	cfg     TCPArgs
}

func (s *tcp) Start() (err error) {
	s.initOutConnPool()
	host, port, _ := net.SplitHostPort(s.cfg.Local)
	p, _ := strconv.Atoi(port)
	s.listen = NewServerChannel(host, p)
	err = s.listen.ListenTCP(s.callback)
	if err != nil {
		return
	}
	return
}

func (s *tcp) Close() {
	if s.inConn != nil {
		CloseConn(s.inConn)
	}
	s.cleanOutPool()
	_ = s.listen.CloseListen()
}

func (s *tcp) cleanOutPool() {
	if s.outPool.Pool != nil {
		s.outPool.Pool.ReleaseAll()
	}
}

func (s *tcp) callback(inConn net.Conn) {
	defer func() {
		if err := recover(); err != nil {
			logger.Sugar().Infof("conn handler crashed with err : %s \nstack: %s", err, string(debug.Stack()))
		}
	}()

	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) (err error) {
	var outConn net.Conn
	var _outConn any
	_outConn, err = s.outPool.Pool.Get()
	if err == nil {
		outConn = _outConn.(net.Conn)
	}
	if err != nil {
		CloseConn(inConn)
		return
	}
	inAddr := (*inConn).RemoteAddr().String()
	inLocalAddr := (*inConn).LocalAddr().String()
	outAddr := outConn.RemoteAddr().String()
	outLocalAddr := outConn.LocalAddr().String()
	IoBind(*inConn, outConn, func(isSrcErr bool, err error) {
		logger.Sugar().Infof("conn %s - %s - %s -%s released", inAddr, inLocalAddr, outLocalAddr, outAddr)
		CloseConn(inConn)
		CloseConn(&outConn)
	}, func(n int, d bool) {}, 0)
	logger.Sugar().Infof("conn %s - %s - %s -%s connected", inAddr, inLocalAddr, outLocalAddr, outAddr)
	return
}

func (s *tcp) initOutConnPool() {
	s.outPool = NewOutPool(s.cfg.CheckParentInterval, s.cfg.Parent, s.cfg.Timeout, s.cfg.PoolSize, s.cfg.PoolSize*2)
}