From dcfd5ef27ab0e6eec2e2d73ef3ef4d15d1516193 Mon Sep 17 00:00:00 2001
From: "Mark D. Roth" <roth@google.com>
Date: Fri, 7 Feb 2025 17:00:58 +0000
Subject: [PATCH] fix use-after-free bug

---
 src/core/xds/xds_client/xds_client.cc | 88 +++++++++++++++------------
 1 file changed, 48 insertions(+), 40 deletions(-)

diff --git a/src/core/xds/xds_client/xds_client.cc b/src/core/xds/xds_client/xds_client.cc
index 9f8a3028ae5d8..e056f9e234921 100644
--- a/src/core/xds/xds_client/xds_client.cc
+++ b/src/core/xds/xds_client/xds_client.cc
@@ -1769,31 +1769,36 @@ 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;
     }
   }
 }
@@ -1801,28 +1806,31 @@ void XdsClient::MaybeRemoveUnsubscribedCacheEntriesLocked(
 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;
     }
   }
 }