Merge pull request #5 from HenryGessau/fix-remove

Fix corner case when top node in ipTree is removed

Author: HenryGessau

ipset_test.go

@@ -197,6 +197,20 @@ func TestIPSetRemoveAll(t *testing.T) {
 	assert.False(t, set.ContainsNet(cidr1))
 }
 
+func TestIPSet_RemoveTop(t *testing.T) {
+	testSet := IPSet{}
+	ip1 := net.ParseIP("10.0.0.1")
+	ip2 := net.ParseIP("10.0.0.2")
+
+	testSet.Insert(ip2) // top
+	testSet.Insert(ip1) // inserted at left
+	testSet.Remove(ip2) // remove top node
+
+	assert.True(t, testSet.Contains(ip1))
+	assert.False(t, testSet.Contains(ip2))
+	assert.Nil(t, testSet.tree.next())
+}
+
 func TestIPSetInsertOverlapping(t *testing.T) {
 	set := IPSet{}
 

iptree.go

@@ -118,6 +118,8 @@ func (t *ipTree) remove() *ipTree {
 			} else {
 				t.up.setRight(newChild)
 			}
+		} else if newChild != nil {
+			newChild.up = t.up
 		}
 		return newChild
 	}