github.com/slackhq/nebula@v1.9.0/cidr/tree6_test.go (about)

     1  package cidr
     2  
     3  import (
     4  	"encoding/binary"
     5  	"net"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestCIDR6Tree_MostSpecificContains(t *testing.T) {
    12  	tree := NewTree6[string]()
    13  	tree.AddCIDR(Parse("1.0.0.0/8"), "1")
    14  	tree.AddCIDR(Parse("2.1.0.0/16"), "2")
    15  	tree.AddCIDR(Parse("3.1.1.0/24"), "3")
    16  	tree.AddCIDR(Parse("4.1.1.1/24"), "4a")
    17  	tree.AddCIDR(Parse("4.1.1.1/30"), "4b")
    18  	tree.AddCIDR(Parse("4.1.1.1/32"), "4c")
    19  	tree.AddCIDR(Parse("254.0.0.0/4"), "5")
    20  	tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/64"), "6a")
    21  	tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/80"), "6b")
    22  	tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/96"), "6c")
    23  
    24  	tests := []struct {
    25  		Found  bool
    26  		Result interface{}
    27  		IP     string
    28  	}{
    29  		{true, "1", "1.0.0.0"},
    30  		{true, "1", "1.255.255.255"},
    31  		{true, "2", "2.1.0.0"},
    32  		{true, "2", "2.1.255.255"},
    33  		{true, "3", "3.1.1.0"},
    34  		{true, "3", "3.1.1.255"},
    35  		{true, "4a", "4.1.1.255"},
    36  		{true, "4b", "4.1.1.2"},
    37  		{true, "4c", "4.1.1.1"},
    38  		{true, "5", "240.0.0.0"},
    39  		{true, "5", "255.255.255.255"},
    40  		{true, "6a", "1:2:0:4:1:1:1:1"},
    41  		{true, "6b", "1:2:0:4:5:1:1:1"},
    42  		{true, "6c", "1:2:0:4:5:0:0:0"},
    43  		{false, "", "239.0.0.0"},
    44  		{false, "", "4.1.2.2"},
    45  	}
    46  
    47  	for _, tt := range tests {
    48  		ok, r := tree.MostSpecificContains(net.ParseIP(tt.IP))
    49  		assert.Equal(t, tt.Found, ok)
    50  		assert.Equal(t, tt.Result, r)
    51  	}
    52  
    53  	tree = NewTree6[string]()
    54  	tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
    55  	tree.AddCIDR(Parse("::/0"), "cool6")
    56  	ok, r := tree.MostSpecificContains(net.ParseIP("0.0.0.0"))
    57  	assert.True(t, ok)
    58  	assert.Equal(t, "cool", r)
    59  
    60  	ok, r = tree.MostSpecificContains(net.ParseIP("255.255.255.255"))
    61  	assert.True(t, ok)
    62  	assert.Equal(t, "cool", r)
    63  
    64  	ok, r = tree.MostSpecificContains(net.ParseIP("::"))
    65  	assert.True(t, ok)
    66  	assert.Equal(t, "cool6", r)
    67  
    68  	ok, r = tree.MostSpecificContains(net.ParseIP("1:2:3:4:5:6:7:8"))
    69  	assert.True(t, ok)
    70  	assert.Equal(t, "cool6", r)
    71  }
    72  
    73  func TestCIDR6Tree_MostSpecificContainsIpV6(t *testing.T) {
    74  	tree := NewTree6[string]()
    75  	tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/64"), "6a")
    76  	tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/80"), "6b")
    77  	tree.AddCIDR(Parse("1:2:0:4:5:0:0:0/96"), "6c")
    78  
    79  	tests := []struct {
    80  		Found  bool
    81  		Result interface{}
    82  		IP     string
    83  	}{
    84  		{true, "6a", "1:2:0:4:1:1:1:1"},
    85  		{true, "6b", "1:2:0:4:5:1:1:1"},
    86  		{true, "6c", "1:2:0:4:5:0:0:0"},
    87  	}
    88  
    89  	for _, tt := range tests {
    90  		ip := net.ParseIP(tt.IP)
    91  		hi := binary.BigEndian.Uint64(ip[:8])
    92  		lo := binary.BigEndian.Uint64(ip[8:])
    93  
    94  		ok, r := tree.MostSpecificContainsIpV6(hi, lo)
    95  		assert.Equal(t, tt.Found, ok)
    96  		assert.Equal(t, tt.Result, r)
    97  	}
    98  }