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

     1  package cidr
     2  
     3  import (
     4  	"net"
     5  	"testing"
     6  
     7  	"github.com/slackhq/nebula/iputil"
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestCIDRTree_List(t *testing.T) {
    12  	tree := NewTree4[string]()
    13  	tree.AddCIDR(Parse("1.0.0.0/16"), "1")
    14  	tree.AddCIDR(Parse("1.0.0.0/8"), "2")
    15  	tree.AddCIDR(Parse("1.0.0.0/16"), "3")
    16  	tree.AddCIDR(Parse("1.0.0.0/16"), "4")
    17  	list := tree.List()
    18  	assert.Len(t, list, 2)
    19  	assert.Equal(t, "1.0.0.0/8", list[0].CIDR.String())
    20  	assert.Equal(t, "2", list[0].Value)
    21  	assert.Equal(t, "1.0.0.0/16", list[1].CIDR.String())
    22  	assert.Equal(t, "4", list[1].Value)
    23  }
    24  
    25  func TestCIDRTree_Contains(t *testing.T) {
    26  	tree := NewTree4[string]()
    27  	tree.AddCIDR(Parse("1.0.0.0/8"), "1")
    28  	tree.AddCIDR(Parse("2.1.0.0/16"), "2")
    29  	tree.AddCIDR(Parse("3.1.1.0/24"), "3")
    30  	tree.AddCIDR(Parse("4.1.1.0/24"), "4a")
    31  	tree.AddCIDR(Parse("4.1.1.1/32"), "4b")
    32  	tree.AddCIDR(Parse("4.1.2.1/32"), "4c")
    33  	tree.AddCIDR(Parse("254.0.0.0/4"), "5")
    34  
    35  	tests := []struct {
    36  		Found  bool
    37  		Result interface{}
    38  		IP     string
    39  	}{
    40  		{true, "1", "1.0.0.0"},
    41  		{true, "1", "1.255.255.255"},
    42  		{true, "2", "2.1.0.0"},
    43  		{true, "2", "2.1.255.255"},
    44  		{true, "3", "3.1.1.0"},
    45  		{true, "3", "3.1.1.255"},
    46  		{true, "4a", "4.1.1.255"},
    47  		{true, "4a", "4.1.1.1"},
    48  		{true, "5", "240.0.0.0"},
    49  		{true, "5", "255.255.255.255"},
    50  		{false, "", "239.0.0.0"},
    51  		{false, "", "4.1.2.2"},
    52  	}
    53  
    54  	for _, tt := range tests {
    55  		ok, r := tree.Contains(iputil.Ip2VpnIp(net.ParseIP(tt.IP)))
    56  		assert.Equal(t, tt.Found, ok)
    57  		assert.Equal(t, tt.Result, r)
    58  	}
    59  
    60  	tree = NewTree4[string]()
    61  	tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
    62  	ok, r := tree.Contains(iputil.Ip2VpnIp(net.ParseIP("0.0.0.0")))
    63  	assert.True(t, ok)
    64  	assert.Equal(t, "cool", r)
    65  
    66  	ok, r = tree.Contains(iputil.Ip2VpnIp(net.ParseIP("255.255.255.255")))
    67  	assert.True(t, ok)
    68  	assert.Equal(t, "cool", r)
    69  }
    70  
    71  func TestCIDRTree_MostSpecificContains(t *testing.T) {
    72  	tree := NewTree4[string]()
    73  	tree.AddCIDR(Parse("1.0.0.0/8"), "1")
    74  	tree.AddCIDR(Parse("2.1.0.0/16"), "2")
    75  	tree.AddCIDR(Parse("3.1.1.0/24"), "3")
    76  	tree.AddCIDR(Parse("4.1.1.0/24"), "4a")
    77  	tree.AddCIDR(Parse("4.1.1.0/30"), "4b")
    78  	tree.AddCIDR(Parse("4.1.1.1/32"), "4c")
    79  	tree.AddCIDR(Parse("254.0.0.0/4"), "5")
    80  
    81  	tests := []struct {
    82  		Found  bool
    83  		Result interface{}
    84  		IP     string
    85  	}{
    86  		{true, "1", "1.0.0.0"},
    87  		{true, "1", "1.255.255.255"},
    88  		{true, "2", "2.1.0.0"},
    89  		{true, "2", "2.1.255.255"},
    90  		{true, "3", "3.1.1.0"},
    91  		{true, "3", "3.1.1.255"},
    92  		{true, "4a", "4.1.1.255"},
    93  		{true, "4b", "4.1.1.2"},
    94  		{true, "4c", "4.1.1.1"},
    95  		{true, "5", "240.0.0.0"},
    96  		{true, "5", "255.255.255.255"},
    97  		{false, "", "239.0.0.0"},
    98  		{false, "", "4.1.2.2"},
    99  	}
   100  
   101  	for _, tt := range tests {
   102  		ok, r := tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP(tt.IP)))
   103  		assert.Equal(t, tt.Found, ok)
   104  		assert.Equal(t, tt.Result, r)
   105  	}
   106  
   107  	tree = NewTree4[string]()
   108  	tree.AddCIDR(Parse("1.1.1.1/0"), "cool")
   109  	ok, r := tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP("0.0.0.0")))
   110  	assert.True(t, ok)
   111  	assert.Equal(t, "cool", r)
   112  
   113  	ok, r = tree.MostSpecificContains(iputil.Ip2VpnIp(net.ParseIP("255.255.255.255")))
   114  	assert.True(t, ok)
   115  	assert.Equal(t, "cool", r)
   116  }
   117  
   118  func TestTree4_GetCIDR(t *testing.T) {
   119  	tree := NewTree4[string]()
   120  	tree.AddCIDR(Parse("1.0.0.0/8"), "1")
   121  	tree.AddCIDR(Parse("2.1.0.0/16"), "2")
   122  	tree.AddCIDR(Parse("3.1.1.0/24"), "3")
   123  	tree.AddCIDR(Parse("4.1.1.0/24"), "4a")
   124  	tree.AddCIDR(Parse("4.1.1.1/32"), "4b")
   125  	tree.AddCIDR(Parse("4.1.2.1/32"), "4c")
   126  	tree.AddCIDR(Parse("254.0.0.0/4"), "5")
   127  
   128  	tests := []struct {
   129  		Found  bool
   130  		Result interface{}
   131  		IPNet  *net.IPNet
   132  	}{
   133  		{true, "1", Parse("1.0.0.0/8")},
   134  		{true, "2", Parse("2.1.0.0/16")},
   135  		{true, "3", Parse("3.1.1.0/24")},
   136  		{true, "4a", Parse("4.1.1.0/24")},
   137  		{true, "4b", Parse("4.1.1.1/32")},
   138  		{true, "4c", Parse("4.1.2.1/32")},
   139  		{true, "5", Parse("254.0.0.0/4")},
   140  		{false, "", Parse("2.0.0.0/8")},
   141  	}
   142  
   143  	for _, tt := range tests {
   144  		ok, r := tree.GetCIDR(tt.IPNet)
   145  		assert.Equal(t, tt.Found, ok)
   146  		assert.Equal(t, tt.Result, r)
   147  	}
   148  }
   149  
   150  func BenchmarkCIDRTree_Contains(b *testing.B) {
   151  	tree := NewTree4[string]()
   152  	tree.AddCIDR(Parse("1.1.0.0/16"), "1")
   153  	tree.AddCIDR(Parse("1.2.1.1/32"), "1")
   154  	tree.AddCIDR(Parse("192.2.1.1/32"), "1")
   155  	tree.AddCIDR(Parse("172.2.1.1/32"), "1")
   156  
   157  	ip := iputil.Ip2VpnIp(net.ParseIP("1.2.1.1"))
   158  	b.Run("found", func(b *testing.B) {
   159  		for i := 0; i < b.N; i++ {
   160  			tree.Contains(ip)
   161  		}
   162  	})
   163  
   164  	ip = iputil.Ip2VpnIp(net.ParseIP("1.2.1.255"))
   165  	b.Run("not found", func(b *testing.B) {
   166  		for i := 0; i < b.N; i++ {
   167  			tree.Contains(ip)
   168  		}
   169  	})
   170  }