From 54be6e87db16d7dbc60d4c4674566ec08b0115c2 Mon Sep 17 00:00:00 2001
From: Nick O'Leary <nick.oleary@gmail.com>
Date: Wed, 18 Jul 2018 11:02:08 +0100
Subject: [PATCH] Check remaining-length encoding is valid

---
 src/PubSubClient.cpp       |  9 +++++++++
 tests/src/receive_spec.cpp | 30 ++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/src/PubSubClient.cpp b/src/PubSubClient.cpp
index 603fcc8..29fbbfa 100755
--- a/src/PubSubClient.cpp
+++ b/src/PubSubClient.cpp
@@ -240,6 +240,12 @@ uint16_t PubSubClient::readPacket(uint8_t* lengthLength) {
     uint8_t start = 0;
 
     do {
+        if (len == 6) {
+            // Invalid remaining length encoding - kill the connection
+            _state = MQTT_DISCONNECTED;
+            _client->stop();
+            return 0;
+        }
         if(!readByte(&digit)) return 0;
         buffer[len++] = digit;
         length += (digit & 127) * multiplier;
@@ -335,6 +341,9 @@ boolean PubSubClient::loop() {
                 } else if (type == MQTTPINGRESP) {
                     pingOutstanding = false;
                 }
+            } else if (!connected()) {
+                // readPacket has closed the connection
+                return false;
             }
         }
         return true;
diff --git a/tests/src/receive_spec.cpp b/tests/src/receive_spec.cpp
index 54a62ee..4ecd439 100644
--- a/tests/src/receive_spec.cpp
+++ b/tests/src/receive_spec.cpp
@@ -160,6 +160,35 @@ int test_receive_oversized_message() {
     END_IT
 }
 
+int test_drop_invalid_remaining_length_message() {
+    IT("drops invalid remaining length message");
+    reset_callback();
+
+    ShimClient shimClient;
+    shimClient.setAllowConnect(true);
+
+    byte connack[] = { 0x20, 0x02, 0x00, 0x00 };
+    shimClient.respond(connack,4);
+
+    PubSubClient client(server, 1883, callback, shimClient);
+    int rc = client.connect((char*)"client_test1");
+    IS_TRUE(rc);
+
+    byte publish[] = {0x30,0x92,0x92,0x92,0x92,0x92,0x92,0x0,0x5,0x74,0x6f,0x70,0x69,0x63,0x70,0x61,0x79,0x6c,0x6f,0x61,0x64};
+    shimClient.respond(publish,21);
+
+    rc = client.loop();
+
+    IS_FALSE(rc);
+
+    IS_FALSE(callback_called);
+
+    IS_FALSE(shimClient.error());
+
+    END_IT
+}
+
+
 int test_receive_oversized_stream_message() {
     IT("drops an oversized message");
     reset_callback();
@@ -241,6 +270,7 @@ int main()
     test_receive_callback();
     test_receive_stream();
     test_receive_max_sized_message();
+    test_drop_invalid_remaining_length_message();
     test_receive_oversized_message();
     test_receive_oversized_stream_message();
     test_receive_qos1();