Refactor latest changes

MQTTnet.Core/Adapter/MqttChannelCommunicationAdapter.cs

@@ -15,7 +15,6 @@ namespace MQTTnet.Core.Adapter
     public class MqttChannelCommunicationAdapter : IMqttCommunicationAdapter
     {
         private readonly IMqttCommunicationChannel _channel;
-        private readonly byte[] _readBuffer = new byte[BufferConstants.Size];
 
         private Task _sendTask = Task.FromResult(0); // this task is used to prevent overlapping write
 
@@ -29,76 +28,105 @@ namespace MQTTnet.Core.Adapter
 
         public Task ConnectAsync(MqttClientOptions options, TimeSpan timeout)
         {
-            return _channel.ConnectAsync(options).TimeoutAfter(timeout);
+            try
+            {
+                return _channel.ConnectAsync(options).TimeoutAfter(timeout);
+            }
+            catch (Exception exception)
+            {
+                throw new MqttCommunicationException(exception);
+            }
         }
 
         public Task DisconnectAsync()
         {
-            return _channel.DisconnectAsync();
+            try
+            {
+                return _channel.DisconnectAsync();
+            }
+            catch (Exception exception)
+            {
+                throw new MqttCommunicationException(exception);
+            }
         }
 
         public async Task SendPacketsAsync(TimeSpan timeout, IEnumerable<MqttBasePacket> packets)
         {
-            lock (_channel)
+            try
             {
-                foreach (var packet in packets)
+                lock (_channel)
                 {
-                    MqttTrace.Information(nameof(MqttChannelCommunicationAdapter), "TX >>> {0} [Timeout={1}]", packet, timeout);
+                    foreach (var packet in packets)
+                    {
+                        MqttTrace.Information(nameof(MqttChannelCommunicationAdapter), "TX >>> {0} [Timeout={1}]", packet, timeout);
 
-                    var writeBuffer = PacketSerializer.Serialize(packet);
-
-                    _sendTask = _sendTask.ContinueWith(p => _channel.SendStream.WriteAsync(writeBuffer, 0, writeBuffer.Length));
+                        var writeBuffer = PacketSerializer.Serialize(packet);
+                        _sendTask = _sendTask.ContinueWith(p => _channel.SendStream.WriteAsync(writeBuffer, 0, writeBuffer.Length));
+                    }
                 }
-            }
 
-            await _sendTask; // configure await false geneates stackoverflow
-            await _channel.SendStream.FlushAsync().TimeoutAfter(timeout).ConfigureAwait(false);
+                await _sendTask; // configure await false geneates stackoverflow
+                await _channel.SendStream.FlushAsync().TimeoutAfter(timeout).ConfigureAwait(false);
+            }
+            catch (Exception exception)
+            {
+                throw new MqttCommunicationException(exception);
+            }
         }
 
         public async Task<MqttBasePacket> ReceivePacketAsync(TimeSpan timeout)
         {
-            Tuple<MqttPacketHeader, MemoryStream> tuple;
-            if (timeout > TimeSpan.Zero)
+            try
             {
-                tuple = await ReceiveAsync(_channel.RawStream).TimeoutAfter(timeout).ConfigureAwait(false);
+                ReceivedMqttPacket receivedMqttPacket;
+                if (timeout > TimeSpan.Zero)
+                {
+                    receivedMqttPacket = await ReceiveAsync(_channel.RawStream).TimeoutAfter(timeout).ConfigureAwait(false);
+                }
+                else
+                {
+                    receivedMqttPacket = await ReceiveAsync(_channel.ReceiveStream).ConfigureAwait(false);
+                }
+
+                var packet = PacketSerializer.Deserialize(receivedMqttPacket);
+                if (packet == null)
+                {
+                    throw new MqttProtocolViolationException("Received malformed packet.");
+                }
+
+                MqttTrace.Information(nameof(MqttChannelCommunicationAdapter), "RX <<< {0}", packet);
+                return packet;
             }
-            else
+            catch (Exception exception)
             {
-                tuple = await ReceiveAsync(_channel.ReceiveStream).ConfigureAwait(false);
+                throw new MqttCommunicationException(exception);
             }
-
-            var packet = PacketSerializer.Deserialize(tuple.Item1, tuple.Item2);
-
-            if (packet == null)
-            {
-                throw new MqttProtocolViolationException("Received malformed packet.");
-            }
-
-            MqttTrace.Information(nameof(MqttChannelCommunicationAdapter), "RX <<< {0}", packet);
-            return packet;
         }
 
-        private async Task<Tuple<MqttPacketHeader, MemoryStream>> ReceiveAsync(Stream stream)
+        private async Task<ReceivedMqttPacket> ReceiveAsync(Stream stream)
         {
             var header = MqttPacketReader.ReadHeaderFromSource(stream);
 
-            MemoryStream body;
-            if (header.BodyLength > 0)
+            if (header.BodyLength == 0)
             {
-                var totalRead = 0;
-                do
-                {
-                    var read = await stream.ReadAsync(_readBuffer, totalRead, header.BodyLength - totalRead).ConfigureAwait(false);
-                    totalRead += read;
-                } while (totalRead < header.BodyLength);
-                body = new MemoryStream(_readBuffer, 0, header.BodyLength);
-            }
-            else
-            {
-                body = new MemoryStream();
+                return new ReceivedMqttPacket(header, new MemoryStream(0));
             }
 
-            return Tuple.Create(header, body);
+            var body = new byte[header.BodyLength];
+
+            var offset = 0;
+            do
+            {
+                var readBytesCount = await stream.ReadAsync(body, offset, body.Length - offset).ConfigureAwait(false);
+                offset += readBytesCount;
+            } while (offset < header.BodyLength);
+
+            if (offset > header.BodyLength)
+            {
+                throw new MqttCommunicationException($"Read more body bytes than required ({offset}/{header.BodyLength}).");
+            }
+
+            return new ReceivedMqttPacket(header, new MemoryStream(body, 0, body.Length));
         }
     }
 }

MQTTnet.Core/Adapter/ReceivedMqttPacket.cs

@@ -0,0 +1,19 @@
+using System;
+using System.IO;
+using MQTTnet.Core.Packets;
+
+namespace MQTTnet.Core.Adapter
+{
+    public class ReceivedMqttPacket
+    {
+        public ReceivedMqttPacket(MqttPacketHeader header, MemoryStream body)
+        {
+            Header = header ?? throw new ArgumentNullException(nameof(header));
+            Body = body ?? throw new ArgumentNullException(nameof(body));
+        }
+
+        public MqttPacketHeader Header { get; }
+
+        public MemoryStream Body { get; }
+    }
+}

MQTTnet.Core/Client/MqttClient.cs

@@ -99,6 +99,11 @@ namespace MQTTnet.Core.Client
 
         public async Task DisconnectAsync()
         {
+            if (!IsConnected)
+            {
+                return;
+            }
+
             try
             {
                 await SendAsync(new MqttDisconnectPacket()).ConfigureAwait(false);

MQTTnet.Core/Client/MqttPacketDispatcher.cs

@@ -11,7 +11,7 @@ namespace MQTTnet.Core.Client
     public class MqttPacketDispatcher
     {
         private readonly ConcurrentDictionary<Type, TaskCompletionSource<MqttBasePacket>> _packetByResponseType = new ConcurrentDictionary<Type, TaskCompletionSource<MqttBasePacket>>();
-        private readonly ConcurrentDictionary<Type, ConcurrentDictionary<ushort,TaskCompletionSource<MqttBasePacket>>> _packetByResponseTypeAndIdentifier = new ConcurrentDictionary<Type, ConcurrentDictionary<ushort,TaskCompletionSource<MqttBasePacket>>>();
+        private readonly ConcurrentDictionary<Type, ConcurrentDictionary<ushort, TaskCompletionSource<MqttBasePacket>>> _packetByResponseTypeAndIdentifier = new ConcurrentDictionary<Type, ConcurrentDictionary<ushort, TaskCompletionSource<MqttBasePacket>>>();
 
         public async Task<MqttBasePacket> WaitForPacketAsync(MqttBasePacket request, Type responseType, TimeSpan timeout)
         {
@@ -24,7 +24,7 @@ namespace MQTTnet.Core.Client
             }
             catch (MqttCommunicationTimedOutException)
             {
-                MqttTrace.Warning(nameof(MqttPacketDispatcher), "Timeout while waiting for packet.");
+                MqttTrace.Warning(nameof(MqttPacketDispatcher), "Timeout while waiting for packet of type '{0}'.", responseType.Name);
                 throw;
             }
             finally
@@ -42,16 +42,20 @@ namespace MQTTnet.Core.Client
             {
                 if (_packetByResponseTypeAndIdentifier.TryGetValue(type, out var byid))
                 {
-                    if (byid.TryRemove( withIdentifier.PacketIdentifier, out var tcs))
+                    if (byid.TryRemove(withIdentifier.PacketIdentifier, out var tcs))
                     {
-                        tcs.TrySetResult( packet );
+                        tcs.TrySetResult(packet);
+                        return;
                     }
                 }
             }
             else if (_packetByResponseType.TryRemove(type, out var tcs))
             {
                 tcs.TrySetResult(packet);
+                return;
             }
+
+            throw new InvalidOperationException($"Packet of type '{type.Name}' not handled or dispatched.");
         }
 
         public void Reset()

MQTTnet.Core/Internal/TaskExtensions.cs

@@ -25,20 +25,22 @@ namespace MQTTnet.Core.Internal
 
                 try
                 {
-                    cancellationTokenSource.CancelAfter(timeout);
+                    #pragma warning disable 4014
                     task.ContinueWith(t =>
-                   {
-                       if (t.IsFaulted)
-                       {
-                           tcs.TrySetException(t.Exception);
-                       }
+                    #pragma warning restore 4014
+                    {
+                        if (t.IsFaulted)
+                        {
+                            tcs.TrySetException(t.Exception);
+                        }
 
-                       if (t.IsCompleted)
-                       {
-                           tcs.TrySetResult(t.Result);
-                       }
-                   }, cancellationTokenSource.Token);
+                        if (t.IsCompleted)
+                        {
+                            tcs.TrySetResult(t.Result);
+                        }
+                    }, cancellationTokenSource.Token);
 
+                    cancellationTokenSource.CancelAfter(timeout);
                     return await tcs.Task;
                 }
                 catch (TaskCanceledException)

MQTTnet.Core/Serializer/IMqttPacketSerializer.cs

@@ -1,4 +1,4 @@
-using System.IO;
+using MQTTnet.Core.Adapter;
 using MQTTnet.Core.Packets;
 
 namespace MQTTnet.Core.Serializer
@@ -9,6 +9,6 @@ namespace MQTTnet.Core.Serializer
 
         byte[] Serialize(MqttBasePacket mqttPacket);
 
-        MqttBasePacket Deserialize(MqttPacketHeader header, MemoryStream stream);
+        MqttBasePacket Deserialize(ReceivedMqttPacket receivedMqttPacket);
     }
 }

MQTTnet.Core/Serializer/MqttPacketReader.cs

@@ -1,6 +1,7 @@
 using System;
 using System.IO;
 using System.Text;
+using MQTTnet.Core.Adapter;
 using MQTTnet.Core.Exceptions;
 using MQTTnet.Core.Protocol;
 using MQTTnet.Core.Packets;
@@ -9,15 +10,29 @@ namespace MQTTnet.Core.Serializer
 {
     public sealed class MqttPacketReader : BinaryReader
     {
-        private readonly MqttPacketHeader _header;
-
-        public MqttPacketReader(Stream stream, MqttPacketHeader header)
-            : base(stream, Encoding.UTF8, true)
+        private readonly ReceivedMqttPacket _receivedMqttPacket;
+        
+        public MqttPacketReader(ReceivedMqttPacket receivedMqttPacket)
+            : base(receivedMqttPacket.Body, Encoding.UTF8, true)
         {
-            _header = header;
+            _receivedMqttPacket = receivedMqttPacket;
         }
 
-        public bool EndOfRemainingData => BaseStream.Position == _header.BodyLength;
+        public bool EndOfRemainingData => BaseStream.Position == _receivedMqttPacket.Header.BodyLength;
+
+        public static MqttPacketHeader ReadHeaderFromSource(Stream stream)
+        {
+            var fixedHeader = (byte)stream.ReadByte();
+            var controlPacketType = (MqttControlPacketType)(fixedHeader >> 4);
+            var bodyLength = ReadBodyLengthFromSource(stream);
+
+            return new MqttPacketHeader
+            {
+                FixedHeader = fixedHeader,
+                ControlPacketType = controlPacketType,
+                BodyLength = bodyLength
+            };
+        }
 
         public override ushort ReadUInt16()
         {
@@ -44,21 +59,7 @@ namespace MQTTnet.Core.Serializer
 
         public byte[] ReadRemainingData()
         {
-            return ReadBytes(_header.BodyLength - (int)BaseStream.Position);
-        }
-
-        public static MqttPacketHeader ReadHeaderFromSource(Stream stream)
-        {
-            var fixedHeader = (byte)stream.ReadByte();
-            var controlPacketType = (MqttControlPacketType)(fixedHeader >> 4);
-            var bodyLength = ReadBodyLengthFromSource(stream);
-
-            return new MqttPacketHeader
-            {
-                FixedHeader = fixedHeader,
-                ControlPacketType = controlPacketType,
-                BodyLength = bodyLength
-            };
+            return ReadBytes(_receivedMqttPacket.Header.BodyLength - (int)BaseStream.Position);
         }
 
         private static int ReadBodyLengthFromSource(Stream stream)
@@ -74,7 +75,7 @@ namespace MQTTnet.Core.Serializer
                 multiplier *= 128;
                 if (multiplier > 128 * 128 * 128)
                 {
-                    throw new MqttProtocolViolationException("Remaining length is ivalid.");
+                    throw new MqttProtocolViolationException("Remaining length is invalid.");
                 }
             } while ((encodedByte & 128) != 0);
             return value;

MQTTnet.Core/Serializer/MqttPacketSerializer.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Text;
+using MQTTnet.Core.Adapter;
 using MQTTnet.Core.Exceptions;
 using MQTTnet.Core.Packets;
 using MQTTnet.Core.Protocol;
@@ -110,14 +111,13 @@ namespace MQTTnet.Core.Serializer
             throw new MqttProtocolViolationException("Packet type invalid.");
         }
 
-        public MqttBasePacket Deserialize(MqttPacketHeader header, MemoryStream body)
+        public MqttBasePacket Deserialize(ReceivedMqttPacket receivedMqttPacket)
         {
-            if (header == null) throw new ArgumentNullException(nameof(header));
-            if (body == null) throw new ArgumentNullException(nameof(body));
+            if (receivedMqttPacket == null) throw new ArgumentNullException(nameof(receivedMqttPacket));
 
-            using (var reader = new MqttPacketReader(body, header))
+            using (var reader = new MqttPacketReader(receivedMqttPacket))
             {
-                return Deserialize(header, reader);
+                return Deserialize(receivedMqttPacket.Header, reader);
             }
         }