From 9e1fed1be72f982a1408f064abd9da788ad17186 Mon Sep 17 00:00:00 2001
From: Hylke van der Schaaf <hylke.vds@gmail.com>
Date: Sat, 17 Aug 2024 18:06:30 +0200
Subject: [PATCH] Fixes 841: Incorrect reference used in compareAndSet in
 CTrie.cleanTomb. (#842)

* Fixes 841: cleanTomb used compareAndSet to update a reference, but
incorrectly re-fetched the 'original' instead of using the version that
was used to make the copy.

* Fixed possible race condition in cleanTomb that could result in the removal of a live node
---
 ChangeLog.txt                                 |  1 +
 .../moquette/broker/subscriptions/CNode.java  |  4 ++--
 .../moquette/broker/subscriptions/CTrie.java  | 19 ++++++++++++-------
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/ChangeLog.txt b/ChangeLog.txt
index 2128e61a3..49043c44b 100644
--- a/ChangeLog.txt
+++ b/ChangeLog.txt
@@ -1,4 +1,5 @@
 Version 0.18-SNAPSHOT:
+   [fix] Incorrect reference used in compareAndSet in CTrie.cleanTomb. (#841)
    [feature] Implement response-information property for request-response flow. (#840)
    [fix] Optimised page file opening for disk-based queues. (#837)
    [feature] Manage payload format indicator property, when set verify payload format. (#826)
diff --git a/broker/src/main/java/io/moquette/broker/subscriptions/CNode.java b/broker/src/main/java/io/moquette/broker/subscriptions/CNode.java
index 00c813636..faa349aa7 100644
--- a/broker/src/main/java/io/moquette/broker/subscriptions/CNode.java
+++ b/broker/src/main/java/io/moquette/broker/subscriptions/CNode.java
@@ -99,9 +99,9 @@ public void add(INode newINode) {
         }
     }
 
-    public void remove(INode node) {
+    public INode remove(INode node) {
         int idx = findIndexForToken(node.mainNode().token);
-        this.children.remove(idx);
+        return this.children.remove(idx);
     }
 
     private List<Subscription> sharedSubscriptions() {
diff --git a/broker/src/main/java/io/moquette/broker/subscriptions/CTrie.java b/broker/src/main/java/io/moquette/broker/subscriptions/CTrie.java
index 874b0f49a..a71d67961 100644
--- a/broker/src/main/java/io/moquette/broker/subscriptions/CTrie.java
+++ b/broker/src/main/java/io/moquette/broker/subscriptions/CTrie.java
@@ -360,10 +360,7 @@ private Action remove(String clientId, Topic topic, INode inode, INode iParent,
             }
         }
         if (cnode instanceof TNode) {
-            // this inode is a tomb, has no clients and should be cleaned up
-            // Because we implemented cleanTomb below, this should be rare, but possible
-            // Consider calling cleanTomb here too
-            return Action.OK;
+            return cleanTomb(inode, iParent);
         }
         if (cnode.containsOnly(clientId) && topic.isEmpty() && cnode.allChildren().isEmpty()) {
             // last client to leave this node, AND there are no downstream children, remove via TNode tomb
@@ -393,9 +390,17 @@ private Action remove(String clientId, Topic topic, INode inode, INode iParent,
      * @return REPEAT if this method wasn't successful or OK.
      */
     private Action cleanTomb(INode inode, INode iParent) {
-        CNode updatedCnode = iParent.mainNode().copy();
-        updatedCnode.remove(inode);
-        return iParent.compareAndSet(iParent.mainNode(), updatedCnode) ? Action.OK : Action.REPEAT;
+        CNode origCnode = iParent.mainNode();
+        CNode updatedCnode = origCnode.copy();
+        INode removed = updatedCnode.remove(inode);
+        if (removed == inode) {
+            return iParent.compareAndSet(origCnode, updatedCnode) ? Action.OK : Action.REPEAT;
+        } else {
+            // The node removed (from the copy!) was not the node we expected to remove.
+            // Probably because another thread replaced the TNode with a live node, so
+            // we don't need to clean it and can return success.
+            return Action.OK;
+        }
     }
 
     public int size() {