From d3a678ec4894a3de127b42ccb9a2450ed7b6aa76 Mon Sep 17 00:00:00 2001
From: Elad Zelingher <dark_leizar@walla.com>
Date: Sun, 7 Sep 2014 23:05:27 +0300
Subject: [PATCH] Handling tasks that return internal types

---
 src/WampSharp.Tests/Api/RpcServerTests.cs     | 44 ++++++++++++++++-
 src/WampSharp/Rpc/Client/TaskExtensions.cs    | 49 +++++++++++++++++--
 src/WampSharp/Rpc/Client/WampRpcSerializer.cs | 15 +-----
 3 files changed, 89 insertions(+), 19 deletions(-)

diff --git a/src/WampSharp.Tests/Api/RpcServerTests.cs b/src/WampSharp.Tests/Api/RpcServerTests.cs
index 6b4ee6e2e..4c3dfaf1e 100644
--- a/src/WampSharp.Tests/Api/RpcServerTests.cs
+++ b/src/WampSharp.Tests/Api/RpcServerTests.cs
@@ -1,4 +1,5 @@
-using Moq;
+using System.Threading.Tasks;
+using Moq;
 using NUnit.Framework;
 using WampSharp.Rpc;
 using WampSharp.Tests.TestHelpers;
@@ -15,6 +16,12 @@ public interface ICalculator
             int Square(int x);
         }
 
+        public interface INumberProcessor
+        {
+            [WampRpcMethod("test/square")]
+            Task ProcessNumber(int x);
+        }
+
         [Test]
         public void RequestContextIsSet()
         {
@@ -44,5 +51,40 @@ public void RequestContextIsSet()
 
             Assert.That(context.SessionId, Is.EqualTo(channel.GetMonitor().SessionId));
         }
+
+#if NET45
+
+        [Test]
+        public void AsyncAwaitTaskWork()
+        {
+            WampPlayground playground = new WampPlayground();
+
+            IWampHost host = playground.Host;
+
+            WampRequestContext context = null;
+
+            Mock<INumberProcessor> mock = new Mock<INumberProcessor>();
+
+            mock.Setup(x => x.ProcessNumber(It.IsAny<int>()))
+                          .Returns(async (int x) =>
+                              {
+                              });
+
+            host.HostService(mock.Object);
+
+            host.Open();
+
+            IWampChannel<MockRaw> channel = playground.CreateNewChannel();
+
+            channel.Open();
+
+            INumberProcessor proxy = channel.GetRpcProxy<INumberProcessor>();
+
+            Task task = proxy.ProcessNumber(4);
+
+            mock.Verify(x => x.ProcessNumber(4));
+        }
+
+#endif
     }
 }
\ No newline at end of file
diff --git a/src/WampSharp/Rpc/Client/TaskExtensions.cs b/src/WampSharp/Rpc/Client/TaskExtensions.cs
index 61df2fb80..2eaedc04f 100644
--- a/src/WampSharp/Rpc/Client/TaskExtensions.cs
+++ b/src/WampSharp/Rpc/Client/TaskExtensions.cs
@@ -1,14 +1,22 @@
 using System;
 using System.Reflection;
 using System.Threading.Tasks;
+using WampSharp.Core.Utilities;
 
 namespace WampSharp.Rpc.Client
 {
     internal static class TaskExtensions
     {
-        private static readonly MethodInfo mCastTask = GetCastTaskMethod();
+        private static readonly MethodInfo mCastTaskToGenericTask = GetCastTaskToGenericTaskMethod();
+        private static readonly MethodInfo mCastToNonGenericTask = GetCastGenericTaskToNonGenericMethod();
 
-        private static MethodInfo GetCastTaskMethod()
+        private static MethodInfo GetCastGenericTaskToNonGenericMethod()
+        {
+            return typeof(TaskExtensions).GetMethod("InnerCastTask",
+                                                     BindingFlags.Static | BindingFlags.NonPublic);
+        }
+
+        private static MethodInfo GetCastTaskToGenericTaskMethod()
         {
             return typeof(TaskExtensions).GetMethod("InternalCastTask",
                                                      BindingFlags.Static | BindingFlags.NonPublic);
@@ -16,7 +24,7 @@ private static MethodInfo GetCastTaskMethod()
 
         public static Task Cast(this Task<object> task, Type taskType)
         {
-            return (Task)mCastTask.MakeGenericMethod(taskType).Invoke(null, new object[] { task });
+            return (Task)mCastTaskToGenericTask.MakeGenericMethod(taskType).Invoke(null, new object[] { task });
         }
 
         private static Task<T> InternalCastTask<T>(Task<object> task)
@@ -39,7 +47,12 @@ public static Task<object> CastTask(this Task task)
             }
             else
             {
-                result = InnerCastTask((dynamic)task);
+                Type underlyingType = UnwrapReturnType(task.GetType());
+
+                MethodInfo method =
+                    mCastToNonGenericTask.MakeGenericMethod(underlyingType);
+
+                result = (Task<object>) method.Invoke(null, new object[] {task});
             }
 
             return result;
@@ -71,5 +84,33 @@ private static TResult ContinueWithSafeCallback<TTask, TResult>(TTask task, Func
 
             return result;
         }
+    
+        /// <summary>
+        /// Unwraps the return type of a given method.
+        /// </summary>
+        /// <param name="returnType">The given return type.</param>
+        /// <returns>The unwrapped return type.</returns>
+        /// <example>
+        /// void, Task -> object
+        /// Task{string} -> string
+        /// int -> int
+        /// </example>
+        public static Type UnwrapReturnType(Type returnType)
+        {
+            if (returnType == typeof(void) || returnType == typeof(Task))
+            {
+                return typeof(object);
+            }
+
+            Type taskType =
+                returnType.GetClosedGenericTypeImplementation(typeof(Task<>));
+
+            if (taskType != null)
+            {
+                return returnType.GetGenericArguments()[0];
+            }
+
+            return returnType;
+        }
     }
 }
\ No newline at end of file
diff --git a/src/WampSharp/Rpc/Client/WampRpcSerializer.cs b/src/WampSharp/Rpc/Client/WampRpcSerializer.cs
index 3585c93d5..96bc60a50 100644
--- a/src/WampSharp/Rpc/Client/WampRpcSerializer.cs
+++ b/src/WampSharp/Rpc/Client/WampRpcSerializer.cs
@@ -38,20 +38,7 @@ public WampRpcCall Serialize(MethodInfo method, object[] arguments)
 
         private Type ExtractReturnType(Type returnType)
         {
-            if (returnType == typeof (void) || returnType == typeof(Task))
-            {
-                return typeof (object);
-            }
-
-            Type taskType = 
-                returnType.GetClosedGenericTypeImplementation(typeof (Task<>));
-
-            if (taskType != null)
-            {
-                return returnType.GetGenericArguments()[0];
-            }
-
-            return returnType;
+            return TaskExtensions.UnwrapReturnType(returnType);
         }
     }
 }
\ No newline at end of file