From 90b3c88fe21a9166c00c4d03d5994df0164b11fc Mon Sep 17 00:00:00 2001
From: ZHANG Dapeng <zdapeng@google.com>
Date: Thu, 3 Oct 2019 15:58:18 -0700
Subject: [PATCH] api: avoid infinite loop in handleResolvedAddresses

If a `LoadBalancer` implementation does not override `handleResolvedAddressGroups()`, or overrides `handleResolvedAddressGroups()` but calls `super.handleResolvedAddressGroups()` at the beginning or the end, it will be trapped in an infinite loop.
---
 api/src/main/java/io/grpc/LoadBalancer.java   | 16 ++++++---
 .../test/java/io/grpc/LoadBalancerTest.java   | 36 +++++++++++++++++++
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/api/src/main/java/io/grpc/LoadBalancer.java b/api/src/main/java/io/grpc/LoadBalancer.java
index 612d25151f..ebae08b23c 100644
--- a/api/src/main/java/io/grpc/LoadBalancer.java
+++ b/api/src/main/java/io/grpc/LoadBalancer.java
@@ -117,6 +117,8 @@ public abstract class LoadBalancer {
   public static final Attributes.Key<Map<String, ?>> ATTR_LOAD_BALANCING_CONFIG =
       Attributes.Key.create("io.grpc.LoadBalancer.loadBalancingConfig");
 
+  private int recursionCount;
+
   /**
    * Handles newly resolved server groups and metadata attributes from name resolution system.
    * {@code servers} contained in {@link EquivalentAddressGroup} should be considered equivalent
@@ -133,8 +135,11 @@ public abstract class LoadBalancer {
   public void handleResolvedAddressGroups(
       List<EquivalentAddressGroup> servers,
       @NameResolver.ResolutionResultAttr Attributes attributes) {
-    handleResolvedAddresses(
-        ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
+    if (recursionCount++ == 0) {
+      handleResolvedAddresses(
+          ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
+    }
+    recursionCount = 0;
   }
 
   /**
@@ -149,8 +154,11 @@ public abstract class LoadBalancer {
    */
   @SuppressWarnings("deprecation")
   public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
-    handleResolvedAddressGroups(
-        resolvedAddresses.getAddresses(), resolvedAddresses.getAttributes());
+    if (recursionCount++ == 0) {
+      handleResolvedAddressGroups(
+          resolvedAddresses.getAddresses(), resolvedAddresses.getAttributes());
+    }
+    recursionCount = 0;
   }
 
   /**
diff --git a/api/src/test/java/io/grpc/LoadBalancerTest.java b/api/src/test/java/io/grpc/LoadBalancerTest.java
index e1a7b9a5d7..6bda67cd4d 100644
--- a/api/src/test/java/io/grpc/LoadBalancerTest.java
+++ b/api/src/test/java/io/grpc/LoadBalancerTest.java
@@ -25,6 +25,7 @@ import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
 import java.net.SocketAddress;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
@@ -360,6 +361,41 @@ public class LoadBalancerTest {
     assertThat(attrsCapture.get()).isEqualTo(attrs);
   }
 
+  @Deprecated
+  @Test
+  public void handleResolvedAddresses_noInfiniteLoop() {
+    final List<List<EquivalentAddressGroup>> serversCapture = new ArrayList<>();
+    final List<Attributes> attrsCapture = new ArrayList<>();
+
+    LoadBalancer balancer = new LoadBalancer() {
+      @Override
+      public void handleResolvedAddressGroups(
+          List<EquivalentAddressGroup> servers, Attributes attrs) {
+        serversCapture.add(servers);
+        attrsCapture.add(attrs);
+        super.handleResolvedAddressGroups(servers, attrs);
+      }
+
+      @Override
+      public void handleNameResolutionError(Status error) {
+      }
+
+      @Override
+      public void shutdown() {
+      }
+    };
+
+    List<EquivalentAddressGroup> servers = Arrays.asList(
+        new EquivalentAddressGroup(new SocketAddress(){}),
+        new EquivalentAddressGroup(new SocketAddress(){}));
+    balancer.handleResolvedAddresses(
+        ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attrs).build());
+    assertThat(serversCapture).hasSize(1);
+    assertThat(attrsCapture).hasSize(1);
+    assertThat(serversCapture.get(0)).isEqualTo(servers);
+    assertThat(attrsCapture.get(0)).isEqualTo(attrs);
+  }
+
   private static class NoopHelper extends LoadBalancer.Helper {
     @Override
     public ManagedChannel createOobChannel(EquivalentAddressGroup eag, String authority) {