github.com/hechain20/hechain@v0.0.0-20220316014945-b544036ba106/common/grpclogging/grpclogging_suite_test.go (about)

     1  /*
     2  Copyright hechain. All Rights Reserved.
     3  
     4  SPDX-License-Identifier: Apache-2.0
     5  */
     6  
     7  package grpclogging_test
     8  
     9  import (
    10  	"bytes"
    11  	"crypto/ecdsa"
    12  	"crypto/elliptic"
    13  	"crypto/rand"
    14  	"crypto/tls"
    15  	"crypto/x509"
    16  	"crypto/x509/pkix"
    17  	"encoding/pem"
    18  	"math/big"
    19  	"net"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/hechain20/hechain/common/grpclogging/testpb"
    24  	. "github.com/onsi/ginkgo"
    25  	. "github.com/onsi/gomega"
    26  )
    27  
    28  //go:generate protoc --proto_path=testpb --go_out=plugins=grpc,paths=source_relative:testpb testpb/echo.proto
    29  
    30  func TestGrpclogging(t *testing.T) {
    31  	RegisterFailHandler(Fail)
    32  	RunSpecs(t, "Grpclogging Suite")
    33  }
    34  
    35  var (
    36  	clientCertWithKey tls.Certificate
    37  	serverCertWithKey tls.Certificate
    38  
    39  	caCertPool      *x509.CertPool
    40  	clientTLSConfig *tls.Config
    41  	serverTLSConfig *tls.Config
    42  )
    43  
    44  var _ = BeforeSuite(func() {
    45  	var err error
    46  	caCert, caKey := generateCA("test-ca", "127.0.0.1")
    47  	clientCert, clientKey := issueCertificate(caCert, caKey, "client", "127.0.0.1")
    48  	clientCertWithKey, err = tls.X509KeyPair(clientCert, clientKey)
    49  	Expect(err).NotTo(HaveOccurred())
    50  	serverCert, serverKey := issueCertificate(caCert, caKey, "server", "127.0.0.1")
    51  	serverCertWithKey, err = tls.X509KeyPair(serverCert, serverKey)
    52  	Expect(err).NotTo(HaveOccurred())
    53  
    54  	caCertPool = x509.NewCertPool()
    55  	added := caCertPool.AppendCertsFromPEM(caCert)
    56  	Expect(added).To(BeTrue())
    57  
    58  	serverTLSConfig = &tls.Config{
    59  		Certificates: []tls.Certificate{serverCertWithKey},
    60  		ClientAuth:   tls.VerifyClientCertIfGiven,
    61  		ClientCAs:    caCertPool,
    62  		RootCAs:      caCertPool,
    63  	}
    64  
    65  	clientTLSConfig = &tls.Config{
    66  		Certificates:       []tls.Certificate{clientCertWithKey},
    67  		RootCAs:            caCertPool,
    68  		ClientSessionCache: tls.NewLRUClientSessionCache(10),
    69  	}
    70  })
    71  
    72  //go:generate counterfeiter -o fakes/echo_service.go --fake-name EchoServiceServer . echoServiceServer
    73  
    74  type echoServiceServer interface {
    75  	testpb.EchoServiceServer
    76  }
    77  
    78  func newTemplate(subjectCN string, hosts ...string) x509.Certificate {
    79  	notBefore := time.Now().Add(-1 * time.Minute)
    80  	notAfter := time.Now().Add(time.Duration(365 * 24 * time.Hour))
    81  
    82  	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
    83  	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
    84  	Expect(err).NotTo(HaveOccurred())
    85  
    86  	template := x509.Certificate{
    87  		Subject:               pkix.Name{CommonName: subjectCN},
    88  		SerialNumber:          serialNumber,
    89  		NotBefore:             notBefore,
    90  		NotAfter:              notAfter,
    91  		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
    92  		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
    93  		BasicConstraintsValid: true,
    94  	}
    95  	for _, h := range hosts {
    96  		if ip := net.ParseIP(h); ip != nil {
    97  			template.IPAddresses = append(template.IPAddresses, ip)
    98  		} else {
    99  			template.DNSNames = append(template.DNSNames, h)
   100  		}
   101  	}
   102  
   103  	return template
   104  }
   105  
   106  func pemEncode(derCert []byte, key *ecdsa.PrivateKey) (pemCert, pemKey []byte) {
   107  	certBuf := &bytes.Buffer{}
   108  	err := pem.Encode(certBuf, &pem.Block{Type: "CERTIFICATE", Bytes: derCert})
   109  	Expect(err).NotTo(HaveOccurred())
   110  
   111  	keyBytes, err := x509.MarshalECPrivateKey(key)
   112  	Expect(err).NotTo(HaveOccurred())
   113  
   114  	keyBuf := &bytes.Buffer{}
   115  	err = pem.Encode(keyBuf, &pem.Block{Type: "EC PRIVATE KEY", Bytes: keyBytes})
   116  	Expect(err).NotTo(HaveOccurred())
   117  
   118  	return certBuf.Bytes(), keyBuf.Bytes()
   119  }
   120  
   121  func generateCA(subjectCN string, hosts ...string) (pemCert, pemKey []byte) {
   122  	privateKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
   123  	Expect(err).NotTo(HaveOccurred())
   124  	publicKey := privateKey.Public()
   125  
   126  	template := newTemplate(subjectCN, hosts...)
   127  	template.KeyUsage |= x509.KeyUsageCertSign
   128  	template.IsCA = true
   129  
   130  	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, publicKey, privateKey)
   131  	Expect(err).NotTo(HaveOccurred())
   132  
   133  	return pemEncode(derBytes, privateKey)
   134  }
   135  
   136  func issueCertificate(caCert, caKey []byte, subjectCN string, hosts ...string) (pemCert, pemKey []byte) {
   137  	tlsCert, err := tls.X509KeyPair(caCert, caKey)
   138  	Expect(err).NotTo(HaveOccurred())
   139  
   140  	ca, err := x509.ParseCertificate(tlsCert.Certificate[0])
   141  	Expect(err).NotTo(HaveOccurred())
   142  
   143  	privateKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
   144  	Expect(err).NotTo(HaveOccurred())
   145  	publicKey := privateKey.Public()
   146  
   147  	template := newTemplate(subjectCN, hosts...)
   148  	derBytes, err := x509.CreateCertificate(rand.Reader, &template, ca, publicKey, tlsCert.PrivateKey)
   149  	Expect(err).NotTo(HaveOccurred())
   150  
   151  	return pemEncode(derBytes, privateKey)
   152  }