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

     1  package grpc
     2  
     3  import (
     4  	context "context"
     5  	"fmt"
     6  	"net"
     7  
     8  	"github.com/Asutorufa/yuhaiin/pkg/log"
     9  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    10  	"github.com/Asutorufa/yuhaiin/pkg/protos/config/listener"
    11  	"github.com/Asutorufa/yuhaiin/pkg/utils/id"
    12  	grpc "google.golang.org/grpc"
    13  )
    14  
    15  type Grpc struct {
    16  	UnimplementedStreamServer
    17  	connChan chan *conn
    18  	id       id.IDGenerator
    19  
    20  	listener net.Listener
    21  	Server   *grpc.Server
    22  }
    23  
    24  func init() {
    25  	listener.RegisterTransport(NewServer)
    26  }
    27  
    28  func NewServer(c *listener.Transport_Grpc) func(netapi.Listener) (netapi.Listener, error) {
    29  	return func(ii netapi.Listener) (netapi.Listener, error) {
    30  		lis, err := ii.Stream(context.TODO())
    31  		if err != nil {
    32  			return nil, err
    33  		}
    34  
    35  		return netapi.PatchStream(NewGrpcNoServer(lis), ii), nil
    36  	}
    37  }
    38  
    39  func NewGrpcNoServer(lis net.Listener) *Grpc {
    40  	s := grpc.NewServer()
    41  
    42  	g := &Grpc{
    43  		connChan: make(chan *conn, 30),
    44  		Server:   s,
    45  		listener: lis,
    46  	}
    47  
    48  	s.RegisterService(&Stream_ServiceDesc, g)
    49  
    50  	if lis != nil {
    51  		go log.IfErr("grpc serve", func() error { return s.Serve(lis) })
    52  	}
    53  
    54  	return g
    55  }
    56  
    57  func (g *Grpc) Addr() net.Addr {
    58  	if g.listener != nil {
    59  		return g.listener.Addr()
    60  	}
    61  
    62  	return netapi.EmptyAddr
    63  }
    64  
    65  func (g *Grpc) Close() error {
    66  	g.Server.Stop()
    67  
    68  	var err error
    69  	if g.listener != nil {
    70  		err = g.listener.Close()
    71  	}
    72  
    73  	return err
    74  }
    75  
    76  func (g *Grpc) Accept() (net.Conn, error) {
    77  	conn, ok := <-g.connChan
    78  	if !ok {
    79  		return nil, net.ErrClosed
    80  	}
    81  
    82  	return conn, nil
    83  }
    84  
    85  func (s *Grpc) Conn(con Stream_ConnServer) error {
    86  	ctx, cancel := context.WithCancel(con.Context())
    87  	s.connChan <- &conn{
    88  		raw:   con,
    89  		raddr: &addr{s.id.Generate()},
    90  		laddr: s.Addr(),
    91  		close: cancel,
    92  	}
    93  
    94  	<-ctx.Done()
    95  
    96  	return nil
    97  }
    98  
    99  type addr struct {
   100  	id uint64
   101  }
   102  
   103  func (addr) Network() string  { return "tcp" }
   104  func (a addr) String() string { return fmt.Sprintf("grpc://%d", a.id) }