From d24d0a4b993587a34ac7d09ac956b0375ee7d1fa Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Tue, 15 May 2018 11:23:12 -0700
Subject: [PATCH] Strip port from server name in grpclb (#2066)

The grpclb server expects server name to not have port number
---
 grpclb.go             | 11 +++++++++++
 grpclb/grpclb_test.go | 15 +++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/grpclb.go b/grpclb.go
index bc2b4452..b313b39d 100644
--- a/grpclb.go
+++ b/grpclb.go
@@ -19,6 +19,7 @@
 package grpc
 
 import (
+	"net"
 	"strconv"
 	"strings"
 	"sync"
@@ -126,6 +127,16 @@ func (b *lbBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) bal
 		target = targetSplitted[1]
 	}
 
+	// Remove port number from target if it exists.
+	if host, _, err := net.SplitHostPort(target); err == nil {
+		// When err is not nil, target remains unchanged.
+		//
+		// Possible non-nil err values:
+		// - missing port in address: no port to strip
+		// - other errors: failed to parse
+		target = host
+	}
+
 	lb := &lbBalancer{
 		cc:              newLBCacheClientConn(cc),
 		target:          target,
diff --git a/grpclb/grpclb_test.go b/grpclb/grpclb_test.go
index a62d5903..78fb599f 100644
--- a/grpclb/grpclb_test.go
+++ b/grpclb/grpclb_test.go
@@ -313,7 +313,7 @@ func newLoadBalancer(numberOfBackends int) (tss *testServers, cleanup func(), er
 	return
 }
 
-func TestGRPCLB(t *testing.T) {
+func testGRPCLBSimple(t *testing.T, target string) {
 	defer leakcheck.Check(t)
 
 	r, cleanup := manual.GenerateAndRegisterManualResolver()
@@ -341,7 +341,7 @@ func TestGRPCLB(t *testing.T) {
 	}
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
+	cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+target,
 		grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
 	if err != nil {
 		t.Fatalf("Failed to dial to the backend %v", err)
@@ -360,6 +360,17 @@ func TestGRPCLB(t *testing.T) {
 	}
 }
 
+func TestGRPCLB(t *testing.T) {
+	testGRPCLBSimple(t, beServerName)
+}
+
+// TestGRPCLBWithPort is same as TestGRPCLB, expect that the dialing target
+// contains fake port number ":8080". The purpose of this test is to make sure
+// that grpclb strips port number from the dialing target in init request.
+func TestGRPCLBWithPort(t *testing.T) {
+	testGRPCLBSimple(t, beServerName+":8080")
+}
+
 // The remote balancer sends response with duplicates to grpclb client.
 func TestGRPCLBWeighted(t *testing.T) {
 	defer leakcheck.Check(t)