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

     1  /*
     2  Copyright hechain. All Rights Reserved.
     3  
     4  SPDX-License-Identifier: Apache-2.0
     5  */
     6  
     7  package util
     8  
     9  import (
    10  	"context"
    11  	"crypto/sha256"
    12  	"crypto/x509"
    13  
    14  	"google.golang.org/grpc/credentials"
    15  	"google.golang.org/grpc/peer"
    16  )
    17  
    18  func ExtractRemoteAddress(ctx context.Context) string {
    19  	var remoteAddress string
    20  	p, ok := peer.FromContext(ctx)
    21  	if !ok {
    22  		return ""
    23  	}
    24  	if address := p.Addr; address != nil {
    25  		remoteAddress = address.String()
    26  	}
    27  	return remoteAddress
    28  }
    29  
    30  // ExtractCertificateHashFromContext extracts the hash of the certificate from the given context.
    31  // If the certificate isn't present, nil is returned
    32  func ExtractCertificateHashFromContext(ctx context.Context) []byte {
    33  	rawCert := ExtractRawCertificateFromContext(ctx)
    34  	if len(rawCert) == 0 {
    35  		return nil
    36  	}
    37  	h := sha256.New()
    38  	h.Write(rawCert)
    39  	return h.Sum(nil)
    40  }
    41  
    42  // ExtractCertificateFromContext returns the TLS certificate (if applicable)
    43  // from the given context of a gRPC stream
    44  func ExtractCertificateFromContext(ctx context.Context) *x509.Certificate {
    45  	pr, extracted := peer.FromContext(ctx)
    46  	if !extracted {
    47  		return nil
    48  	}
    49  
    50  	authInfo := pr.AuthInfo
    51  	if authInfo == nil {
    52  		return nil
    53  	}
    54  
    55  	tlsInfo, isTLSConn := authInfo.(credentials.TLSInfo)
    56  	if !isTLSConn {
    57  		return nil
    58  	}
    59  	certs := tlsInfo.State.PeerCertificates
    60  	if len(certs) == 0 {
    61  		return nil
    62  	}
    63  	return certs[0]
    64  }
    65  
    66  // ExtractRawCertificateFromContext returns the raw TLS certificate (if applicable)
    67  // from the given context of a gRPC stream
    68  func ExtractRawCertificateFromContext(ctx context.Context) []byte {
    69  	cert := ExtractCertificateFromContext(ctx)
    70  	if cert == nil {
    71  		return nil
    72  	}
    73  	return cert.Raw
    74  }