Add ability to send disconnect packets with code

Wiebe Cazemier
1 parent 4d957a3d
Showing 6 changed files with 76 additions and 2 deletions
exceptions.h
mqttpacket.cpp
mqttpacket.h
threadloop.cpp
types.cpp
types.h
@@ -22,10 +22,18 @@ License along with FlashMQ. If not, see &lt;https://www.gnu.org/licenses/&gt;.
 #include <stdexcept>
 #include <sstream>
  
+#include "types.h"
+
 class ProtocolError : public std::runtime_error
 {
 public:
-    ProtocolError(const std::string &msg) : std::runtime_error(msg) {}
+    const ReasonCodes reasonCode;
+
+    ProtocolError(const std::string &msg, ReasonCodes reasonCode = ReasonCodes::UnspecifiedError) : std::runtime_error(msg),
+        reasonCode(reasonCode)
+    {
+
+    }
 };
  
 class NotImplementedException : public std::runtime_error
@@ -34,7 +34,7 @@ MqttPacket::MqttPacket(CirBuf &amp;buf, size_t packet_len, size_t fixed_header_lengt
  
     if (packet_len > sender->getMaxIncomingPacketSize())
     {
-        throw ProtocolError("Incoming packet size exceeded. TODO: DISCONNECT WITH CODE 0x95");
+        throw ProtocolError("Incoming packet size exceeded.", ReasonCodes::PacketTooLarge);
     }
  
     buf.read(bites.data(), packet_len);
@@ -204,6 +204,19 @@ MqttPacket::MqttPacket(const PubResponse &amp;pubAck) :
     }
 }
  
+MqttPacket::MqttPacket(const Disconnect &disconnect) :
+    bites(disconnect.getLengthWithoutFixedHeader())
+{
+    this->protocolVersion = ProtocolVersion::Mqtt5;
+
+    packetType = PacketType::DISCONNECT;
+    first_byte = static_cast<char>(packetType) << 4;
+
+    writeByte(static_cast<uint8_t>(disconnect.reasonCode));
+    writeProperties(disconnect.propertyBuilder);
+    calculateRemainingLength();
+}
+
 void MqttPacket::bufferToMqttPackets(CirBuf &buf, std::vector<MqttPacket> &packetQueueIn, std::shared_ptr<Client> &sender)
 {
     while (buf.usedBytes() >= MQTT_HEADER_LENGH)
@@ -630,6 +643,9 @@ void MqttPacket::handleConnect()
  
 void MqttPacket::handleDisconnect()
 {
+    if (first_byte & 0b1111)
+        throw ProtocolError("Disconnect packet first 4 bits should be 0.", ReasonCodes::MalformedPacket);
+
     logger->logf(LOG_NOTICE, "Client '%s' cleanly disconnecting", sender->repr().c_str());
     sender->setDisconnectReason("MQTT Disconnect received.");
     sender->markAsDisconnecting();
@@ -100,6 +100,7 @@ public:
     MqttPacket(const UnsubAck &unsubAck);
     MqttPacket(const ProtocolVersion protocolVersion, Publish &_publish);
     MqttPacket(const PubResponse &pubAck);
+    MqttPacket(const Disconnect &disconnect);
  
     static void bufferToMqttPackets(CirBuf &buf, std::vector<MqttPacket> &packetQueueIn, std::shared_ptr<Client> &sender);
  
@@ -135,6 +135,22 @@ void do_thread_work(ThreadData *threadData)
                             }
                         }
                     }
+                    catch (ProtocolError &ex)
+                    {
+                        client->setDisconnectReason(ex.what());
+                        if (client->getProtocolVersion() >= ProtocolVersion::Mqtt5 && client->hasConnectPacketSeen())
+                        {
+                            Disconnect d(client->getProtocolVersion(), ex.reasonCode);
+                            MqttPacket p(d);
+                            client->writeMqttPacket(p);
+                            client->setReadyForDisconnect();
+                        }
+                        else
+                        {
+                            logger->logf(LOG_ERR, "Protocol error: %s. Removing client.", ex.what());
+                            threadData->removeClient(client);
+                        }
+                    }
                     catch(std::exception &ex)
                     {
                         client->setDisconnectReason(ex.what());
@@ -289,3 +289,27 @@ size_t UnsubAck::getLengthWithoutFixedHeader() const
  
     return result;
 }
+
+Disconnect::Disconnect(const ProtocolVersion protVersion, ReasonCodes reason_code) :
+    reasonCode(reason_code)
+{
+    assert(protVersion >= ProtocolVersion::Mqtt5);
+
+
+}
+
+size_t Disconnect::getLengthWithoutFixedHeader() const
+{
+    size_t result = 1;
+    const size_t proplen = propertyBuilder ? propertyBuilder->getLength() : 1;
+    result += proplen;
+    return result;
+}
+
+
+
+
+
+
+
+
@@ -249,4 +249,13 @@ public:
     bool needsReasonCode() const;
 };
  
+class Disconnect
+{
+public:
+    ReasonCodes reasonCode;
+    std::shared_ptr<Mqtt5PropertyBuilder> propertyBuilder;
+    Disconnect(const ProtocolVersion protVersion, ReasonCodes reason_code);
+    size_t getLengthWithoutFixedHeader() const;
+};
+
 #endif // TYPES_H