Fix from review

Author: embhorn

src/mqtt_broker.c

@@ -3417,6 +3417,7 @@ WOLFMQTT_LOCAL void BrokerOrphan_DropFull(MqttBroker* broker,
                     WOLFMQTT_FREE(sp->client_id);
                 }
                 WOLFMQTT_FREE(sp);
+                broker->subs_gen++;
             }
             else {
                 prev = sp;
@@ -3902,6 +3903,7 @@ static void BrokerSubs_RemoveClient(MqttBroker* broker, BrokerClient* bc)
                 WOLFMQTT_FREE(cur->client_id);
             }
             WOLFMQTT_FREE(cur);
+            broker->subs_gen++;
         }
         else {
             prev = cur;
@@ -4053,6 +4055,11 @@ static int BrokerSubs_Add(MqttBroker* broker, BrokerClient* bc,
         sub->filter[filter_len] = '\0';
         sub->next = broker->subs;
         broker->subs = sub;
+        /* Bump on add too so subs_gen reflects every structural change to the
+         * list, not only removals. Also narrows the ABA window where a free and
+         * a same-address add in one fan-out iteration could otherwise pass the
+         * successor-still-linked check. */
+        broker->subs_gen++;
     }
     else if (sub != NULL) {
         WOLFMQTT_FREE(sub);
@@ -4138,6 +4145,7 @@ static void BrokerSubs_Remove(MqttBroker* broker, BrokerClient* bc,
                 WOLFMQTT_FREE(cur->client_id);
             }
             WOLFMQTT_FREE(cur);
+            broker->subs_gen++;
             if (bc->sub_count > 0) {
                 bc->sub_count--;
             }
@@ -4263,6 +4271,7 @@ static void BrokerSubs_RemoveByClientId(MqttBroker* broker,
                 WOLFMQTT_FREE(cur->client_id);
             }
             WOLFMQTT_FREE(cur);
+            broker->subs_gen++;
         }
         else {
             prev = cur;
@@ -5394,6 +5403,7 @@ static void BrokerClient_PublishWillImmediate(MqttBroker* broker,
 #else
     BrokerSub* sub;
     BrokerSub* next_sub = NULL;
+    word32 subs_gen_snapshot = 0;
 #endif
 
     if (broker == NULL || topic == NULL) {
@@ -5426,8 +5436,10 @@ static void BrokerClient_PublishWillImmediate(MqttBroker* broker,
         /* Snapshot the successor before any MqttPacket_Write: a WS fan-out
          * write can drive an lws_service spin whose re-entrant CLOSED frees
          * this client's BrokerSub nodes, so reading sub->next afterwards
-         * would dereference a freed node. */
+         * would dereference a freed node. Also snapshot the subs generation so
+         * next_sub is only re-validated when a free actually happened. */
         next_sub = sub->next;
+        subs_gen_snapshot = broker->subs_gen;
 #endif
         if (sub->client != NULL && sub->client->protocol_level != 0 &&
 #ifdef WOLFMQTT_STATIC_MEMORY
@@ -5534,6 +5546,13 @@ static void BrokerClient_PublishWillImmediate(MqttBroker* broker,
             {
                 MqttPublish out_pub;
                 int enc_rc, wr_rc;
+                /* Cache the client before the write: MqttPacket_Write can drive
+                 * a re-entrant WS close that frees this subscriber's BrokerSub
+                 * nodes (this `sub`), while the BrokerClient itself survives via
+                 * deferred removal. Reaching it through the freed `sub` after
+                 * the write would be a use-after-free; go through wc instead,
+                 * mirroring BrokerHandle_Publish. */
+                BrokerClient* wc = sub->client;
                 XMEMSET(&out_pub, 0, sizeof(out_pub));
                 out_pub.topic_name = (char*)topic;
                 out_pub.qos = eff_qos;
@@ -5545,16 +5564,16 @@ static void BrokerClient_PublishWillImmediate(MqttBroker* broker,
                     out_pub.packet_id = BrokerNextPacketId(broker);
                 }
 #ifdef WOLFMQTT_V5
-                out_pub.protocol_level = sub->client->protocol_level;
+                out_pub.protocol_level = wc->protocol_level;
 #endif
-                enc_rc = MqttEncode_Publish(sub->client->tx_buf,
-                    BROKER_CLIENT_TX_SZ(sub->client), &out_pub, 0);
+                enc_rc = MqttEncode_Publish(wc->tx_buf,
+                    BROKER_CLIENT_TX_SZ(wc), &out_pub, 0);
                 if (enc_rc > 0) {
-                    wr_rc = MqttPacket_Write(&sub->client->client,
-                            sub->client->tx_buf, enc_rc);
+                    wr_rc = MqttPacket_Write(&wc->client,
+                            wc->tx_buf, enc_rc);
                     /* Scrub tx_buf unless still in-progress (CONTINUE). */
                     if (wr_rc != MQTT_CODE_CONTINUE) {
-                        BROKER_FORCE_ZERO(sub->client->tx_buf, enc_rc);
+                        BROKER_FORCE_ZERO(wc->tx_buf, enc_rc);
                     }
                 }
             }
@@ -5586,9 +5605,11 @@ static void BrokerClient_PublishWillImmediate(MqttBroker* broker,
         }
 #endif
 #ifndef WOLFMQTT_STATIC_MEMORY
-        /* The write above can drive a re-entrant WS close that frees next_sub;
-         * stop the walk if that snapshot is no longer linked. */
-        if (next_sub != NULL && !BrokerSubs_StillLinked(broker, next_sub)) {
+        /* The write above can drive a re-entrant WS close that frees next_sub.
+         * Only re-validate it when a subscription was actually removed during
+         * the write (generation changed), so the common case stays O(1). */
+        if (next_sub != NULL && broker->subs_gen != subs_gen_snapshot &&
+                !BrokerSubs_StillLinked(broker, next_sub)) {
             break;
         }
         sub = next_sub;
@@ -7177,6 +7198,7 @@ static int BrokerHandle_Publish(BrokerClient* bc, int rx_len,
 #ifndef WOLFMQTT_STATIC_MEMORY
         BrokerSub* sub = broker->subs;
         BrokerSub* next_sub = NULL;
+        word32 subs_gen_snapshot = 0;
 #endif
         /* Fan out to matching subscribers */
 #ifdef WOLFMQTT_STATIC_MEMORY
@@ -7188,8 +7210,11 @@ static int BrokerHandle_Publish(BrokerClient* bc, int rx_len,
             /* Snapshot the successor before any MqttPacket_Write: a fan-out
              * write can drive an lws_service spin that frees this client's
              * BrokerSub nodes re-entrantly (LWS_CALLBACK_CLOSED), so reading
-             * sub->next afterwards would dereference a freed node. */
+             * sub->next afterwards would dereference a freed node. Also snapshot
+             * the subs generation so next_sub is only re-validated when a free
+             * actually happened. */
             next_sub = sub->next;
+            subs_gen_snapshot = broker->subs_gen;
 #endif
             if (sub->client != NULL &&
                 sub->client->protocol_level != 0 &&
@@ -7422,8 +7447,11 @@ static int BrokerHandle_Publish(BrokerClient* bc, int rx_len,
                 }
             }
             /* The write above can drive a re-entrant WS close that frees
-             * next_sub; stop the walk if that snapshot is no longer linked. */
-            if (next_sub != NULL && !BrokerSubs_StillLinked(broker, next_sub)) {
+             * next_sub. Only re-validate it when a subscription was actually
+             * removed during the write (generation changed), so the common
+             * case stays O(1). */
+            if (next_sub != NULL && broker->subs_gen != subs_gen_snapshot &&
+                    !BrokerSubs_StillLinked(broker, next_sub)) {
                 break;
             }
             sub = next_sub;
@@ -8564,6 +8592,7 @@ static void BrokerSubs_FreeAll(MqttBroker* broker)
             WOLFMQTT_FREE(broker->subs->client_id);
         }
         WOLFMQTT_FREE(broker->subs);
+        broker->subs_gen++;
         broker->subs = next;
     }
 #endif

src/mqtt_client.c

@@ -610,6 +610,11 @@ void MqttClient_RespList_Remove(MqttClient *client, MqttPendResp *rmResp)
         }
     }
     if (tmpResp) {
+    #ifdef WOLFMQTT_V5
+        /* Drop the reserved-quota back-reference so a recycled publish object
+         * cannot carry a stale pointer into a future pending response. */
+        tmpResp->recvQuotaStat = NULL;
+    #endif
         /* Fix up the first and last pointers */
         if (client->firstPendResp == tmpResp) {
             client->firstPendResp = tmpResp->next;
@@ -2916,7 +2921,14 @@ static int MqttPublishMsg(MqttClient *client, MqttPublish *publish,
                         publish->packet_id);
                 #ifndef WOLFMQTT_NONBLOCK
                     if (rc == MQTT_CODE_CONTINUE) {
-                        /* mark success, let other thread handle response */
+                        /* No non-blocking re-entry will complete this ack, so
+                         * report success now. Release the reserved quota since
+                         * the ack will no longer be tracked, and the pending
+                         * response is unlinked below so no reference to the
+                         * caller's publish object remains. */
+                    #ifdef WOLFMQTT_V5
+                        MqttClient_RestoreRecvQuota(client, publish);
+                    #endif
                         rc = MQTT_CODE_SUCCESS;
                     }
                 #endif
@@ -2968,13 +2980,13 @@ static int MqttPublishMsg(MqttClient *client, MqttPublish *publish,
                     break;
             #endif
             #ifdef WOLFMQTT_MULTITHREAD
-                /* Leave a write-only publish's pending response in the list: a
-                 * build without WOLFMQTT_NONBLOCK reports success here without
-                 * the ack, and unlinking it now would prevent the reading thread
-                 * from completing it and releasing the reserved Receive Maximum
-                 * unit, leaking quota. The reader removes it when the ack
-                 * arrives. */
-                if (!writeOnly && wm_SemLock(&client->lockClient) == 0) {
+                /* Remove the pending response before returning: a caller told
+                 * SUCCESS may free its publish object, so no reference to it may
+                 * remain in the list. A write-only publish under
+                 * WOLFMQTT_NONBLOCK returned CONTINUE above and broke out before
+                 * here, leaving its entry for the reading thread to complete and
+                 * remove. */
+                if (wm_SemLock(&client->lockClient) == 0) {
                     MqttClient_RespList_Remove(client, &publish->pendResp);
                     wm_SemUnlock(&client->lockClient);
                 }
@@ -3867,6 +3879,15 @@ int MqttClient_CancelMessage(MqttClient *client, MqttObject* msg)
                 MqttPacket_TypeDesc(tmpResp->packet_type),
                 tmpResp->packet_type, tmpResp->packet_id,
                 tmpResp->packetProcessing, tmpResp->packetDone);
+        #endif
+        #ifdef WOLFMQTT_V5
+            /* Release any Receive Maximum unit this response reserved so
+             * cancelling a QoS>0 (e.g. write-only) publish does not leak it.
+             * lockClient is held; idempotent via recvQuotaHeld. */
+            if (tmpResp->recvQuotaStat != NULL) {
+                MqttClient_RecvQuotaRelease_Locked(client,
+                    tmpResp->recvQuotaStat);
+            }
         #endif
             MqttClient_RespList_Remove(client, tmpResp);
             break;
@@ -3972,6 +3993,15 @@ int MqttClient_NetDisconnect(MqttClient *client)
                 MqttPacket_TypeDesc(tmpResp->packet_type),
                 tmpResp->packet_type, tmpResp->packet_id,
                 tmpResp->packetProcessing, tmpResp->packetDone);
+        #endif
+        #ifdef WOLFMQTT_V5
+            /* Release any reserved Receive Maximum unit as the list is purged
+             * on disconnect, so it is not left held. Idempotent via
+             * recvQuotaHeld; lockClient is held. */
+            if (tmpResp->recvQuotaStat != NULL) {
+                MqttClient_RecvQuotaRelease_Locked(client,
+                    tmpResp->recvQuotaStat);
+            }
         #endif
             MqttClient_RespList_Remove(client, tmpResp);
         }

tests/test_broker_connect.c

@@ -1239,6 +1239,74 @@ static int first_disconnect_reason(const byte* buf, size_t len)
  *
  * Pre-fix the broker fanned out twice and the subscriber would see two
  * forwarded PUBLISHes, breaking exactly-once delivery. */
+#ifndef WOLFMQTT_STATIC_MEMORY
+/* The fan-out loop re-validates a snapshotted successor only when the broker's
+ * subscription generation changed during the write, keeping the common case
+ * O(1). This pins two things: an ordinary PUBLISH is still delivered to a
+ * matching subscriber, and removing a subscription (UNSUBSCRIBE) bumps the
+ * generation, so a re-entrant free during a later fan-out is still detected. */
+TEST(fanout_subs_generation_bumped_on_unsubscribe)
+{
+    MqttBroker broker;
+    MqttBrokerNet net;
+    int i;
+    int sub_pubs;
+    word32 gen_after_sub;
+    /* CONNECT subscriber "A". */
+    static const byte connect_sub[] = {
+        0x10, 0x0D, 0x00, 0x04, 'M', 'Q', 'T', 'T', 0x04, 0x02, 0x00, 0x3C,
+        0x00, 0x01, 'A'
+    };
+    /* CONNECT publisher "B". */
+    static const byte connect_pub[] = {
+        0x10, 0x0D, 0x00, 0x04, 'M', 'Q', 'T', 'T', 0x04, 0x02, 0x00, 0x3C,
+        0x00, 0x01, 'B'
+    };
+    /* SUBSCRIBE packet_id=1, filter "x", QoS 0. */
+    static const byte subscribe_x[] = {
+        0x82, 0x06, 0x00, 0x01, 0x00, 0x01, 'x', 0x00
+    };
+    /* PUBLISH QoS 0, topic "x", payload "hi". remain = 2+1+2 = 5. */
+    static const byte publish_x[] = {
+        0x30, 0x05, 0x00, 0x01, 'x', 'h', 'i'
+    };
+    /* UNSUBSCRIBE packet_id=2, filter "x". remain = 2 + (2+1) = 5. */
+    static const byte unsubscribe_x[] = {
+        0xA2, 0x05, 0x00, 0x02, 0x00, 0x01, 'x'
+    };
+
+    install_mock_net(&net);
+    XMEMSET(&broker, 0, sizeof(broker));
+    ASSERT_EQ(MQTT_CODE_SUCCESS, MqttBroker_Init(&broker, &net));
+    ASSERT_EQ(MQTT_CODE_SUCCESS, MqttBroker_Start(&broker));
+
+    reset_mock_clients(2);
+    mock_client_input_append(0, connect_sub, sizeof(connect_sub));
+    mock_client_input_append(0, subscribe_x, sizeof(subscribe_x));
+    mock_client_input_append(1, connect_pub, sizeof(connect_pub));
+    mock_client_input_append(1, publish_x, sizeof(publish_x));
+    for (i = 0; i < 32; i++) {
+        MqttBroker_Step(&broker);
+    }
+
+    /* The common-case fan-out delivered the PUBLISH to the subscriber. */
+    sub_pubs = count_packets_of_type(g_clients[0].out_buf,
+        g_clients[0].out_len, MQTT_PACKET_TYPE_PUBLISH);
+    ASSERT_EQ(1, sub_pubs);
+    gen_after_sub = broker.subs_gen;
+
+    /* UNSUBSCRIBE removes the BrokerSub and must bump the generation. */
+    mock_client_input_append(0, unsubscribe_x, sizeof(unsubscribe_x));
+    for (i = 0; i < 16; i++) {
+        MqttBroker_Step(&broker);
+    }
+    ASSERT_TRUE(broker.subs_gen != gen_after_sub);
+
+    MqttBroker_Stop(&broker);
+    MqttBroker_Free(&broker);
+}
+#endif /* !WOLFMQTT_STATIC_MEMORY */
+
 TEST(qos2_duplicate_publish_dedup)
 {
     MqttBroker broker;
@@ -4976,6 +5044,9 @@ int main(int argc, char** argv)
 #ifdef WOLFMQTT_V5
     RUN_TEST(connect_v5_emptyid_assigned_id_emitted);
     RUN_TEST(connect_v5_emptyid_clean0_accepted);
+#endif
+#ifndef WOLFMQTT_STATIC_MEMORY
+    RUN_TEST(fanout_subs_generation_bumped_on_unsubscribe);
 #endif
     RUN_TEST(qos2_duplicate_publish_dedup);
     RUN_TEST(qos2_phantom_dup_publish_is_fresh);