From 83d44936ef61fbe16d96bc60304ccc60aa6b02a4 Mon Sep 17 00:00:00 2001
From: Rueian <rueiancsie@gmail.com>
Date: Tue, 30 Apr 2024 00:00:18 +0800
Subject: [PATCH] Merge pull request #540 from redis/fix-cluster-move-ipv6

fix: ipv6 handling in MOVED and ASK messages of cluster mode
Signed-off-by: Rueian <rueiancsie@gmail.com>
---
 message.go      | 14 ++++++++++++--
 message_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/message.go b/message.go
index 32d39ef..081d325 100644
--- a/message.go
+++ b/message.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net"
 	"strconv"
 	"strings"
 	"time"
@@ -55,7 +56,7 @@ func (r *ValkeyError) IsNil() bool {
 // IsMoved checks if it is a valkey MOVED message and returns moved address.
 func (r *ValkeyError) IsMoved() (addr string, ok bool) {
 	if ok = strings.HasPrefix(r.string, "MOVED"); ok {
-		addr = strings.Split(r.string, " ")[2]
+		addr = fixIPv6HostPort(strings.Split(r.string, " ")[2])
 	}
 	return
 }
@@ -63,11 +64,20 @@ func (r *ValkeyError) IsMoved() (addr string, ok bool) {
 // IsAsk checks if it is a valkey ASK message and returns ask address.
 func (r *ValkeyError) IsAsk() (addr string, ok bool) {
 	if ok = strings.HasPrefix(r.string, "ASK"); ok {
-		addr = strings.Split(r.string, " ")[2]
+		addr = fixIPv6HostPort(strings.Split(r.string, " ")[2])
 	}
 	return
 }
 
+func fixIPv6HostPort(addr string) string {
+	if strings.IndexByte(addr, '.') < 0 && len(addr) > 0 && addr[0] != '[' { // skip ipv4 and enclosed ipv6
+		if i := strings.LastIndexByte(addr, ':'); i >= 0 {
+			return net.JoinHostPort(addr[:i], addr[i+1:])
+		}
+	}
+	return addr
+}
+
 // IsTryAgain checks if it is a valkey TRYAGAIN message and returns ask address.
 func (r *ValkeyError) IsTryAgain() bool {
 	return strings.HasPrefix(r.string, "TRYAGAIN")
diff --git a/message_test.go b/message_test.go
index 9b87978..f66195f 100644
--- a/message_test.go
+++ b/message_test.go
@@ -58,6 +58,38 @@ func TestIsValkeyErr(t *testing.T) {
 	}
 }
 
+func TestValkeyErrorIsMoved(t *testing.T) {
+	for _, c := range []struct {
+		err  string
+		addr string
+	}{
+		{err: "MOVED 1 127.0.0.1:1", addr: "127.0.0.1:1"},
+		{err: "MOVED 1 [::1]:1", addr: "[::1]:1"},
+		{err: "MOVED 1 ::1:1", addr: "[::1]:1"},
+	} {
+		e := ValkeyError{typ: '-', string: c.err}
+		if addr, ok := e.IsMoved(); !ok || addr != c.addr {
+			t.Fail()
+		}
+	}
+}
+
+func TestValkeyErrorIsAsk(t *testing.T) {
+	for _, c := range []struct {
+		err  string
+		addr string
+	}{
+		{err: "ASK 1 127.0.0.1:1", addr: "127.0.0.1:1"},
+		{err: "ASK 1 [::1]:1", addr: "[::1]:1"},
+		{err: "ASK 1 ::1:1", addr: "[::1]:1"},
+	} {
+		e := ValkeyError{typ: '-', string: c.err}
+		if addr, ok := e.IsAsk(); !ok || addr != c.addr {
+			t.Fail()
+		}
+	}
+}
+
 func TestIsValkeyBusyGroup(t *testing.T) {
 	err := errors.New("other")
 	if IsValkeyBusyGroup(err) {