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 }