Make CancellationToken available in call credentials interceptor

src/Grpc.AspNetCore.Server/Internal/HttpContextServerCallContext.cs

@@ -77,12 +77,7 @@ protected override string PeerCore
         get
         {
             // Follows the standard at https://github.com/grpc/grpc/blob/master/doc/naming.md
-            if (_peer == null)
-            {
-                _peer = BuildPeer();
-            }
-
-            return _peer;
+            return _peer ??= BuildPeer();
         }
     }
 
@@ -291,10 +286,7 @@ private void EndCallCore()
 
     private void LogCallEnd()
     {
-        if (_activity != null)
-        {
-            _activity.AddTag(GrpcServerConstants.ActivityStatusCodeTag, _status.StatusCode.ToTrailerString());
-        }
+        _activity?.AddTag(GrpcServerConstants.ActivityStatusCodeTag, _status.StatusCode.ToTrailerString());
         if (_status.StatusCode != StatusCode.OK)
         {
             if (GrpcEventSource.Log.IsEnabled())
@@ -387,10 +379,7 @@ protected override Task WriteResponseHeadersAsyncCore(Metadata responseHeaders)
     public void Initialize(ISystemClock? clock = null)
     {
         _activity = GetHostActivity();
-        if (_activity != null)
-        {
-            _activity.AddTag(GrpcServerConstants.ActivityMethodTag, MethodCore);
-        }
+        _activity?.AddTag(GrpcServerConstants.ActivityMethodTag, MethodCore);
 
         if (GrpcEventSource.Log.IsEnabled())
         {

src/Grpc.Core.Api/AsyncAuthInterceptor.cs

@@ -16,6 +16,7 @@
 
 #endregion
 
+using System.Threading;
 using System.Threading.Tasks;
 using Grpc.Core.Utils;
 
@@ -34,16 +35,25 @@ namespace Grpc.Core;
 /// </summary>
 public class AuthInterceptorContext
 {
-    readonly string serviceUrl;
-    readonly string methodName;
+    private readonly string serviceUrl;
+    private readonly string methodName;
+    private readonly CancellationToken cancellationToken;
 
     /// <summary>
     /// Initializes a new instance of <c>AuthInterceptorContext</c>.
     /// </summary>
-    public AuthInterceptorContext(string serviceUrl, string methodName)
+    public AuthInterceptorContext(string serviceUrl, string methodName) : this(serviceUrl, methodName, CancellationToken.None)
+    {
+    }
+
+    /// <summary>
+    /// Initializes a new instance of <c>AuthInterceptorContext</c>.
+    /// </summary>
+    public AuthInterceptorContext(string serviceUrl, string methodName, CancellationToken cancellationToken)
     {
         this.serviceUrl = GrpcPreconditions.CheckNotNull(serviceUrl, nameof(serviceUrl));
         this.methodName = GrpcPreconditions.CheckNotNull(methodName, nameof(methodName));
+        this.cancellationToken = cancellationToken;
     }
 
     /// <summary>
@@ -61,4 +71,12 @@ public string MethodName
     {
         get { return methodName; }
     }
+
+    /// <summary>
+    /// The cancellation token of the RPC being called.
+    /// </summary>
+    public CancellationToken CancellationToken
+    {
+        get { return cancellationToken; }
+    }
 }

src/Grpc.Net.Client/Internal/DefaultCallCredentialsConfigurator.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -20,15 +20,19 @@
 
 namespace Grpc.Net.Client.Internal;
 
-internal class DefaultCallCredentialsConfigurator : CallCredentialsConfiguratorBase
+internal sealed class DefaultCallCredentialsConfigurator : CallCredentialsConfiguratorBase
 {
     public AsyncAuthInterceptor? Interceptor { get; private set; }
-    public IReadOnlyList<CallCredentials>? Credentials { get; private set; }
+    public IReadOnlyList<CallCredentials>? CompositeCredentials { get; private set; }
 
-    public void Reset()
+    // A place to cache the context to avoid creating a new instance for each auth interceptor call.
+    // It's ok not to reset this state because the context is only used for the lifetime of the call.
+    public AuthInterceptorContext? CachedContext { get; set; }
+
+    public void ResetPerCallCredentialState()
     {
         Interceptor = null;
-        Credentials = null;
+        CompositeCredentials = null;
     }
 
     public override void SetAsyncAuthInterceptorCredentials(object? state, AsyncAuthInterceptor interceptor)
@@ -38,6 +42,6 @@ public override void SetAsyncAuthInterceptorCredentials(object? state, AsyncAuth
 
     public override void SetCompositeCredentials(object? state, IReadOnlyList<CallCredentials> credentials)
     {
-        Credentials = credentials;
+        CompositeCredentials = credentials;
     }
 }

src/Grpc.Net.Client/Internal/GrpcCall.cs

@@ -955,13 +955,13 @@ private async Task ReadCredentials(HttpRequestMessage request)
 
             if (Options.Credentials != null)
             {
-                await GrpcProtocolHelpers.ReadCredentialMetadata(configurator, Channel, request, Method, Options.Credentials).ConfigureAwait(false);
+                await GrpcProtocolHelpers.ReadCredentialMetadata(configurator, Channel, request, Method, Options.Credentials, _callCts.Token).ConfigureAwait(false);
             }
             if (Channel.CallCredentials?.Count > 0)
             {
                 foreach (var credentials in Channel.CallCredentials)
                 {
-                    await GrpcProtocolHelpers.ReadCredentialMetadata(configurator, Channel, request, Method, credentials).ConfigureAwait(false);
+                    await GrpcProtocolHelpers.ReadCredentialMetadata(configurator, Channel, request, Method, credentials, _callCts.Token).ConfigureAwait(false);
                 }
             }
         }

src/Grpc.Net.Client/Internal/GrpcProtocolHelpers.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -121,13 +121,34 @@ internal static bool ShouldSkipHeader(string name)
     /* round an integer up to the next value with three significant figures */
     private static long TimeoutRoundUpToThreeSignificantFigures(long x)
     {
-        if (x < 1000) return x;
-        if (x < 10000) return RoundUp(x, 10);
-        if (x < 100000) return RoundUp(x, 100);
-        if (x < 1000000) return RoundUp(x, 1000);
-        if (x < 10000000) return RoundUp(x, 10000);
-        if (x < 100000000) return RoundUp(x, 100000);
-        if (x < 1000000000) return RoundUp(x, 1000000);
+        if (x < 1000)
+        {
+            return x;
+        }
+        if (x < 10000)
+        {
+            return RoundUp(x, 10);
+        }
+        if (x < 100000)
+        {
+            return RoundUp(x, 100);
+        }
+        if (x < 1000000)
+        {
+            return RoundUp(x, 1000);
+        }
+        if (x < 10000000)
+        {
+            return RoundUp(x, 10000);
+        }
+        if (x < 100000000)
+        {
+            return RoundUp(x, 100000);
+        }
+        if (x < 1000000000)
+        {
+            return RoundUp(x, 1000000);
+        }
         return RoundUp(x, 10000000);
 
         static long RoundUp(long x, long divisor)
@@ -235,7 +256,7 @@ internal static bool CanWriteCompressed(WriteOptions? writeOptions)
         return canCompress;
     }
 
-    internal static AuthInterceptorContext CreateAuthInterceptorContext(Uri baseAddress, IMethod method)
+    internal static AuthInterceptorContext CreateAuthInterceptorContext(Uri baseAddress, IMethod method, CancellationToken cancellationToken)
     {
         var authority = baseAddress.Authority;
         if (baseAddress.Scheme == Uri.UriSchemeHttps && authority.EndsWith(":443", StringComparison.Ordinal))
@@ -252,38 +273,44 @@ internal static AuthInterceptorContext CreateAuthInterceptorContext(Uri baseAddr
             serviceUrl += "/";
         }
         serviceUrl += method.ServiceName;
-        return new AuthInterceptorContext(serviceUrl, method.Name);
+        return new AuthInterceptorContext(serviceUrl, method.Name, cancellationToken);
     }
 
     internal static async Task ReadCredentialMetadata(
         DefaultCallCredentialsConfigurator configurator,
         GrpcChannel channel,
         HttpRequestMessage message,
         IMethod method,
-        CallCredentials credentials)
+        CallCredentials credentials,
+        CancellationToken cancellationToken)
     {
         credentials.InternalPopulateConfiguration(configurator, null);
 
         if (configurator.Interceptor != null)
         {
-            var authInterceptorContext = GrpcProtocolHelpers.CreateAuthInterceptorContext(channel.Address, method);
+            // Multiple auth interceptors can be called for a gRPC call.
+            // These all have the same data: address, method and cancellation token.
+            // Lazily allocate the context if it is needed.
+            // Stored on the configurator instead of a ref parameter because ref parameters are not supported in async methods.
+            configurator.CachedContext ??= CreateAuthInterceptorContext(channel.Address, method, cancellationToken);
+
             var metadata = new Metadata();
-            await configurator.Interceptor(authInterceptorContext, metadata).ConfigureAwait(false);
+            await configurator.Interceptor(configurator.CachedContext, metadata).ConfigureAwait(false);
 
             foreach (var entry in metadata)
             {
                 AddHeader(message.Headers, entry);
             }
         }
 
-        if (configurator.Credentials != null)
+        if (configurator.CompositeCredentials != null)
         {
             // Copy credentials locally. ReadCredentialMetadata will update it.
-            var callCredentials = configurator.Credentials;
-            foreach (var c in callCredentials)
+            var compositeCredentials = configurator.CompositeCredentials;
+            foreach (var callCredentials in compositeCredentials)
             {
-                configurator.Reset();
-                await ReadCredentialMetadata(configurator, channel, message, method, c).ConfigureAwait(false);
+                configurator.ResetPerCallCredentialState();
+                await ReadCredentialMetadata(configurator, channel, message, method, callCredentials, cancellationToken).ConfigureAwait(false);
             }
         }
     }

test/Grpc.Net.Client.Tests/CallCredentialTests.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -18,6 +18,7 @@
 
 using System.Net;
 using System.Net.Http.Headers;
+using System.Threading;
 using Greet;
 using Grpc.Core;
 using Grpc.Net.Client.Tests.Infrastructure;
@@ -79,19 +80,71 @@ public async Task CallCredentialsWithHttps_MetadataOnRequest()
         var invoker = HttpClientCallInvokerFactory.Create(httpClient);
 
         // Act
+        var syncPoint = new SyncPoint(runContinuationsAsynchronously: true);
         var callCredentials = CallCredentials.FromInterceptor(async (context, metadata) =>
         {
-            // The operation is asynchronous to ensure delegate is awaited
-            await Task.Delay(50);
+            // The operation is asynchronous to ensure auth interceptor is awaited.
+            // Sending the request and returning a response is blocked until the auth interceptor completes.
+            await syncPoint.WaitToContinue();
+
+            // Set header.
             metadata.Add("authorization", "SECRET_TOKEN");
         });
         var call = invoker.AsyncUnaryCall<HelloRequest, HelloReply>(ClientTestHelpers.ServiceMethod, string.Empty, new CallOptions(credentials: callCredentials), new HelloRequest());
-        await call.ResponseAsync.DefaultTimeout();
+        var responseTask = call.ResponseAsync;
+
+        await syncPoint.WaitForSyncPoint().DefaultTimeout();
+
+        // Response task should be blocked waiting for the auth interceptor to complete.
+        Assert.False(responseTask.IsCompleted);
+        // Sending the request should be blocked waiting for the auth interceptor to complete.
+        Assert.Null(authorizationValue);
+
+        syncPoint.Continue();
+        await responseTask.DefaultTimeout();
 
         // Assert
         Assert.AreEqual("SECRET_TOKEN", authorizationValue);
     }
 
+    [Test]
+    public async Task CallCredentialsWithHttps_CancellationToken()
+    {
+        // Arrange
+        string? authorizationValue = null;
+        var httpClient = ClientTestHelpers.CreateTestClient(async request =>
+        {
+            authorizationValue = request.Headers.GetValues("authorization").Single();
+
+            var reply = new HelloReply { Message = "Hello world" };
+            var streamContent = await ClientTestHelpers.CreateResponseContent(reply).DefaultTimeout();
+            return ResponseUtils.CreateResponse(HttpStatusCode.OK, streamContent);
+        });
+        var invoker = HttpClientCallInvokerFactory.Create(httpClient);
+
+        // Act
+        var unreachableAuthInterceptorSection = false;
+        var callCredentials = CallCredentials.FromInterceptor(async (context, metadata) =>
+        {
+            var tcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+            context.CancellationToken.Register(s => ((TaskCompletionSource<object?>)s!).SetCanceled(), tcs);
+
+            await tcs.Task;
+
+            unreachableAuthInterceptorSection = true;
+        });
+        var call = invoker.AsyncUnaryCall<HelloRequest, HelloReply>(ClientTestHelpers.ServiceMethod, string.Empty, new CallOptions(credentials: callCredentials), new HelloRequest());
+        var responseTask = call.ResponseAsync;
+
+        call.Dispose();
+
+        var ex = await ExceptionAssert.ThrowsAsync<RpcException>(() => responseTask).DefaultTimeout();
+        Assert.AreEqual(StatusCode.Cancelled, ex.StatusCode);
+
+        // Assert
+        Assert.False(unreachableAuthInterceptorSection);
+    }
+
     [Test]
     public async Task CallCredentialsWithHttp_NoMetadataOnRequest()
     {