tcp.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. package proxy
  2. import (
  3. "log"
  4. "net"
  5. "runtime/debug"
  6. "strconv"
  7. )
  8. var _ TCP = (*tcp)(nil)
  9. type TCPArgs struct {
  10. Local string //监听地址
  11. Parent string //被代理地址
  12. Timeout int //拨号超时(毫秒)
  13. PoolSize int //代理池数
  14. CheckParentInterval int //检查被代理连接间隔(秒)
  15. OutCallback func() bool //回调 //是否允许代理
  16. }
  17. type TCP interface {
  18. Start() (err error)
  19. Close()
  20. cleanOutPool()
  21. callback(inConn net.Conn)
  22. outToTCP(inConn *net.Conn) (err error)
  23. initOutConnPool()
  24. }
  25. type tcp struct {
  26. inConn *net.Conn
  27. outPool OutPool
  28. listen ServerChannel
  29. cfg TCPArgs
  30. }
  31. func (s *tcp) Start() (err error) {
  32. s.initOutConnPool()
  33. host, port, _ := net.SplitHostPort(s.cfg.Local)
  34. p, _ := strconv.Atoi(port)
  35. s.listen = NewServerChannel(host, p)
  36. err = s.listen.ListenTCP(s.callback)
  37. if err != nil {
  38. return
  39. }
  40. return
  41. }
  42. func (s *tcp) Close() {
  43. if s.inConn != nil {
  44. CloseConn(s.inConn)
  45. }
  46. s.cleanOutPool()
  47. _ = s.listen.CloseListen()
  48. }
  49. func (s *tcp) cleanOutPool() {
  50. if s.outPool.Pool != nil {
  51. s.outPool.Pool.ReleaseAll()
  52. }
  53. }
  54. func (s *tcp) callback(inConn net.Conn) {
  55. defer func() {
  56. if err := recover(); err != nil {
  57. log.Printf("conn handler crashed with err : %s \nstack: %s", err, string(debug.Stack()))
  58. }
  59. }()
  60. if s.cfg.OutCallback != nil {
  61. if !s.cfg.OutCallback() {
  62. CloseConn(&inConn)
  63. return
  64. }
  65. }
  66. err := s.outToTCP(&inConn)
  67. if err != nil {
  68. CloseConn(&inConn)
  69. }
  70. s.inConn = &inConn
  71. }
  72. func (s *tcp) outToTCP(inConn *net.Conn) (err error) {
  73. var outConn net.Conn
  74. var _outConn interface{}
  75. _outConn, err = s.outPool.Pool.Get()
  76. if err == nil {
  77. outConn = _outConn.(net.Conn)
  78. }
  79. if err != nil {
  80. CloseConn(inConn)
  81. return
  82. }
  83. inAddr := (*inConn).RemoteAddr().String()
  84. inLocalAddr := (*inConn).LocalAddr().String()
  85. outAddr := outConn.RemoteAddr().String()
  86. outLocalAddr := outConn.LocalAddr().String()
  87. IoBind(*inConn, outConn, func(isSrcErr bool, err error) {
  88. log.Printf("conn %s - %s - %s -%s released", inAddr, inLocalAddr, outLocalAddr, outAddr)
  89. CloseConn(inConn)
  90. CloseConn(&outConn)
  91. }, func(n int, d bool) {}, 0)
  92. log.Printf("conn %s - %s - %s -%s connected", inAddr, inLocalAddr, outLocalAddr, outAddr)
  93. return
  94. }
  95. func (s *tcp) initOutConnPool() {
  96. s.outPool = NewOutPool(s.cfg.CheckParentInterval, s.cfg.Parent, s.cfg.Timeout, s.cfg.PoolSize, s.cfg.PoolSize*2)
  97. }