From 1ab4adf22dbaf4eb6cfd01eaa711612ffcf87890 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Tue, 13 Jun 2017 13:43:09 -0700
Subject: [PATCH] Support ipv6 addresses in grpclb (#1303)

Add square brackets to ipv6 addresses, otherwise net.Dial() and
net.SplitHostPort() will return too many colons error.
---
 grpclb.go             |  9 ++++++++-
 grpclb/grpclb_test.go | 11 ++++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/grpclb.go b/grpclb.go
index bb2ada5fa..c5bd65225 100644
--- a/grpclb.go
+++ b/grpclb.go
@@ -235,8 +235,15 @@ func (b *balancer) processServerList(l *lbpb.ServerList, seq int) {
 	)
 	for _, s := range servers {
 		md := metadata.Pairs("lb-token", s.LoadBalanceToken)
+		ip := net.IP(s.IpAddress)
+		ipStr := ip.String()
+		if ip.To4() == nil {
+			// Add square brackets to ipv6 addresses, otherwise net.Dial() and
+			// net.SplitHostPort() will return too many colons error.
+			ipStr = fmt.Sprintf("[%s]", ipStr)
+		}
 		addr := Address{
-			Addr:     fmt.Sprintf("%s:%d", net.IP(s.IpAddress), s.Port),
+			Addr:     fmt.Sprintf("%s:%d", ipStr, s.Port),
 			Metadata: &md,
 		}
 		sl = append(sl, &grpclbAddrInfo{
diff --git a/grpclb/grpclb_test.go b/grpclb/grpclb_test.go
index ed3dd9029..4948c7d44 100644
--- a/grpclb/grpclb_test.go
+++ b/grpclb/grpclb_test.go
@@ -129,8 +129,9 @@ func (r *testNameResolver) inject(updates []*naming.Update) {
 }
 
 type serverNameCheckCreds struct {
-	expected string
+	mu       sync.Mutex
 	sn       string
+	expected string
 }
 
 func (c *serverNameCheckCreds) ServerHandshake(rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
@@ -141,6 +142,8 @@ func (c *serverNameCheckCreds) ServerHandshake(rawConn net.Conn) (net.Conn, cred
 	return rawConn, nil, nil
 }
 func (c *serverNameCheckCreds) ClientHandshake(ctx context.Context, addr string, rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	b := make([]byte, len(c.expected))
 	if _, err := rawConn.Read(b); err != nil {
 		fmt.Printf("Failed to read the server name from the server %v", err)
@@ -153,14 +156,20 @@ func (c *serverNameCheckCreds) ClientHandshake(ctx context.Context, addr string,
 	return rawConn, nil, nil
 }
 func (c *serverNameCheckCreds) Info() credentials.ProtocolInfo {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	return credentials.ProtocolInfo{}
 }
 func (c *serverNameCheckCreds) Clone() credentials.TransportCredentials {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	return &serverNameCheckCreds{
 		expected: c.expected,
 	}
 }
 func (c *serverNameCheckCreds) OverrideServerName(s string) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	c.expected = s
 	return nil
 }