From 4876d85a9c84acbefd869b402c8b34dd961cec72 Mon Sep 17 00:00:00 2001 From: Senrian <47714364+Senrian@users.noreply.github.com> Date: Mon, 30 Mar 2026 17:21:47 +0800 Subject: [PATCH] Fix TOCTOU race condition in brokerVersionTable access Issue: #10214 Replace non-atomic check-then-act pattern with thread-safe operations: - sendHeartbeatToBroker(): use computeIfAbsent instead of containsKey+put - sendHeartbeatToBrokerV2(): same fix - findBrokerVersion(): use local variable to avoid NPE if entry removed between calls --- .../client/impl/factory/MQClientInstance.java | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/client/src/main/java/org/apache/rocketmq/client/impl/factory/MQClientInstance.java b/client/src/main/java/org/apache/rocketmq/client/impl/factory/MQClientInstance.java index e0b28fef646..4e92e918bc1 100644 --- a/client/src/main/java/org/apache/rocketmq/client/impl/factory/MQClientInstance.java +++ b/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 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