You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/zrpc/internal/client.go

198 lines
5.1 KiB
Go

4 years ago
package internal
4 years ago
import (
"context"
"crypto/tls"
"crypto/x509"
"errors"
4 years ago
"fmt"
"io/ioutil"
"log"
"strings"
4 years ago
"time"
"github.com/tal-tech/go-zero/zrpc/internal/balancer/p2c"
"github.com/tal-tech/go-zero/zrpc/internal/clientinterceptors"
"github.com/tal-tech/go-zero/zrpc/internal/resolver"
4 years ago
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
4 years ago
)
const (
dialTimeout = time.Second * 3
separator = '/'
)
4 years ago
func init() {
resolver.RegisterResolver()
}
4 years ago
type (
// Client interface wraps the Conn method.
Client interface {
Conn() *grpc.ClientConn
}
// A ClientOptions is a client options.
4 years ago
ClientOptions struct {
NonBlock bool
4 years ago
Timeout time.Duration
Secure bool
Retry bool
4 years ago
DialOptions []grpc.DialOption
}
// ClientOption defines the method to customize a ClientOptions.
4 years ago
ClientOption func(options *ClientOptions)
client struct {
conn *grpc.ClientConn
}
4 years ago
)
// NewClient returns a Client.
func NewClient(target string, opts ...ClientOption) (Client, error) {
var cli client
opts = append([]ClientOption{WithDialOption(grpc.WithBalancerName(p2c.Name))}, opts...)
if err := cli.dial(target, opts...); err != nil {
return nil, err
}
return &cli, nil
}
func (c *client) Conn() *grpc.ClientConn {
return c.conn
4 years ago
}
func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
var cliOpts ClientOptions
4 years ago
for _, opt := range opts {
opt(&cliOpts)
4 years ago
}
var options []grpc.DialOption
if !cliOpts.Secure {
options = append([]grpc.DialOption(nil), grpc.WithInsecure())
}
if !cliOpts.NonBlock {
options = append(options, grpc.WithBlock())
}
options = append(options,
4 years ago
WithUnaryClientInterceptors(
clientinterceptors.UnaryTracingInterceptor,
4 years ago
clientinterceptors.DurationInterceptor,
clientinterceptors.PrometheusInterceptor,
3 years ago
clientinterceptors.BreakerInterceptor,
clientinterceptors.TimeoutInterceptor(cliOpts.Timeout),
clientinterceptors.RetryInterceptor(cliOpts.Retry),
),
WithStreamClientInterceptors(
clientinterceptors.StreamTracingInterceptor,
4 years ago
),
)
4 years ago
return append(options, cliOpts.DialOptions...)
4 years ago
}
func (c *client) dial(server string, opts ...ClientOption) error {
options := c.buildDialOptions(opts...)
4 years ago
timeCtx, cancel := context.WithTimeout(context.Background(), dialTimeout)
defer cancel()
conn, err := grpc.DialContext(timeCtx, server, options...)
if err != nil {
service := server
if errors.Is(err, context.DeadlineExceeded) {
pos := strings.LastIndexByte(server, separator)
// len(server) - 1 is the index of last char
if 0 < pos && pos < len(server)-1 {
service = server[pos+1:]
}
}
return fmt.Errorf("rpc dial: %s, error: %s, make sure rpc service %q is already started",
server, err.Error(), service)
4 years ago
}
c.conn = conn
return nil
}
// WithDialOption returns a func to customize a ClientOptions with given dial option.
func WithDialOption(opt grpc.DialOption) ClientOption {
return func(options *ClientOptions) {
options.DialOptions = append(options.DialOptions, opt)
}
}
// WithNonBlock sets the dialing to be nonblock.
func WithNonBlock() ClientOption {
return func(options *ClientOptions) {
options.NonBlock = true
}
}
// WithTimeout returns a func to customize a ClientOptions with given timeout.
func WithTimeout(timeout time.Duration) ClientOption {
return func(options *ClientOptions) {
options.Timeout = timeout
}
4 years ago
}
// WithRetry returns a func to customize a ClientOptions with auto retry.
func WithRetry() ClientOption {
return func(options *ClientOptions) {
options.Retry = true
}
}
// WithUnaryClientInterceptor returns a func to customize a ClientOptions with given interceptor.
func WithUnaryClientInterceptor(interceptor grpc.UnaryClientInterceptor) ClientOption {
return func(options *ClientOptions) {
options.DialOptions = append(options.DialOptions, WithUnaryClientInterceptors(interceptor))
}
}
// WithTlsClientFromUnilateral return a func to customize a ClientOptions Verify with Unilateralism authentication.
func WithTlsClientFromUnilateral(crt, domainName string) ClientOption {
return func(options *ClientOptions) {
c, err := credentials.NewClientTLSFromFile(crt, domainName)
if err != nil {
log.Fatalf("credentials.NewClientTLSFromFile err: %v", err)
}
options.Secure = true
options.DialOptions = append(options.DialOptions, grpc.WithTransportCredentials(c))
}
}
// WithTlsClientFromMutual return a func to customize a ClientOptions Verify with mutual authentication.
func WithTlsClientFromMutual(crtFile, keyFile, caFile string) ClientOption {
return func(options *ClientOptions) {
cert, err := tls.LoadX509KeyPair(crtFile, keyFile)
if err != nil {
log.Fatalf("tls.LoadX509KeyPair err: %v", err)
}
certPool := x509.NewCertPool()
ca, err := ioutil.ReadFile(caFile)
if err != nil {
log.Fatalf("credentials: failed to ReadFile CA certificates err: %v", err)
}
if !certPool.AppendCertsFromPEM(ca) {
log.Fatalf("credentials: failed to append certificates err: %v", err)
}
config := &tls.Config{
Certificates: []tls.Certificate{cert},
RootCAs: certPool,
}
options.Secure = true
options.DialOptions = append(options.DialOptions,
grpc.WithTransportCredentials(credentials.NewTLS(config)))
}
}