From 415d04381a341bd7c872a572b5b126cc86ce6e58 Mon Sep 17 00:00:00 2001
From: Steven Rhodes <steven.rhodes@snowflake.com>
Date: Fri, 17 Nov 2023 13:18:34 -0800
Subject: [PATCH] Change how peer info is cached (#375)

---
 auth/opa/rpcauth/input.go   | 26 +++++++++++++++++++-------
 auth/opa/rpcauth/rpcauth.go | 14 --------------
 2 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/auth/opa/rpcauth/input.go b/auth/opa/rpcauth/input.go
index 81d2bc82..4b20dd2b 100644
--- a/auth/opa/rpcauth/input.go
+++ b/auth/opa/rpcauth/input.go
@@ -21,6 +21,7 @@ import (
 	"crypto/x509/pkix"
 	"encoding/json"
 	"net"
+	"reflect"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -177,20 +178,31 @@ func AddPeerToContext(ctx context.Context, p *PeerAuthInput) context.Context {
 // PeerInputFromContext populates peer information from the supplied
 // context, if available.
 func PeerInputFromContext(ctx context.Context) *PeerAuthInput {
-	// If this runs after rpcauth hooks, we can return richer data that includes
-	// information added by the hooks.
-	cached, ok := ctx.Value(peerInfoKey{}).(*PeerAuthInput)
-	if ok {
-		return cached
-	}
+	cached, _ := ctx.Value(peerInfoKey{}).(*PeerAuthInput)
 
 	out := &PeerAuthInput{}
 	p, ok := peer.FromContext(ctx)
 	if !ok {
-		return nil
+		// If there's no peer info, returned cached data so that invocations
+		// of AddPeerToContext can work outside of RPC contexts.
+		return cached
 	}
+
 	out.Net = NetInputFromAddr(p.Addr)
 	out.Cert = CertInputFrom(p.AuthInfo)
+
+	// If this runs after rpcauth hooks, we can return richer data that includes
+	// information added by the hooks.
+	// We need to compare cached data to peer info because we might be calling
+	// PeerInputFromContext on the context of a client stream, which has a peer
+	// of the target being called and may have the cached value from an earlier
+	// server authorization.
+	if cached != nil && cached.Principal != nil && reflect.DeepEqual(out.Net, cached.Net) {
+		out.Principal = &PrincipalAuthInput{
+			ID:     cached.Principal.ID,
+			Groups: cached.Principal.Groups,
+		}
+	}
 	return out
 }
 
diff --git a/auth/opa/rpcauth/rpcauth.go b/auth/opa/rpcauth/rpcauth.go
index f41075ef..01eb86ae 100644
--- a/auth/opa/rpcauth/rpcauth.go
+++ b/auth/opa/rpcauth/rpcauth.go
@@ -189,7 +189,6 @@ func (g *Authorizer) AuthorizeClient(ctx context.Context, method string, req, re
 	if err := g.Eval(ctx, authInput); err != nil {
 		return err
 	}
-	ctx = AddPeerToContext(ctx, authInput.Peer)
 	return invoker(ctx, method, req, reply, cc, opts...)
 }
 
@@ -212,16 +211,6 @@ type wrappedClientStream struct {
 	grpc.ClientStream
 	method string
 	authz  *Authorizer
-
-	peerMu            sync.Mutex
-	lastPeerAuthInput *PeerAuthInput
-}
-
-func (e *wrappedClientStream) Context() context.Context {
-	e.peerMu.Lock()
-	ctx := AddPeerToContext(e.ClientStream.Context(), e.lastPeerAuthInput)
-	e.peerMu.Unlock()
-	return ctx
 }
 
 // see: grpc.ClientStream.SendMsg
@@ -238,9 +227,6 @@ func (e *wrappedClientStream) SendMsg(req interface{}) error {
 	if err := e.authz.Eval(ctx, authInput); err != nil {
 		return err
 	}
-	e.peerMu.Lock()
-	e.lastPeerAuthInput = authInput.Peer
-	e.peerMu.Unlock()
 	return e.ClientStream.SendMsg(req)
 }