From 23eb35bf0c0178aa8254e4d6886f713b62cc3a48 Mon Sep 17 00:00:00 2001
From: Bing Han <h.bing612@gmail.com>
Date: Thu, 14 Sep 2023 16:30:29 +0800
Subject: [PATCH] fix(sd): support ipv6 for consul and eureka

---
 sd/consul/instancer.go      |  4 +++-
 sd/consul/instancer_test.go | 42 +++++++++++++++++++++++++++++++++++++
 sd/eureka/instancer.go      |  5 +++--
 sd/eureka/instancer_test.go | 16 ++++++++++++++
 sd/eureka/registrar.go      |  4 +++-
 5 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/sd/consul/instancer.go b/sd/consul/instancer.go
index 3215695eb..a764ff5e8 100644
--- a/sd/consul/instancer.go
+++ b/sd/consul/instancer.go
@@ -3,6 +3,8 @@ package consul
 import (
 	"errors"
 	"fmt"
+	"net"
+	"strconv"
 	"time"
 
 	consul "github.com/hashicorp/consul/api"
@@ -181,7 +183,7 @@ func makeInstances(entries []*consul.ServiceEntry) []string {
 		if entry.Service.Address != "" {
 			addr = entry.Service.Address
 		}
-		instances[i] = fmt.Sprintf("%s:%d", addr, entry.Service.Port)
+		instances[i] = net.JoinHostPort(addr, strconv.Itoa(entry.Service.Port))
 	}
 	return instances
 }
diff --git a/sd/consul/instancer_test.go b/sd/consul/instancer_test.go
index 84737d187..9389839f9 100644
--- a/sd/consul/instancer_test.go
+++ b/sd/consul/instancer_test.go
@@ -61,6 +61,21 @@ var consulState = []*consul.ServiceEntry{
 			},
 		},
 	},
+	{
+		Node: &consul.Node{
+			Address: "10.0.0.1",
+			Node:    "app01.local",
+		},
+		Service: &consul.AgentService{
+			Address: "2001:db8:1::ab9:C0A8:102",
+			ID:      "search-db-1",
+			Port:    9000,
+			Service: "search",
+			Tags: []string{
+				"ipv6",
+			},
+		},
+	},
 }
 
 func TestInstancer(t *testing.T) {
@@ -135,6 +150,33 @@ func TestInstancerAddressOverride(t *testing.T) {
 	}
 }
 
+func TestInstancerAddressIpv6(t *testing.T) {
+	s := NewInstancer(newTestClient(consulState), log.NewNopLogger(), "search", []string{"ipv6"}, true)
+	defer s.Stop()
+
+	state := s.cache.State()
+	if want, have := 1, len(state.Instances); want != have {
+		t.Fatalf("want %d, have %d", want, have)
+	}
+
+	endpoint, closer, err := testFactory(state.Instances[0])
+	if err != nil {
+		t.Fatal(err)
+	}
+	if closer != nil {
+		defer closer.Close()
+	}
+
+	response, err := endpoint(context.Background(), struct{}{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if want, have := "[2001:db8:1::ab9:C0A8:102]:9000", response.(string); want != have {
+		t.Errorf("want %q, have %q", want, have)
+	}
+}
+
 type eofTestClient struct {
 	client *testClient
 	eofSig chan bool
diff --git a/sd/eureka/instancer.go b/sd/eureka/instancer.go
index 2849cf438..0fb82ef77 100644
--- a/sd/eureka/instancer.go
+++ b/sd/eureka/instancer.go
@@ -1,7 +1,8 @@
 package eureka
 
 import (
-	"fmt"
+	"net"
+	"strconv"
 
 	"github.com/hudl/fargo"
 
@@ -84,7 +85,7 @@ func (s *Instancer) getInstances() ([]string, error) {
 func convertFargoAppToInstances(app *fargo.Application) []string {
 	instances := make([]string, len(app.Instances))
 	for i, inst := range app.Instances {
-		instances[i] = fmt.Sprintf("%s:%d", inst.IPAddr, inst.Port)
+		instances[i] = net.JoinHostPort(inst.IPAddr, strconv.Itoa(inst.Port))
 	}
 	return instances
 }
diff --git a/sd/eureka/instancer_test.go b/sd/eureka/instancer_test.go
index 9363e5055..1f82fff43 100644
--- a/sd/eureka/instancer_test.go
+++ b/sd/eureka/instancer_test.go
@@ -90,3 +90,19 @@ func TestBadInstancerScheduleUpdates(t *testing.T) {
 		t.Errorf("want %d, have %d", want, have)
 	}
 }
+
+func TestConvertIpv6Instance(t *testing.T) {
+	app := &fargo.Application{
+		Instances: []*fargo.Instance{
+			{IPAddr: "10.0.0.10", Port: 8000},
+			{IPAddr: "2001:db8:1::ab9:C0A8:102", Port: 8000},
+		},
+	}
+	expect := []string{"10.0.0.10:8000", "[2001:db8:1::ab9:C0A8:102]:8000"}
+	for i, inst := range convertFargoAppToInstances(app) {
+		if inst != expect[i] {
+			t.Fatalf("instance %s converting is wrong", expect[i])
+		}
+	}
+
+}
diff --git a/sd/eureka/registrar.go b/sd/eureka/registrar.go
index 2d84a7c80..621f932f4 100644
--- a/sd/eureka/registrar.go
+++ b/sd/eureka/registrar.go
@@ -2,7 +2,9 @@ package eureka
 
 import (
 	"fmt"
+	"net"
 	"net/http"
+	"strconv"
 	"sync"
 	"time"
 
@@ -51,7 +53,7 @@ func NewRegistrar(conn fargoConnection, instance *fargo.Instance, logger log.Log
 	return &Registrar{
 		conn:     conn,
 		instance: instance,
-		logger:   log.With(logger, "service", instance.App, "address", fmt.Sprintf("%s:%d", instance.IPAddr, instance.Port)),
+		logger:   log.With(logger, "service", instance.App, "address", net.JoinHostPort(instance.IPAddr, strconv.Itoa(instance.Port))),
 	}
 }