From a4d7da49735fd51c7a17a44b84630c64268106e6 Mon Sep 17 00:00:00 2001
From: David Pichler <david.a.pichler@gmail.com>
Date: Tue, 23 Jul 2024 02:39:43 -0500
Subject: [PATCH] feat: add ipv6 support (#32960)

related with #34917
This PR updates the func utils get IP address function to support
getting a IPv6 address if it is available and return it. I've tested
that this works locally when running inside a docker container that has
an IPv6 network address.

This PR should address
https://github.com/milvus-io/milvus/discussions/20217

Thank you for review!

Signed-off-by: David Pichler <david.pichler@grainger.com>
Co-authored-by: David Pichler <david.pichler@grainger.com>
---
 pkg/util/funcutil/func.go      | 27 ++++++++++++++++++++++-----
 pkg/util/funcutil/func_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/pkg/util/funcutil/func.go b/pkg/util/funcutil/func.go
index b928290e4e7a1..89da8c3877fb4 100644
--- a/pkg/util/funcutil/func.go
+++ b/pkg/util/funcutil/func.go
@@ -64,16 +64,33 @@ func GetIP(ip string) string {
 func GetLocalIP() string {
 	addrs, err := net.InterfaceAddrs()
 	if err == nil {
-		for _, addr := range addrs {
-			ipaddr, ok := addr.(*net.IPNet)
-			if ok && ipaddr.IP.IsGlobalUnicast() && ipaddr.IP.To4() != nil {
-				return ipaddr.IP.String()
-			}
+		ip := GetValidLocalIP(addrs)
+		if len(ip) != 0 {
+			return ip
 		}
 	}
 	return "127.0.0.1"
 }
 
+// GetValidLocalIP return the first valid local ip address
+func GetValidLocalIP(addrs []net.Addr) string {
+	// Search for valid ipv4 addresses
+	for _, addr := range addrs {
+		ipaddr, ok := addr.(*net.IPNet)
+		if ok && ipaddr.IP.IsGlobalUnicast() && ipaddr.IP.To4() != nil {
+			return ipaddr.IP.String()
+		}
+	}
+	// Search for valid ipv6 addresses
+	for _, addr := range addrs {
+		ipaddr, ok := addr.(*net.IPNet)
+		if ok && ipaddr.IP.IsGlobalUnicast() && ipaddr.IP.To16() != nil && ipaddr.IP.To4() == nil {
+			return "[" + ipaddr.IP.String() + "]"
+		}
+	}
+	return ""
+}
+
 // JSONToMap parse the jsonic index parameters to map
 func JSONToMap(mStr string) (map[string]string, error) {
 	buffer := make(map[string]any)
diff --git a/pkg/util/funcutil/func_test.go b/pkg/util/funcutil/func_test.go
index ded4184f0f673..d9437e0a30191 100644
--- a/pkg/util/funcutil/func_test.go
+++ b/pkg/util/funcutil/func_test.go
@@ -21,6 +21,7 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"net"
 	"reflect"
 	"strconv"
 	"testing"
@@ -54,6 +55,35 @@ func Test_CheckGrpcReady(t *testing.T) {
 	cancel()
 }
 
+func Test_GetValidLocalIPNoValid(t *testing.T) {
+	addrs := make([]net.Addr, 0, 1)
+	addrs = append(addrs, &net.IPNet{IP: net.IPv4(127, 1, 1, 1), Mask: net.IPv4Mask(255, 255, 255, 255)})
+	ip := GetValidLocalIP(addrs)
+	assert.Equal(t, "", ip)
+}
+
+func Test_GetValidLocalIPIPv4(t *testing.T) {
+	addrs := make([]net.Addr, 0, 1)
+	addrs = append(addrs, &net.IPNet{IP: net.IPv4(100, 1, 1, 1), Mask: net.IPv4Mask(255, 255, 255, 255)})
+	ip := GetValidLocalIP(addrs)
+	assert.Equal(t, "100.1.1.1", ip)
+}
+
+func Test_GetValidLocalIPIPv6(t *testing.T) {
+	addrs := make([]net.Addr, 0, 1)
+	addrs = append(addrs, &net.IPNet{IP: net.IP{8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, Mask: net.IPMask{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}})
+	ip := GetValidLocalIP(addrs)
+	assert.Equal(t, "[800::]", ip)
+}
+
+func Test_GetValidLocalIPIPv4Priority(t *testing.T) {
+	addrs := make([]net.Addr, 0, 1)
+	addrs = append(addrs, &net.IPNet{IP: net.IP{8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, Mask: net.IPMask{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}})
+	addrs = append(addrs, &net.IPNet{IP: net.IPv4(100, 1, 1, 1), Mask: net.IPv4Mask(255, 255, 255, 255)})
+	ip := GetValidLocalIP(addrs)
+	assert.Equal(t, "100.1.1.1", ip)
+}
+
 func Test_GetLocalIP(t *testing.T) {
 	ip := GetLocalIP()
 	assert.NotNil(t, ip)