MQTTnet/Source/MQTTnet.Server/Internal/MqttClientSubscriptionsManager.cs

// Licensed to the .NET Foundation under one or more agreements.
// 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.Packets;
using MQTTnet.Protocol;

namespace MQTTnet.Server.Internal
{
    public sealed class MqttClientSubscriptionsManager : IDisposable
    {
        static readonly List<uint> EmptySubscriptionIdentifiers = new List<uint>();

        readonly MqttServerEventContainer _eventContainer;
        readonly Dictionary<ulong, HashSet<MqttSubscription>> _noWildcardSubscriptionsByTopicHash = new Dictionary<ulong, HashSet<MqttSubscription>>();
        readonly MqttRetainedMessagesManager _retainedMessagesManager;

        readonly MqttSession _session;

        // Callback to maintain list of subscriber clients
        readonly ISubscriptionChangedNotification _subscriptionChangedNotification;

        // Subscriptions are stored in various dictionaries and use a "topic hash"; see the MqttSubscription object for a detailed explanation.
        // The additional lock is important to coordinate complex update logic with multiple steps, checks and interceptors.
        readonly Dictionary<string, MqttSubscription> _subscriptions = new Dictionary<string, MqttSubscription>();

        // Use subscription lock to maintain consistency across subscriptions and topic hash dictionaries
        readonly ReaderWriterLockSlim _subscriptionsLock = new ReaderWriterLockSlim();
        readonly Dictionary<ulong, TopicHashMaskSubscriptions> _wildcardSubscriptionsByTopicHash = new Dictionary<ulong, TopicHashMaskSubscriptions>();

        public MqttClientSubscriptionsManager(
            MqttSession session,
            MqttServerEventContainer eventContainer,
            MqttRetainedMessagesManager retainedMessagesManager,
            ISubscriptionChangedNotification subscriptionChangedNotification)
        {
            _session = session ?? throw new ArgumentNullException(nameof(session));
            _eventContainer = eventContainer ?? throw new ArgumentNullException(nameof(eventContainer));
            _retainedMessagesManager = retainedMessagesManager ?? throw new ArgumentNullException(nameof(retainedMessagesManager));
            _subscriptionChangedNotification = subscriptionChangedNotification;
        }

        public CheckSubscriptionsResult CheckSubscriptions(string topic, ulong topicHash, MqttQualityOfServiceLevel qualityOfServiceLevel, string senderId)
        {
            var possibleSubscriptions = new List<MqttSubscription>();

            // Check for possible subscriptions. They might have collisions but this is fine.
            _subscriptionsLock.EnterReadLock();
            try
            {
                if (_noWildcardSubscriptionsByTopicHash.TryGetValue(topicHash, out var noWildcardSubscriptions))
                {
                    possibleSubscriptions.AddRange(noWildcardSubscriptions);
                }

                foreach (var wcs in _wildcardSubscriptionsByTopicHash)
                {
                    var subscriptionHash = wcs.Key;
                    var subscriptionsByHashMask = wcs.Value.SubscriptionsByHashMask;
                    foreach (var shm in subscriptionsByHashMask)
                    {
                        var subscriptionHashMask = shm.Key;
                        if ((topicHash & subscriptionHashMask) == subscriptionHash)
                        {
                            var subscriptions = shm.Value;
                            possibleSubscriptions.AddRange(subscriptions);
                        }
                    }
                }
            }
            finally
            {
                _subscriptionsLock.ExitReadLock();
            }

            // The pre check has evaluated that nothing is subscribed.
            // If there were some possible candidates they get checked below
            // again to avoid collisions.
            if (possibleSubscriptions.Count == 0)
            {
                return CheckSubscriptionsResult.NotSubscribed;
            }

            var senderIsReceiver = string.Equals(senderId, _session.Id);
            var maxQoSLevel = -1; // Not subscribed.

            HashSet<uint> subscriptionIdentifiers = null;
            var retainAsPublished = false;

            foreach (var subscription in possibleSubscriptions)
            {
                if (subscription.NoLocal && senderIsReceiver)
                {
                    // This is a MQTTv5 feature!
                    continue;
                }

                if (MqttTopicFilterComparer.Compare(topic, subscription.Topic) != MqttTopicFilterCompareResult.IsMatch)
                {
                    continue;
                }

                if (subscription.RetainAsPublished)
                {
                    // This is a MQTTv5 feature!
                    retainAsPublished = true;
                }

                if ((int)subscription.GrantedQualityOfServiceLevel > maxQoSLevel)
                {
                    maxQoSLevel = (int)subscription.GrantedQualityOfServiceLevel;
                }

                if (subscription.Identifier > 0)
                {
                    if (subscriptionIdentifiers == null)
                    {
                        subscriptionIdentifiers = new HashSet<uint>();
                    }

                    subscriptionIdentifiers.Add(subscription.Identifier);
                }
            }

            if (maxQoSLevel == -1)
            {
                return CheckSubscriptionsResult.NotSubscribed;
            }

            var result = new CheckSubscriptionsResult
            {
                IsSubscribed = true,
                RetainAsPublished = retainAsPublished,
                SubscriptionIdentifiers = subscriptionIdentifiers?.ToList() ?? EmptySubscriptionIdentifiers,

                // Start with the same QoS as the publisher.
                QualityOfServiceLevel = qualityOfServiceLevel
            };

            // Now downgrade if required.
            //
            // If a subscribing Client has been granted maximum QoS 1 for a particular Topic Filter, then a QoS 0 Application Message matching the filter is delivered
            // to the Client at QoS 0. This means that at most one copy of the message is received by the Client. On the other hand, a QoS 2 Message published to
            // the same topic is downgraded by the Server to QoS 1 for delivery to the Client, so that Client might receive duplicate copies of the Message.

            // Subscribing to a Topic Filter at QoS 2 is equivalent to saying "I would like to receive Messages matching this filter at the QoS with which they were published".
            // This means a publisher is responsible for determining the maximum QoS a Message can be delivered at, but a subscriber is able to require that the Server
            // downgrades the QoS to one more suitable for its usage.
            if (maxQoSLevel < (int)qualityOfServiceLevel)
            {
                result.QualityOfServiceLevel = (MqttQualityOfServiceLevel)maxQoSLevel;
            }

            return result;
        }

        public void Dispose()
        {
            _subscriptionsLock?.Dispose();
        }

        public async Task<SubscribeResult> Subscribe(MqttSubscribePacket subscribePacket, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(subscribePacket);

            var retainedApplicationMessages = await _retainedMessagesManager.GetMessages().ConfigureAwait(false);
            var result = new SubscribeResult(subscribePacket.TopicFilters.Count);

            var addedSubscriptions = new List<string>();
            var finalTopicFilters = new List<MqttTopicFilter>();

            // The topic filters are order by its QoS so that the higher QoS will win over a
            // lower one.
            foreach (var topicFilterItem in subscribePacket.TopicFilters.OrderByDescending(f => f.QualityOfServiceLevel))
            {
                var interceptorEventArgs = await InterceptSubscribe(topicFilterItem, subscribePacket.UserProperties, cancellationToken).ConfigureAwait(false);
                var topicFilter = interceptorEventArgs.TopicFilter;
                var processSubscription = interceptorEventArgs.ProcessSubscription && interceptorEventArgs.Response.ReasonCode <= MqttSubscribeReasonCode.GrantedQoS2;

                result.UserProperties = interceptorEventArgs.UserProperties;
                result.ReasonString = interceptorEventArgs.ReasonString;
                result.ReasonCodes.Add(interceptorEventArgs.Response.ReasonCode);

                if (interceptorEventArgs.CloseConnection)
                {
                    // When any of the interceptor calls leads to a connection close the connection
                    // must be closed. So do not revert to false!
                    result.CloseConnection = true;
                }

                if (!processSubscription || string.IsNullOrEmpty(topicFilter.Topic))
                {
                    continue;
                }

                var createSubscriptionResult = CreateSubscription(topicFilter, subscribePacket.SubscriptionIdentifier, interceptorEventArgs.Response.ReasonCode);

                addedSubscriptions.Add(topicFilter.Topic);
                finalTopicFilters.Add(topicFilter);

                FilterRetainedApplicationMessages(retainedApplicationMessages, createSubscriptionResult, result);
            }

            // This call will add the new subscription to the internal storage.
            // So the event _ClientSubscribedTopicEvent_ must be called afterwards.
            _subscriptionChangedNotification?.OnSubscriptionsAdded(_session, addedSubscriptions);

            if (_eventContainer.ClientSubscribedTopicEvent.HasHandlers)
            {
                foreach (var finalTopicFilter in finalTopicFilters)
                {
                    var eventArgs = new ClientSubscribedTopicEventArgs(_session.Id, _session.UserName, finalTopicFilter, _session.Items);
                    await _eventContainer.ClientSubscribedTopicEvent.InvokeAsync(eventArgs).ConfigureAwait(false);
                }
            }

            return result;
        }

        public async Task<UnsubscribeResult> Unsubscribe(MqttUnsubscribePacket unsubscribePacket, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(unsubscribePacket);

            var result = new UnsubscribeResult();

            var removedSubscriptions = new List<string>();

            _subscriptionsLock.EnterWriteLock();
            try
            {
                foreach (var topicFilter in unsubscribePacket.TopicFilters)
                {
                    _subscriptions.TryGetValue(topicFilter, out var existingSubscription);

                    var interceptorEventArgs = await InterceptUnsubscribe(topicFilter, existingSubscription, unsubscribePacket.UserProperties, cancellationToken).ConfigureAwait(false);
                    var acceptUnsubscription = interceptorEventArgs.Response.ReasonCode == MqttUnsubscribeReasonCode.Success;

                    result.UserProperties = interceptorEventArgs.UserProperties;
                    result.ReasonCodes.Add(interceptorEventArgs.Response.ReasonCode);

                    if (interceptorEventArgs.CloseConnection)
                    {
                        // When any of the interceptor calls leads to a connection close the connection
                        // must be closed. So do not revert to false!
                        result.CloseConnection = true;
                    }

                    if (!acceptUnsubscription)
                    {
                        continue;
                    }

                    if (interceptorEventArgs.ProcessUnsubscription)
                    {
                        _subscriptions.Remove(topicFilter);

                        // must remove subscription object from topic hash dictionary also
                        if (existingSubscription != null)
                        {
                            var topicHash = existingSubscription.TopicHash;

                            if (existingSubscription.TopicHasWildcard)
                            {
                                if (_wildcardSubscriptionsByTopicHash.TryGetValue(topicHash, out var subscriptions))
                                {
                                    subscriptions.RemoveSubscription(existingSubscription);
                                    if (subscriptions.SubscriptionsByHashMask.Count == 0)
                                    {
                                        _wildcardSubscriptionsByTopicHash.Remove(topicHash);
                                    }
                                }
                            }
                            else
                            {
                                if (_noWildcardSubscriptionsByTopicHash.TryGetValue(topicHash, out var subscriptions))
                                {
                                    subscriptions.Remove(existingSubscription);
                                    if (subscriptions.Count == 0)
                                    {
                                        _noWildcardSubscriptionsByTopicHash.Remove(topicHash);
                                    }
                                }
                            }
                        }

                        removedSubscriptions.Add(topicFilter);
                    }
                }
            }
            finally
            {
                _subscriptionsLock.ExitWriteLock();
                _subscriptionChangedNotification?.OnSubscriptionsRemoved(_session, removedSubscriptions);
            }

            if (_eventContainer.ClientUnsubscribedTopicEvent.HasHandlers)
            {
                foreach (var topicFilter in unsubscribePacket.TopicFilters)
                {
                    var eventArgs = new ClientUnsubscribedTopicEventArgs(_session.Id, _session.UserName, topicFilter, _session.Items);
                    await _eventContainer.ClientUnsubscribedTopicEvent.InvokeAsync(eventArgs).ConfigureAwait(false);
                }
            }

            return result;
        }

        CreateSubscriptionResult CreateSubscription(MqttTopicFilter topicFilter, uint subscriptionIdentifier, MqttSubscribeReasonCode reasonCode)
        {
            MqttQualityOfServiceLevel grantedQualityOfServiceLevel;

            if (reasonCode == MqttSubscribeReasonCode.GrantedQoS0)
            {
                grantedQualityOfServiceLevel = MqttQualityOfServiceLevel.AtMostOnce;
            }
            else if (reasonCode == MqttSubscribeReasonCode.GrantedQoS1)
            {
                grantedQualityOfServiceLevel = MqttQualityOfServiceLevel.AtLeastOnce;
            }
            else if (reasonCode == MqttSubscribeReasonCode.GrantedQoS2)
            {
                grantedQualityOfServiceLevel = MqttQualityOfServiceLevel.ExactlyOnce;
            }
            else
            {
                throw new InvalidOperationException();
            }

            var subscription = new MqttSubscription(
                topicFilter.Topic,
                topicFilter.NoLocal,
                topicFilter.RetainHandling,
                topicFilter.RetainAsPublished,
                grantedQualityOfServiceLevel,
                subscriptionIdentifier);

            bool isNewSubscription;

            // Add to subscriptions and maintain topic hash dictionaries

            _subscriptionsLock.EnterWriteLock();
            try
            {
                MqttTopicHash.Calculate(topicFilter.Topic, out var topicHash, out _, out var hasWildcard);

                if (_subscriptions.TryGetValue(topicFilter.Topic, out var existingSubscription))
                {
                    // must remove object from topic hash dictionary first
                    if (hasWildcard)
                    {
                        if (_wildcardSubscriptionsByTopicHash.TryGetValue(topicHash, out var subscriptions))
                        {
                            subscriptions.RemoveSubscription(existingSubscription);
                            // no need to remove empty entry because we'll be adding subscription again below
                        }
                    }
                    else
                    {
                        if (_noWildcardSubscriptionsByTopicHash.TryGetValue(topicHash, out var subscriptions))
                        {
                            subscriptions.Remove(existingSubscription);
                            // no need to remove empty entry because we'll be adding subscription again below
                        }
                    }
                }

                isNewSubscription = existingSubscription == null;
                _subscriptions[topicFilter.Topic] = subscription;

                // Add or re-add to topic hash dictionary
                if (hasWildcard)
                {
                    if (!_wildcardSubscriptionsByTopicHash.TryGetValue(topicHash, out var subscriptions))
                    {
                        subscriptions = new TopicHashMaskSubscriptions();
                        _wildcardSubscriptionsByTopicHash.Add(topicHash, subscriptions);
                    }

                    subscriptions.AddSubscription(subscription);
                }
                else
                {
                    if (!_noWildcardSubscriptionsByTopicHash.TryGetValue(topicHash, out var subscriptions))
                    {
                        subscriptions = new HashSet<MqttSubscription>();
                        _noWildcardSubscriptionsByTopicHash.Add(topicHash, subscriptions);
                    }

                    subscriptions.Add(subscription);
                }
            }
            finally
            {
                _subscriptionsLock.ExitWriteLock();
            }

            return new CreateSubscriptionResult
            {
                IsNewSubscription = isNewSubscription,
                Subscription = subscription
            };
        }

        static void FilterRetainedApplicationMessages(
            IList<MqttApplicationMessage> retainedMessages,
            CreateSubscriptionResult createSubscriptionResult,
            SubscribeResult subscribeResult)
        {
            if (createSubscriptionResult.Subscription.RetainHandling == MqttRetainHandling.DoNotSendOnSubscribe)
            {
                // This is a MQTT V5+ feature.
                return;
            }

            if (createSubscriptionResult.Subscription.RetainHandling == MqttRetainHandling.SendAtSubscribeIfNewSubscriptionOnly && !createSubscriptionResult.IsNewSubscription)
            {
                // This is a MQTT V5+ feature.
                return;
            }

            for (var index = retainedMessages.Count - 1; index >= 0; index--)
            {
                var retainedMessage = retainedMessages[index];
                if (retainedMessage == null)
                {
                    continue;
                }

                if (MqttTopicFilterComparer.Compare(retainedMessage.Topic, createSubscriptionResult.Subscription.Topic) != MqttTopicFilterCompareResult.IsMatch)
                {
                    continue;
                }

                var retainedMessageMatch = new MqttRetainedMessageMatch(retainedMessage, createSubscriptionResult.Subscription.GrantedQualityOfServiceLevel);
                if (retainedMessageMatch.SubscriptionQualityOfServiceLevel > retainedMessageMatch.ApplicationMessage.QualityOfServiceLevel)
                {
                    // UPGRADING the QoS is not allowed!
                    // From MQTT spec: Subscribing to a Topic Filter at QoS 2 is equivalent to saying
                    // "I would like to receive Messages matching this filter at the QoS with which they were published".
                    // This means a publisher is responsible for determining the maximum QoS a Message can be delivered at,
                    // but a subscriber is able to require that the Server downgrades the QoS to one more suitable for its usage.
                    retainedMessageMatch.SubscriptionQualityOfServiceLevel = retainedMessageMatch.ApplicationMessage.QualityOfServiceLevel;
                }

                if (subscribeResult.RetainedMessages == null)
                {
                    subscribeResult.RetainedMessages = new List<MqttRetainedMessageMatch>();
                }

                subscribeResult.RetainedMessages.Add(retainedMessageMatch);

                // Clear the retained message from the list because the client should receive every message only
                // one time even if multiple subscriptions affect them.
                retainedMessages[index] = null;
            }
        }

        async Task<InterceptingSubscriptionEventArgs> InterceptSubscribe(
            MqttTopicFilter topicFilter,
            List<MqttUserProperty> userProperties,
            CancellationToken cancellationToken)
        {
            var eventArgs = new InterceptingSubscriptionEventArgs(cancellationToken, _session.Id, _session.UserName, new MqttSessionStatus(_session), topicFilter, userProperties);

            if (topicFilter.QualityOfServiceLevel == MqttQualityOfServiceLevel.AtMostOnce)
            {
                eventArgs.Response.ReasonCode = MqttSubscribeReasonCode.GrantedQoS0;
            }
            else if (topicFilter.QualityOfServiceLevel == MqttQualityOfServiceLevel.AtLeastOnce)
            {
                eventArgs.Response.ReasonCode = MqttSubscribeReasonCode.GrantedQoS1;
            }
            else if (topicFilter.QualityOfServiceLevel == MqttQualityOfServiceLevel.ExactlyOnce)
            {
                eventArgs.Response.ReasonCode = MqttSubscribeReasonCode.GrantedQoS2;
            }

            if (topicFilter.Topic.StartsWith("$share/"))
            {
                eventArgs.Response.ReasonCode = MqttSubscribeReasonCode.SharedSubscriptionsNotSupported;
            }
            else
            {
                await _eventContainer.InterceptingSubscriptionEvent.InvokeAsync(eventArgs).ConfigureAwait(false);
            }

            return eventArgs;
        }

        async Task<InterceptingUnsubscriptionEventArgs> InterceptUnsubscribe(
            string topicFilter,
            MqttSubscription mqttSubscription,
            List<MqttUserProperty> userProperties,
            CancellationToken cancellationToken)
        {
            var clientUnsubscribingTopicEventArgs = new InterceptingUnsubscriptionEventArgs(cancellationToken, _session.Id, _session.UserName, _session.Items, topicFilter, userProperties)
            {
                Response =
                {
                    ReasonCode = mqttSubscription == null ? MqttUnsubscribeReasonCode.NoSubscriptionExisted : MqttUnsubscribeReasonCode.Success
                }
            };

            await _eventContainer.InterceptingUnsubscriptionEvent.InvokeAsync(clientUnsubscribingTopicEventArgs).ConfigureAwait(false);

            return clientUnsubscribingTopicEventArgs;
        }

        sealed class CreateSubscriptionResult
        {
            public bool IsNewSubscription { get; set; }

            public MqttSubscription Subscription { get; set; }
        }
    }
}