Merge implementation of IsAllowPacketFragmentation.

xljiulang · Dec 8, 2024 · e01e5f3 · e01e5f3
1 parent 9a7a8bd
commit e01e5f3
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/Source/MQTTnet.AspnetCore/Features/PacketFragmentationFeature.cs b/Source/MQTTnet.AspnetCore/Features/PacketFragmentationFeature.cs
@@ -3,12 +3,32 @@
 // See the LICENSE file in the project root for more information.
 
 using MQTTnet.Adapter;
+using MQTTnet.Server;
 using System;
 
 namespace MQTTnet.AspNetCore
 {
     sealed class PacketFragmentationFeature(Func<IMqttChannelAdapter, bool> allowPacketFragmentationSelector)
     {
-        public Func<IMqttChannelAdapter, bool> AllowPacketFragmentationSelector { get; } = allowPacketFragmentationSelector;       
+        public Func<IMqttChannelAdapter, bool> AllowPacketFragmentationSelector { get; } = allowPacketFragmentationSelector;
+
+        public static bool IsAllowPacketFragmentation(IMqttChannelAdapter channelAdapter, MqttServerTcpEndpointBaseOptions? endpointOptions)
+        {
+            //if (endpointOptions != null && endpointOptions.AllowPacketFragmentationSelector != null)
+            //{
+            //    return endpointOptions.AllowPacketFragmentationSelector(channelAdapter);
+            //}
+
+            // In the AspNetCore environment, we need to exclude WebSocket before AllowPacketFragmentation.
+            if (channelAdapter is MqttServerChannelAdapter serverChannelAdapter)
+            {
+                if (serverChannelAdapter.IsWebSocketConnection)
+                {
+                    return false;
+                }
+            }
+
+            return endpointOptions == null || endpointOptions.AllowPacketFragmentation;
+        }
     }
 }
diff --git a/Source/MQTTnet.AspnetCore/Internal/MqttServerChannelAdapter.cs b/Source/MQTTnet.AspnetCore/Internal/MqttServerChannelAdapter.cs
@@ -24,18 +24,18 @@ private void SetAllowPacketFragmentation(ConnectionContext connection, HttpConte
         // When connection is from MapMqtt(),
         // the PacketFragmentationFeature instance is copied from kestrel's ConnectionContext.Features to HttpContext.Features,
         // but no longer from HttpContext.Features to connection.Features.     
-        var feature = httpContext == null
+        var packetFragmentationFeature = httpContext == null
             ? connection.Features.Get<PacketFragmentationFeature>()
             : httpContext.Features.Get<PacketFragmentationFeature>();
 
-        if (feature == null)
+        if (packetFragmentationFeature == null)
         {
-            var value = !IsWebSocketConnection;
+            var value = PacketFragmentationFeature.IsAllowPacketFragmentation(this, null);
             SetAllowPacketFragmentation(value);
         }
         else
         {
-            var value = feature.AllowPacketFragmentationSelector(this);
+            var value = packetFragmentationFeature.AllowPacketFragmentationSelector(this);
             SetAllowPacketFragmentation(value);
         }
     }

diff --git a/Source/MQTTnet.AspnetCore/KestrelServerOptionsExtensions.cs b/Source/MQTTnet.AspnetCore/KestrelServerOptionsExtensions.cs
@@ -6,7 +6,6 @@
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Https;
 using Microsoft.Extensions.DependencyInjection;
-using MQTTnet.Adapter;
 using MQTTnet.Exceptions;
 using MQTTnet.Server;
 using System;
@@ -91,19 +90,7 @@ void UseMiddleware(ListenOptions listenOptions)
                             tlsConfigure?.Invoke(httpsOptions);
                         });
                     }
-                    listenOptions.UseMqtt(protocols, AllowPacketFragmentationSelector);
-                }
-
-                bool AllowPacketFragmentationSelector(IMqttChannelAdapter channelAdapter)
-                {
-                    if (channelAdapter is MqttServerChannelAdapter serverChannelAdapter)
-                    {
-                        if (serverChannelAdapter.IsWebSocketConnection)
-                        {
-                            return false;
-                        }
-                    }
-                    return endpoint.AllowPacketFragmentation;
+                    listenOptions.UseMqtt(protocols, channelAdapter => PacketFragmentationFeature.IsAllowPacketFragmentation(channelAdapter, endpoint));
                 }
             }
         }