From e67dcab05fc72730b27b3d1e0e853d035de9ac22 Mon Sep 17 00:00:00 2001
From: Greg Medding <gregmedd@users.noreply.github.com>
Date: Wed, 24 Apr 2024 13:01:45 -0700
Subject: [PATCH] Fixing locking around queryMap_ (#42)

The queryMap_ member is a std::map that can be accessed from multiple
threads. Locking was added in previous commits to guard modify
operations, but read-only operations were not guarded. This resulted in
unpredictable behavior.
---
 lib/src/zenohUTransport.cpp | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/lib/src/zenohUTransport.cpp b/lib/src/zenohUTransport.cpp
index 90aa911..0b80894 100644
--- a/lib/src/zenohUTransport.cpp
+++ b/lib/src/zenohUTransport.cpp
@@ -185,12 +185,19 @@ UCode ZenohUTransport::sendQueryable(const UMessage &message) noexcept {
 
     auto uuidStr = UuidSerializer::serializeToString(message.attributes().reqid());
 
-    if (queryMap_.find(uuidStr) == queryMap_.end()) {
-        spdlog::error("failed to find UUID = {}", uuidStr);
-        return UCode::UNAVAILABLE;
-    }
+    z_owned_query_t query;
+    {
+        std::unique_lock<std::mutex> lock(queryMapMutex_);
 
-    auto query = queryMap_[uuidStr];
+        if (auto query_it = queryMap_.find(uuidStr);
+                query_it == queryMap_.end()) {
+            spdlog::error("failed to find UUID = {}", uuidStr);
+            return UCode::UNAVAILABLE;
+
+        } else {
+            query = query_it->second;
+        }
+    }
 
     z_query_reply_options_t options = z_query_reply_options_default();
 
@@ -228,9 +235,10 @@ UCode ZenohUTransport::sendQueryable(const UMessage &message) noexcept {
 
     spdlog::debug("replied on query with uid = {}", uuidStr);
     /* once replied remove the uuid from the map, as it cannot be reused */
-    std::unique_lock<std::mutex> lock(queryMapMutex_);
-    queryMap_.erase(uuidStr);
-    lock.unlock();
+    {
+        std::unique_lock<std::mutex> lock(queryMapMutex_);
+        queryMap_.erase(uuidStr);
+    }
 
     z_drop(z_move(map));
 
@@ -459,16 +467,15 @@ void ZenohUTransport::QueryHandler(const z_query_t *query, void *arg) {
 
     auto uuidStr = UuidSerializer::serializeToString(attributes.id());
 
-    instance->queryMap_[uuidStr] = z_query_clone(query);
-
     if (UMessageType::UMESSAGE_TYPE_REQUEST != attributes.type()) {
        spdlog::error("Wrong message type = {}", static_cast<int>(attributes.type()));
        return;
     }
 
-    std::unique_lock<std::mutex> lock(instance->queryMapMutex_);
-    instance->queryMap_[uuidStr] = z_query_clone(query);
-    lock.unlock();
+    {
+        std::unique_lock<std::mutex> lock(instance->queryMapMutex_);
+        instance->queryMap_[uuidStr] = z_query_clone(query);
+    }
 
     UMessage message(payload, attributes);