base-golang/pkg/proxy/tcp.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
}