From 10873b30bf245efc365a957b425da2e15f47a560 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Wed, 22 Nov 2017 13:59:20 -0800
Subject: [PATCH] Fix panics on balancer and resolver updates (#1684)

 - NewAddress with empty list (addrConn with an empty address list)
 - NewServiceConfig to switch balancer before the first balancer is built
---
 balancer_conn_wrappers.go |  8 ++++++++
 clientconn.go             |  8 +++++++-
 clientconn_test.go        | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/balancer_conn_wrappers.go b/balancer_conn_wrappers.go
index ebfee4a8..c747e92b 100644
--- a/balancer_conn_wrappers.go
+++ b/balancer_conn_wrappers.go
@@ -19,6 +19,7 @@
 package grpc
 
 import (
+	"fmt"
 	"sync"
 
 	"google.golang.org/grpc/balancer"
@@ -183,6 +184,9 @@ func (ccb *ccBalancerWrapper) handleResolvedAddrs(addrs []resolver.Address, err
 }
 
 func (ccb *ccBalancerWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	if len(addrs) <= 0 {
+		return nil, fmt.Errorf("grpc: cannot create SubConn with empty address list")
+	}
 	ac, err := ccb.cc.newAddrConn(addrs)
 	if err != nil {
 		return nil, err
@@ -223,6 +227,10 @@ type acBalancerWrapper struct {
 func (acbw *acBalancerWrapper) UpdateAddresses(addrs []resolver.Address) {
 	acbw.mu.Lock()
 	defer acbw.mu.Unlock()
+	if len(addrs) <= 0 {
+		acbw.ac.tearDown(errConnDrain)
+		return
+	}
 	if !acbw.ac.tryUpdateAddrs(addrs) {
 		cc := acbw.ac.cc
 		acbw.ac.mu.Lock()
diff --git a/clientconn.go b/clientconn.go
index ae605bc3..9d68df9f 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -642,6 +642,8 @@ func (cc *ClientConn) handleResolvedAddrs(addrs []resolver.Address, err error) {
 }
 
 // switchBalancer starts the switching from current balancer to the balancer with name.
+//
+// Caller must hold cc.mu.
 func (cc *ClientConn) switchBalancer(name string) {
 	if cc.conns == nil {
 		return
@@ -659,7 +661,9 @@ func (cc *ClientConn) switchBalancer(name string) {
 
 	// TODO(bar switching) change this to two steps: drain and close.
 	// Keep track of sc in wrapper.
-	cc.balancerWrapper.close()
+	if cc.balancerWrapper != nil {
+		cc.balancerWrapper.close()
+	}
 
 	builder := balancer.Get(name)
 	if builder == nil {
@@ -684,6 +688,8 @@ func (cc *ClientConn) handleSubConnStateChange(sc balancer.SubConn, s connectivi
 }
 
 // newAddrConn creates an addrConn for addrs and adds it to cc.conns.
+//
+// Caller needs to make sure len(addrs) > 0.
 func (cc *ClientConn) newAddrConn(addrs []resolver.Address) (*addrConn, error) {
 	ac := &addrConn{
 		cc:    cc,
diff --git a/clientconn_test.go b/clientconn_test.go
index c0b0ba43..54caf015 100644
--- a/clientconn_test.go
+++ b/clientconn_test.go
@@ -30,6 +30,8 @@ import (
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/naming"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
 	_ "google.golang.org/grpc/resolver/passthrough"
 	"google.golang.org/grpc/test/leakcheck"
 	"google.golang.org/grpc/testdata"
@@ -329,6 +331,41 @@ func TestNonblockingDialWithEmptyBalancer(t *testing.T) {
 	}
 }
 
+func TestResolverServiceConfigBeforeAddressNotPanic(t *testing.T) {
+	defer leakcheck.Check(t)
+	r, rcleanup := manual.GenerateAndRegisterManualResolver()
+	defer rcleanup()
+
+	cc, err := Dial(r.Scheme()+":///test.server", WithInsecure())
+	if err != nil {
+		t.Fatalf("failed to dial: %v", err)
+	}
+	defer cc.Close()
+
+	// SwitchBalancer before NewAddress. There was no balancer created, this
+	// makes sure we don't call close on nil balancerWrapper.
+	r.NewServiceConfig(`{"loadBalancingPolicy": "round_robin"}`) // This should not panic.
+
+	time.Sleep(time.Second) // Sleep to make sure the service config is handled by ClientConn.
+}
+
+func TestResolverEmptyUpdateNotPanic(t *testing.T) {
+	defer leakcheck.Check(t)
+	r, rcleanup := manual.GenerateAndRegisterManualResolver()
+	defer rcleanup()
+
+	cc, err := Dial(r.Scheme()+":///test.server", WithInsecure())
+	if err != nil {
+		t.Fatalf("failed to dial: %v", err)
+	}
+	defer cc.Close()
+
+	// This make sure we don't create addrConn with empty address list.
+	r.NewAddress([]resolver.Address{}) // This should not panic.
+
+	time.Sleep(time.Second) // Sleep to make sure the service config is handled by ClientConn.
+}
+
 func TestClientUpdatesParamsAfterGoAway(t *testing.T) {
 	defer leakcheck.Check(t)
 	lis, err := net.Listen("tcp", "localhost:0")