github.com/ipfans/trojan-go@v0.11.0/tunnel/adapter/server.go (about)

     1  package adapter
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"sync"
     7  
     8  	"github.com/ipfans/trojan-go/common"
     9  	"github.com/ipfans/trojan-go/config"
    10  	"github.com/ipfans/trojan-go/log"
    11  	"github.com/ipfans/trojan-go/tunnel"
    12  	"github.com/ipfans/trojan-go/tunnel/freedom"
    13  	"github.com/ipfans/trojan-go/tunnel/http"
    14  	"github.com/ipfans/trojan-go/tunnel/socks"
    15  )
    16  
    17  type Server struct {
    18  	tcpListener net.Listener
    19  	udpListener net.PacketConn
    20  	socksConn   chan tunnel.Conn
    21  	httpConn    chan tunnel.Conn
    22  	socksLock   sync.RWMutex
    23  	nextSocks   bool
    24  	ctx         context.Context
    25  	cancel      context.CancelFunc
    26  }
    27  
    28  func (s *Server) acceptConnLoop() {
    29  	for {
    30  		conn, err := s.tcpListener.Accept()
    31  		if err != nil {
    32  			select {
    33  			case <-s.ctx.Done():
    34  				log.Debug("exiting")
    35  				return
    36  			default:
    37  				continue
    38  			}
    39  		}
    40  		rewindConn := common.NewRewindConn(conn)
    41  		rewindConn.SetBufferSize(16)
    42  		buf := [3]byte{}
    43  		_, err = rewindConn.Read(buf[:])
    44  		rewindConn.Rewind()
    45  		rewindConn.StopBuffering()
    46  		if err != nil {
    47  			log.Error(common.NewError("failed to detect proxy protocol type").Base(err))
    48  			continue
    49  		}
    50  		s.socksLock.RLock()
    51  		if buf[0] == 5 && s.nextSocks {
    52  			s.socksLock.RUnlock()
    53  			log.Debug("socks5 connection")
    54  			s.socksConn <- &freedom.Conn{
    55  				Conn: rewindConn,
    56  			}
    57  		} else {
    58  			s.socksLock.RUnlock()
    59  			log.Debug("http connection")
    60  			s.httpConn <- &freedom.Conn{
    61  				Conn: rewindConn,
    62  			}
    63  		}
    64  	}
    65  }
    66  
    67  func (s *Server) AcceptConn(overlay tunnel.Tunnel) (tunnel.Conn, error) {
    68  	if _, ok := overlay.(*http.Tunnel); ok {
    69  		select {
    70  		case conn := <-s.httpConn:
    71  			return conn, nil
    72  		case <-s.ctx.Done():
    73  			return nil, common.NewError("adapter closed")
    74  		}
    75  	} else if _, ok := overlay.(*socks.Tunnel); ok {
    76  		s.socksLock.Lock()
    77  		s.nextSocks = true
    78  		s.socksLock.Unlock()
    79  		select {
    80  		case conn := <-s.socksConn:
    81  			return conn, nil
    82  		case <-s.ctx.Done():
    83  			return nil, common.NewError("adapter closed")
    84  		}
    85  	} else {
    86  		panic("invalid overlay")
    87  	}
    88  }
    89  
    90  func (s *Server) AcceptPacket(tunnel.Tunnel) (tunnel.PacketConn, error) {
    91  	return &freedom.PacketConn{
    92  		UDPConn: s.udpListener.(*net.UDPConn),
    93  	}, nil
    94  }
    95  
    96  func (s *Server) Close() error {
    97  	s.cancel()
    98  	s.tcpListener.Close()
    99  	return s.udpListener.Close()
   100  }
   101  
   102  func NewServer(ctx context.Context, _ tunnel.Server) (*Server, error) {
   103  	cfg := config.FromContext(ctx, Name).(*Config)
   104  	var cancel context.CancelFunc
   105  	ctx, cancel = context.WithCancel(ctx)
   106  
   107  	addr := tunnel.NewAddressFromHostPort("tcp", cfg.LocalHost, cfg.LocalPort)
   108  	tcpListener, err := net.Listen("tcp", addr.String())
   109  	if err != nil {
   110  		cancel()
   111  		return nil, common.NewError("adapter failed to create tcp listener").Base(err)
   112  	}
   113  	udpListener, err := net.ListenPacket("udp", addr.String())
   114  	if err != nil {
   115  		cancel()
   116  		return nil, common.NewError("adapter failed to create tcp listener").Base(err)
   117  	}
   118  	server := &Server{
   119  		tcpListener: tcpListener,
   120  		udpListener: udpListener,
   121  		socksConn:   make(chan tunnel.Conn, 32),
   122  		httpConn:    make(chan tunnel.Conn, 32),
   123  		ctx:         ctx,
   124  		cancel:      cancel,
   125  	}
   126  	log.Info("adapter listening on tcp/udp:", addr)
   127  	go server.acceptConnLoop()
   128  	return server, nil
   129  }