fix use-after-free bug

markdroth · Feb 7, 2025 · dcfd5ef · dcfd5ef
1 parent e6cde8c
commit dcfd5ef
Showing 1 changed file with 48 additions and 40 deletions.
diff --git a/src/core/xds/xds_client/xds_client.cc b/src/core/xds/xds_client/xds_client.cc
@@ -1769,60 +1769,68 @@ void XdsClient::CancelResourceWatch(const XdsResourceType* type,
 void XdsClient::MaybeRemoveUnsubscribedCacheEntriesLocked(
     XdsChannel* xds_channel) {
   // Look at each authority for which xds_channel is the last channel.
-  for (auto& [_, authority_state] : authority_state_map_) {
-    if (authority_state.xds_channels.back() != xds_channel) continue;
-    // Look at each resource type.
-    for (auto type_it = authority_state.type_map.begin();
-         type_it != authority_state.type_map.end();) {
-      auto& resource_map = type_it->second;
-      // Remove the cache entry for any resource without watchers.
-      for (auto resource_it = resource_map.begin();
-           resource_it != resource_map.end();) {
-        ResourceState& resource_state = resource_it->second;
-        if (!resource_state.HasWatchers()) {
-          resource_map.erase(resource_it++);
+  for (auto authority_it = authority_state_map_.begin();
+       authority_it != authority_state_map_.end();) {
+    AuthorityState& authority_state = authority_it->second;
+    if (authority_state.xds_channels.back() == xds_channel) {
+      // Look at each resource type.
+      for (auto type_it = authority_state.type_map.begin();
+           type_it != authority_state.type_map.end();) {
+        auto& resource_map = type_it->second;
+        // Remove the cache entry for any resource without watchers.
+        for (auto resource_it = resource_map.begin();
+             resource_it != resource_map.end();) {
+          ResourceState& resource_state = resource_it->second;
+          if (!resource_state.HasWatchers()) {
+            resource_map.erase(resource_it++);
+          } else {
+            ++resource_it;
+          }
+        }
+        // Clean up empty entries in the map.
+        if (resource_map.empty()) {
+          authority_state.type_map.erase(type_it++);
         } else {
-          ++resource_it;
+          ++type_it;
         }
       }
-      // Clean up empty entries in the map.
-      if (resource_map.empty()) {
-        authority_state.type_map.erase(type_it++);
-      } else {
-        ++type_it;
-      }
     }
     if (authority_state.type_map.empty()) {
-      authority_state.xds_channels.clear();
+      authority_state_map_.erase(authority_it++);
+    } else {
+      ++authority_it;
     }
   }
 }
 
 void XdsClient::MaybeRemoveUnsubscribedCacheEntriesForTypeLocked(
     XdsChannel* xds_channel, const XdsResourceType* type) {
   // Look at each authority for which xds_channel is the last channel.
-  for (auto& [_, authority_state] : authority_state_map_) {
-    if (authority_state.xds_channels.back() != xds_channel) continue;
-    // Find type map.
-    auto type_it = authority_state.type_map.find(type);
-    if (type_it == authority_state.type_map.end()) return;
-    auto& resource_map = type_it->second;
-    // Remove the cache entry for any resource without watchers.
-    for (auto resource_it = resource_map.begin();
-         resource_it != resource_map.end();) {
-      ResourceState& resource_state = resource_it->second;
-      if (!resource_state.HasWatchers()) {
-        resource_map.erase(resource_it++);
-      } else {
-        ++resource_it;
+  for (auto authority_it = authority_state_map_.begin();
+       authority_it != authority_state_map_.end();) {
+    AuthorityState& authority_state = authority_it->second;
+    if (authority_state.xds_channels.back() == xds_channel) {
+      // Find type map.
+      auto type_it = authority_state.type_map.find(type);
+      if (type_it == authority_state.type_map.end()) return;
+      auto& resource_map = type_it->second;
+      // Remove the cache entry for any resource without watchers.
+      for (auto resource_it = resource_map.begin();
+           resource_it != resource_map.end();) {
+        ResourceState& resource_state = resource_it->second;
+        if (!resource_state.HasWatchers()) {
+          resource_map.erase(resource_it++);
+        } else {
+          ++resource_it;
+        }
       }
+      // Clean up empty entries in the map.
+      if (resource_map.empty()) authority_state.type_map.erase(type_it);
     }
-    // Clean up empty entries in the map.
-    if (resource_map.empty()) {
-      authority_state.type_map.erase(type_it);
-      if (authority_state.type_map.empty()) {
-        authority_state.xds_channels.clear();
-      }
+    if (authority_state.type_map.empty()) {
+      authority_state_map_.erase(authority_it++);
+    } else {
+      ++authority_it;
     }
   }
 }