github.com/hechain20/hechain@v0.0.0-20220316014945-b544036ba106/cmd/common/comm/client.go (about)

     1  /*
     2  Copyright hechain. All Rights Reserved.
     3  
     4  SPDX-License-Identifier: Apache-2.0
     5  */
     6  
     7  package comm
     8  
     9  import (
    10  	"time"
    11  
    12  	"github.com/hechain20/hechain/common/crypto/tlsgen"
    13  	"github.com/hechain20/hechain/common/util"
    14  	"github.com/hechain20/hechain/internal/pkg/comm"
    15  	"github.com/pkg/errors"
    16  	"google.golang.org/grpc"
    17  )
    18  
    19  const defaultTimeout = time.Second * 5
    20  
    21  // Client deals with TLS connections
    22  // to the discovery server
    23  type Client struct {
    24  	config      comm.ClientConfig
    25  	TLSCertHash []byte
    26  }
    27  
    28  // NewClient creates a new comm client out of the given configuration
    29  func NewClient(conf Config) (*Client, error) {
    30  	if conf.Timeout == time.Duration(0) {
    31  		conf.Timeout = defaultTimeout
    32  	}
    33  	sop, err := conf.ToSecureOptions(newSelfSignedTLSCert)
    34  	if err != nil {
    35  		return nil, errors.WithStack(err)
    36  	}
    37  	cc := comm.ClientConfig{
    38  		SecOpts:     sop,
    39  		DialTimeout: conf.Timeout,
    40  	}
    41  	return &Client{config: cc, TLSCertHash: util.ComputeSHA256(sop.Certificate)}, nil
    42  }
    43  
    44  // NewDialer creates a new dialer from the given endpoint
    45  func (c *Client) NewDialer(endpoint string) func() (*grpc.ClientConn, error) {
    46  	return func() (*grpc.ClientConn, error) {
    47  		conn, err := c.config.Dial(endpoint)
    48  		if err != nil {
    49  			return nil, errors.WithStack(err)
    50  		}
    51  		return conn, nil
    52  	}
    53  }
    54  
    55  func newSelfSignedTLSCert() (*tlsgen.CertKeyPair, error) {
    56  	ca, err := tlsgen.NewCA()
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	return ca.NewClientCertKeyPair()
    61  }