[ISSUE #10214] Fix TOCTOU race condition in MQClientInstance.brokerVersionTable

Fixes #10214

Author: daguimu

client/src/main/java/org/apache/rocketmq/client/impl/factory/MQClientInstance.java

@@ -654,10 +654,8 @@ public boolean sendHeartbeatToBroker(long id, String brokerName, String addr, bo
     private boolean sendHeartbeatToBroker(long id, String brokerName, String addr, HeartbeatData heartbeatData) {
         try {
             int version = this.mQClientAPIImpl.sendHeartbeat(addr, heartbeatData, clientConfig.getMqClientApiTimeout());
-            if (!this.brokerVersionTable.containsKey(brokerName)) {
-                this.brokerVersionTable.put(brokerName, new ConcurrentHashMap<>(4));
-            }
-            this.brokerVersionTable.get(brokerName).put(addr, version);
+            this.brokerVersionTable.computeIfAbsent(brokerName, k -> new ConcurrentHashMap<>(4))
+                .put(addr, version);
             long times = this.sendHeartbeatTimesTotal.getAndIncrement();
             if (times % 20 == 0) {
                 log.info("send heart beat to broker[{} {} {}] success", brokerName, id, addr);
@@ -734,10 +732,8 @@ private boolean sendHeartbeatToBrokerV2(long id, String brokerName, String addr,
                 log.info("sendHeartbeatToAllBrokerV2 normal brokerName: {} subChange: {} brokerAddrHeartbeatFingerprintTable: {}", brokerName, heartbeatV2Result.isSubChange(), JSON.toJSONString(brokerAddrHeartbeatFingerprintTable));
             }
             version = heartbeatV2Result.getVersion();
-            if (!this.brokerVersionTable.containsKey(brokerName)) {
-                this.brokerVersionTable.put(brokerName, new ConcurrentHashMap<>(4));
-            }
-            this.brokerVersionTable.get(brokerName).put(addr, version);
+            this.brokerVersionTable.computeIfAbsent(brokerName, k -> new ConcurrentHashMap<>(4))
+                .put(addr, version);
             long times = this.sendHeartbeatTimesTotal.getAndIncrement();
             if (times % 20 == 0) {
                 log.info("send heart beat to broker[{} {} {}] success", brokerName, id, addr);
@@ -1301,9 +1297,11 @@ public FindBrokerResult findBrokerAddressInSubscribe(
     }
 
     private int findBrokerVersion(String brokerName, String brokerAddr) {
-        if (this.brokerVersionTable.containsKey(brokerName)) {
-            if (this.brokerVersionTable.get(brokerName).containsKey(brokerAddr)) {
-                return this.brokerVersionTable.get(brokerName).get(brokerAddr);
+        ConcurrentHashMap<String, Integer> brokerVersions = this.brokerVersionTable.get(brokerName);
+        if (brokerVersions != null) {
+            Integer version = brokerVersions.get(brokerAddr);
+            if (version != null) {
+                return version;
             }
         }
         //To do need to fresh the version

client/src/test/java/org/apache/rocketmq/client/impl/factory/MQClientInstanceTest.java

@@ -554,4 +554,74 @@ public void testSendHeartbeatToAllBrokerConcurrently() {
             fail("failed: " + e.getMessage());
         }
     }
+
+    @Test
+    public void testFindBrokerVersionWhenBrokerNameNotExist() throws Exception {
+        ConcurrentMap<String, ConcurrentHashMap<String, Integer>> brokerVersionTable =
+            (ConcurrentMap<String, ConcurrentHashMap<String, Integer>>) FieldUtils.readDeclaredField(
+                mqClientInstance, "brokerVersionTable", true);
+        brokerVersionTable.clear();
+
+        FindBrokerResult result = mqClientInstance.findBrokerAddressInSubscribe("nonExistentBroker", 0, false);
+        assertNull(result);
+    }
+
+    @Test
+    public void testFindBrokerVersionWhenBrokerAddrNotExist() throws Exception {
+        ConcurrentMap<String, ConcurrentHashMap<String, Integer>> brokerVersionTable =
+            (ConcurrentMap<String, ConcurrentHashMap<String, Integer>>) FieldUtils.readDeclaredField(
+                mqClientInstance, "brokerVersionTable", true);
+        brokerVersionTable.clear();
+        brokerVersionTable.put("broker-a", new ConcurrentHashMap<>());
+
+        FindBrokerResult result = mqClientInstance.findBrokerAddressInSubscribe("broker-a", 0, false);
+        assertNull(result);
+    }
+
+    @Test
+    public void testFindBrokerVersionConcurrentRemoval() throws Exception {
+        ConcurrentMap<String, ConcurrentHashMap<String, Integer>> brokerVersionTable =
+            (ConcurrentMap<String, ConcurrentHashMap<String, Integer>>) FieldUtils.readDeclaredField(
+                mqClientInstance, "brokerVersionTable", true);
+        brokerVersionTable.clear();
+
+        String brokerName = "broker-a";
+        ConcurrentHashMap<String, Integer> versionMap = new ConcurrentHashMap<>(4);
+        versionMap.put("127.0.0.1:10911", 401);
+        brokerVersionTable.put(brokerName, versionMap);
+
+        // Simulate concurrent removal between containsKey and get
+        // With the fix using local variable caching, this should not cause NPE
+        brokerVersionTable.remove(brokerName);
+
+        // findBrokerVersion is private, but it's called via findBrokerAddressInSubscribe
+        // The key test is that no NPE is thrown when brokerName is absent
+        FindBrokerResult result = mqClientInstance.findBrokerAddressInSubscribe(brokerName, 0, false);
+        assertNull(result);
+    }
+
+    @Test
+    public void testBrokerVersionTableComputeIfAbsent() throws Exception {
+        ConcurrentMap<String, ConcurrentHashMap<String, Integer>> brokerVersionTable =
+            (ConcurrentMap<String, ConcurrentHashMap<String, Integer>>) FieldUtils.readDeclaredField(
+                mqClientInstance, "brokerVersionTable", true);
+        brokerVersionTable.clear();
+
+        String brokerName = "broker-a";
+
+        // Use computeIfAbsent as the fix does
+        ConcurrentHashMap<String, Integer> inner = brokerVersionTable.computeIfAbsent(
+            brokerName, k -> new ConcurrentHashMap<>(4));
+        inner.put("127.0.0.1:10911", 401);
+
+        // Verify the entry was created atomically
+        assertNotNull(brokerVersionTable.get(brokerName));
+        assertEquals(Integer.valueOf(401), brokerVersionTable.get(brokerName).get("127.0.0.1:10911"));
+
+        // Calling computeIfAbsent again should return the same map, not create a new one
+        ConcurrentHashMap<String, Integer> inner2 = brokerVersionTable.computeIfAbsent(
+            brokerName, k -> new ConcurrentHashMap<>(4));
+        assertEquals(inner, inner2);
+        assertEquals(Integer.valueOf(401), inner2.get("127.0.0.1:10911"));
+    }
 }