Start of retained messages

Also materializes some concepts about MqttPacket.

Start of retained messages
Also materializes some concepts about MqttPacket.
Wiebe Cazemier
1 parent e3ad153c
Showing 9 changed files with 215 additions and 24 deletions
client.cpp
mqttpacket.cpp
mqttpacket.h
rwlockguard.cpp
rwlockguard.h
subscriptionstore.cpp
subscriptionstore.h
types.cpp
types.h
@@ -83,11 +83,22 @@ void Client::writeMqttPacket(const MqttPacket &amp;packet)
     if (packet.packetType == PacketType::PUBLISH && wwi > CLIENT_MAX_BUFFER_SIZE)
         return;
  
-    if (packet.getSize() > getWriteBufMaxWriteSize())
-        growWriteBuffer(packet.getSize());
+    if (packet.getSizeIncludingNonPresentHeader() > getWriteBufMaxWriteSize())
+        growWriteBuffer(packet.getSizeIncludingNonPresentHeader());
  
-    std::memcpy(&writebuf[wwi], &packet.getBites()[0], packet.getSize());
-    wwi += packet.getSize();
+    if (!packet.containsFixedHeader())
+    {
+        writebuf[wwi++] = packet.getFirstByte();
+        RemainingLength r = packet.getRemainingLength();
+        std::memcpy(&writebuf[wwi], r.bytes, r.len);
+        wwi += r.len;
+    }
+
+    std::memcpy(&writebuf[wwi], &packet.getBites()[0], packet.getBites().size());
+    wwi += packet.getBites().size();
+
+    assert(wwi >= static_cast<int>(packet.getSizeIncludingNonPresentHeader()));
+    assert(wwi <= static_cast<int>(writeBufsize));
  
     setReadyForWriting(true);
 }
@@ -2,6 +2,12 @@
 #include <cstring>
 #include <iostream>
 #include <list>
+#include <cassert>
+
+RemainingLength::RemainingLength()
+{
+    memset(bytes, 0, 4);
+}
  
 // constructor for parsing incoming packets
 MqttPacket::MqttPacket(char *buf, size_t len, size_t fixed_header_length, Client_p &sender) :
@@ -16,9 +22,11 @@ MqttPacket::MqttPacket(char *buf, size_t len, size_t fixed_header_length, Client
     pos += fixed_header_length;
 }
  
+// This constructor cheats and doesn't use calculateRemainingLength, because it's always the same. It allocates enough space in the vector.
 MqttPacket::MqttPacket(const ConnAck &connAck) :
-    bites(4)
+    bites(connAck.getLength() + 2)
 {
+    fixed_header_length = 2;
     packetType = PacketType::CONNACK;
     char first_byte = static_cast<char>(packetType) << 4;
     writeByte(first_byte);
@@ -31,6 +39,7 @@ MqttPacket::MqttPacket(const ConnAck &amp;connAck) :
 MqttPacket::MqttPacket(const SubAck &subAck) :
     bites(3)
 {
+    fixed_header_length = 2; // TODO: this is wrong, pending implementation of the new method in SubAck
     packetType = PacketType::SUBACK;
     char first_byte = static_cast<char>(packetType) << 4;
     writeByte(first_byte);
@@ -48,11 +57,17 @@ MqttPacket::MqttPacket(const SubAck &amp;subAck) :
 }
  
 MqttPacket::MqttPacket(const Publish &publish) :
-    bites(publish.topic.length() + publish.payload.length() + 2 + 2) // TODO: same as above
+    bites(publish.getLength())
 {
+    if (publish.topic.length() > 0xFFFF)
+    {
+        throw ProtocolError("Topic path too long.");
+    }
+
     packetType = PacketType::PUBLISH;
-    char first_byte = static_cast<char>(packetType) << 4;
-    writeByte(first_byte);
+    first_byte = static_cast<char>(packetType) << 4;
+    first_byte |= (publish.qos << 1);
+    first_byte |= (static_cast<char>(publish.retain) & 0b00000001);
  
     char topicLenMSB = (publish.topic.length() & 0xF0) >> 8;
     char topicLenLSB = publish.topic.length() & 0x0F;
@@ -60,9 +75,13 @@ MqttPacket::MqttPacket(const Publish &amp;publish) :
     writeByte(topicLenLSB);
     writeBytes(publish.topic.c_str(), publish.topic.length());
  
-    writeBytes(publish.payload.c_str(), publish.payload.length());
+    if (publish.qos)
+    {
+        throw NotImplementedException("I would write two bytes containing the packet id here, but QoS is not done yet.");
+    }
  
-    // TODO: untested. May be unnecessary, because a received packet can be resent just like that.
+    writeBytes(publish.payload.c_str(), publish.payload.length());
+    calculateRemainingLength();
 }
  
 void MqttPacket::handle(std::shared_ptr<SubscriptionStore> &subscriptionStore)
@@ -201,23 +220,77 @@ void MqttPacket::handlePublish(std::shared_ptr&lt;SubscriptionStore&gt; &amp;subscriptionS
     if (variable_header_length == 0)
         throw ProtocolError("Empty publish topic");
  
+    bool retain = (first_byte & 0b00000001);
+    bool dup = !!(first_byte & 0b00001000);
     char qos = (first_byte & 0b00000110) >> 1;
  
+    if (qos == 3)
+        throw ProtocolError("QoS 3 is a protocol violation.");
+
+    // TODO: validate UTF8.
     std::string topic(readBytes(variable_header_length), variable_header_length);
  
     if (qos)
     {
         throw ProtocolError("Qos not implemented.");
-        //uint16_t packet_id = readTwoBytesToUInt16();
+        uint16_t packet_id = readTwoBytesToUInt16();
     }
  
-    // TODO: validate UTF8.
-    size_t payload_length = remainingAfterPos();
-    std::string payload(readBytes(payload_length), payload_length);
+    if (retain)
+    {
+        size_t payload_length = remainingAfterPos();
+        std::string payload(readBytes(payload_length), payload_length);
  
+        subscriptionStore->setRetainedMessage(topic, payload, qos);
+    }
+
+    // Set dup flag to 0, because that must not be propagated [MQTT-3.3.1-3].
+    // Existing subscribers don't get retain=1. [MQTT-3.3.1-9]
+    bites[0] &= 0b11110110;
+
+    // For the existing clients, we can just write the same packet back out, with our small alterations.
     subscriptionStore->queuePacketAtSubscribers(topic, *this, sender);
 }
  
+void MqttPacket::calculateRemainingLength()
+{
+    assert(fixed_header_length == 0); // because you're not supposed to call this on packet that we already know the length of.
+
+    size_t x = bites.size();
+
+    do
+    {
+        if (remainingLength.len > 4)
+            throw std::runtime_error("Calculated remaining length is longer than 4 bytes.");
+
+        char encodedByte = x % 128;
+        x = x / 128;
+        if (x > 0)
+            encodedByte = encodedByte | 128;
+        remainingLength.bytes[remainingLength.len++] = encodedByte;
+    }
+    while(x > 0);
+}
+
+RemainingLength MqttPacket::getRemainingLength() const
+{
+    assert(remainingLength.len > 0);
+    return remainingLength;
+}
+
+size_t MqttPacket::getSizeIncludingNonPresentHeader() const
+{
+    size_t total = bites.size();
+
+    if (fixed_header_length == 0)
+    {
+        total++;
+        total += remainingLength.len;
+    }
+
+    return total;
+}
+
  
 Client_p MqttPacket::getSender() const
 {
@@ -229,6 +302,16 @@ void MqttPacket::setSender(const Client_p &amp;value)
     sender = value;
 }
  
+bool MqttPacket::containsFixedHeader() const
+{
+    return fixed_header_length > 0;
+}
+
+char MqttPacket::getFirstByte() const
+{
+    return first_byte;
+}
+
 char *MqttPacket::readBytes(size_t length)
 {
     if (pos + length > bites.size())
@@ -262,6 +345,7 @@ void MqttPacket::writeBytes(const char *b, size_t len)
         throw ProtocolError("Exceeding packet size");
  
     memcpy(&bites[pos], b, len);
+    pos += len;
 }
  
 uint16_t MqttPacket::readTwoBytesToUInt16()
@@ -287,3 +371,5 @@ size_t MqttPacket::remainingAfterPos()
  
  
  
+
+
@@ -13,13 +13,21 @@
 #include "types.h"
 #include "subscriptionstore.h"
  
+struct RemainingLength
+{
+    char bytes[4];
+    int len = 0;
+public:
+    RemainingLength();
+};
  
 class MqttPacket
 {
     std::vector<char> bites;
-    size_t fixed_header_length = 0;
+    size_t fixed_header_length = 0; // if 0, this packet does not contain the bytes of the fixed header.
+    RemainingLength remainingLength;
     Client_p sender;
-    char first_byte;
+    char first_byte = 0;
     size_t pos = 0;
     ProtocolVersion protocolVersion = ProtocolVersion::None;
  
@@ -30,12 +38,13 @@ class MqttPacket
     uint16_t readTwoBytesToUInt16();
     size_t remainingAfterPos();
  
+    void calculateRemainingLength();
+
 public:
     PacketType packetType = PacketType::Reserved;
-    MqttPacket(char *buf, size_t len, size_t fixed_header_length, Client_p &sender);
+    MqttPacket(char *buf, size_t len, size_t fixed_header_length, Client_p &sender); // Constructor for parsing incoming packets.
  
-    // TODO: not constructors, but static functions that return all the stuff after the fixed header, then a constructor with vector.
-    // Or, I can not have the fixed header, and calculate that on write-to-buf.
+    // 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 Publish &publish);
@@ -46,11 +55,15 @@ public:
     void handlePing();
     void handlePublish(std::shared_ptr<SubscriptionStore> &subscriptionStore);
  
-    size_t getSize() const { return bites.size(); }
+    size_t getSizeIncludingNonPresentHeader() const;
     const std::vector<char> &getBites() const { return bites; }
  
     Client_p getSender() const;
     void setSender(const Client_p &value);
+
+    bool containsFixedHeader() const;
+    char getFirstByte() const;
+    RemainingLength getRemainingLength() const;
 };
  
 #endif // MQTTPACKET_H
@@ -10,7 +10,7 @@ RWLockGuard::RWLockGuard(pthread_rwlock_t *rwlock) :
  
 RWLockGuard::~RWLockGuard()
 {
-    pthread_rwlock_unlock(rwlock);
+    unlock();
 }
  
 void RWLockGuard::wrlock()
@@ -24,3 +24,12 @@ void RWLockGuard::rdlock()
     if (pthread_rwlock_wrlock(rwlock) != 0)
         throw std::runtime_error("rdlock failed.");
 }
+
+void RWLockGuard::unlock()
+{
+    if (rwlock != NULL)
+    {
+        pthread_rwlock_unlock(rwlock);
+        rwlock = NULL;
+    }
+}
@@ -12,6 +12,7 @@ public:
     ~RWLockGuard();
     void wrlock();
     void rdlock();
+    void unlock();
 };
  
 #endif // RWLOCKGUARD_H
@@ -44,6 +44,9 @@ void SubscriptionStore::addSubscription(Client_p &amp;client, const std::string &amp;top
     }
  
     clients_by_id[client->getClientId()] = client;
+    lock_guard.unlock();
+
+    giveClientRetainedMessage(client, topic); // TODO: wildcards
 }
  
 void SubscriptionStore::removeClient(const Client_p &client)
@@ -119,4 +122,44 @@ void SubscriptionStore::queuePacketAtSubscribers(const std::string &amp;topic, const
     publishRecursively(subtopics.begin(), subtopics.end(), root, packet);
 }
  
+void SubscriptionStore::giveClientRetainedMessage(Client_p &client, const std::string &topic)
+{
+    RWLockGuard locker(&retainedMessagesRwlock);
+    locker.rdlock();
+
+    auto retained_ptr = retainedMessages.find(topic);
+
+    if (retained_ptr == retainedMessages.end())
+        return;
+
+    const RetainedPayload &m = retained_ptr->second;
+
+    Publish publish(topic, m.payload, m.qos);
+    publish.retain = true;
+    const MqttPacket packet(publish);
+    client->writeMqttPacket(packet);
+}
+
+void SubscriptionStore::setRetainedMessage(const std::string &topic, const std::string &payload, char qos)
+{
+    RWLockGuard locker(&retainedMessagesRwlock);
+    locker.wrlock();
+
+    auto retained_ptr = retainedMessages.find(topic);
+    bool retained_found = retained_ptr != retainedMessages.end();
+
+    if (!retained_found && payload.empty())
+        return;
+
+    if (retained_found && payload.empty())
+    {
+        retainedMessages.erase(topic);
+        return;
+    }
+
+    RetainedPayload &m = retainedMessages[topic];
+    m.payload = payload;
+    m.qos = qos;
+}
+
  
@@ -12,6 +12,12 @@
 #include "client.h"
 #include "utils.h"
  
+struct RetainedPayload
+{
+    std::string payload;
+    char qos;
+};
+
 class SubscriptionNode
 {
     std::string subtopic;
@@ -32,6 +38,9 @@ class SubscriptionStore
     std::unordered_map<std::string, Client_p> clients_by_id;
     const std::unordered_map<std::string, Client_p> &clients_by_id_const;
  
+    pthread_rwlock_t retainedMessagesRwlock = PTHREAD_RWLOCK_INITIALIZER;
+    std::unordered_map<std::string, RetainedPayload> retainedMessages;
+
     bool publishNonRecursively(const MqttPacket &packet, const std::forward_list<std::string> &subscribers) const;
     bool publishRecursively(std::list<std::string>::const_iterator cur_subtopic_it, std::list<std::string>::const_iterator end,
                             std::unique_ptr<SubscriptionNode> &next, const MqttPacket &packet) const;
@@ -42,6 +51,9 @@ public:
     void removeClient(const Client_p &client);
  
     void queuePacketAtSubscribers(const std::string &topic, const MqttPacket &packet, const Client_p &sender);
+    void giveClientRetainedMessage(Client_p &client, const std::string &topic);
+
+    void setRetainedMessage(const std::string &topic, const std::string &payload, char qos);
 };
  
 #endif // SUBSCRIPTIONSTORE_H
@@ -16,9 +16,21 @@ SubAck::SubAck(uint16_t packet_id, const std::list&lt;std::string&gt; &amp;subs) :
     }
 }
  
-Publish::Publish(std::string &topic, std::string payload) :
+Publish::Publish(const std::string &topic, const std::string payload, char qos) :
     topic(topic),
-    payload(payload)
+    payload(payload),
+    qos(qos)
 {
  
 }
+
+// Length starting at the variable header, not the fixed header.
+size_t Publish::getLength() const
+{
+    int result = topic.length() + payload.length() + 2;
+
+    if (qos)
+        result += 2;
+
+    return result;
+}
@@ -48,6 +48,7 @@ class ConnAck
 public:
     ConnAck(ConnAckReturnCodes return_code);
     ConnAckReturnCodes return_code;
+    size_t getLength() const { return 2;} // size of connack is always the same
 };
  
 enum class SubAckReturnCodes
@@ -71,7 +72,10 @@ class Publish
 public:
     std::string topic;
     std::string payload;
-    Publish(std::string &topic, std::string payload);
+    char qos = 0;
+    bool retain = false; // Note: existing subscribers don't get publishes of retained messages with retain=1. [MQTT-3.3.1-9]
+    Publish(const std::string &topic, const std::string payload, char qos);
+    size_t getLength() const;
 };
  
 #endif // TYPES_H