From d6f44fda1728520d216e1d49ad81f97ef062e07f Mon Sep 17 00:00:00 2001
From: LPeter1997 <lenkefi.peti@gmail.com>
Date: Sat, 30 Sep 2023 21:13:22 +0200
Subject: [PATCH] Fixes

---
 .../IJsonRpcMethodHandlerAttribute.cs         | 23 ++++++++++++++++
 src/Draco.JsonRpc/JsonRpcConnection.cs        |  8 ++++--
 src/Draco.JsonRpc/JsonRpcMethodHandler.cs     | 26 ++++++++++++++++---
 3 files changed, 51 insertions(+), 6 deletions(-)
 create mode 100644 src/Draco.JsonRpc/IJsonRpcMethodHandlerAttribute.cs

diff --git a/src/Draco.JsonRpc/IJsonRpcMethodHandlerAttribute.cs b/src/Draco.JsonRpc/IJsonRpcMethodHandlerAttribute.cs
new file mode 100644
index 0000000000..96e71b7f44
--- /dev/null
+++ b/src/Draco.JsonRpc/IJsonRpcMethodHandlerAttribute.cs
@@ -0,0 +1,23 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Draco.JsonRpc;
+
+/// <summary>
+/// Implemented by attributes that mark JSON RPC method handles.
+/// </summary>
+public interface IJsonRpcMethodHandlerAttribute
+{
+    /// <summary>
+    /// The name of the JSON RPC method to handle.
+    /// </summary>
+    public string Method { get; }
+
+    /// <summary>
+    /// True, if this method mutates the state.
+    /// </summary>
+    public bool Mutating { get; }
+}
diff --git a/src/Draco.JsonRpc/JsonRpcConnection.cs b/src/Draco.JsonRpc/JsonRpcConnection.cs
index abc628384a..0db4503dde 100644
--- a/src/Draco.JsonRpc/JsonRpcConnection.cs
+++ b/src/Draco.JsonRpc/JsonRpcConnection.cs
@@ -48,11 +48,15 @@ public abstract class JsonRpcConnection<TMessage>
     /// <summary>
     /// Registers an RPC handler method.
     /// </summary>
+    /// <typeparam name="TRequest">The request attribute type.</typeparam>
+    /// <typeparam name="TNotification">The notification attribute type.</typeparam>
     /// <param name="handlerMethod">The handler method info.</param>
     /// <param name="target">The method invocation target.</param>
-    public void AddRpcMethod(MethodInfo handlerMethod, object? target)
+    public void AddRpcMethod<TRequest, TNotification>(MethodInfo handlerMethod, object? target)
+        where TRequest : Attribute, IJsonRpcMethodHandlerAttribute
+        where TNotification : Attribute, IJsonRpcMethodHandlerAttribute
     {
-        var handler = new JsonRpcMethodHandler(handlerMethod, target);
+        var handler = JsonRpcMethodHandler.Create<TRequest, TNotification>(handlerMethod, target);
         this.methodHandlers.Add(handler.MethodName, handler);
     }
 
diff --git a/src/Draco.JsonRpc/JsonRpcMethodHandler.cs b/src/Draco.JsonRpc/JsonRpcMethodHandler.cs
index 95390b7b01..effad7ebcd 100644
--- a/src/Draco.JsonRpc/JsonRpcMethodHandler.cs
+++ b/src/Draco.JsonRpc/JsonRpcMethodHandler.cs
@@ -14,6 +14,23 @@ namespace Draco.JsonRpc;
 /// </summary>
 internal sealed class JsonRpcMethodHandler
 {
+    /// <summary>
+    /// Constructs a new <see cref="JsonRpcMethodHandler"/>.
+    /// </summary>
+    /// <typeparam name="TRequest">The request attribute type.</typeparam>
+    /// <typeparam name="TNotification">The notification attribute type.</typeparam>
+    /// <param name="handlerMethod">The handler method.</param>
+    /// <param name="target">The invocation target.</param>
+    /// <returns>The constructed handler object.</returns>
+    public static JsonRpcMethodHandler Create<TRequest, TNotification>(MethodInfo handlerMethod, object? target)
+        where TRequest : Attribute, IJsonRpcMethodHandlerAttribute
+        where TNotification : Attribute, IJsonRpcMethodHandlerAttribute
+    {
+        var requestAttr = handlerMethod.GetCustomAttribute<TRequest>();
+        var notificationAttr = handlerMethod.GetCustomAttribute<TNotification>();
+        return new(handlerMethod, target, requestAttr, notificationAttr);
+    }
+
     /// <summary>
     /// The handler <see cref="MethodInfo"/>.
     /// </summary>
@@ -60,7 +77,11 @@ internal sealed class JsonRpcMethodHandler
     /// </summary>
     public Type DeclaredReturnType { get; }
 
-    public JsonRpcMethodHandler(MethodInfo handlerMethod, object? target)
+    private JsonRpcMethodHandler(
+        MethodInfo handlerMethod,
+        object? target,
+        IJsonRpcMethodHandlerAttribute? requestAttr,
+        IJsonRpcMethodHandlerAttribute? notificationAttr)
     {
         this.HandlerMethod = handlerMethod;
         this.Target = target;
@@ -104,9 +125,6 @@ public JsonRpcMethodHandler(MethodInfo handlerMethod, object? target)
         }
 
         // Verify attributes
-        var requestAttr = handlerMethod.GetCustomAttribute<RequestAttribute>();
-        var notificationAttr = handlerMethod.GetCustomAttribute<NotificationAttribute>();
-
         if (requestAttr is null && notificationAttr is null)
         {
             throw new ArgumentException($"Handler must be marked as either a request or notification handler.", nameof(handlerMethod));