Fix CancellationToken serialization in static commands - inject from request context

Co-authored-by: tomasherceg <5599524+tomasherceg@users.noreply.github.com>
Fixed-by: Standa Lukeš

Author: 2 people

src/Framework/Framework/Compilation/Binding/StaticCommandBindingCompiler.cs

@@ -165,6 +165,7 @@ public override string ToString() =>
                 StaticCommandParameterType.Argument => "?",
                 StaticCommandParameterType.Inject => $"service({((Type)Arg!).ToCode(stripNamespace: true)})",
                 StaticCommandParameterType.Invocation => Arg!.ToString()!,
+                StaticCommandParameterType.CurrentCancellationToken => "cancellationToken",
                 _ => "...invalid argument..."
             };
     }
@@ -174,6 +175,7 @@ public enum StaticCommandParameterType : byte
         Inject,
         Constant,
         DefaultValue,
-        Invocation
+        Invocation,
+        CurrentCancellationToken
     }
 }

src/Framework/Framework/Compilation/Binding/StaticCommandExecutionPlanSerializer.cs

@@ -55,6 +55,10 @@ public static JsonNode SerializePlan(StaticCommandInvocationPlan plan)
                 {
                     array.Add(SerializePlan((StaticCommandInvocationPlan)arg.Arg!));
                 }
+                else if (arg.Type == StaticCommandParameterType.CurrentCancellationToken)
+                {
+                    array.Add(null);
+                }
                 else throw new NotSupportedException(arg.Type.ToString());
             }
             return array;
@@ -160,6 +164,8 @@ public static StaticCommandInvocationPlan DeserializePlan(ref Utf8JsonReader jso
                         new StaticCommandParameterPlan(type, methodParameters[i]!.DefaultValue),
                     StaticCommandParameterType.Invocation =>
                         new StaticCommandParameterPlan(type, DeserializePlan(ref json)),
+                    StaticCommandParameterType.CurrentCancellationToken =>
+                        new StaticCommandParameterPlan(type, null),
                     _ => throw new NotSupportedException(type.ToString())
                 };
                 json.AssertRead();

src/Framework/Framework/Compilation/Binding/StaticCommandMethodTranslator.cs

@@ -75,6 +75,9 @@ ValidationPathFormatter validationPathFormatter
             var argPlans = allArguments.Select((arg, index) => {
                 if (arg.OriginalExpression.GetParameterAnnotation() is BindingParameterAnnotation { ExtensionParameter:  InjectedServiceExtensionParameter service })
                     return new StaticCommandParameterPlan(StaticCommandParameterType.Inject, ResolvedTypeDescriptor.ToSystemType(service.ParameterType));
+                else if (arg.OriginalExpression.Type == typeof(System.Threading.CancellationToken))
+                    // CancellationToken cannot be serialized from the client - always provide from the request context
+                    return new StaticCommandParameterPlan(StaticCommandParameterType.CurrentCancellationToken, null);
                 else if (arg.OriginalExpression is ConstantExpression constant)
                 {
                     if (constant.Value == method.GetParameters()[index - (method.IsStatic ? 0 : 1)].DefaultValue)

src/Framework/Framework/Hosting/StaticCommandExecutor.cs

@@ -80,6 +80,8 @@ IDotvvmRequestContext context
 #pragma warning disable CS0618
                         (serviceLoader.GetStaticCommandService((Type)a.Arg!, context), null),
 #pragma warning restore CS0618
+                    StaticCommandParameterType.CurrentCancellationToken =>
+                        (context.RequestAborted, null),
                     StaticCommandParameterType.Invocation =>
                         (await Execute((StaticCommandInvocationPlan)a.Arg!, arguments, argumentValidationPaths, context), null),
                     _ => throw new NotSupportedException("" + a.Type)

src/Tests/Binding/StaticCommandCompilationTests.cs

@@ -147,6 +147,14 @@ public void StaticCommandCompilation_DateTimeResultAssignment()
             Assert.AreEqual("{options.viewModel.DateFrom(dotvvm.serialization.serializeDate(await dotvvm.staticCommandPostback(\"XXXX\",[],options),false));}", result);
         }
 
+        [TestMethod]
+        public void StaticCommandCompilation_AsyncMethodWithCancellationToken_CancellationTokenIsNotInClientArgs()
+        {
+            // CancellationToken should be provided by the server from the request context, not sent from the client
+            var result = CompileBinding("StaticCommands.GetLengthAsync(StringProp)", niceMode: false, typeof(TestViewModel));
+            Assert.AreEqual("await dotvvm.staticCommandPostback(\"XXXX\",[options.viewModel.StringProp.state],options)", result);
+        }
+
         [TestMethod]
         public void StaticCommandCompilation_DateTimeAssignment()
         {
@@ -461,6 +469,9 @@ public static class StaticCommands
 
         [AllowStaticCommand]
         public static DateTime GetDate() => DateTime.UtcNow;
+
+        [AllowStaticCommand]
+        public static Task<int> GetLengthAsync(string str, System.Threading.CancellationToken cancellationToken = default) => Task.FromResult(str.Length);
     }
 
     public abstract class TestInnerService<TOutput>

src/Tests/Binding/StaticCommandExecutorTests.cs

@@ -5,6 +5,7 @@
 using System.Linq.Expressions;
 using System.Reflection;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using DotVVM.Framework.Compilation.Binding;
 using DotVVM.Framework.Compilation.Javascript;
@@ -247,5 +248,44 @@ public async Task Validation_CustomPrimitives()
         internal static void CustomPrimitivesValidation(TestViewModel viewModel, [RegularExpression(@"\d\d7", ErrorMessage = "Vehicle must have lucky number.")] VehicleNumber vehicle)
         {
         }
+
+        [TestMethod]
+        public async Task CancellationToken_IsInjectedFromContext()
+        {
+            var plan = new StaticCommandInvocationPlan(
+                ((Func<CancellationToken, bool>)MethodWithCancellationToken).Method,
+                [ new StaticCommandParameterPlan(StaticCommandParameterType.CurrentCancellationToken, null) ]
+            );
+            var result = await Invoke(plan);
+            // The method should have been called successfully (not thrown an exception)
+            Assert.IsTrue((bool)result);
+        }
+
+        [TestMethod]
+        public async Task CancellationToken_OptionalParameter_IsInjectedFromContext()
+        {
+            var plan = new StaticCommandInvocationPlan(
+                ((Func<string, CancellationToken, Task<string>>)MethodWithOptionalCancellationToken).Method,
+                new[] {
+                    new StaticCommandParameterPlan(StaticCommandParameterType.Argument, typeof(string)),
+                    new StaticCommandParameterPlan(StaticCommandParameterType.CurrentCancellationToken, null)
+                }
+            );
+            var result = await Invoke(plan, ("hello", "/Input"));
+            Assert.AreEqual("hello", result);
+        }
+
+        [AllowStaticCommand]
+        internal static bool MethodWithCancellationToken(CancellationToken cancellationToken)
+        {
+            return !cancellationToken.IsCancellationRequested;
+        }
+
+        [AllowStaticCommand]
+        internal static async Task<string> MethodWithOptionalCancellationToken(string input, CancellationToken cancellationToken = default)
+        {
+            await Task.Yield();
+            return input;
+        }
     }
 }

src/Tests/Binding/StaticCommandPlanSerializationTests.cs

@@ -115,6 +115,25 @@ public static void Method(float arg) { }
             public static void NotOverloadedMethod(int arg) { }
 
             public static void MethodNotUsableInStaticCommand() { }
+
+            [AllowStaticCommand]
+            public static Task MethodWithCancellationToken(int arg, System.Threading.CancellationToken cancellationToken) => Task.CompletedTask;
+        }
+
+        [TestMethod]
+        public void StaticCommandPlanSerialization_CancellationToken_DeserializedPlanIsIdentical()
+        {
+            var plan = MakeInvocationPlan(() => StaticCommandMethodCollection.MethodWithCancellationToken(0, default),
+                new StaticCommandParameterPlan(StaticCommandParameterType.Constant, 42),
+                new StaticCommandParameterPlan(StaticCommandParameterType.CurrentCancellationToken, null));
+            var json = StaticCommandExecutionPlanSerializer.SerializePlan(plan);
+            var deserializedPlan = Deserialize(json);
+
+            Assert.AreEqual(plan.Method, deserializedPlan.Method);
+            Assert.AreEqual(2, deserializedPlan.Arguments.Length);
+            Assert.AreEqual(StaticCommandParameterType.Constant, deserializedPlan.Arguments[0].Type);
+            Assert.AreEqual(StaticCommandParameterType.CurrentCancellationToken, deserializedPlan.Arguments[1].Type);
+            Assert.IsNull(deserializedPlan.Arguments[1].Arg);
         }
     }
 }