base-golang/pkg/proxy/tcp.go

package proxy

import (
	"fmt"
	"log"
	"net"
	"runtime/debug"
	"strconv"
	"sync"
)

var ServicesMap = new(sync.Map)

var _ TCP = (*tcp)(nil)

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

type TCP interface {
	Clean()

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

type tcp struct {
	outPool OutPool
	cfg     TCPArgs
}

type ServiceItem struct {
	S    TCP
	Name string
}

func Regist(name string, args TCPArgs) *sync.Map {
	item := &ServiceItem{
		S:    &tcp{cfg: args},
		Name: name,
	}
	ServicesMap.LoadOrStore(name, item)
	return ServicesMap
}

func Run(name string) (err error) {
	item, ok := ServicesMap.Load(name)
	service := item.(*ServiceItem)
	if ok {
		go func() {
			defer func() {
				recoverErr := recover()
				if recoverErr != nil {
					log.Fatalf("%s servcie crashed, ERR: %s\ntrace:%s", name, recoverErr, string(debug.Stack()))
				}
			}()
			startErr := service.S.start()
			if startErr != nil {
				log.Fatalf("%s servcie fail, ERR: %s", name, startErr)
			}
		}()
	}
	if !ok {
		err = fmt.Errorf("service %s not found", name)
	}
	return
}

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

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

func (s *tcp) callback(inConn net.Conn) {
	defer func() {
		if err := recover(); err != nil {
			log.Printf("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)
	}
}

func (s *tcp) outToTCP(inConn *net.Conn) (err error) {
	var outConn net.Conn
	var _outConn interface{}
	_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) {
		log.Printf("conn %s - %s - %s -%s released", inAddr, inLocalAddr, outLocalAddr, outAddr)
		CloseConn(inConn)
		CloseConn(&outConn)
	}, func(n int, d bool) {}, 0)
	log.Printf("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)
}