From 6a720386a7065a0145cfcc30516e0ac364b3ef05 Mon Sep 17 00:00:00 2001
From: Yulin Chen <yulin@uber.com>
Date: Thu, 21 Sep 2017 19:10:18 -0700
Subject: [PATCH] Fix index out of range panic when finding ip from trie with
 single ip network, referencing #4

---
 trie.go      |  3 +++
 trie_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/trie.go b/trie.go
index 16792d4..2743014 100644
--- a/trie.go
+++ b/trie.go
@@ -149,6 +149,9 @@ func (p *prefixTrie) containingNetworks(number rnet.NetworkNumber) ([]RangerEntr
 	if p.hasEntry() {
 		results = []RangerEntry{p.entry}
 	}
+	if p.targetBitPosition() < 0 {
+		return results, nil
+	}
 	bit, err := p.targetBitFromIP(number)
 	if err != nil {
 		return nil, err
diff --git a/trie_test.go b/trie_test.go
index 6794da1..44ca547 100644
--- a/trie_test.go
+++ b/trie_test.go
@@ -193,6 +193,52 @@ func TestPrefixTrieRemove(t *testing.T) {
 	}
 }
 
+func TestToReplicateIssue(t *testing.T) {
+	cases := []struct {
+		version  rnet.IPVersion
+		inserts  []string
+		ip       net.IP
+		networks []string
+		name     string
+	}{
+		{
+			rnet.IPv4,
+			[]string{"192.168.0.1/32"},
+			net.ParseIP("192.168.0.1"),
+			[]string{"192.168.0.1/32"},
+			"basic containing network for /32 mask",
+		},
+		{
+			rnet.IPv6,
+			[]string{"a::1/128"},
+			net.ParseIP("a::1"),
+			[]string{"a::1/128"},
+			"basic containing network for /128 mask",
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			trie := newPrefixTree(tc.version)
+			for _, insert := range tc.inserts {
+				_, network, _ := net.ParseCIDR(insert)
+				err := trie.Insert(NewBasicRangerEntry(*network))
+				assert.NoError(t, err)
+			}
+			expectedEntries := []RangerEntry{}
+			for _, network := range tc.networks {
+				_, net, _ := net.ParseCIDR(network)
+				expectedEntries = append(expectedEntries, NewBasicRangerEntry(*net))
+			}
+			contains, err := trie.Contains(tc.ip)
+			assert.NoError(t, err)
+			assert.True(t, contains)
+			networks, err := trie.ContainingNetworks(tc.ip)
+			assert.NoError(t, err)
+			assert.Equal(t, expectedEntries, networks)
+		})
+	}
+}
+
 type expectedIPRange struct {
 	start net.IP
 	end   net.IP
-- 
GitLab