github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/simple/server.go (about)

     1  package simple
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"net"
     7  	"sync"
     8  
     9  	"github.com/Asutorufa/yuhaiin/pkg/net/dialer"
    10  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    11  	"github.com/Asutorufa/yuhaiin/pkg/protos/config/listener"
    12  )
    13  
    14  type Server struct {
    15  	net.Listener
    16  	net.PacketConn
    17  
    18  	control listener.TcpUdpControl
    19  	host    string
    20  	pmu     sync.Mutex
    21  	smu     sync.Mutex
    22  }
    23  
    24  func (s *Server) Close() error {
    25  	var err error
    26  
    27  	if s.Listener != nil {
    28  		if er := s.Listener.Close(); er != nil {
    29  			err = errors.Join(err, er)
    30  		}
    31  	}
    32  
    33  	if s.PacketConn != nil {
    34  		if er := s.PacketConn.Close(); er != nil {
    35  			err = errors.Join(err, er)
    36  		}
    37  	}
    38  
    39  	return err
    40  }
    41  
    42  func (s *Server) initPacketConn() error {
    43  	if s.PacketConn != nil {
    44  		return nil
    45  	}
    46  
    47  	s.pmu.Lock()
    48  	defer s.pmu.Unlock()
    49  
    50  	if s.PacketConn != nil {
    51  		return nil
    52  	}
    53  
    54  	p, err := dialer.ListenPacket("udp", s.host)
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	s.PacketConn = p
    60  
    61  	return nil
    62  }
    63  
    64  func (s *Server) initStream() error {
    65  	if s.Listener != nil {
    66  		return nil
    67  	}
    68  
    69  	s.smu.Lock()
    70  	defer s.smu.Unlock()
    71  
    72  	if s.Listener != nil {
    73  		return nil
    74  	}
    75  
    76  	lis, err := dialer.ListenContext(context.TODO(), "tcp", s.host)
    77  	if err != nil {
    78  		return err
    79  	}
    80  
    81  	s.Listener = lis
    82  
    83  	return nil
    84  }
    85  
    86  func (s *Server) Packet(ctx context.Context) (net.PacketConn, error) {
    87  	if s.control == listener.TcpUdpControl_disable_udp {
    88  		return nil, errors.ErrUnsupported
    89  	}
    90  
    91  	if err := s.initPacketConn(); err != nil {
    92  		return nil, err
    93  	}
    94  
    95  	return s.PacketConn, nil
    96  }
    97  
    98  func (s *Server) Stream(ctx context.Context) (net.Listener, error) {
    99  	if s.control == listener.TcpUdpControl_disable_tcp {
   100  		return nil, errors.ErrUnsupported
   101  	}
   102  
   103  	if err := s.initStream(); err != nil {
   104  		return nil, err
   105  	}
   106  
   107  	return s.Listener, nil
   108  }
   109  
   110  func NewServer(c *listener.Inbound_Tcpudp) (netapi.Listener, error) {
   111  	return &Server{
   112  		host:    c.Tcpudp.Host,
   113  		control: c.Tcpudp.Control,
   114  	}, nil
   115  }
   116  
   117  func init() {
   118  	listener.RegisterNetwork(NewServer)
   119  }