Set web socket boundaries correctly (#1514)

Source/MQTTnet.Benchmarks/MqttTcpChannelBenchmark.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System;
 using BenchmarkDotNet.Attributes;
 using MQTTnet.Channel;
 using MQTTnet.Implementations;
@@ -70,8 +71,8 @@ namespace MQTTnet.Benchmarks
 
             while (read < expected)
             {
-                var readresult = await _clientChannel.ReadAsync(new byte[size], 0, size, CancellationToken.None).ConfigureAwait(false);
-                read += readresult;
+                var readResult = await _clientChannel.ReadAsync(new byte[size], 0, size, CancellationToken.None).ConfigureAwait(false);
+                read += readResult;
             }
         }
 
@@ -79,9 +80,11 @@ namespace MQTTnet.Benchmarks
         {
             await Task.Yield();
 
+            var buffer = new ArraySegment<byte>(new byte[size]);
+            
             for (var i = 0; i < iterations; i++)
             {
-                await _serverChannel.WriteAsync(new byte[size], 0, size, CancellationToken.None).ConfigureAwait(false);
+                await _serverChannel.WriteAsync(buffer, true, CancellationToken.None).ConfigureAwait(false);
             }
         }
     }

Source/MQTTnet.Benchmarks/SerializerBenchmark.cs

@@ -104,7 +104,7 @@ namespace MQTTnet.Benchmarks
                 return Task.FromResult(count);
             }
 
-            public Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+            public Task WriteAsync(ArraySegment<byte> buffer, bool isEndOfPacket, CancellationToken cancellationToken)
             {
                 throw new NotSupportedException();
             }

Source/MQTTnet.Extensions.WebSocket4Net/WebSocket4NetMqttChannel.cs

@@ -154,9 +154,9 @@ namespace MQTTnet.Extensions.WebSocket4Net
             return Task.FromResult(readBytes);
         }
 
-        public Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public Task WriteAsync(ArraySegment<byte> buffer, bool isEndOfPacket, CancellationToken cancellationToken)
         {
-            _webSocket.Send(buffer, offset, count);
+            _webSocket.Send(buffer.Array, buffer.Offset, buffer.Count);
             return Task.FromResult(0);
         }
 

Source/MQTTnet.Tests/Mockups/MemoryMqttChannel.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System;
 using System.IO;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
@@ -45,9 +46,9 @@ namespace MQTTnet.Tests.Mockups
             return _stream.ReadAsync(buffer, offset, count, cancellationToken);
         }
 
-        public Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public Task WriteAsync(ArraySegment<byte> buffer, bool isEndOfPacket, CancellationToken cancellationToken)
         {
-            return _stream.WriteAsync(buffer, offset, count, cancellationToken);
+            return _stream.WriteAsync(buffer.Array, buffer.Offset, buffer.Count, cancellationToken);
         }
 
         public void Dispose()

Source/MQTTnet/Adapter/MqttChannelAdapter.cs

@@ -188,11 +188,14 @@ namespace MQTTnet.Adapter
 
                     _logger.Verbose("TX ({0} bytes) >>> {1}", packetBuffer.Length, packet);
 
-                    await _channel.WriteAsync(packetBuffer.Packet.Array, packetBuffer.Packet.Offset, packetBuffer.Packet.Count, cancellationToken).ConfigureAwait(false);
-
                     if (packetBuffer.Payload.Count > 0)
                     {
-                        await _channel.WriteAsync(packetBuffer.Payload.Array, packetBuffer.Payload.Offset, packetBuffer.Payload.Count, cancellationToken).ConfigureAwait(false);
+                        await _channel.WriteAsync(packetBuffer.Packet, false, cancellationToken).ConfigureAwait(false);
+                        await _channel.WriteAsync(packetBuffer.Payload, true, cancellationToken).ConfigureAwait(false);
+                    }
+                    else
+                    {
+                        await _channel.WriteAsync(packetBuffer.Packet, true, cancellationToken).ConfigureAwait(false);
                     }
 
                     Interlocked.Add(ref _bytesReceived, packetBuffer.Length);

Source/MQTTnet/Channel/IMqttChannel.cs

@@ -12,13 +12,17 @@ namespace MQTTnet.Channel
     public interface IMqttChannel : IDisposable
     {
         string Endpoint { get; }
+        
         bool IsSecureConnection { get; }
+        
         X509Certificate2 ClientCertificate { get; }
 
         Task ConnectAsync(CancellationToken cancellationToken);
+        
         Task DisconnectAsync(CancellationToken cancellationToken);
 
         Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken);
-        Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken);
+        
+        Task WriteAsync(ArraySegment<byte> buffer, bool isEndOfPacket, CancellationToken cancellationToken);
     }
 }

Source/MQTTnet/Implementations/MqttTcpChannel.Uwp.cs

@@ -108,12 +108,12 @@ namespace MQTTnet.Implementations
             return _readStream.ReadAsync(buffer, offset, count, cancellationToken);
         }
 
-        public Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public Task WriteAsync(ArraySegment<byte> buffer, bool isEndOfPacket, CancellationToken cancellationToken)
         {
             // In the write method only the internal buffer will be filled. So here is no
             // async/await required. The real network transmit is done when calling the
             // Flush method.
-            _writeStream.Write(buffer, offset, count);
+            _writeStream.Write(buffer.Array, buffer.Offset, buffer.Count);
             return _writeStream.FlushAsync(cancellationToken);
         }
 

Source/MQTTnet/Implementations/MqttTcpChannel.cs

@@ -225,7 +225,7 @@ namespace MQTTnet.Implementations
             }
         }
 
-        public async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public async Task WriteAsync(ArraySegment<byte> buffer, bool isEndOfPacket, CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
 
@@ -239,12 +239,12 @@ namespace MQTTnet.Implementations
                 }
 
 #if NETCOREAPP3_0_OR_GREATER || NETSTANDARD2_1_OR_GREATER
-                await stream.WriteAsync(buffer.AsMemory(offset, count), cancellationToken).ConfigureAwait(false);
+                await stream.WriteAsync(buffer.AsMemory(), cancellationToken).ConfigureAwait(false);
 #else
                 // Workaround for: https://github.com/dotnet/corefx/issues/24430
                 using (cancellationToken.Register(_disposeAction))
                 {
-                    await stream.WriteAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
+                    await stream.WriteAsync(buffer.Array, buffer.Offset, buffer.Count, cancellationToken).ConfigureAwait(false);
                 }
 #endif
             }

Source/MQTTnet/Implementations/MqttWebSocketChannel.cs

@@ -2,15 +2,15 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using MQTTnet.Channel;
-using MQTTnet.Internal;
 using System;
 using System.Net;
 using System.Net.WebSockets;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
 using System.Threading.Tasks;
+using MQTTnet.Channel;
 using MQTTnet.Client;
+using MQTTnet.Internal;
 
 namespace MQTTnet.Implementations
 {
@@ -35,12 +35,12 @@ namespace MQTTnet.Implementations
             ClientCertificate = clientCertificate;
         }
 
+        public X509Certificate2 ClientCertificate { get; }
+
         public string Endpoint { get; }
 
         public bool IsSecureConnection { get; private set; }
 
-        public X509Certificate2 ClientCertificate { get; }
-
         public async Task ConnectAsync(CancellationToken cancellationToken)
         {
             var uri = _options.Uri;
@@ -89,14 +89,26 @@ namespace MQTTnet.Implementations
             Cleanup();
         }
 
+        public void Dispose()
+        {
+            Cleanup();
+        }
+
         public async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
             var response = await _webSocket.ReceiveAsync(new ArraySegment<byte>(buffer, offset, count), cancellationToken).ConfigureAwait(false);
             return response.Count;
         }
 
-        public async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public async Task WriteAsync(ArraySegment<byte> buffer, bool isEndOfPacket, CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
+            
+#if NET5_0_OR_GREATER
+            // MQTT Control Packets MUST be sent in WebSocket binary data frames. If any other type of data frame is received the recipient MUST close the Network Connection [MQTT-6.0.0-1].
+            // A single WebSocket data frame can contain multiple or partial MQTT Control Packets. The receiver MUST NOT assume that MQTT Control Packets are aligned on WebSocket frame boundaries [MQTT-6.0.0-2].
+            await _webSocket.SendAsync(buffer, WebSocketMessageType.Binary, isEndOfPacket, cancellationToken).ConfigureAwait(false);
+#else
             // The lock is required because the client will throw an exception if _SendAsync_ is 
             // called from multiple threads at the same time. But this issue only happens with several
             // framework versions.
@@ -107,86 +119,9 @@ namespace MQTTnet.Implementations
 
             using (await _sendLock.WaitAsync(cancellationToken).ConfigureAwait(false))
             {
-                await _webSocket.SendAsync(new ArraySegment<byte>(buffer, offset, count), WebSocketMessageType.Binary, true, cancellationToken).ConfigureAwait(false);
+                await _webSocket.SendAsync(buffer, WebSocketMessageType.Binary, isEndOfPacket, cancellationToken).ConfigureAwait(false);
             }
-        }
-
-        public void Dispose()
-        {
-            Cleanup();
-        }
-
-        void SetupClientWebSocket(ClientWebSocket clientWebSocket)
-        {
-            if (_options.ProxyOptions != null)
-            {
-                clientWebSocket.Options.Proxy = CreateProxy();
-            }
-
-            if (_options.RequestHeaders != null)
-            {
-                foreach (var requestHeader in _options.RequestHeaders)
-                {
-                    clientWebSocket.Options.SetRequestHeader(requestHeader.Key, requestHeader.Value);
-                }
-            }
-
-            if (_options.SubProtocols != null)
-            {
-                foreach (var subProtocol in _options.SubProtocols)
-                {
-                    clientWebSocket.Options.AddSubProtocol(subProtocol);
-                }
-            }
-
-            if (_options.CookieContainer != null)
-            {
-                clientWebSocket.Options.Cookies = _options.CookieContainer;
-            }
-
-            if (_options.TlsOptions?.UseTls == true && _options.TlsOptions?.Certificates != null)
-            {
-                clientWebSocket.Options.ClientCertificates = new X509CertificateCollection();
-                foreach (var certificate in _options.TlsOptions.Certificates)
-                {
-#if WINDOWS_UWP
-                    clientWebSocket.Options.ClientCertificates.Add(new X509Certificate(certificate));
-#else
-                    clientWebSocket.Options.ClientCertificates.Add(certificate);
 #endif
-
-                }
-            }
-
-            var certificateValidationHandler = _options.TlsOptions?.CertificateValidationHandler;
-            if (certificateValidationHandler != null)
-            {
-#if NETSTANDARD1_3
-                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'netstandard1.3'.");
-#elif NETSTANDARD2_0
-                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'netstandard2.0'.");
-#elif WINDOWS_UWP
-                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'uap10.0'.");
-#elif NET452
-                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'net452'.");
-#elif NET461
-                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'net461'.");
-#else
-                clientWebSocket.Options.RemoteCertificateValidationCallback = (sender, certificate, chain, sslPolicyErrors) =>
-                {
-                    // TODO: Find a way to add client options to same callback. Problem is that they have a different type.
-                    var context = new MqttClientCertificateValidationEventArgs
-                    {
-                        Certificate = certificate,
-                        Chain = chain,
-                        SslPolicyErrors = sslPolicyErrors,
-                        ClientOptions = _options
-                    };
-
-                    return certificateValidationHandler(context);
-                };
-#endif
-            }
         }
 
         void Cleanup()
@@ -230,5 +165,77 @@ namespace MQTTnet.Implementations
             return new WebProxy(proxyUri, _options.ProxyOptions.BypassOnLocal, _options.ProxyOptions.BypassList);
 #endif
         }
+
+        void SetupClientWebSocket(ClientWebSocket clientWebSocket)
+        {
+            if (_options.ProxyOptions != null)
+            {
+                clientWebSocket.Options.Proxy = CreateProxy();
+            }
+
+            if (_options.RequestHeaders != null)
+            {
+                foreach (var requestHeader in _options.RequestHeaders)
+                {
+                    clientWebSocket.Options.SetRequestHeader(requestHeader.Key, requestHeader.Value);
+                }
+            }
+
+            if (_options.SubProtocols != null)
+            {
+                foreach (var subProtocol in _options.SubProtocols)
+                {
+                    clientWebSocket.Options.AddSubProtocol(subProtocol);
+                }
+            }
+
+            if (_options.CookieContainer != null)
+            {
+                clientWebSocket.Options.Cookies = _options.CookieContainer;
+            }
+
+            if (_options.TlsOptions?.UseTls == true && _options.TlsOptions?.Certificates != null)
+            {
+                clientWebSocket.Options.ClientCertificates = new X509CertificateCollection();
+                foreach (var certificate in _options.TlsOptions.Certificates)
+                {
+#if WINDOWS_UWP
+                    clientWebSocket.Options.ClientCertificates.Add(new X509Certificate(certificate));
+#else
+                    clientWebSocket.Options.ClientCertificates.Add(certificate);
+#endif
+                }
+            }
+
+            var certificateValidationHandler = _options.TlsOptions?.CertificateValidationHandler;
+            if (certificateValidationHandler != null)
+            {
+#if NETSTANDARD1_3
+                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'netstandard1.3'.");
+#elif NETSTANDARD2_0
+                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'netstandard2.0'.");
+#elif WINDOWS_UWP
+                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'uap10.0'.");
+#elif NET452
+                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'net452'.");
+#elif NET461
+                throw new NotSupportedException("Remote certificate validation callback is not supported when using 'net461'.");
+#else
+                clientWebSocket.Options.RemoteCertificateValidationCallback = (sender, certificate, chain, sslPolicyErrors) =>
+                {
+                    // TODO: Find a way to add client options to same callback. Problem is that they have a different type.
+                    var context = new MqttClientCertificateValidationEventArgs
+                    {
+                        Certificate = certificate,
+                        Chain = chain,
+                        SslPolicyErrors = sslPolicyErrors,
+                        ClientOptions = _options
+                    };
+
+                    return certificateValidationHandler(context);
+                };
+#endif
+            }
+        }
     }
 }