credentials: expose NewContextWithRequestInfo publicly (#8198)

Author: vadimsht

credentials/credentials.go

@@ -207,14 +207,32 @@ type RequestInfo struct {
 	AuthInfo AuthInfo
 }
 
+// requestInfoKey is a struct to be used as the key to store RequestInfo in a
+// context.
+type requestInfoKey struct{}
+
 // RequestInfoFromContext extracts the RequestInfo from the context if it exists.
 //
 // This API is experimental.
 func RequestInfoFromContext(ctx context.Context) (ri RequestInfo, ok bool) {
-	ri, ok = icredentials.RequestInfoFromContext(ctx).(RequestInfo)
+	ri, ok = ctx.Value(requestInfoKey{}).(RequestInfo)
 	return ri, ok
 }
 
+// NewContextWithRequestInfo creates a new context from ctx and attaches ri to it.
+//
+// This RequestInfo will be accessible via RequestInfoFromContext.
+//
+// Intended to be used from tests for PerRPCCredentials implementations (that
+// often need to check connection's SecurityLevel). Should not be used from
+// non-test code: the gRPC client already prepares a context with the correct
+// RequestInfo attached when calling PerRPCCredentials.GetRequestMetadata.
+//
+// This API is experimental.
+func NewContextWithRequestInfo(ctx context.Context, ri RequestInfo) context.Context {
+	return context.WithValue(ctx, requestInfoKey{}, ri)
+}
+
 // ClientHandshakeInfo holds data to be passed to ClientHandshake. This makes
 // it possible to pass arbitrary data to the handshaker from gRPC, resolver,
 // balancer etc. Individual credential implementations control the actual

credentials/google/google_test.go

@@ -255,7 +255,7 @@ func TestDefaultCredentialsWithOptions(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			bundle := NewDefaultCredentialsWithOptions(tc.defaultCredsOpts)
 			ri := credentials.RequestInfo{AuthInfo: tc.authInfo}
-			ctx := icredentials.NewRequestInfoContext(ctx, ri)
+			ctx := credentials.NewContextWithRequestInfo(ctx, ri)
 			got, err := bundle.PerRPCCredentials().GetRequestMetadata(ctx, "uri")
 			if err != nil {
 				t.Fatalf("Bundle's PerRPCCredentials().GetRequestMetadata() unexpected error = %v", err)

credentials/sts/sts_test.go

@@ -35,7 +35,6 @@ import (
 	"github.com/google/go-cmp/cmp"
 
 	"google.golang.org/grpc/credentials"
-	icredentials "google.golang.org/grpc/internal/credentials"
 	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/testutils"
 )
@@ -102,7 +101,7 @@ func createTestContext(ctx context.Context, s credentials.SecurityLevel) context
 		Method:   "testInfo",
 		AuthInfo: auth,
 	}
-	return icredentials.NewRequestInfoContext(ctx, ri)
+	return credentials.NewContextWithRequestInfo(ctx, ri)
 }
 
 // errReader implements the io.Reader interface and returns an error from the

internal/credentials/credentials.go

@@ -20,20 +20,6 @@ import (
 	"context"
 )
 
-// requestInfoKey is a struct to be used as the key to store RequestInfo in a
-// context.
-type requestInfoKey struct{}
-
-// NewRequestInfoContext creates a context with ri.
-func NewRequestInfoContext(ctx context.Context, ri any) context.Context {
-	return context.WithValue(ctx, requestInfoKey{}, ri)
-}
-
-// RequestInfoFromContext extracts the RequestInfo from ctx.
-func RequestInfoFromContext(ctx context.Context) any {
-	return ctx.Value(requestInfoKey{})
-}
-
 // clientHandshakeInfoKey is a struct used as the key to store
 // ClientHandshakeInfo in a context.
 type clientHandshakeInfoKey struct{}

internal/transport/http2_client.go

@@ -545,7 +545,7 @@ func (t *http2Client) createHeaderFields(ctx context.Context, callHdr *CallHdr)
 		Method:   callHdr.Method,
 		AuthInfo: t.authInfo,
 	}
-	ctxWithRequestInfo := icredentials.NewRequestInfoContext(ctx, ri)
+	ctxWithRequestInfo := credentials.NewContextWithRequestInfo(ctx, ri)
 	authData, err := t.getTrAuthData(ctxWithRequestInfo, aud)
 	if err != nil {
 		return nil, err