github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/tunnel/adapter/server.go (about)

     1  package adapter
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"errors"
     7  	"net"
     8  
     9  	"github.com/laof/lite-speed-test/common"
    10  	"github.com/laof/lite-speed-test/log"
    11  	"github.com/laof/lite-speed-test/tunnel"
    12  	"github.com/laof/lite-speed-test/tunnel/freedom"
    13  	"github.com/laof/lite-speed-test/tunnel/http"
    14  	"github.com/laof/lite-speed-test/tunnel/socks"
    15  	"github.com/laof/lite-speed-test/utils"
    16  )
    17  
    18  type Server struct {
    19  	tcpListener net.Listener
    20  	udpListener net.PacketConn
    21  	socksConn   chan tunnel.Conn
    22  	httpConn    chan tunnel.Conn
    23  	nextSocks   bool
    24  	ctx         context.Context
    25  	cancel      context.CancelFunc
    26  }
    27  
    28  // TODO: https socks4
    29  func (s *Server) acceptConnLoop(tcpListener net.Listener) {
    30  	for {
    31  		conn, err := tcpListener.Accept()
    32  		if err != nil {
    33  			select {
    34  			case <-s.ctx.Done():
    35  				log.D("exiting")
    36  				return
    37  			default:
    38  				continue
    39  			}
    40  		}
    41  
    42  		br := bufio.NewReader(conn)
    43  		b, err := br.Peek(1)
    44  		if err != nil {
    45  			log.Error(common.NewError("failed to detect proxy protocol type").Base(err))
    46  			conn.Close()
    47  			continue
    48  		}
    49  		cc := &common.BufferdConn{Conn: conn, Br: br}
    50  		switch b[0] {
    51  		case 4:
    52  			log.Error(common.NewError("not support proxy protocol type").Base(err))
    53  			conn.Close()
    54  			continue
    55  		case 5:
    56  			s.socksConn <- &freedom.Conn{
    57  				Conn: cc,
    58  			}
    59  		default:
    60  			s.httpConn <- &freedom.Conn{
    61  				Conn: cc,
    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, errors.New("adapter closed")
    74  		}
    75  	} else if _, ok := overlay.(*socks.Tunnel); ok {
    76  		s.nextSocks = true
    77  		select {
    78  		case conn := <-s.socksConn:
    79  			return conn, nil
    80  		case <-s.ctx.Done():
    81  			return nil, errors.New("adapter closed")
    82  		}
    83  	} else {
    84  		panic("invalid overlay")
    85  	}
    86  }
    87  
    88  func (s *Server) AcceptPacket(tunnel.Tunnel) (tunnel.PacketConn, error) {
    89  	return &freedom.PacketConn{
    90  		UDPConn: s.udpListener.(*net.UDPConn),
    91  	}, nil
    92  }
    93  
    94  func (s *Server) Close() error {
    95  	s.cancel()
    96  	s.tcpListener.Close()
    97  	return s.udpListener.Close()
    98  }
    99  
   100  func NewServer(ctx context.Context, _ tunnel.Server) (*Server, error) {
   101  	var cancel context.CancelFunc
   102  	ctx, cancel = context.WithCancel(ctx)
   103  	localHost := ctx.Value("LocalHost").(string)
   104  	localPort := ctx.Value("LocalPort").(int)
   105  	addr := tunnel.NewAddressFromHostPort("tcp", localHost, localPort)
   106  	// tcpListener, err := utils.Listen(ctx, "tcp", addr.String())
   107  	// if err != nil {
   108  	// 	return nil, common.NewError("adapter failed to create tcp listener").Base(err)
   109  	// }
   110  	tcpListeners, err := utils.GetListens(ctx, "tcp", addr.String())
   111  	if err != nil || len(tcpListeners) < 1 {
   112  		return nil, common.NewError("adapter failed to create tcp listener").Base(err)
   113  	}
   114  	udpListener, err := net.ListenPacket("udp", addr.String())
   115  	if err != nil {
   116  		return nil, common.NewError("adapter failed to create tcp listener").Base(err)
   117  	}
   118  	server := &Server{
   119  		tcpListener: tcpListeners[0],
   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  	for _, v := range tcpListeners {
   127  		go server.acceptConnLoop(v)
   128  	}
   129  	return server, nil
   130  }