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

     1  package netapi
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/Asutorufa/yuhaiin/pkg/utils/pool"
     9  	"github.com/Asutorufa/yuhaiin/pkg/utils/system"
    10  )
    11  
    12  type Server interface {
    13  	io.Closer
    14  }
    15  
    16  type Listener interface {
    17  	Stream(context.Context) (net.Listener, error)
    18  	Packet(context.Context) (net.PacketConn, error)
    19  	Server
    20  }
    21  
    22  type Accepter interface {
    23  	Server
    24  	AcceptStream() (*StreamMeta, error)
    25  	AcceptPacket() (*Packet, error)
    26  }
    27  
    28  type StreamMeta struct {
    29  	Source      net.Addr
    30  	Destination net.Addr
    31  	Inbound     net.Addr
    32  
    33  	Src     net.Conn
    34  	Address Address
    35  }
    36  
    37  type WriteBack func(b []byte, addr net.Addr) (int, error)
    38  
    39  type Packet struct {
    40  	Src       net.Addr
    41  	Dst       Address
    42  	WriteBack WriteBack
    43  	Payload   *pool.Bytes
    44  }
    45  
    46  type DNSServer interface {
    47  	Server
    48  	HandleUDP(context.Context, net.PacketConn) error
    49  	HandleTCP(context.Context, net.Conn) error
    50  	Do(context.Context, *pool.Bytes, func([]byte) error) error
    51  }
    52  
    53  var EmptyDNSServer DNSServer = &emptyHandler{}
    54  
    55  type emptyHandler struct{}
    56  
    57  func (e *emptyHandler) Close() error                                    { return nil }
    58  func (e *emptyHandler) HandleUDP(context.Context, net.PacketConn) error { return io.EOF }
    59  func (e *emptyHandler) HandleTCP(context.Context, net.Conn) error       { return io.EOF }
    60  func (e *emptyHandler) Do(_ context.Context, b *pool.Bytes, _ func([]byte) error) error {
    61  	b.Free()
    62  	return io.EOF
    63  }
    64  
    65  type ChannelListener struct {
    66  	ctx     context.Context
    67  	cancel  context.CancelFunc
    68  	channel chan net.Conn
    69  	addr    net.Addr
    70  }
    71  
    72  func NewChannelListener(addr net.Addr) *ChannelListener {
    73  	ctx, cancel := context.WithCancel(context.Background())
    74  	return &ChannelListener{
    75  		addr:    addr,
    76  		ctx:     ctx,
    77  		cancel:  cancel,
    78  		channel: make(chan net.Conn, system.Procs)}
    79  }
    80  
    81  func (c *ChannelListener) Accept() (net.Conn, error) {
    82  	select {
    83  	case <-c.ctx.Done():
    84  		return nil, c.ctx.Err()
    85  
    86  	case conn := <-c.channel:
    87  		return conn, nil
    88  	}
    89  }
    90  
    91  func (c *ChannelListener) NewConn(conn net.Conn) {
    92  	select {
    93  	case <-c.ctx.Done():
    94  		conn.Close()
    95  	case c.channel <- conn:
    96  	}
    97  }
    98  
    99  func (c *ChannelListener) Close() error {
   100  	c.cancel()
   101  	return nil
   102  }
   103  
   104  func (c *ChannelListener) Addr() net.Addr { return c.addr }
   105  
   106  type ListenerPatch struct {
   107  	Listener
   108  	lis net.Listener
   109  }
   110  
   111  func PatchStream(lis net.Listener, inbound Listener) *ListenerPatch {
   112  	return &ListenerPatch{
   113  		Listener: inbound,
   114  		lis:      lis,
   115  	}
   116  }
   117  
   118  func (w *ListenerPatch) Stream(ctx context.Context) (net.Listener, error) { return w.lis, nil }
   119  
   120  func (w *ListenerPatch) Close() error {
   121  	w.lis.Close()
   122  	return w.Listener.Close()
   123  }
   124  
   125  type ChannelServer struct {
   126  	packetChan chan *Packet
   127  	streamChan chan *StreamMeta
   128  	ctx        context.Context
   129  	cancel     context.CancelFunc
   130  }
   131  
   132  func NewChannelServer() *ChannelServer {
   133  	ctx, cancel := context.WithCancel(context.Background())
   134  	return &ChannelServer{
   135  		packetChan: make(chan *Packet, 100),
   136  		streamChan: make(chan *StreamMeta, 100),
   137  		ctx:        ctx,
   138  		cancel:     cancel,
   139  	}
   140  }
   141  
   142  func (s *ChannelServer) AcceptPacket() (*Packet, error) {
   143  	select {
   144  	case <-s.ctx.Done():
   145  		return nil, s.ctx.Err()
   146  	case p := <-s.packetChan:
   147  		return p, nil
   148  	}
   149  }
   150  
   151  func (s *ChannelServer) AcceptStream() (*StreamMeta, error) {
   152  	select {
   153  	case <-s.ctx.Done():
   154  		return nil, s.ctx.Err()
   155  	case p := <-s.streamChan:
   156  		return p, nil
   157  	}
   158  }
   159  
   160  func (s *ChannelServer) Close() error {
   161  	s.cancel()
   162  	return nil
   163  }
   164  
   165  func (s *ChannelServer) SendPacket(packet *Packet) error {
   166  	select {
   167  	case <-s.ctx.Done():
   168  		return s.ctx.Err()
   169  	case s.packetChan <- packet:
   170  		return nil
   171  	}
   172  }
   173  
   174  func (s *ChannelServer) SendStream(stream *StreamMeta) error {
   175  	select {
   176  	case <-s.ctx.Done():
   177  		return s.ctx.Err()
   178  	case s.streamChan <- stream:
   179  		return nil
   180  	}
   181  }
   182  
   183  func (s *ChannelServer) Context() context.Context { return s.ctx }