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

     1  package simplesocks
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"github.com/ipfans/trojan-go/common"
     8  	"github.com/ipfans/trojan-go/log"
     9  	"github.com/ipfans/trojan-go/tunnel"
    10  	"github.com/ipfans/trojan-go/tunnel/trojan"
    11  )
    12  
    13  // Server is a simplesocks server
    14  type Server struct {
    15  	underlay   tunnel.Server
    16  	connChan   chan tunnel.Conn
    17  	packetChan chan tunnel.PacketConn
    18  	ctx        context.Context
    19  	cancel     context.CancelFunc
    20  }
    21  
    22  func (s *Server) Close() error {
    23  	s.cancel()
    24  	return s.underlay.Close()
    25  }
    26  
    27  func (s *Server) acceptLoop() {
    28  	for {
    29  		conn, err := s.underlay.AcceptConn(&Tunnel{})
    30  		if err != nil {
    31  			log.Error(common.NewError("simplesocks failed to accept connection from underlying tunnel").Base(err))
    32  			select {
    33  			case <-s.ctx.Done():
    34  				return
    35  			default:
    36  			}
    37  			continue
    38  		}
    39  		metadata := new(tunnel.Metadata)
    40  		if err := metadata.ReadFrom(conn); err != nil {
    41  			log.Error(common.NewError("simplesocks server faield to read header").Base(err))
    42  			conn.Close()
    43  			continue
    44  		}
    45  		switch metadata.Command {
    46  		case Connect:
    47  			s.connChan <- &Conn{
    48  				metadata: metadata,
    49  				Conn:     conn,
    50  			}
    51  		case Associate:
    52  			s.packetChan <- &PacketConn{
    53  				PacketConn: trojan.PacketConn{
    54  					Conn: conn,
    55  				},
    56  			}
    57  		default:
    58  			log.Error(common.NewError(fmt.Sprintf("simplesocks unknown command %d", metadata.Command)))
    59  			conn.Close()
    60  		}
    61  	}
    62  }
    63  
    64  func (s *Server) AcceptConn(tunnel.Tunnel) (tunnel.Conn, error) {
    65  	select {
    66  	case conn := <-s.connChan:
    67  		return conn, nil
    68  	case <-s.ctx.Done():
    69  		return nil, common.NewError("simplesocks server closed")
    70  	}
    71  }
    72  
    73  func (s *Server) AcceptPacket(tunnel.Tunnel) (tunnel.PacketConn, error) {
    74  	select {
    75  	case packetConn := <-s.packetChan:
    76  		return packetConn, nil
    77  	case <-s.ctx.Done():
    78  		return nil, common.NewError("simplesocks server closed")
    79  	}
    80  }
    81  
    82  func NewServer(ctx context.Context, underlay tunnel.Server) (*Server, error) {
    83  	ctx, cancel := context.WithCancel(ctx)
    84  	server := &Server{
    85  		underlay:   underlay,
    86  		ctx:        ctx,
    87  		connChan:   make(chan tunnel.Conn, 32),
    88  		packetChan: make(chan tunnel.PacketConn, 32),
    89  		cancel:     cancel,
    90  	}
    91  	go server.acceptLoop()
    92  	log.Debug("simplesocks server created")
    93  	return server, nil
    94  }