Merge pull request #205 from Snowflake-Labs/elinardi-log-peer

Add peer identity to proxy and server logs
Snowflake-Labs · Mar 22, 2023 · 78d8162 · 78d8162
2 parents c5d60b4 + 12def00
commit 78d8162
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 9 deletions.
diff --git a/auth/opa/rpcauth/input.go b/auth/opa/rpcauth/input.go
@@ -159,7 +159,7 @@ func PeerInputFromContext(ctx context.Context) *PeerAuthInput {
 	out := &PeerAuthInput{}
 	p, ok := peer.FromContext(ctx)
 	if !ok {
-		return out
+		return nil
 	}
 	out.Net = NetInputFromAddr(p.Addr)
 	out.Cert = CertInputFrom(p.AuthInfo)

diff --git a/auth/opa/rpcauth/rpcauth_test.go b/auth/opa/rpcauth/rpcauth_test.go
@@ -403,7 +403,7 @@ func TestRpcAuthInput(t *testing.T) {
 			method: "/AMethod",
 			compare: &RPCAuthInput{
 				Method: "/AMethod",
-				Peer:   &PeerAuthInput{},
+				Peer:   nil,
 			},
 		},
 		{
@@ -413,7 +413,7 @@ func TestRpcAuthInput(t *testing.T) {
 			compare: &RPCAuthInput{
 				Method:   "/AMethod",
 				Metadata: md,
-				Peer:     &PeerAuthInput{},
+				Peer:     nil,
 			},
 		},
 		{
@@ -425,7 +425,7 @@ func TestRpcAuthInput(t *testing.T) {
 				Method:      "/AMethod",
 				Message:     json.RawMessage{0x7b, 0x7d},
 				MessageType: "google.protobuf.Empty",
-				Peer:        &PeerAuthInput{},
+				Peer:        nil,
 			},
 		},
 		{

diff --git a/proxy/server/target.go b/proxy/server/target.go
@@ -274,6 +274,11 @@ func (s *TargetStream) Run(nonce uint32, replyChan chan *pb.ProxyReply) {
 				return err
 			}
 			streamPeerInfo := s.PeerAuthInfo()
+			if streamPeerInfo == nil {
+				err := status.Errorf(codes.Internal, "peer auth info cannot be nil")
+				s.CloseWith(err)
+				return err
+			}
 			authinput.Host = &rpcauth.HostAuthInput{
 				Net:       streamPeerInfo.Net,
 				Cert:      streamPeerInfo.Cert,

diff --git a/telemetry/telemetry.go b/telemetry/telemetry.go
@@ -23,10 +23,10 @@ import (
 	"io"
 	"strings"
 
+	"github.com/Snowflake-Labs/sansshell/auth/opa/rpcauth"
 	"github.com/go-logr/logr"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/peer"
 )
 
 const (
@@ -155,8 +155,9 @@ func (l *loggedClientStream) CloseSend() error {
 func UnaryServerLogInterceptor(logger logr.Logger) grpc.UnaryServerInterceptor {
 	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 		l := logger.WithValues("method", info.FullMethod)
-		if p, ok := peer.FromContext(ctx); ok {
-			l = l.WithValues("peer-address", p.Addr)
+		p := rpcauth.PeerInputFromContext(ctx)
+		if p != nil {
+			l = l.WithValues("peer", p)
 		}
 		l = logMetadata(ctx, l)
 		l.Info("new request")
@@ -177,8 +178,9 @@ func UnaryServerLogInterceptor(logger logr.Logger) grpc.UnaryServerInterceptor {
 func StreamServerLogInterceptor(logger logr.Logger) grpc.StreamServerInterceptor {
 	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 		l := logger.WithValues("method", info.FullMethod)
-		if p, ok := peer.FromContext(ss.Context()); ok {
-			l = l.WithValues("peer-address", p.Addr)
+		p := rpcauth.PeerInputFromContext(ss.Context())
+		if p != nil {
+			l = l.WithValues("peer", p)
 		}
 		l = logMetadata(ss.Context(), l)
 		l.Info("new stream")