From 8257369ec1ed8c341f35b3e161a6fd460c3e8ac5 Mon Sep 17 00:00:00 2001
From: Kahbazi <A.Kahbazi@gmail.com>
Date: Fri, 13 Sep 2019 01:58:46 +0430
Subject: [PATCH] Add support for cancellation token in typed client hub
 (#13816)

---
 .../Core/src/Internal/TypedClientBuilder.cs   | 22 ++++++++--
 .../test/Internal/TypedClientBuilderTests.cs  | 41 +++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/src/SignalR/server/Core/src/Internal/TypedClientBuilder.cs b/src/SignalR/server/Core/src/Internal/TypedClientBuilder.cs
index 511280d636..a333c600a8 100644
--- a/src/SignalR/server/Core/src/Internal/TypedClientBuilder.cs
+++ b/src/SignalR/server/Core/src/Internal/TypedClientBuilder.cs
@@ -132,6 +132,14 @@ namespace Microsoft.AspNetCore.SignalR.Internal
                 methodBuilder.DefineGenericParameters(genericTypeNames);
             }
 
+            // Check to see if the last parameter of the method is a CancellationToken
+            bool hasCancellationToken = paramTypes.LastOrDefault() == typeof(CancellationToken);
+            if (hasCancellationToken)
+            {
+                // remove CancellationToken from input paramTypes
+                paramTypes = paramTypes.Take(paramTypes.Length - 1).ToArray();
+            }
+
             var generator = methodBuilder.GetILGenerator();
 
             // Declare local variable to store the arguments to IClientProxy.SendCoreAsync
@@ -145,7 +153,7 @@ namespace Microsoft.AspNetCore.SignalR.Internal
             generator.Emit(OpCodes.Ldstr, interfaceMethodInfo.Name);
 
             // Create an new object array to hold all the parameters to this method
-            generator.Emit(OpCodes.Ldc_I4, parameters.Length); // Stack: 
+            generator.Emit(OpCodes.Ldc_I4, paramTypes.Length); // Stack: 
             generator.Emit(OpCodes.Newarr, typeof(object)); // allocate object array
             generator.Emit(OpCodes.Stloc_0);
 
@@ -162,8 +170,16 @@ namespace Microsoft.AspNetCore.SignalR.Internal
             // Load parameter array on to the stack.
             generator.Emit(OpCodes.Ldloc_0);
 
-            // Get 'CancellationToken.None' and put it on the stack, since we don't support CancellationToken right now
-            generator.Emit(OpCodes.Call, CancellationTokenNoneProperty.GetMethod);
+            if (hasCancellationToken)
+            {
+                // Get CancellationToken from input argument and put it on the stack
+                generator.Emit(OpCodes.Ldarg, paramTypes.Length + 1);
+            }
+            else
+            {
+                // Get 'CancellationToken.None' and put it on the stack, for when method does not have CancellationToken
+                generator.Emit(OpCodes.Call, CancellationTokenNoneProperty.GetMethod);
+            }
 
             // Send!
             generator.Emit(OpCodes.Callvirt, invokeMethod);
diff --git a/src/SignalR/server/SignalR/test/Internal/TypedClientBuilderTests.cs b/src/SignalR/server/SignalR/test/Internal/TypedClientBuilderTests.cs
index 4f68f6fe74..9f412af8f4 100644
--- a/src/SignalR/server/SignalR/test/Internal/TypedClientBuilderTests.cs
+++ b/src/SignalR/server/SignalR/test/Internal/TypedClientBuilderTests.cs
@@ -75,6 +75,41 @@ namespace Microsoft.AspNetCore.SignalR.Tests.Internal
             await task2.OrTimeout();
         }
 
+        [Fact]
+        public async Task SupportsCancellationToken()
+        {
+            var clientProxy = new MockProxy();
+            var typedProxy = TypedClientBuilder<ICancellationTokenMethod>.Build(clientProxy);
+            CancellationTokenSource cts1 = new CancellationTokenSource();
+            var task1 = typedProxy.Method("foo", cts1.Token);
+            Assert.False(task1.IsCompleted);
+
+            CancellationTokenSource cts2 = new CancellationTokenSource();
+            var task2 = typedProxy.NoArgumentMethod(cts2.Token);
+            Assert.False(task2.IsCompleted);
+
+            Assert.Collection(clientProxy.Sends,
+                send1 =>
+                {
+                    Assert.Equal("Method", send1.Method);
+                    Assert.Equal(1, send1.Arguments.Length);
+                    Assert.Collection(send1.Arguments,
+                        arg1 => Assert.Equal("foo", arg1));
+                    Assert.Equal(cts1.Token, send1.CancellationToken);
+                    send1.Complete();
+                },
+                send2 =>
+                {
+                    Assert.Equal("NoArgumentMethod", send2.Method);
+                    Assert.Equal(0, send2.Arguments.Length);
+                    Assert.Equal(cts2.Token, send2.CancellationToken);
+                    send2.Complete();
+                });
+
+            await task1.OrTimeout();
+            await task2.OrTimeout();
+        }
+
         [Fact]
         public void ThrowsIfProvidedAClass()
         {
@@ -179,6 +214,12 @@ namespace Microsoft.AspNetCore.SignalR.Tests.Internal
             Task SubMethod(string foo);
         }
 
+        public interface ICancellationTokenMethod
+        {
+            Task Method(string foo, CancellationToken cancellationToken);
+            Task NoArgumentMethod(CancellationToken cancellationToken);
+        }
+
         public interface IPropertiesClient
         {
             string Property { get; }