From 3dcbbb786587f4016e71babad365a03de46d52ea Mon Sep 17 00:00:00 2001
From: liwen <liwen@youzan.com>
Date: Thu, 25 Feb 2021 21:16:58 +0800
Subject: [PATCH] fix: iterator lower bound may panic if search with the same
 prefix on keys

---
 iradix_test.go | 35 ++++++++++++++++++++++++++++++++++-
 iter.go        | 13 ++++++++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/iradix_test.go b/iradix_test.go
index c65995f..8252ea2 100644
--- a/iradix_test.go
+++ b/iradix_test.go
@@ -1551,7 +1551,10 @@ func TestIterateLowerBound(t *testing.T) {
 		"zap",
 		"zip",
 	}
-
+	specialKeys := []string{
+		string([]byte{1, 0, 1, 0, 0}),
+		string([]byte{1, 0, 1, 0, 1, 0}),
+	}
 	type exp struct {
 		keys   []string
 		search string
@@ -1665,6 +1668,36 @@ func TestIterateLowerBound(t *testing.T) {
 			"cbacb",
 			[]string{"cbbaa", "cbbab", "cbbbc", "cbcbb", "cbcbc", "cbcca", "ccaaa", "ccabc", "ccaca", "ccacc", "ccbac", "cccaa", "cccac", "cccca"},
 		},
+		{
+			specialKeys,
+			"",
+			specialKeys,
+		},
+		{
+			specialKeys,
+			specialKeys[0][:4],
+			specialKeys,
+		},
+		{
+			specialKeys,
+			specialKeys[0],
+			specialKeys,
+		},
+		{
+			specialKeys,
+			specialKeys[1][:5],
+			[]string{specialKeys[1]},
+		},
+		{
+			specialKeys,
+			specialKeys[1],
+			[]string{specialKeys[1]},
+		},
+		{
+			specialKeys,
+			specialKeys[1] + "0",
+			[]string{},
+		},
 	}
 
 	for idx, test := range cases {
diff --git a/iter.go b/iter.go
index cd16d3b..7f1bcda 100644
--- a/iter.go
+++ b/iter.go
@@ -20,7 +20,7 @@ func (i *Iterator) SeekPrefixWatch(prefix []byte) (watch <-chan struct{}) {
 	watch = n.mutateCh
 	search := prefix
 	for {
-		// Check for key exhaution
+		// Check for key exhaustion
 		if len(search) == 0 {
 			i.node = n
 			return
@@ -132,6 +132,17 @@ func (i *Iterator) SeekLowerBound(key []byte) {
 			search = search[len(n.prefix):]
 		}
 
+		// Check for key exhaustion
+		if len(search) == 0 {
+			// Prefix is same with the search key and not leaf node, that means the lower bound is greater than the search
+			// and from now on we need to follow the minimum path to the smallest
+			// leaf under this subtree.
+			n = i.recurseMin(n)
+			if n != nil {
+				found(n)
+			}
+			return
+		}
 		// Otherwise, take the lower bound next edge.
 		idx, lbNode := n.getLowerBoundEdge(search[0])
 		if lbNode == nil {