Add interceptor for client subscriptions

MQTTnet.Core/Server/MqttApplicationMessageInterceptorContext.cs

@@ -0,0 +1,7 @@
+namespace MQTTnet.Core.Server
+{
+    public class MqttApplicationMessageInterceptorContext
+    {
+        public MqttApplicationMessage ApplicationMessage { get; set; }
+    }
+}

MQTTnet.Core/Server/MqttClientSession.cs

@@ -124,7 +124,7 @@ namespace MQTTnet.Core.Server
                 while (!cancellationToken.IsCancellationRequested)
                 {
                     var packet = await adapter.ReceivePacketAsync(TimeSpan.Zero, cancellationToken).ConfigureAwait(false);
-                    await ProcessReceivedPacketAsync(adapter, packet).ConfigureAwait(false);
+                    await ProcessReceivedPacketAsync(adapter, packet, cancellationToken).ConfigureAwait(false);
                 }
             }
             catch (OperationCanceledException)
@@ -142,28 +142,35 @@ namespace MQTTnet.Core.Server
             }
         }
 
-        private async Task ProcessReceivedPacketAsync(IMqttCommunicationAdapter adapter, MqttBasePacket packet)
+        private async Task ProcessReceivedPacketAsync(IMqttCommunicationAdapter adapter, MqttBasePacket packet, CancellationToken cancellationToken)
         {
             if (packet is MqttSubscribePacket subscribePacket)
             {
-                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, _cancellationTokenSource.Token, _subscriptionsManager.Subscribe(subscribePacket));
+                var subscribeResult = _subscriptionsManager.Subscribe(subscribePacket);
+                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, subscribeResult.ResponsePacket);
                 EnqueueRetainedMessages(subscribePacket);
+
+                if (subscribeResult.CloseConnection)
+                {
+                    await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new MqttDisconnectPacket());
+                    Stop();
+                }
             }
             else if (packet is MqttUnsubscribePacket unsubscribePacket)
             {
-                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, _cancellationTokenSource.Token, _subscriptionsManager.Unsubscribe(unsubscribePacket));
+                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, _subscriptionsManager.Unsubscribe(unsubscribePacket));
             }
             else if (packet is MqttPublishPacket publishPacket)
             {
-                await HandleIncomingPublishPacketAsync(adapter, publishPacket);
+                await HandleIncomingPublishPacketAsync(adapter, publishPacket, cancellationToken);
             }
             else if (packet is MqttPubRelPacket pubRelPacket)
             {
-                await HandleIncomingPubRelPacketAsync(adapter, pubRelPacket);
+                await HandleIncomingPubRelPacketAsync(adapter, pubRelPacket, cancellationToken);
             }
             else if (packet is MqttPubRecPacket pubRecPacket)
             {
-                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, _cancellationTokenSource.Token, pubRecPacket.CreateResponse<MqttPubRelPacket>());
+                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, pubRecPacket.CreateResponse<MqttPubRelPacket>());
             }
             else if (packet is MqttPubAckPacket || packet is MqttPubCompPacket)
             {
@@ -171,7 +178,7 @@ namespace MQTTnet.Core.Server
             }
             else if (packet is MqttPingReqPacket)
             {
-                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, _cancellationTokenSource.Token, new MqttPingRespPacket());
+                await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new MqttPingRespPacket());
             }
             else if (packet is MqttDisconnectPacket || packet is MqttConnectPacket)
             {
@@ -193,10 +200,17 @@ namespace MQTTnet.Core.Server
             }
         }
 
-        private async Task HandleIncomingPublishPacketAsync(IMqttCommunicationAdapter adapter, MqttPublishPacket publishPacket)
+        private async Task HandleIncomingPublishPacketAsync(IMqttCommunicationAdapter adapter, MqttPublishPacket publishPacket, CancellationToken cancellationToken)
         {
             var applicationMessage = publishPacket.ToApplicationMessage();
-            _options.ApplicationMessageInterceptor?.Invoke(applicationMessage);
+
+            var interceptorContext = new MqttApplicationMessageInterceptorContext
+            {
+                ApplicationMessage = applicationMessage
+            };
+
+            _options.ApplicationMessageInterceptor?.Invoke(interceptorContext);
+            applicationMessage = interceptorContext.ApplicationMessage;
 
             if (applicationMessage.Retain)
             {
@@ -214,7 +228,7 @@ namespace MQTTnet.Core.Server
                     {
                         _sessionsManager.DispatchApplicationMessage(this, applicationMessage);
 
-                        await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, _cancellationTokenSource.Token,
+                        await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken,
                             new MqttPubAckPacket { PacketIdentifier = publishPacket.PacketIdentifier });
 
                         return;
@@ -229,7 +243,7 @@ namespace MQTTnet.Core.Server
 
                         _sessionsManager.DispatchApplicationMessage(this, applicationMessage);
 
-                        await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, _cancellationTokenSource.Token,
+                        await adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken,
                             new MqttPubRecPacket { PacketIdentifier = publishPacket.PacketIdentifier });
 
                         return;
@@ -239,14 +253,14 @@ namespace MQTTnet.Core.Server
             }
         }
 
-        private Task HandleIncomingPubRelPacketAsync(IMqttCommunicationAdapter adapter, MqttPubRelPacket pubRelPacket)
+        private Task HandleIncomingPubRelPacketAsync(IMqttCommunicationAdapter adapter, MqttPubRelPacket pubRelPacket, CancellationToken cancellationToken)
         {
             lock (_unacknowledgedPublishPackets)
             {
                 _unacknowledgedPublishPackets.Remove(pubRelPacket.PacketIdentifier);
             }
 
-            return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, _cancellationTokenSource.Token, new MqttPubCompPacket { PacketIdentifier = pubRelPacket.PacketIdentifier });
+            return adapter.SendPacketsAsync(_options.DefaultCommunicationTimeout, cancellationToken, new MqttPubCompPacket { PacketIdentifier = pubRelPacket.PacketIdentifier });
         }
     }
 }

MQTTnet.Core/Server/MqttClientSubscribeResult.cs

@@ -0,0 +1,11 @@
+using MQTTnet.Core.Packets;
+
+namespace MQTTnet.Core.Server
+{
+    public class MqttClientSubscribeResult
+    {
+        public MqttSubAckPacket ResponsePacket { get; set; }
+
+        public bool CloseConnection { get; set; }
+    }
+}

MQTTnet.Core/Server/MqttClientSubscriptionsManager.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using Microsoft.Extensions.Options;
 using MQTTnet.Core.Packets;
 using MQTTnet.Core.Protocol;
 
@@ -8,30 +9,45 @@ namespace MQTTnet.Core.Server
     public sealed class MqttClientSubscriptionsManager
     {
         private readonly Dictionary<string, MqttQualityOfServiceLevel> _subscribedTopics = new Dictionary<string, MqttQualityOfServiceLevel>();
+        private readonly MqttServerOptions _options;
 
-        public MqttClientSubscriptionsManager()
+        public MqttClientSubscriptionsManager(IOptions<MqttServerOptions> options)
         {
-            
+            _options = options?.Value ?? throw new ArgumentNullException(nameof(options));
         }
 
-        public MqttSubAckPacket Subscribe(MqttSubscribePacket subscribePacket)
+        public MqttClientSubscribeResult Subscribe(MqttSubscribePacket subscribePacket)
         {
             if (subscribePacket == null) throw new ArgumentNullException(nameof(subscribePacket));
 
             var responsePacket = subscribePacket.CreateResponse<MqttSubAckPacket>();
+            var closeConnection = false;
 
             lock (_subscribedTopics)
             {
                 foreach (var topicFilter in subscribePacket.TopicFilters)
                 {
+                    var interceptorContext = new MqttSubscriptionInterceptorContext("", topicFilter);
+                    _options.SubscriptionsInterceptor?.Invoke(interceptorContext);
+                    responsePacket.SubscribeReturnCodes.Add(interceptorContext.AcceptSubscription ? MqttSubscribeReturnCode.SuccessMaximumQoS1 : MqttSubscribeReturnCode.Failure);
+                    
+                    if (interceptorContext.CloseConnection)
+                    {
+                        closeConnection = true;
+                    }
 
-
-                    _subscribedTopics[topicFilter.Topic] = topicFilter.QualityOfServiceLevel;
-                    responsePacket.SubscribeReturnCodes.Add(MqttSubscribeReturnCode.SuccessMaximumQoS1); // TODO: Add support for QoS 2.
+                    if (interceptorContext.AcceptSubscription)
+                    {
+                        _subscribedTopics[topicFilter.Topic] = topicFilter.QualityOfServiceLevel;
+                    }
                 }
             }
 
-            return responsePacket;
+            return new MqttClientSubscribeResult
+            {
+                ResponsePacket = responsePacket,
+                CloseConnection = closeConnection
+            };
         }
 
         public MqttUnsubAckPacket Unsubscribe(MqttUnsubscribePacket unsubscribePacket)

MQTTnet.Core/Server/MqttServer.cs

@@ -56,8 +56,14 @@ namespace MQTTnet.Core.Server
 
             foreach (var applicationMessage in applicationMessages)
             {
-                _options.ApplicationMessageInterceptor?.Invoke(applicationMessage);
-                _clientSessionsManager.DispatchApplicationMessage(null, applicationMessage);
+                var interceptorContext = new MqttApplicationMessageInterceptorContext
+                {
+                    ApplicationMessage = applicationMessage
+                };
+
+                _options.ApplicationMessageInterceptor?.Invoke(interceptorContext);
+                
+                _clientSessionsManager.DispatchApplicationMessage(null, interceptorContext.ApplicationMessage);
             }
         }
 

MQTTnet.Core/Server/MqttServerOptions.cs

@@ -16,7 +16,9 @@ namespace MQTTnet.Core.Server
 
         public Func<MqttConnectPacket, MqttConnectReturnCode> ConnectionValidator { get; set; }
 
-        public Func<MqttApplicationMessage, MqttApplicationMessage> ApplicationMessageInterceptor { get; set; }
+        public Action<MqttApplicationMessageInterceptorContext> ApplicationMessageInterceptor { get; set; }
+
+        public Action<MqttSubscriptionInterceptorContext> SubscriptionsInterceptor { get; set; }
 
         public IMqttServerStorage Storage { get; set; }
     }

MQTTnet.Core/Server/MqttSubscriptionInterceptorContext.cs

@@ -0,0 +1,19 @@
+namespace MQTTnet.Core.Server
+{
+    public class MqttSubscriptionInterceptorContext
+    {
+        public MqttSubscriptionInterceptorContext(string clientId, TopicFilter topicFilter)
+        {
+            ClientId = clientId;
+            TopicFilter = topicFilter;
+        }
+
+        public string ClientId { get; }
+
+        public TopicFilter TopicFilter { get; }
+
+        public bool AcceptSubscription { get; set; } = true;
+
+        public bool CloseConnection { get; set; }
+    }
+}