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

     1  package grpc
     2  
     3  import (
     4  	context "context"
     5  	"crypto/tls"
     6  	"net"
     7  	"sync"
     8  	"sync/atomic"
     9  	"time"
    10  
    11  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    12  	"github.com/Asutorufa/yuhaiin/pkg/protos/node/point"
    13  	"github.com/Asutorufa/yuhaiin/pkg/protos/node/protocol"
    14  	grpc "google.golang.org/grpc"
    15  	"google.golang.org/grpc/backoff"
    16  	"google.golang.org/grpc/credentials"
    17  	"google.golang.org/grpc/credentials/insecure"
    18  )
    19  
    20  type client struct {
    21  	netapi.Proxy
    22  
    23  	clientConn *grpc.ClientConn
    24  	client     StreamClient
    25  
    26  	tlsConfig *tls.Config
    27  
    28  	count     *atomic.Int64
    29  	stopTimer *time.Timer
    30  	mu        sync.Mutex
    31  }
    32  
    33  func init() {
    34  	point.RegisterProtocol(NewClient)
    35  }
    36  
    37  func NewClient(config *protocol.Protocol_Grpc) point.WrapProxy {
    38  	return func(p netapi.Proxy) (netapi.Proxy, error) {
    39  		return &client{
    40  			Proxy:     p,
    41  			count:     &atomic.Int64{},
    42  			tlsConfig: point.ParseTLSConfig(config.Grpc.Tls),
    43  		}, nil
    44  	}
    45  }
    46  
    47  func (c *client) initClient() error {
    48  	c.mu.Lock()
    49  	defer c.mu.Unlock()
    50  
    51  	if c.clientConn != nil {
    52  		c.clientCountAdd()
    53  		return nil
    54  	}
    55  
    56  	var tlsOption grpc.DialOption
    57  	if c.tlsConfig == nil {
    58  		tlsOption = grpc.WithTransportCredentials(insecure.NewCredentials())
    59  	} else {
    60  		tlsOption = grpc.WithTransportCredentials(credentials.NewTLS(c.tlsConfig))
    61  	}
    62  
    63  	clientConn, err := grpc.Dial("",
    64  		grpc.WithConnectParams(grpc.ConnectParams{
    65  			Backoff: backoff.Config{
    66  				BaseDelay:  500 * time.Millisecond,
    67  				Multiplier: 1.5,
    68  				Jitter:     0.2,
    69  				MaxDelay:   19 * time.Second,
    70  			},
    71  			MinConnectTimeout: 5 * time.Second,
    72  		}),
    73  		grpc.WithInitialWindowSize(65536),
    74  		tlsOption,
    75  		grpc.WithContextDialer(func(ctx context.Context, s string) (net.Conn, error) {
    76  			return c.Proxy.Conn(ctx, netapi.EmptyAddr)
    77  		}))
    78  	if err != nil {
    79  		return err
    80  	}
    81  
    82  	c.clientConn = clientConn
    83  	c.client = NewStreamClient(clientConn)
    84  	c.clientCountAdd()
    85  
    86  	return nil
    87  }
    88  
    89  func (c *client) clientCountAdd() {
    90  	if c.count.Add(1) == 1 && c.stopTimer != nil {
    91  		c.stopTimer.Stop()
    92  	}
    93  }
    94  
    95  func (c *client) clientCountSub() {
    96  	c.mu.Lock()
    97  	defer c.mu.Unlock()
    98  
    99  	if c.count.Add(-1) != 0 {
   100  		return
   101  	}
   102  
   103  	if c.stopTimer == nil {
   104  		c.stopTimer = time.AfterFunc(time.Minute, c.close)
   105  	} else {
   106  		c.stopTimer.Reset(time.Minute)
   107  	}
   108  }
   109  
   110  func (c *client) reconnect() error {
   111  	c.close()
   112  	return c.initClient()
   113  }
   114  
   115  func (c *client) close() {
   116  	c.mu.Lock()
   117  	if c.clientConn != nil {
   118  		c.clientConn.Close()
   119  		c.clientConn = nil
   120  		c.client = nil
   121  	}
   122  	c.mu.Unlock()
   123  }
   124  
   125  func (c *client) Conn(ctx context.Context, addr netapi.Address) (net.Conn, error) {
   126  	if err := c.initClient(); err != nil {
   127  		return nil, err
   128  	}
   129  	var retried bool
   130  
   131  _retry:
   132  	ctx, cancel := context.WithCancel(ctx)
   133  	con, err := c.client.Conn(ctx)
   134  	if err != nil {
   135  		cancel()
   136  		if !retried {
   137  			if er := c.reconnect(); er != nil {
   138  				return nil, er
   139  			}
   140  			retried = true
   141  			goto _retry
   142  		}
   143  		return nil, err
   144  	}
   145  
   146  	return &conn{
   147  		raw:   con,
   148  		laddr: caddr{},
   149  		close: func() {
   150  			cancel()
   151  			_ = con.CloseSend()
   152  			c.clientCountSub()
   153  		},
   154  		raddr: addr,
   155  	}, nil
   156  }
   157  
   158  type caddr struct{}
   159  
   160  func (caddr) Network() string { return "tcp" }
   161  func (caddr) String() string  { return "GRPC" }