Implement unsubscribe handling

Wiebe Cazemier
1 parent c3961e7f
Showing 6 changed files with 121 additions and 0 deletions
mqttpacket.cpp
mqttpacket.h
subscriptionstore.cpp
subscriptionstore.h
types.cpp
types.h
@@ -66,6 +66,16 @@ MqttPacket::MqttPacket(const SubAck &amp;subAck) :
     calculateRemainingLength();
 }
+MqttPacket::MqttPacket(const UnsubAck &unsubAck) :
+    bites(unsubAck.getLengthWithoutFixedHeader())
+{
+    packetType = PacketType::SUBACK;
+    first_byte = static_cast<char>(packetType) << 4;
+    writeByte((unsubAck.packet_id & 0xFF00) >> 8);
+    writeByte(unsubAck.packet_id & 0x00FF);
+    calculateRemainingLength();
+}
+
 MqttPacket::MqttPacket(const Publish &publish) :
     bites(publish.getLengthWithoutFixedHeader())
 {
@@ -136,6 +146,8 @@ void MqttPacket::handle()
         sender->writePingResp();
     else if (packetType == PacketType::SUBSCRIBE)
         handleSubscribe();
+    else if (packetType == PacketType::UNSUBSCRIBE)
+        handleUnsubscribe();
     else if (packetType == PacketType::PUBLISH)
         handlePublish();
     else if (packetType == PacketType::PUBACK)
@@ -358,6 +370,32 @@ void MqttPacket::handleSubscribe()
     sender->writeMqttPacket(response);
 }
+void MqttPacket::handleUnsubscribe()
+{
+    const char firstByteFirstNibble = (first_byte & 0x0F);
+
+    if (firstByteFirstNibble != 2)
+        throw ProtocolError("First LSB of first byte is wrong value for subscribe packet.");
+
+    uint16_t packet_id = readTwoBytesToUInt16();
+
+    while (remainingAfterPos() > 0)
+    {
+        uint16_t topicLength = readTwoBytesToUInt16();
+        std::string topic(readBytes(topicLength), topicLength);
+
+        if (topic.empty() || !isValidUtf8(topic))
+            throw ProtocolError("Subscribe topic not valid UTF-8.");
+
+        sender->getThreadData()->getSubscriptionStore()->removeSubscription(sender, topic);
+        logger->logf(LOG_INFO, "Client '%s' unsubscribed to '%s'", sender->repr().c_str(), topic.c_str());
+    }
+
+    UnsubAck unsubAck(packet_id);
+    MqttPacket response(unsubAck);
+    sender->writeMqttPacket(response);
+}
+
 void MqttPacket::handlePublish()
 {
     uint16_t variable_header_length = readTwoBytesToUInt16();
@@ -59,6 +59,7 @@ public:
     // Constructor for outgoing packets. These may not allocate room for the fixed header, because we don't (always) know the length in advance.
     MqttPacket(const ConnAck &connAck);
     MqttPacket(const SubAck &subAck);
+    MqttPacket(const UnsubAck &unsubAck);
     MqttPacket(const Publish &publish);
     MqttPacket(const PubAck &pubAck);
@@ -66,6 +67,7 @@ public:
     void handleConnect();
     void handleDisconnect();
     void handleSubscribe();
+    void handleUnsubscribe();
     void handlePing();
     void handlePublish();
     void handlePubAck();
@@ -30,6 +30,20 @@ void SubscriptionNode::addSubscriber(const std::shared_ptr&lt;Session&gt; &amp;subscriber,
     }
 }
+void SubscriptionNode::removeSubscriber(const std::shared_ptr<Session> &subscriber)
+{
+    Subscription sub;
+    sub.session = subscriber;
+    sub.qos = 0;
+
+    auto it = std::find(subscribers.begin(), subscribers.end(), sub);
+
+    if (it != subscribers.end())
+    {
+        subscribers.erase(it);
+    }
+}
+
 SubscriptionStore::SubscriptionStore() :
     root(new SubscriptionNode("root")),
@@ -84,6 +98,52 @@ void SubscriptionStore::addSubscription(Client_p &amp;client, const std::string &amp;top
 }
+void SubscriptionStore::removeSubscription(Client_p &client, const std::string &topic)
+{
+    const std::list<std::string> subtopics = split(topic, '/');
+
+    RWLockGuard lock_guard(&subscriptionsRwlock);
+    lock_guard.wrlock();
+
+    // TODO: because it's so similar to adding a subscription, make a function to retrieve the deepest node?
+    SubscriptionNode *deepestNode = root.get();
+    for(const std::string &subtopic : subtopics)
+    {
+        std::unique_ptr<SubscriptionNode> *selectedChildren = nullptr;
+
+        if (subtopic == "#")
+            selectedChildren = &deepestNode->childrenPound;
+        else if (subtopic == "+")
+            selectedChildren = &deepestNode->childrenPlus;
+        else
+            selectedChildren = &deepestNode->children[subtopic];
+
+        std::unique_ptr<SubscriptionNode> &node = *selectedChildren;
+
+        if (!node)
+        {
+            return;
+        }
+        deepestNode = node.get();
+    }
+
+    assert(deepestNode);
+
+    if (deepestNode)
+    {
+        auto session_it = sessionsByIdConst.find(client->getClientId());
+        if (session_it != sessionsByIdConst.end())
+        {
+            const std::shared_ptr<Session> &ses = session_it->second;
+            deepestNode->removeSubscriber(ses);
+        }
+    }
+
+    lock_guard.unlock();
+
+
+}
+
 // Removes an existing client when it already exists [MQTT-3.1.4-2].
 void SubscriptionStore::registerClientAndKickExistingOne(Client_p &client)
 {
@@ -42,6 +42,7 @@ public:
     std::vector<Subscription> &getSubscribers();
     void addSubscriber(const std::shared_ptr<Session> &subscriber, char qos);
+    void removeSubscriber(const std::shared_ptr<Session> &subscriber);
     std::unordered_map<std::string, std::unique_ptr<SubscriptionNode>> children;
     std::unique_ptr<SubscriptionNode> childrenPlus;
     std::unique_ptr<SubscriptionNode> childrenPound;
@@ -69,6 +70,7 @@ public:
     SubscriptionStore();
     void addSubscription(Client_p &client, const std::string &topic, char qos);
+    void removeSubscription(Client_p &client, const std::string &topic);
     void registerClientAndKickExistingOne(Client_p &client);
     bool sessionPresent(const std::string &clientid);
@@ -54,3 +54,14 @@ size_t PubAck::getLengthWithoutFixedHeader() const
 {
     return 2;
 }
+
+UnsubAck::UnsubAck(uint16_t packet_id) :
+    packet_id(packet_id)
+{
+
+}
+
+size_t UnsubAck::getLengthWithoutFixedHeader() const
+{
+    return 2;
+}
@@ -69,6 +69,14 @@ public:
     size_t getLengthWithoutFixedHeader() const;
 };
+class UnsubAck
+{
+public:
+    uint16_t packet_id;
+    UnsubAck(uint16_t packet_id);
+    size_t getLengthWithoutFixedHeader() const;
+};
+
 class Publish
 {
 public: