Deny context values from flowing across disjoint requests (#21863)

This prevents things like custom HTTP headers from showing up in an
authentication request.
Removed duplicate import from a test.

Author: jhendrixMSFT

sdk/azcore/CHANGELOG.md

@@ -17,6 +17,7 @@
 * Passing a `nil` credential value will no longer cause a panic. Instead, the authentication is skipped.
 * Calling `Error` on a zero-value `azcore.ResponseError` will no longer panic.
 * Fixed an issue in `fake.PagerResponder[T]` that would cause a trailing error to be omitted when iterating over pages.
+* Context values created by `azcore` will no longer flow across disjoint HTTP requests.
 
 ### Other Changes
 

sdk/azcore/arm/runtime/policy_bearer_token_test.go

@@ -15,7 +15,6 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	armpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/shared"
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	azpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/errorinfo"
@@ -292,7 +291,7 @@ func TestBearerTokenPolicyRequiresHTTPS(t *testing.T) {
 	srv, close := mock.NewServer()
 	defer close()
 	b := NewBearerTokenPolicy(mockCredential{}, nil)
-	pl := newTestPipeline(&policy.ClientOptions{Transport: srv, PerRetryPolicies: []policy.Policy{b}})
+	pl := newTestPipeline(&azpolicy.ClientOptions{Transport: srv, PerRetryPolicies: []azpolicy.Policy{b}})
 	req, err := runtime.NewRequest(context.Background(), "GET", srv.URL())
 	require.NoError(t, err)
 	_, err = pl.Do(req)

sdk/azcore/arm/runtime/policy_register_rp.go

@@ -126,12 +126,13 @@ func (r *rpRegistrationPolicy) Do(req *azpolicy.Request) (*http.Response, error)
 			u:     r.endpoint,
 			subID: subID,
 		}
-		if _, err = rpOps.Register(req.Raw().Context(), rp); err != nil {
+		if _, err = rpOps.Register(&shared.ContextWithDeniedValues{Context: req.Raw().Context()}, rp); err != nil {
 			logRegistrationExit(err)
 			return resp, err
 		}
+
 		// RP was registered, however we need to wait for the registration to complete
-		pollCtx, pollCancel := context.WithTimeout(req.Raw().Context(), r.options.PollingDuration)
+		pollCtx, pollCancel := context.WithTimeout(&shared.ContextWithDeniedValues{Context: req.Raw().Context()}, r.options.PollingDuration)
 		var lastRegState string
 		for {
 			// get the current registration state

sdk/azcore/internal/shared/shared.go

@@ -18,6 +18,9 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/errorinfo"
 )
 
+// NOTE: when adding a new context key type, it likely needs to be
+// added to the deny-list of key types in ContextWithDeniedValues
+
 // CtxWithHTTPHeaderKey is used as a context key for adding/retrieving http.Header.
 type CtxWithHTTPHeaderKey struct{}
 
@@ -110,6 +113,25 @@ func ExtractModuleName(clientName string) (string, string, error) {
 	return matches[3], matches[2], nil
 }
 
+// ContextWithDeniedValues wraps an existing [context.Context], denying access to certain context values.
+// Pipeline policies that create new requests to be sent down their own pipeline MUST wrap the caller's
+// context with an instance of this type. This is to prevent context values from flowing across disjoint
+// requests which can have unintended side-effects.
+type ContextWithDeniedValues struct {
+	context.Context
+}
+
+// Value implements part of the [context.Context] interface.
+// It acts as a deny-list for certain context keys.
+func (c *ContextWithDeniedValues) Value(key any) any {
+	switch key.(type) {
+	case CtxAPINameKey, CtxWithCaptureResponse, CtxWithHTTPHeaderKey, CtxWithRetryOptionsKey, CtxWithTracingTracer:
+		return nil
+	default:
+		return c.Context.Value(key)
+	}
+}
+
 // NonRetriableError marks the specified error as non-retriable.
 func NonRetriableError(err error) error {
 	return &nonRetriableError{err}

sdk/azcore/internal/shared/shared_test.go

@@ -136,3 +136,23 @@ func TestExtractModuleName(t *testing.T) {
 	require.Empty(t, mod)
 	require.Empty(t, client)
 }
+
+func TestContextWithDeniedValues(t *testing.T) {
+	type testKey struct{}
+	const value = "value"
+
+	ctx := context.WithValue(context.Background(), testKey{}, value)
+	ctx = context.WithValue(ctx, CtxAPINameKey{}, value)
+	ctx = context.WithValue(ctx, CtxWithCaptureResponse{}, value)
+	ctx = context.WithValue(ctx, CtxWithHTTPHeaderKey{}, value)
+	ctx = context.WithValue(ctx, CtxWithRetryOptionsKey{}, value)
+	ctx = context.WithValue(ctx, CtxWithTracingTracer{}, value)
+	ctx = &ContextWithDeniedValues{Context: ctx}
+
+	require.Nil(t, ctx.Value(CtxAPINameKey{}))
+	require.Nil(t, ctx.Value(CtxWithCaptureResponse{}))
+	require.Nil(t, ctx.Value(CtxWithHTTPHeaderKey{}))
+	require.Nil(t, ctx.Value(CtxWithRetryOptionsKey{}))
+	require.Nil(t, ctx.Value(CtxWithTracingTracer{}))
+	require.NotNil(t, ctx.Value(testKey{}))
+}

sdk/azcore/runtime/policy_bearer_token.go

@@ -34,7 +34,7 @@ type acquiringResourceState struct {
 // acquire acquires or updates the resource; only one
 // thread/goroutine at a time ever calls this function
 func acquire(state acquiringResourceState) (newResource exported.AccessToken, newExpiration time.Time, err error) {
-	tk, err := state.p.cred.GetToken(state.req.Raw().Context(), state.tro)
+	tk, err := state.p.cred.GetToken(&shared.ContextWithDeniedValues{Context: state.req.Raw().Context()}, state.tro)
 	if err != nil {
 		return exported.AccessToken{}, time.Time{}, err
 	}

sdk/azcore/runtime/policy_http_trace.go

@@ -110,6 +110,10 @@ func StartSpan(ctx context.Context, name string, tracer tracing.Tracer, options
 	if !tracer.Enabled() {
 		return ctx, func(err error) {}
 	}
+
+	// we MUST propagate the active tracer before returning so that the trace policy can access it
+	ctx = context.WithValue(ctx, shared.CtxWithTracingTracer{}, tracer)
+
 	const newSpanKind = tracing.SpanKindInternal
 	if activeSpan := ctx.Value(ctxActiveSpan{}); activeSpan != nil {
 		// per the design guidelines, if a SDK method Foo() calls SDK method Bar(),
@@ -125,7 +129,6 @@ func StartSpan(ctx context.Context, name string, tracer tracing.Tracer, options
 	ctx, span := tracer.Start(ctx, name, &tracing.SpanOptions{
 		Kind: newSpanKind,
 	})
-	ctx = context.WithValue(ctx, shared.CtxWithTracingTracer{}, tracer)
 	ctx = context.WithValue(ctx, ctxActiveSpan{}, newSpanKind)
 	return ctx, func(err error) {
 		if err != nil {

sdk/azcore/runtime/policy_http_trace_test.go

@@ -205,7 +205,6 @@ func TestStartSpansDontNest(t *testing.T) {
 
 	barMethod := func(ctx context.Context) {
 		ourCtx, endSpan := StartSpan(ctx, "BarMethod", tr, nil)
-		require.Same(t, ctx, ourCtx)
 		defer endSpan(nil)
 		req, err := exported.NewRequest(ourCtx, http.MethodGet, srv.URL()+"/bar")
 		require.NoError(t, err)