trie_test.go 14.5 KB
Newer Older
1
package cidranger
2 3

import (
4 5
	"encoding/binary"
	"math/rand"
6
	"net"
7
	"runtime"
8
	"testing"
9
	"time"
10 11

	"github.com/stretchr/testify/assert"
Steven Allen's avatar
Steven Allen committed
12
	rnet "github.com/libp2p/go-cidranger/net"
13 14
)

15 16 17 18 19 20 21
func getAllByVersion(version rnet.IPVersion) *net.IPNet {
	if version == rnet.IPv6 {
		return AllIPv6
	}
	return AllIPv4
}

22
func TestPrefixTrieInsert(t *testing.T) {
23
	cases := []struct {
24
		version                      rnet.IPVersion
25 26 27 28
		inserts                      []string
		expectedNetworksInDepthOrder []string
		name                         string
	}{
29
		{rnet.IPv4, []string{"192.168.0.1/24"}, []string{"192.168.0.1/24"}, "basic insert"},
30 31 32 33 34 35 36 37 38 39 40 41
		{
			rnet.IPv4,
			[]string{"1.2.3.4/32", "1.2.3.5/32"},
			[]string{"1.2.3.4/32", "1.2.3.5/32"},
			"single ip IPv4 network insert",
		},
		{
			rnet.IPv6,
			[]string{"0::1/128", "0::2/128"},
			[]string{"0::1/128", "0::2/128"},
			"single ip IPv6 network insert",
		},
42
		{
43
			rnet.IPv4,
44 45 46 47
			[]string{"192.168.0.1/16", "192.168.0.1/24"},
			[]string{"192.168.0.1/16", "192.168.0.1/24"},
			"in order insert",
		},
48 49 50 51 52 53
		{
			rnet.IPv4,
			[]string{"192.168.0.1/32", "192.168.0.1/32"},
			[]string{"192.168.0.1/32"},
			"duplicate network insert",
		},
54
		{
55
			rnet.IPv4,
56 57 58 59 60
			[]string{"192.168.0.1/24", "192.168.0.1/16"},
			[]string{"192.168.0.1/16", "192.168.0.1/24"},
			"reverse insert",
		},
		{
61
			rnet.IPv4,
62 63 64 65 66
			[]string{"192.168.0.1/24", "192.168.1.1/24"},
			[]string{"192.168.0.1/24", "192.168.1.1/24"},
			"branch insert",
		},
		{
67
			rnet.IPv4,
68 69 70 71 72 73 74
			[]string{"192.168.0.1/24", "192.168.1.1/24", "192.168.1.1/30"},
			[]string{"192.168.0.1/24", "192.168.1.1/24", "192.168.1.1/30"},
			"branch inserts",
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
75
			trie := newPrefixTree(tc.version).(*prefixTrie)
76 77
			for _, insert := range tc.inserts {
				_, network, _ := net.ParseCIDR(insert)
78
				err := trie.Insert(NewBasicRangerEntry(*network))
79 80
				assert.NoError(t, err)
			}
Yulin Chen's avatar
Yulin Chen committed
81 82 83

			assert.Equal(t, len(tc.expectedNetworksInDepthOrder), trie.Len(), "trie size should match")

84 85 86 87
			allNetworks, err := trie.CoveredNetworks(*getAllByVersion(tc.version))
			assert.Nil(t, err)
			assert.Equal(t, len(allNetworks), trie.Len(), "trie size should match")

88 89
			walk := trie.walkDepth()
			for _, network := range tc.expectedNetworksInDepthOrder {
90 91
				_, ipnet, _ := net.ParseCIDR(network)
				expected := NewBasicRangerEntry(*ipnet)
92
				actual := <-walk
93
				assert.Equal(t, expected, actual)
94 95 96 97 98 99 100 101 102 103
			}

			// Ensure no unexpected elements in trie.
			for network := range walk {
				assert.Nil(t, network)
			}
		})
	}
}

104
func TestPrefixTrieString(t *testing.T) {
Yulin Chen's avatar
Yulin Chen committed
105
	inserts := []string{"192.168.0.1/24", "192.168.1.1/24", "192.168.1.1/30"}
106
	trie := newPrefixTree(rnet.IPv4).(*prefixTrie)
Yulin Chen's avatar
Yulin Chen committed
107 108
	for _, insert := range inserts {
		_, network, _ := net.ParseCIDR(insert)
109
		trie.Insert(NewBasicRangerEntry(*network))
Yulin Chen's avatar
Yulin Chen committed
110 111 112 113 114 115 116 117 118
	}
	expected := `0.0.0.0/0 (target_pos:31:has_entry:false)
| 1--> 192.168.0.0/23 (target_pos:8:has_entry:false)
| | 0--> 192.168.0.0/24 (target_pos:7:has_entry:true)
| | 1--> 192.168.1.0/24 (target_pos:7:has_entry:true)
| | | 0--> 192.168.1.0/30 (target_pos:1:has_entry:true)`
	assert.Equal(t, expected, trie.String())
}

119
func TestPrefixTrieRemove(t *testing.T) {
120
	cases := []struct {
121
		version                      rnet.IPVersion
122 123 124 125
		inserts                      []string
		removes                      []string
		expectedRemoves              []string
		expectedNetworksInDepthOrder []string
126
		expectedTrieString           string
127 128 129
		name                         string
	}{
		{
130
			rnet.IPv4,
131 132 133 134
			[]string{"192.168.0.1/24"},
			[]string{"192.168.0.1/24"},
			[]string{"192.168.0.1/24"},
			[]string{},
135
			"0.0.0.0/0 (target_pos:31:has_entry:false)",
136 137
			"basic remove",
		},
138 139 140 141 142 143
		{
			rnet.IPv4,
			[]string{"1.2.3.4/32", "1.2.3.5/32"},
			[]string{"1.2.3.5/32"},
			[]string{"1.2.3.5/32"},
			[]string{"1.2.3.4/32"},
144 145
			`0.0.0.0/0 (target_pos:31:has_entry:false)
| 0--> 1.2.3.4/32 (target_pos:-1:has_entry:true)`,
146 147 148 149 150 151 152 153
			"single ip IPv4 network remove",
		},
		{
			rnet.IPv4,
			[]string{"0::1/128", "0::2/128"},
			[]string{"0::2/128"},
			[]string{"0::2/128"},
			[]string{"0::1/128"},
154 155
			`0.0.0.0/0 (target_pos:31:has_entry:false)
| 0--> ::1/128 (target_pos:-1:has_entry:true)`,
156 157
			"single ip IPv6 network remove",
		},
158
		{
159
			rnet.IPv4,
160 161 162 163
			[]string{"192.168.0.1/24", "192.168.0.1/25", "192.168.0.1/26"},
			[]string{"192.168.0.1/25"},
			[]string{"192.168.0.1/25"},
			[]string{"192.168.0.1/24", "192.168.0.1/26"},
164 165 166
			`0.0.0.0/0 (target_pos:31:has_entry:false)
| 1--> 192.168.0.0/24 (target_pos:7:has_entry:true)
| | 0--> 192.168.0.0/26 (target_pos:5:has_entry:true)`,
167 168 169
			"remove path prefix",
		},
		{
170
			rnet.IPv4,
171 172 173 174
			[]string{"192.168.0.1/24", "192.168.0.1/25", "192.168.0.64/26", "192.168.0.1/26"},
			[]string{"192.168.0.1/25"},
			[]string{"192.168.0.1/25"},
			[]string{"192.168.0.1/24", "192.168.0.1/26", "192.168.0.64/26"},
175 176 177 178 179
			`0.0.0.0/0 (target_pos:31:has_entry:false)
| 1--> 192.168.0.0/24 (target_pos:7:has_entry:true)
| | 0--> 192.168.0.0/25 (target_pos:6:has_entry:false)
| | | 0--> 192.168.0.0/26 (target_pos:5:has_entry:true)
| | | 1--> 192.168.0.64/26 (target_pos:5:has_entry:true)`,
180 181 182
			"remove path prefix with more than 1 children",
		},
		{
183
			rnet.IPv4,
184 185 186 187
			[]string{"192.168.0.1/24", "192.168.0.1/25"},
			[]string{"192.168.0.1/26"},
			[]string{""},
			[]string{"192.168.0.1/24", "192.168.0.1/25"},
188 189 190
			`0.0.0.0/0 (target_pos:31:has_entry:false)
| 1--> 192.168.0.0/24 (target_pos:7:has_entry:true)
| | 0--> 192.168.0.0/25 (target_pos:6:has_entry:true)`,
Yulin Chen's avatar
Yulin Chen committed
191
			"remove non existent",
192 193 194 195 196
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
197
			trie := newPrefixTree(tc.version).(*prefixTrie)
198 199
			for _, insert := range tc.inserts {
				_, network, _ := net.ParseCIDR(insert)
200
				err := trie.Insert(NewBasicRangerEntry(*network))
201 202 203 204 205 206 207
				assert.NoError(t, err)
			}
			for i, remove := range tc.removes {
				_, network, _ := net.ParseCIDR(remove)
				removed, err := trie.Remove(*network)
				assert.NoError(t, err)
				if str := tc.expectedRemoves[i]; str != "" {
208 209 210
					_, ipnet, _ := net.ParseCIDR(str)
					expected := NewBasicRangerEntry(*ipnet)
					assert.Equal(t, expected, removed)
211 212 213 214
				} else {
					assert.Nil(t, removed)
				}
			}
Yulin Chen's avatar
Yulin Chen committed
215 216 217

			assert.Equal(t, len(tc.expectedNetworksInDepthOrder), trie.Len(), "trie size should match after revmoval")

218 219 220 221
			allNetworks, err := trie.CoveredNetworks(*getAllByVersion(tc.version))
			assert.Nil(t, err)
			assert.Equal(t, len(allNetworks), trie.Len(), "trie size should match")

222 223
			walk := trie.walkDepth()
			for _, network := range tc.expectedNetworksInDepthOrder {
224 225
				_, ipnet, _ := net.ParseCIDR(network)
				expected := NewBasicRangerEntry(*ipnet)
226
				actual := <-walk
227
				assert.Equal(t, expected, actual)
228 229 230 231 232 233
			}

			// Ensure no unexpected elements in trie.
			for network := range walk {
				assert.Nil(t, network)
			}
234 235

			assert.Equal(t, tc.expectedTrieString, trie.String())
236 237 238 239
		})
	}
}

240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285
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)
		})
	}
}

286 287 288 289 290
type expectedIPRange struct {
	start net.IP
	end   net.IP
}

291
func TestPrefixTrieContains(t *testing.T) {
292
	cases := []struct {
293
		version     rnet.IPVersion
294 295 296 297 298
		inserts     []string
		expectedIPs []expectedIPRange
		name        string
	}{
		{
299
			rnet.IPv4,
300 301
			[]string{"192.168.0.0/24"},
			[]expectedIPRange{
Yulin Chen's avatar
Yulin Chen committed
302
				{net.ParseIP("192.168.0.0"), net.ParseIP("192.168.1.0")},
303 304 305 306
			},
			"basic contains",
		},
		{
307
			rnet.IPv4,
308 309
			[]string{"192.168.0.0/24", "128.168.0.0/24"},
			[]expectedIPRange{
Yulin Chen's avatar
Yulin Chen committed
310 311
				{net.ParseIP("192.168.0.0"), net.ParseIP("192.168.1.0")},
				{net.ParseIP("128.168.0.0"), net.ParseIP("128.168.1.0")},
312 313 314 315 316 317 318
			},
			"multiple ranges contains",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
319
			trie := newPrefixTree(tc.version)
320 321
			for _, insert := range tc.inserts {
				_, network, _ := net.ParseCIDR(insert)
322
				err := trie.Insert(NewBasicRangerEntry(*network))
323 324 325 326 327
				assert.NoError(t, err)
			}
			for _, expectedIPRange := range tc.expectedIPs {
				var contains bool
				var err error
328 329 330
				start := expectedIPRange.start
				for ; !expectedIPRange.end.Equal(start); start = rnet.NextIP(start) {
					contains, err = trie.Contains(start)
331 332 333 334 335
					assert.NoError(t, err)
					assert.True(t, contains)
				}

				// Check out of bounds ips on both ends
336
				contains, err = trie.Contains(rnet.PreviousIP(expectedIPRange.start))
337 338
				assert.NoError(t, err)
				assert.False(t, contains)
339
				contains, err = trie.Contains(rnet.NextIP(expectedIPRange.end))
340 341 342 343 344 345 346
				assert.NoError(t, err)
				assert.False(t, contains)
			}
		})
	}
}

347
func TestPrefixTrieContainingNetworks(t *testing.T) {
348
	cases := []struct {
349
		version  rnet.IPVersion
350 351 352 353 354 355
		inserts  []string
		ip       net.IP
		networks []string
		name     string
	}{
		{
356
			rnet.IPv4,
357 358 359 360 361 362
			[]string{"192.168.0.0/24"},
			net.ParseIP("192.168.0.1"),
			[]string{"192.168.0.0/24"},
			"basic containing networks",
		},
		{
363
			rnet.IPv4,
364 365 366 367 368 369 370 371
			[]string{"192.168.0.0/24", "192.168.0.0/25"},
			net.ParseIP("192.168.0.1"),
			[]string{"192.168.0.0/24", "192.168.0.0/25"},
			"inclusive networks",
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
372
			trie := newPrefixTree(tc.version)
373 374
			for _, insert := range tc.inserts {
				_, network, _ := net.ParseCIDR(insert)
375
				err := trie.Insert(NewBasicRangerEntry(*network))
376 377
				assert.NoError(t, err)
			}
378
			expectedEntries := []RangerEntry{}
379 380
			for _, network := range tc.networks {
				_, net, _ := net.ParseCIDR(network)
381
				expectedEntries = append(expectedEntries, NewBasicRangerEntry(*net))
382 383 384
			}
			networks, err := trie.ContainingNetworks(tc.ip)
			assert.NoError(t, err)
385
			assert.Equal(t, expectedEntries, networks)
386 387 388
		})
	}
}
Rob Adams's avatar
Rob Adams committed
389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450

type coveredNetworkTest struct {
	version  rnet.IPVersion
	inserts  []string
	search   string
	networks []string
	name     string
}

var coveredNetworkTests = []coveredNetworkTest{
	{
		rnet.IPv4,
		[]string{"192.168.0.0/24"},
		"192.168.0.0/16",
		[]string{"192.168.0.0/24"},
		"basic covered networks",
	},
	{
		rnet.IPv4,
		[]string{"192.168.0.0/24"},
		"10.1.0.0/16",
		nil,
		"nothing",
	},
	{
		rnet.IPv4,
		[]string{"192.168.0.0/24", "192.168.0.0/25"},
		"192.168.0.0/16",
		[]string{"192.168.0.0/24", "192.168.0.0/25"},
		"multiple networks",
	},
	{
		rnet.IPv4,
		[]string{"192.168.0.0/24", "192.168.0.0/25", "192.168.0.1/32"},
		"192.168.0.0/16",
		[]string{"192.168.0.0/24", "192.168.0.0/25", "192.168.0.1/32"},
		"multiple networks 2",
	},
	{
		rnet.IPv4,
		[]string{"192.168.1.1/32"},
		"192.168.0.0/16",
		[]string{"192.168.1.1/32"},
		"leaf",
	},
	{
		rnet.IPv4,
		[]string{"0.0.0.0/0", "192.168.1.1/32"},
		"192.168.0.0/16",
		[]string{"192.168.1.1/32"},
		"leaf with root",
	},
	{
		rnet.IPv4,
		[]string{
			"0.0.0.0/0", "192.168.0.0/24", "192.168.1.1/32",
			"10.1.0.0/16", "10.1.1.0/24",
		},
		"192.168.0.0/16",
		[]string{"192.168.0.0/24", "192.168.1.1/32"},
		"path not taken",
	},
451 452 453 454 455 456 457 458 459
	{
		rnet.IPv4,
		[]string{
			"192.168.0.0/15",
		},
		"192.168.0.0/16",
		nil,
		"only masks different",
	},
Rob Adams's avatar
Rob Adams committed
460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483
}

func TestPrefixTrieCoveredNetworks(t *testing.T) {
	for _, tc := range coveredNetworkTests {
		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)
			}
			var expectedEntries []RangerEntry
			for _, network := range tc.networks {
				_, net, _ := net.ParseCIDR(network)
				expectedEntries = append(expectedEntries,
					NewBasicRangerEntry(*net))
			}
			_, snet, _ := net.ParseCIDR(tc.search)
			networks, err := trie.CoveredNetworks(*snet)
			assert.NoError(t, err)
			assert.Equal(t, expectedEntries, networks)
		})
	}
}
484 485 486 487 488 489

func TestTrieMemUsage(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping memory test in `-short` mode")
	}
	numIPs := 100000
Yulin Chen's avatar
Yulin Chen committed
490
	runs := 10
491 492 493 494 495 496 497 498 499

	// Avg heap allocation over all runs should not be more than the heap allocation of first run multiplied
	// by threshold, picking 1% as sane number for detecting memory leak.
	thresh := 1.01

	trie := newPrefixTree(rnet.IPv4)

	var baseLineHeap, totalHeapAllocOverRuns uint64
	for i := 0; i < runs; i++ {
Yulin Chen's avatar
Yulin Chen committed
500
		t.Logf("Executing Run %d of %d", i+1, runs)
501 502 503 504 505

		// Insert networks.
		for n := 0; n < numIPs; n++ {
			trie.Insert(NewBasicRangerEntry(GenLeafIPNet(GenIPV4())))
		}
Yulin Chen's avatar
Yulin Chen committed
506 507 508
		t.Logf("Inserted All (%d networks)", trie.Len())
		assert.Less(t, 0, trie.Len(), "Len should > 0")
		assert.LessOrEqualf(t, trie.Len(), numIPs, "Len should <= %d", numIPs)
509

510 511 512 513
		allNetworks, err := trie.CoveredNetworks(*getAllByVersion(rnet.IPv4))
		assert.Nil(t, err)
		assert.Equal(t, len(allNetworks), trie.Len(), "trie size should match")

514 515 516 517 518 519
		// Remove networks.
		_, all, _ := net.ParseCIDR("0.0.0.0/0")
		ll, _ := trie.CoveredNetworks(*all)
		for i := 0; i < len(ll); i++ {
			trie.Remove(ll[i].Network())
		}
Yulin Chen's avatar
Yulin Chen committed
520 521
		t.Logf("Removed All (%d networks)", len(ll))
		assert.Equal(t, 0, trie.Len(), "Len after removal should == 0")
522 523 524 525 526 527 528

		// Perform GC
		runtime.GC()

		// Get HeapAlloc stats.
		heapAlloc := GetHeapAllocation()
		totalHeapAllocOverRuns += heapAlloc
Yulin Chen's avatar
Yulin Chen committed
529
		if i == 0 {
530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562
			baseLineHeap = heapAlloc
		}
	}

	// Assert that heap allocation from first loop is within set threshold of avg over all runs.
	assert.Less(t, uint64(0), baseLineHeap)
	assert.LessOrEqual(t, float64(baseLineHeap), float64(totalHeapAllocOverRuns/uint64(runs))*thresh)
}

func GenLeafIPNet(ip net.IP) net.IPNet {
	return net.IPNet{
		IP:   ip,
		Mask: net.CIDRMask(32, 32),
	}
}

// GenIPV4 generates an IPV4 address
func GenIPV4() net.IP {
	rand.Seed(time.Now().UnixNano())
	var min, max int
	min = 1
	max = 4294967295
	nn := rand.Intn(max-min) + min
	ip := make(net.IP, 4)
	binary.BigEndian.PutUint32(ip, uint32(nn))
	return ip
}

func GetHeapAllocation() uint64 {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	return m.HeapAlloc
}