From 07a62470e8ccf1c8604fa585ec44e8636eb7e9e8 Mon Sep 17 00:00:00 2001
From: Takatoshi Kondo <redboltz@gmail.com>
Date: Fri, 20 Oct 2023 14:31:36 +0900
Subject: [PATCH] Fixed store_ management.

Improved protocol error checking.
---
 include/mqtt/endpoint.hpp         | 39 +++++++++++---
 include/mqtt/store.hpp            |  2 +-
 test/system/st_async_pubsub_2.cpp | 88 +++----------------------------
 test/system/st_pubsub_1.cpp       | 87 +++---------------------------
 4 files changed, 44 insertions(+), 172 deletions(-)

diff --git a/include/mqtt/endpoint.hpp b/include/mqtt/endpoint.hpp
index 280121cd1..7488c4bf8 100644
--- a/include/mqtt/endpoint.hpp
+++ b/include/mqtt/endpoint.hpp
@@ -8140,7 +8140,16 @@ class endpoint : public std::enable_shared_from_this<endpoint<Mutex, LockGuard,
                         ep_.pid_man_.release_id(packet_id_);
                         return true;
                     } ();
-                if (erased) ep_.on_serialize_remove(packet_id_);
+                if (erased) {
+                    ep_.on_serialize_remove(packet_id_);
+                }
+                else {
+                    MQTT_LOG("mqtt_impl", error)
+                        << MQTT_ADD_VALUE(address, &ep_)
+                        << "invalid puback received. packet_id:" << packet_id_;
+                    ep_.call_protocol_error_handlers();
+                    return;
+                }
                 switch (ep_.version_) {
                 case protocol_version::v3_1_1:
                     if (ep_.on_puback(packet_id_)) {
@@ -8288,6 +8297,13 @@ class endpoint : public std::enable_shared_from_this<endpoint<Mutex, LockGuard,
                         if (is_error(reason_code_)) ep_.pid_man_.release_id(packet_id_);
                         return true;
                     } ();
+                if (!erased) {
+                    MQTT_LOG("mqtt_impl", error)
+                        << MQTT_ADD_VALUE(address, &ep_)
+                        << "invalid pubrec received. packet_id:" << packet_id_;
+                    ep_.call_protocol_error_handlers();
+                    return;
+                }
                 {
                     auto res =
                         [&] {
@@ -8647,7 +8663,16 @@ class endpoint : public std::enable_shared_from_this<endpoint<Mutex, LockGuard,
                         ep_.pid_man_.release_id(packet_id_);
                         return true;
                     } ();
-                if (erased) ep_.on_serialize_remove(packet_id_);
+                if (erased) {
+                    ep_.on_serialize_remove(packet_id_);
+                }
+                else {
+                    MQTT_LOG("mqtt_impl", error)
+                        << MQTT_ADD_VALUE(address, &ep_)
+                        << "invalid pubcomp received. packet_id:" << packet_id_;
+                    ep_.call_protocol_error_handlers();
+                    return;
+                }
                 switch (ep_.version_) {
                 case protocol_version::v3_1_1:
                     if (ep_.on_pubcomp(packet_id_)) {
@@ -9993,9 +10018,8 @@ class endpoint : public std::enable_shared_from_this<endpoint<Mutex, LockGuard,
                     // publish store is erased when pubrec is received.
                     // pubrel store is erased when pubcomp is received.
                     // If invalid client send pubrec twice with the same packet id,
-                    // then send corresponding pubrel twice is a possible client/server
-                    // implementation.
-                    // In this case, overwrite store_.
+                    // then send disconnect with protocol_error reason_code (v5), or
+                    // simply close the socket (v3.1.1).
                     if (store_.insert_or_update(
                             packet_id,
                             control_packet_type::pubcomp,
@@ -10794,9 +10818,8 @@ class endpoint : public std::enable_shared_from_this<endpoint<Mutex, LockGuard,
                         // publish store is erased when pubrec is received.
                         // pubrel store is erased when pubcomp is received.
                         // If invalid client send pubrec twice with the same packet id,
-                        // then send corresponding pubrel twice is a possible client/server
-                        // implementation.
-                        // In this case, overwrite store_.
+                        // then send disconnect with protocol_error reason_code (v5), or
+                        // simply close the socket (v3.1.1).
                         MQTT_LOG("mqtt_impl", warning)
                             << MQTT_ADD_VALUE(address, this)
                             << "overwrite pubrel"
diff --git a/include/mqtt/store.hpp b/include/mqtt/store.hpp
index 717ad4f44..797675968 100644
--- a/include/mqtt/store.hpp
+++ b/include/mqtt/store.hpp
@@ -179,7 +179,7 @@ class store {
                     >
                 >
             >,
-            mi::ordered_non_unique<
+            mi::ordered_unique<
                 mi::tag<tag_packet_id>,
                 mi::const_mem_fun<
                     elem_t, packet_id_t,
diff --git a/test/system/st_async_pubsub_2.cpp b/test/system/st_async_pubsub_2.cpp
index e5dd1279b..4817f98e5 100644
--- a/test/system/st_async_pubsub_2.cpp
+++ b/test/system/st_async_pubsub_2.cpp
@@ -621,23 +621,12 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
             cont("h_suback"),
             // publish topic1 QoS2
             cont("h_pubrec"),
-            cont("h_pubcomp"),
             deps("h_publish", "h_suback"),
             // pubrec send twice
-            cont("h_pubrel1"),
-            cont("h_pubrel2"),
-            deps("h_unsuback", "h_pubcomp", "h_pubrel2"),
             // disconnect
-            cont("h_close"),
+            cont("h_error"),
         };
 
-        auto g = MQTT_NS::shared_scope_guard(
-            [&c] {
-                auto unsub_pid = c->acquire_unique_packet_id();
-                c->async_unsubscribe(unsub_pid, "topic1");
-            }
-        );
-
         switch (c->get_protocol_version()) {
         case MQTT_NS::protocol_version::v3_1_1:
             c->set_connack_handler(
@@ -663,13 +652,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->async_pubrel(packet_id);
                     return true;
                 });
-            c->set_pubcomp_handler(
-                [&chk, g]
-                (packet_id_t) mutable {
-                    MQTT_CHK("h_pubcomp");
-                    g.reset();
-                    return true;
-                });
             c->set_suback_handler(
                 [&chk, &c]
                 (packet_id_t, std::vector<MQTT_NS::suback_return_code> results) {
@@ -680,13 +662,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->async_publish(pid_pub, "topic1", "topic1_contents", MQTT_NS::qos::exactly_once);
                     return true;
                 });
-            c->set_unsuback_handler(
-                [&chk, &c]
-                (packet_id_t) {
-                    MQTT_CHK("h_unsuback");
-                    c->async_disconnect();
-                    return true;
-                });
             c->set_publish_handler(
                 [&chk, &c]
                 (MQTT_NS::optional<packet_id_t> packet_id,
@@ -705,23 +680,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->async_pubrec(*packet_id);
                     return true;
                 });
-            c->set_pubrel_handler(
-                [&chk, &c, g]
-                (packet_id_t packet_id) mutable {
-                    auto ret = MQTT_ORDERED(
-                        [&] {
-                            MQTT_CHK("h_pubrel1");
-                            c->async_pubcomp(packet_id);
-                        },
-                        [&] () {
-                            MQTT_CHK("h_pubrel2");
-                            c->async_pubcomp(packet_id);
-                            g.reset();
-                        }
-                    );
-                    BOOST_TEST(ret);
-                    return true;
-                });
             break;
         case MQTT_NS::protocol_version::v5:
             c->set_v5_connack_handler(
@@ -747,13 +705,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->async_pubrel(packet_id);
                     return true;
                 });
-            c->set_v5_pubcomp_handler(
-                [&chk, g]
-                (packet_id_t, MQTT_NS::v5::pubcomp_reason_code, MQTT_NS::v5::properties /*props*/) mutable {
-                    MQTT_CHK("h_pubcomp");
-                    g.reset();
-                    return true;
-                });
             c->set_v5_suback_handler(
                 [&chk, &c]
                 (packet_id_t, std::vector<MQTT_NS::v5::suback_reason_code> reasons, MQTT_NS::v5::properties /*props*/) {
@@ -764,15 +715,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->async_publish(pid_pub, "topic1", "topic1_contents", MQTT_NS::qos::exactly_once);
                     return true;
                 });
-            c->set_v5_unsuback_handler(
-                [&chk, &c]
-                (packet_id_t, std::vector<MQTT_NS::v5::unsuback_reason_code> reasons, MQTT_NS::v5::properties /*props*/) {
-                    MQTT_CHK("h_unsuback");
-                    BOOST_TEST(reasons.size() == 1U);
-                    BOOST_TEST(reasons[0] == MQTT_NS::v5::unsuback_reason_code::success);
-                    c->async_disconnect();
-                    return true;
-                });
             c->set_v5_publish_handler(
                 [&chk, &c]
                 (MQTT_NS::optional<packet_id_t> packet_id,
@@ -792,23 +734,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->async_pubrec(*packet_id);
                     return true;
                 });
-            c->set_v5_pubrel_handler(
-                [&chk, &c, g]
-                (packet_id_t packet_id, MQTT_NS::v5::pubrel_reason_code, MQTT_NS::v5::properties /*props*/) mutable {
-                    auto ret = MQTT_ORDERED(
-                        [&] {
-                            MQTT_CHK("h_pubrel1");
-                            c->async_pubcomp(packet_id);
-                        },
-                        [&] () {
-                            MQTT_CHK("h_pubrel2");
-                            c->async_pubcomp(packet_id);
-                            g.reset();
-                        }
-                    );
-                    BOOST_TEST(ret);
-                    return true;
-                });
             break;
         default:
             BOOST_CHECK(false);
@@ -816,17 +741,16 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
         }
 
         c->set_close_handler(
-            [&chk, &finish]
+            []
             () {
-                MQTT_CHK("h_close");
-                finish();
+                BOOST_CHECK(false);
             });
         c->set_error_handler(
-            []
+            [&chk, &finish]
             (MQTT_NS::error_code) {
-                BOOST_CHECK(false);
+                MQTT_CHK("h_error");
+                finish();
             });
-        g.reset();
         c->async_connect();
         ioc.run();
         BOOST_TEST(chk.all());
diff --git a/test/system/st_pubsub_1.cpp b/test/system/st_pubsub_1.cpp
index 859f40a26..90336f4a8 100644
--- a/test/system/st_pubsub_1.cpp
+++ b/test/system/st_pubsub_1.cpp
@@ -1757,22 +1757,12 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
             cont("h_suback"),
             // publish topic1 QoS2
             cont("h_pubrec"),
-            cont("h_pubcomp"),
             deps("h_publish", "h_suback"),
             // pubrec send twice
-            cont("h_pubrel1"),
-            cont("h_pubrel2"),
-            deps("h_unsuback", "h_pubcomp", "h_pubrel2"),
             // disconnect
-            cont("h_close"),
+            cont("h_error"),
         };
 
-        auto g = MQTT_NS::shared_scope_guard(
-            [&c] {
-                c->unsubscribe("topic1");
-            }
-        );
-
         switch (c->get_protocol_version()) {
         case MQTT_NS::protocol_version::v3_1_1:
             c->set_connack_handler(
@@ -1797,13 +1787,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->pubrel(packet_id);
                     return true;
                 });
-            c->set_pubcomp_handler(
-                [&chk, g]
-                (packet_id_t) mutable {
-                    MQTT_CHK("h_pubcomp");
-                    g.reset();
-                    return true;
-                });
             c->set_suback_handler(
                 [&chk, &c]
                 (packet_id_t, std::vector<MQTT_NS::suback_return_code> results) {
@@ -1813,13 +1796,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->publish("topic1", "topic1_contents", MQTT_NS::qos::exactly_once);
                     return true;
                 });
-            c->set_unsuback_handler(
-                [&chk, &c]
-                (packet_id_t) {
-                    MQTT_CHK("h_unsuback");
-                    c->disconnect();
-                    return true;
-                });
             c->set_publish_handler(
                 [&chk, &c]
                 (MQTT_NS::optional<packet_id_t> packet_id,
@@ -1838,23 +1814,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->pubrec(*packet_id);
                     return true;
                 });
-            c->set_pubrel_handler(
-                [&chk, &c, g]
-                (packet_id_t packet_id) mutable {
-                    auto ret = MQTT_ORDERED(
-                        [&] {
-                            MQTT_CHK("h_pubrel1");
-                            c->pubcomp(packet_id);
-                        },
-                        [&] () {
-                            MQTT_CHK("h_pubrel2");
-                            c->pubcomp(packet_id);
-                            g.reset();
-                        }
-                    );
-                    BOOST_TEST(ret);
-                    return true;
-                });
             break;
         case MQTT_NS::protocol_version::v5:
             c->set_v5_connack_handler(
@@ -1879,13 +1838,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->pubrel(packet_id);
                     return true;
                 });
-            c->set_v5_pubcomp_handler(
-                [&chk, g]
-                (packet_id_t, MQTT_NS::v5::pubcomp_reason_code, MQTT_NS::v5::properties /*props*/) mutable {
-                    MQTT_CHK("h_pubcomp");
-                    g.reset();
-                    return true;
-                });
             c->set_v5_suback_handler(
                 [&chk, &c]
                 (packet_id_t, std::vector<MQTT_NS::v5::suback_reason_code> reasons, MQTT_NS::v5::properties /*props*/) {
@@ -1895,15 +1847,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->publish("topic1", "topic1_contents", MQTT_NS::qos::exactly_once);
                     return true;
                 });
-            c->set_v5_unsuback_handler(
-                [&chk, &c]
-                (packet_id_t, std::vector<MQTT_NS::v5::unsuback_reason_code> reasons, MQTT_NS::v5::properties /*props*/) {
-                    MQTT_CHK("h_unsuback");
-                    BOOST_TEST(reasons.size() == 1U);
-                    BOOST_TEST(reasons[0] == MQTT_NS::v5::unsuback_reason_code::success);
-                    c->disconnect();
-                    return true;
-                });
             c->set_v5_publish_handler(
                 [&chk, &c]
                 (MQTT_NS::optional<packet_id_t> packet_id,
@@ -1923,23 +1866,6 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
                     c->pubrec(*packet_id);
                     return true;
                 });
-            c->set_v5_pubrel_handler(
-                [&chk, &c, g]
-                (packet_id_t packet_id, MQTT_NS::v5::pubrel_reason_code, MQTT_NS::v5::properties /*props*/) mutable {
-                    auto ret = MQTT_ORDERED(
-                        [&] {
-                            MQTT_CHK("h_pubrel1");
-                            c->pubcomp(packet_id);
-                        },
-                        [&] () {
-                            MQTT_CHK("h_pubrel2");
-                            c->pubcomp(packet_id);
-                            g.reset();
-                        }
-                    );
-                    BOOST_TEST(ret);
-                    return true;
-                });
             break;
         default:
             BOOST_CHECK(false);
@@ -1947,18 +1873,17 @@ BOOST_AUTO_TEST_CASE( pub_qos2_sub_qos2_protocol_error_resend_pubrec ) {
         }
 
         c->set_close_handler(
-            [&chk, &finish]
+            []
             () {
-                MQTT_CHK("h_close");
-                finish();
+                BOOST_CHECK(false);
             });
         c->set_error_handler(
-            []
+            [&chk, &finish]
             (MQTT_NS::error_code) {
-                BOOST_CHECK(false);
+                MQTT_CHK("h_error");
+                finish();
             });
 
-        g.reset();
         c->connect();
         ioc.run();
         BOOST_TEST(chk.all());