From 120b7fb7d1e2bac5cfdce5c7d5ec741780dee5eb Mon Sep 17 00:00:00 2001
From: Christian <christian.kratky@googlemail.com>
Date: Thu, 18 Feb 2021 18:56:15 +0100
Subject: [PATCH] Fix RPC issue #1093 by calling both message handlers.

---
 .../MQTTnet.Extensions.Rpc/MqttRpcClient.cs   | 22 +++++++++++--------
 ...cAwareApplicationMessageReceivedHandler.cs |  9 ++++----
 ...MqttApplicationMessageReceivedEventArgs.cs |  8 +++++++
 3 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs b/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs
index c225230e..3fbec759 100644
--- a/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs
+++ b/Source/MQTTnet.Extensions.Rpc/MqttRpcClient.cs
@@ -8,6 +8,7 @@ using System.Collections.Concurrent;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using MQTTnet.Implementations;
 
 namespace MQTTnet.Extensions.Rpc
 {
@@ -88,8 +89,8 @@ namespace MQTTnet.Extensions.Rpc
 
             try
             {
-                var tcs = new TaskCompletionSource<byte[]>();
-                if (!_waitingCalls.TryAdd(responseTopic, tcs))
+                var promise = new TaskCompletionSource<byte[]>();
+                if (!_waitingCalls.TryAdd(responseTopic, promise))
                 {
                     throw new InvalidOperationException();
                 }
@@ -102,15 +103,15 @@ namespace MQTTnet.Extensions.Rpc
                 {
                     linkedCts.Token.Register(() =>
                     {
-                        if (!tcs.Task.IsCompleted && !tcs.Task.IsFaulted && !tcs.Task.IsCanceled)
+                        if (!promise.Task.IsCompleted && !promise.Task.IsFaulted && !promise.Task.IsCanceled)
                         {
-                            tcs.TrySetCanceled();
+                            promise.TrySetCanceled();
                         }
                     });
 
                     try
                     {
-                        var result = await tcs.Task.ConfigureAwait(false);
+                        var result = await promise.Task.ConfigureAwait(false);
                         timeoutCts.Cancel(false);
                         return result;
                     }
@@ -136,14 +137,17 @@ namespace MQTTnet.Extensions.Rpc
 
         Task HandleApplicationMessageReceivedAsync(MqttApplicationMessageReceivedEventArgs eventArgs)
         {
-            if (!_waitingCalls.TryRemove(eventArgs.ApplicationMessage.Topic, out var tcs))
+            if (!_waitingCalls.TryRemove(eventArgs.ApplicationMessage.Topic, out var promise))
             {
-                return Task.FromResult(0);
+                return PlatformAbstractionLayer.CompletedTask;
             }
 
-            tcs.TrySetResult(eventArgs.ApplicationMessage.Payload);
+            promise.TrySetResult(eventArgs.ApplicationMessage.Payload);
 
-            return Task.FromResult(0);
+            // Set this message to handled to that other code can avoid execution etc.
+            eventArgs.IsHandled = true;
+            
+            return PlatformAbstractionLayer.CompletedTask;
         }
 
         public void Dispose()
diff --git a/Source/MQTTnet.Extensions.Rpc/RpcAwareApplicationMessageReceivedHandler.cs b/Source/MQTTnet.Extensions.Rpc/RpcAwareApplicationMessageReceivedHandler.cs
index d7b75fc9..5d43eb68 100644
--- a/Source/MQTTnet.Extensions.Rpc/RpcAwareApplicationMessageReceivedHandler.cs
+++ b/Source/MQTTnet.Extensions.Rpc/RpcAwareApplicationMessageReceivedHandler.cs
@@ -18,14 +18,15 @@ namespace MQTTnet.Extensions.Rpc
 
         public IMqttApplicationMessageReceivedHandler OriginalHandler { get; }
 
-        public Task HandleApplicationMessageReceivedAsync(MqttApplicationMessageReceivedEventArgs eventArgs)
+        public async Task HandleApplicationMessageReceivedAsync(MqttApplicationMessageReceivedEventArgs eventArgs)
         {
+            // First try to check if there is a pending RPC call! 
+            await _handleReceivedApplicationMessageAsync(eventArgs).ConfigureAwait(false);
+            
             if (OriginalHandler != null)
             {
-                return OriginalHandler.HandleApplicationMessageReceivedAsync(eventArgs);
+                await OriginalHandler.HandleApplicationMessageReceivedAsync(eventArgs).ConfigureAwait(false);
             }
-
-            return _handleReceivedApplicationMessageAsync(eventArgs);
         }
     }
 }
diff --git a/Source/MQTTnet/MqttApplicationMessageReceivedEventArgs.cs b/Source/MQTTnet/MqttApplicationMessageReceivedEventArgs.cs
index 533f4df6..c44f1793 100644
--- a/Source/MQTTnet/MqttApplicationMessageReceivedEventArgs.cs
+++ b/Source/MQTTnet/MqttApplicationMessageReceivedEventArgs.cs
@@ -21,5 +21,13 @@ namespace MQTTnet
         public bool ProcessingFailed { get; set; }
 
         public MqttApplicationMessageReceivedReasonCode ReasonCode { get; set; } = MqttApplicationMessageReceivedReasonCode.Success;
+        
+        /// <summary>
+        /// Gets or sets whether this message was handled.
+        /// This value can be used in user code for custom control flow.
+        /// </summary>
+        public bool IsHandled { get; set; }
+        
+        public object Tag { get; set; }
     }
 }