From b527059b722dd093b65fb9bb7cdde6a4517bf010 Mon Sep 17 00:00:00 2001
From: Christian Kratky <christian.kratky@googlemail.com>
Date: Thu, 27 Dec 2018 23:05:57 +0100
Subject: [PATCH] Fix session takeover.

---
 Source/MQTTnet/Server/MqttClientSession.cs    | 34 ++++++++++++-------
 .../Server/MqttClientSessionsManager.cs       | 10 +++---
 Source/MQTTnet/Server/MqttServer.cs           |  2 +-
 3 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/Source/MQTTnet/Server/MqttClientSession.cs b/Source/MQTTnet/Server/MqttClientSession.cs
index 9eaab4f1..9d327b3c 100644
--- a/Source/MQTTnet/Server/MqttClientSession.cs
+++ b/Source/MQTTnet/Server/MqttClientSession.cs
@@ -30,7 +30,7 @@ namespace MQTTnet.Server
         private MqttApplicationMessage _willMessage;
         private bool _wasCleanDisconnect;
         private Task _workerTask;
-        private IDisposable _cleanupHandle;
+        //private IDisposable _cleanupHandle;
 
         private string _adapterEndpoint;
         private MqttProtocolVersion? _adapterProtocolVersion;
@@ -73,11 +73,11 @@ namespace MQTTnet.Server
             status.LastNonKeepAlivePacketReceived = _keepAliveMonitor.LastNonKeepAlivePacketReceived;
         }
 
-        public Task RunAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
-        {
-            _workerTask = RunInternalAsync(connectPacket, adapter);
-            return _workerTask;
-        }
+        //public Task RunAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
+        //{
+        //    _workerTask = RunInternalAsync(connectPacket, adapter);
+        //    return _workerTask;
+        //}
 
         public Task StopAsync(MqttClientDisconnectType type)
         {
@@ -121,16 +121,22 @@ namespace MQTTnet.Server
             _cancellationTokenSource = null;
         }
 
+        public Task RunAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
+        {
+            _workerTask = RunInternalAsync(connectPacket, adapter);
+            return _workerTask;
+        }
+
         private async Task RunInternalAsync(MqttConnectPacket connectPacket, IMqttChannelAdapter adapter)
         {
             if (adapter == null) throw new ArgumentNullException(nameof(adapter));
 
             try
             {
-                if (_cancellationTokenSource != null)
-                {
-                    await StopAsync(MqttClientDisconnectType.Clean, true).ConfigureAwait(false);
-                }
+                //if (_cancellationTokenSource != null)
+                //{
+                //    await StopAsync(MqttClientDisconnectType.Clean, true).ConfigureAwait(false);
+                //}
 
                 adapter.ReadingPacketStarted += OnAdapterReadingPacketStarted;
                 adapter.ReadingPacketCompleted += OnAdapterReadingPacketCompleted;
@@ -203,12 +209,14 @@ namespace MQTTnet.Server
                 //await TryDisconnectAdapterAsync(adapter).ConfigureAwait(false);
                 //TryDisposeAdapter(adapter);
 
-                _cleanupHandle?.Dispose();
-                _cleanupHandle = null;
+                //_cleanupHandle?.Dispose();
+                //_cleanupHandle = null;
 
                 _cancellationTokenSource?.Cancel(false);
                 _cancellationTokenSource?.Dispose();
                 _cancellationTokenSource = null;
+
+                _workerTask = null;
             }
         }
 
@@ -264,6 +272,8 @@ namespace MQTTnet.Server
                         await _workerTask.ConfigureAwait(false);
                     }
                 }
+                
+                await Task.FromResult(0);
             }
             finally
             {
diff --git a/Source/MQTTnet/Server/MqttClientSessionsManager.cs b/Source/MQTTnet/Server/MqttClientSessionsManager.cs
index 356782e0..9707e580 100644
--- a/Source/MQTTnet/Server/MqttClientSessionsManager.cs
+++ b/Source/MQTTnet/Server/MqttClientSessionsManager.cs
@@ -60,9 +60,9 @@ namespace MQTTnet.Server
             }
         }
 
-        public Task StartSession(IMqttChannelAdapter clientAdapter)
+        public Task HandleConnectionAsync(IMqttChannelAdapter clientAdapter)
         {
-            return Task.Run(() => RunSessionAsync(clientAdapter, _cancellationToken), _cancellationToken);
+            return Task.Run(() => HandleConnectionAsync(clientAdapter, _cancellationToken), _cancellationToken);
         }
 
         public async Task<IList<IMqttClientSessionStatus>> GetClientStatusAsync()
@@ -207,7 +207,7 @@ namespace MQTTnet.Server
             }
         }
 
-        private async Task RunSessionAsync(IMqttChannelAdapter clientAdapter, CancellationToken cancellationToken)
+        private async Task HandleConnectionAsync(IMqttChannelAdapter clientAdapter, CancellationToken cancellationToken)
         {
             var clientId = string.Empty;
             
@@ -295,11 +295,11 @@ namespace MQTTnet.Server
                 var isSessionPresent = _sessions.TryGetValue(connectPacket.ClientId, out var clientSession);
                 if (isSessionPresent)
                 {
+                    await clientSession.StopAsync(MqttClientDisconnectType.Clean).ConfigureAwait(false);
+
                     if (connectPacket.CleanSession)
                     {
                         _sessions.Remove(connectPacket.ClientId);
-
-                        await clientSession.StopAsync(MqttClientDisconnectType.Clean).ConfigureAwait(false);
                         clientSession.Dispose();
                         clientSession = null;
 
diff --git a/Source/MQTTnet/Server/MqttServer.cs b/Source/MQTTnet/Server/MqttServer.cs
index aae37e6a..1b5fb3be 100644
--- a/Source/MQTTnet/Server/MqttServer.cs
+++ b/Source/MQTTnet/Server/MqttServer.cs
@@ -148,7 +148,7 @@ namespace MQTTnet.Server
 
         private void OnClientAccepted(object sender, MqttServerAdapterClientAcceptedEventArgs eventArgs)
         {
-            eventArgs.SessionTask = _clientSessionsManager.StartSession(eventArgs.Client);
+            eventArgs.SessionTask = _clientSessionsManager.HandleConnectionAsync(eventArgs.Client);
         }
     }
 }