...

Source file src/go4.org/netipx/ipset_test.go

Documentation: go4.org/netipx

     1  // Copyright 2020 The Inet.Af AUTHORS. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package netipx
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"log"
    11  	"math/rand"
    12  	"net/netip"
    13  	"reflect"
    14  	"testing"
    15  )
    16  
    17  func buildIPSet(b *IPSetBuilder) *IPSet {
    18  	ret, err := b.IPSet()
    19  	if err != nil {
    20  		panic(err)
    21  	}
    22  	return ret
    23  }
    24  
    25  func TestIPSet(t *testing.T) {
    26  	tests := []struct {
    27  		name         string
    28  		f            func(s *IPSetBuilder)
    29  		wantRanges   []IPRange
    30  		wantPrefixes []IPPrefix      // non-nil to test
    31  		wantContains map[string]bool // optional non-exhaustive IPs to test for in resulting set
    32  	}{
    33  		{
    34  			name: "mix_family",
    35  			f: func(s *IPSetBuilder) {
    36  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
    37  				s.AddPrefix(mustIPPrefix("::/0"))
    38  				s.RemovePrefix(mustIPPrefix("10.2.0.0/16"))
    39  			},
    40  			wantRanges: []IPRange{
    41  				{mustIP("10.0.0.0"), mustIP("10.1.255.255")},
    42  				{mustIP("10.3.0.0"), mustIP("10.255.255.255")},
    43  				{mustIP("::"), mustIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")},
    44  			},
    45  		},
    46  		{
    47  			name: "merge_adjacent",
    48  			f: func(s *IPSetBuilder) {
    49  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
    50  				s.AddPrefix(mustIPPrefix("11.0.0.0/8"))
    51  			},
    52  			wantRanges: []IPRange{
    53  				{mustIP("10.0.0.0"), mustIP("11.255.255.255")},
    54  			},
    55  			wantPrefixes: pxv("10.0.0.0/7"),
    56  		},
    57  		{
    58  			name: "remove_32",
    59  			f: func(s *IPSetBuilder) {
    60  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
    61  				s.RemovePrefix(mustIPPrefix("10.1.2.3/32"))
    62  			},
    63  			wantRanges: []IPRange{
    64  				{mustIP("10.0.0.0"), mustIP("10.1.2.2")},
    65  				{mustIP("10.1.2.4"), mustIP("10.255.255.255")},
    66  			},
    67  			wantPrefixes: pxv(
    68  				"10.0.0.0/16",
    69  				"10.1.0.0/23",
    70  				"10.1.2.0/31",
    71  				"10.1.2.2/32",
    72  				"10.1.2.4/30",
    73  				"10.1.2.8/29",
    74  				"10.1.2.16/28",
    75  				"10.1.2.32/27",
    76  				"10.1.2.64/26",
    77  				"10.1.2.128/25",
    78  				"10.1.3.0/24",
    79  				"10.1.4.0/22",
    80  				"10.1.8.0/21",
    81  				"10.1.16.0/20",
    82  				"10.1.32.0/19",
    83  				"10.1.64.0/18",
    84  				"10.1.128.0/17",
    85  				"10.2.0.0/15",
    86  				"10.4.0.0/14",
    87  				"10.8.0.0/13",
    88  				"10.16.0.0/12",
    89  				"10.32.0.0/11",
    90  				"10.64.0.0/10",
    91  				"10.128.0.0/9",
    92  			),
    93  		},
    94  		{
    95  			name: "remove_32_and_first_16",
    96  			f: func(s *IPSetBuilder) {
    97  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
    98  				s.RemovePrefix(mustIPPrefix("10.1.2.3/32"))
    99  				s.RemovePrefix(mustIPPrefix("10.0.0.0/16"))
   100  			},
   101  			wantRanges: []IPRange{
   102  				{mustIP("10.1.0.0"), mustIP("10.1.2.2")},
   103  				{mustIP("10.1.2.4"), mustIP("10.255.255.255")},
   104  			},
   105  			wantPrefixes: pxv(
   106  				"10.1.0.0/23",
   107  				"10.1.2.0/31",
   108  				"10.1.2.2/32",
   109  				"10.1.2.4/30",
   110  				"10.1.2.8/29",
   111  				"10.1.2.16/28",
   112  				"10.1.2.32/27",
   113  				"10.1.2.64/26",
   114  				"10.1.2.128/25",
   115  				"10.1.3.0/24",
   116  				"10.1.4.0/22",
   117  				"10.1.8.0/21",
   118  				"10.1.16.0/20",
   119  				"10.1.32.0/19",
   120  				"10.1.64.0/18",
   121  				"10.1.128.0/17",
   122  				"10.2.0.0/15",
   123  				"10.4.0.0/14",
   124  				"10.8.0.0/13",
   125  				"10.16.0.0/12",
   126  				"10.32.0.0/11",
   127  				"10.64.0.0/10",
   128  				"10.128.0.0/9",
   129  			),
   130  		},
   131  		{
   132  			name: "add_dup",
   133  			f: func(s *IPSetBuilder) {
   134  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
   135  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
   136  			},
   137  			wantRanges: []IPRange{
   138  				{mustIP("10.0.0.0"), mustIP("10.255.255.255")},
   139  			},
   140  		},
   141  		{
   142  			name: "add_dup_subet",
   143  			f: func(s *IPSetBuilder) {
   144  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
   145  				s.AddPrefix(mustIPPrefix("10.0.0.0/16"))
   146  			},
   147  			wantRanges: []IPRange{
   148  				{mustIP("10.0.0.0"), mustIP("10.255.255.255")},
   149  			},
   150  		},
   151  		{
   152  			name: "add_remove_add",
   153  			f: func(s *IPSetBuilder) {
   154  				s.AddPrefix(mustIPPrefix("10.0.0.0/8"))
   155  				s.RemovePrefix(mustIPPrefix("10.1.2.3/32"))
   156  				s.AddPrefix(mustIPPrefix("10.1.0.0/16")) // undoes prior line
   157  			},
   158  			wantRanges: []IPRange{
   159  				{mustIP("10.0.0.0"), mustIP("10.255.255.255")},
   160  			},
   161  		},
   162  		{
   163  			name: "remove_then_add",
   164  			f: func(s *IPSetBuilder) {
   165  				s.RemovePrefix(mustIPPrefix("1.2.3.4/32")) // no-op
   166  				s.AddPrefix(mustIPPrefix("1.2.3.4/32"))
   167  			},
   168  			wantRanges: []IPRange{
   169  				{mustIP("1.2.3.4"), mustIP("1.2.3.4")},
   170  			},
   171  		},
   172  		{
   173  			name: "remove_end_on_add_start",
   174  			f: func(s *IPSetBuilder) {
   175  				s.AddRange(IPRange{mustIP("0.0.0.38"), mustIP("0.0.0.177")})
   176  				s.RemoveRange(IPRange{mustIP("0.0.0.18"), mustIP("0.0.0.38")})
   177  			},
   178  			wantRanges: []IPRange{
   179  				{mustIP("0.0.0.39"), mustIP("0.0.0.177")},
   180  			},
   181  		},
   182  		{
   183  			name: "fuzz_fail_2",
   184  			f: func(s *IPSetBuilder) {
   185  				s.AddRange(IPRange{mustIP("0.0.0.143"), mustIP("0.0.0.185")})
   186  				s.AddRange(IPRange{mustIP("0.0.0.84"), mustIP("0.0.0.174")})
   187  				s.AddRange(IPRange{mustIP("0.0.0.51"), mustIP("0.0.0.61")})
   188  				s.RemoveRange(IPRange{mustIP("0.0.0.66"), mustIP("0.0.0.146")})
   189  				s.AddRange(IPRange{mustIP("0.0.0.22"), mustIP("0.0.0.207")})
   190  				s.RemoveRange(IPRange{mustIP("0.0.0.198"), mustIP("0.0.0.203")})
   191  				s.RemoveRange(IPRange{mustIP("0.0.0.23"), mustIP("0.0.0.69")})
   192  				s.AddRange(IPRange{mustIP("0.0.0.64"), mustIP("0.0.0.105")})
   193  				s.AddRange(IPRange{mustIP("0.0.0.151"), mustIP("0.0.0.203")})
   194  				s.AddRange(IPRange{mustIP("0.0.0.138"), mustIP("0.0.0.160")})
   195  				s.RemoveRange(IPRange{mustIP("0.0.0.64"), mustIP("0.0.0.161")})
   196  			},
   197  			wantRanges: []IPRange{
   198  				{mustIP("0.0.0.22"), mustIP("0.0.0.22")},
   199  				{mustIP("0.0.0.162"), mustIP("0.0.0.207")},
   200  			},
   201  			wantContains: map[string]bool{
   202  				"0.0.0.22": true,
   203  			},
   204  		},
   205  		{
   206  			name: "single_ips",
   207  			f: func(s *IPSetBuilder) {
   208  				s.Add(mustIP("10.0.0.0"))
   209  				s.Add(mustIP("10.0.0.1"))
   210  				s.Add(mustIP("10.0.0.2"))
   211  				s.Add(mustIP("10.0.0.3"))
   212  				s.Add(mustIP("10.0.0.4"))
   213  				s.Remove(mustIP("10.0.0.4"))
   214  				s.Add(mustIP("10.0.0.255"))
   215  			},
   216  			wantRanges: []IPRange{
   217  				{mustIP("10.0.0.0"), mustIP("10.0.0.3")},
   218  				{mustIP("10.0.0.255"), mustIP("10.0.0.255")},
   219  			},
   220  			wantPrefixes: pxv("10.0.0.0/30", "10.0.0.255/32"),
   221  		},
   222  		{
   223  			// regression test for a bug where Ranges returned invalid IPRanges.
   224  			name: "single_ip_removal",
   225  			f: func(s *IPSetBuilder) {
   226  				s.Add(mustIP("10.0.0.0"))
   227  				s.Add(mustIP("10.0.0.1"))
   228  				s.Add(mustIP("10.0.0.2"))
   229  				s.Add(mustIP("10.0.0.3"))
   230  				s.Add(mustIP("10.0.0.4"))
   231  				s.Remove(mustIP("10.0.0.4"))
   232  			},
   233  			wantRanges: []IPRange{
   234  				{mustIP("10.0.0.0"), mustIP("10.0.0.3")},
   235  			},
   236  			wantPrefixes: pxv("10.0.0.0/30"),
   237  		},
   238  		{
   239  			name: "invert_empty",
   240  			f: func(s *IPSetBuilder) {
   241  				s.Complement()
   242  			},
   243  			wantRanges: []IPRange{
   244  				{mustIP("0.0.0.0"), mustIP("255.255.255.255")},
   245  				{mustIP("::"), mustIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")},
   246  			},
   247  			wantPrefixes: pxv("0.0.0.0/0", "::/0"),
   248  		},
   249  		{
   250  			name: "invert_full",
   251  			f: func(s *IPSetBuilder) {
   252  				s.AddPrefix(mustIPPrefix("0.0.0.0/0"))
   253  				s.AddPrefix(mustIPPrefix("::/0"))
   254  				s.Complement()
   255  			},
   256  			wantRanges:   []IPRange{},
   257  			wantPrefixes: pxv(),
   258  		},
   259  		{
   260  			name: "invert_partial",
   261  			f: func(s *IPSetBuilder) {
   262  				s.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("2.2.2.2")})
   263  				s.Add(mustIP("3.3.3.3"))
   264  				s.AddPrefix(mustIPPrefix("4.4.4.0/24"))
   265  				s.Add(mustIP("1::1"))
   266  				s.Complement()
   267  			},
   268  			wantRanges: []IPRange{
   269  				{mustIP("0.0.0.0"), mustIP("1.1.1.0")},
   270  				{mustIP("2.2.2.3"), mustIP("3.3.3.2")},
   271  				{mustIP("3.3.3.4"), mustIP("4.4.3.255")},
   272  				{mustIP("4.4.5.0"), mustIP("255.255.255.255")},
   273  				{mustIP("::"), mustIP("1::")},
   274  				{mustIP("1::2"), mustIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")},
   275  			},
   276  		},
   277  		{
   278  			name: "intersect",
   279  			f: func(s *IPSetBuilder) {
   280  				var t IPSetBuilder
   281  				t.AddRange(IPRange{mustIP("2.2.2.2"), mustIP("3.3.3.3")})
   282  
   283  				s.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("4.4.4.4")})
   284  				s.Intersect(buildIPSet(&t))
   285  			},
   286  			wantRanges: []IPRange{
   287  				{mustIP("2.2.2.2"), mustIP("3.3.3.3")},
   288  			},
   289  		},
   290  		{
   291  			name: "intersect_disjoint",
   292  			f: func(s *IPSetBuilder) {
   293  				var t IPSetBuilder
   294  				t.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("2.2.2.2")})
   295  
   296  				s.AddRange(IPRange{mustIP("3.3.3.3"), mustIP("4.4.4.4")})
   297  				s.Intersect(buildIPSet(&t))
   298  			},
   299  			wantRanges: []IPRange{},
   300  		},
   301  		{
   302  			name: "intersect_partial",
   303  			f: func(s *IPSetBuilder) {
   304  				var t IPSetBuilder
   305  				t.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("3.3.3.3")})
   306  
   307  				s.AddRange(IPRange{mustIP("2.2.2.2"), mustIP("4.4.4.4")})
   308  				s.Intersect(buildIPSet(&t))
   309  			},
   310  			wantRanges: []IPRange{
   311  				{mustIP("2.2.2.2"), mustIP("3.3.3.3")},
   312  			},
   313  		},
   314  	}
   315  	for _, tt := range tests {
   316  		t.Run(tt.name, func(t *testing.T) {
   317  			debugf = t.Logf
   318  			defer func() { debugf = discardf }()
   319  			var build IPSetBuilder
   320  			tt.f(&build)
   321  			s := buildIPSet(&build)
   322  			got := s.Ranges()
   323  			t.Run("ranges", func(t *testing.T) {
   324  				for _, v := range got {
   325  					if !v.IsValid() {
   326  						t.Errorf("invalid IPRange in result: %s -> %s", v.From(), v.To())
   327  					}
   328  				}
   329  				if reflect.DeepEqual(got, tt.wantRanges) {
   330  					return
   331  				}
   332  				t.Error("failed. got:\n")
   333  				for _, v := range got {
   334  					t.Errorf("  %s -> %s", v.From(), v.To())
   335  				}
   336  				t.Error("want:\n")
   337  				for _, v := range tt.wantRanges {
   338  					t.Errorf("  %s -> %s", v.From(), v.To())
   339  				}
   340  			})
   341  			if tt.wantPrefixes != nil {
   342  				t.Run("prefixes", func(t *testing.T) {
   343  					got := s.Prefixes()
   344  					if got == nil {
   345  						got = []IPPrefix{}
   346  					}
   347  					if reflect.DeepEqual(got, tt.wantPrefixes) {
   348  						return
   349  					}
   350  					t.Error("failed. got:\n")
   351  					for _, v := range got {
   352  						t.Errorf("  %v", v)
   353  					}
   354  					t.Error("want:\n")
   355  					for _, v := range tt.wantPrefixes {
   356  						t.Errorf("  %v", v)
   357  					}
   358  				})
   359  			}
   360  			if len(tt.wantContains) > 0 {
   361  				for ipStr, want := range tt.wantContains {
   362  					got := s.Contains(mustIP(ipStr))
   363  					if got != want {
   364  						t.Errorf("Contains(%q) = %v; want %v", s, got, want)
   365  					}
   366  				}
   367  			}
   368  		})
   369  	}
   370  }
   371  
   372  func TestIPSetRemoveFreePrefix(t *testing.T) {
   373  	pfx := mustIPPrefix
   374  	tests := []struct {
   375  		name         string
   376  		f            func(s *IPSetBuilder)
   377  		b            uint8
   378  		wantPrefix   IPPrefix
   379  		wantPrefixes []IPPrefix
   380  		wantOK       bool
   381  	}{
   382  		{
   383  			name: "cut in half",
   384  			f: func(s *IPSetBuilder) {
   385  				s.AddPrefix(pfx("10.0.0.0/8"))
   386  			},
   387  			b:            9,
   388  			wantPrefix:   pfx("10.0.0.0/9"),
   389  			wantPrefixes: pxv("10.128.0.0/9"),
   390  			wantOK:       true,
   391  		},
   392  		{
   393  			name: "on prefix left",
   394  			f: func(s *IPSetBuilder) {
   395  				s.AddPrefix(pfx("10.0.0.0/8"))
   396  				s.RemovePrefix(pfx("10.0.0.0/9"))
   397  			},
   398  			b:            9,
   399  			wantPrefix:   pfx("10.128.0.0/9"),
   400  			wantPrefixes: []IPPrefix{},
   401  			wantOK:       true,
   402  		},
   403  	}
   404  	for _, tt := range tests {
   405  		t.Run(tt.name, func(t *testing.T) {
   406  			debugf = t.Logf
   407  			var build IPSetBuilder
   408  			tt.f(&build)
   409  			s := buildIPSet(&build)
   410  			gotPrefix, gotSet, ok := s.RemoveFreePrefix(tt.b)
   411  			if ok != tt.wantOK {
   412  				t.Errorf("extractPrefix() ok = %t, wantOK %t", ok, tt.wantOK)
   413  				return
   414  			}
   415  			if !reflect.DeepEqual(gotPrefix, tt.wantPrefix) {
   416  				t.Errorf("extractPrefix() = %v, want %v", gotPrefix, tt.wantPrefix)
   417  			}
   418  			if !reflect.DeepEqual(gotSet.Prefixes(), tt.wantPrefixes) {
   419  				t.Errorf("extractPrefix() = %v, want %v", gotSet.Prefixes(), tt.wantPrefixes)
   420  			}
   421  		})
   422  	}
   423  }
   424  
   425  func mustIPSet(ranges ...string) *IPSet {
   426  	var ret IPSetBuilder
   427  	for _, r := range ranges {
   428  		ipr, err := ParseIPRange(r[1:])
   429  		if err != nil {
   430  			panic(err)
   431  		}
   432  		switch r[0] {
   433  		case '+':
   434  			ret.AddRange(ipr)
   435  		case '-':
   436  			ret.RemoveRange(ipr)
   437  		default:
   438  			panic(fmt.Sprintf("unknown command %q", r[0]))
   439  		}
   440  	}
   441  	return buildIPSet(&ret)
   442  }
   443  
   444  func TestIPSetOverlaps(t *testing.T) {
   445  	tests := []struct {
   446  		a, b *IPSet
   447  		want bool
   448  	}{
   449  		{
   450  			mustIPSet(),
   451  			mustIPSet(),
   452  			false,
   453  		},
   454  		{
   455  			mustIPSet("+10.0.0.0-10.0.0.5"),
   456  			mustIPSet("+10.0.0.0-10.0.0.5"),
   457  			true, // exact match
   458  		},
   459  		{
   460  			mustIPSet("+10.0.0.0-10.0.0.5"),
   461  			mustIPSet("+10.0.0.5-10.0.0.10"),
   462  			true, // overlap on edge
   463  		},
   464  		{
   465  			mustIPSet("+10.0.0.0-10.0.0.5"),
   466  			mustIPSet("+10.0.0.3-10.0.0.7"),
   467  			true, // overlap in middle
   468  		},
   469  		{
   470  			mustIPSet("+10.0.0.0-10.0.0.5"),
   471  			mustIPSet("+10.0.0.2-10.0.0.3"),
   472  			true, // one inside other
   473  		},
   474  		{
   475  			mustIPSet("+10.0.0.0-10.0.0.5", "+10.1.0.0-10.1.0.5"),
   476  			mustIPSet("+10.1.0.1-10.1.0.2"),
   477  			true, // overlap in non-first
   478  		},
   479  		{
   480  			mustIPSet("+10.0.0.0-10.0.0.10", "-10.0.0.5-10.0.0.10"),
   481  			mustIPSet("+10.0.0.7-10.0.0.8"),
   482  			false, // removal cancels overlap
   483  		},
   484  		{
   485  			mustIPSet("+10.0.0.0-10.0.0.10", "-10.0.0.5-10.0.0.10", "+10.0.0.5-10.0.0.10"),
   486  			mustIPSet("+10.0.0.7-10.0.0.8"),
   487  			true, // removal+readd restores overlap
   488  		},
   489  	}
   490  
   491  	for _, test := range tests {
   492  		got := test.a.Overlaps(test.b)
   493  		if got != test.want {
   494  			t.Errorf("(%s).Overlaps(%s) = %v, want %v", test.a, test.b, got, test.want)
   495  		}
   496  		got = test.b.Overlaps(test.a)
   497  		if got != test.want {
   498  			t.Errorf("(%s).Overlaps(%s) = %v, want %v", test.b, test.a, got, test.want)
   499  		}
   500  	}
   501  }
   502  
   503  func TestIPSetContains(t *testing.T) {
   504  	var build IPSetBuilder
   505  	build.AddPrefix(mustIPPrefix("10.0.0.0/8"))
   506  	build.AddPrefix(mustIPPrefix("1.2.3.4/32"))
   507  	build.AddPrefix(mustIPPrefix("fc00::/7"))
   508  	s := buildIPSet(&build)
   509  
   510  	tests := []struct {
   511  		ip   string
   512  		want bool
   513  	}{
   514  		{"0.0.0.0", false},
   515  		{"::", false},
   516  
   517  		{"1.2.3.3", false},
   518  		{"1.2.3.4", true},
   519  		{"1.2.3.5", false},
   520  
   521  		{"9.255.255.255", false},
   522  		{"10.0.0.0", true},
   523  		{"10.1.2.3", true},
   524  		{"10.255.255.255", true},
   525  		{"11.0.0.0", false},
   526  
   527  		{"::", false},
   528  		{"fc00::", true},
   529  		{"fc00::1", true},
   530  		{"fd00::1", true},
   531  		{"ff00::1", false},
   532  
   533  		{"fd00::%a", false},
   534  		{"fd00::1%a", false},
   535  	}
   536  	for _, tt := range tests {
   537  		got := s.Contains(mustIP(tt.ip))
   538  		if got != tt.want {
   539  			t.Errorf("contains(%q) = %v; want %v", tt.ip, got, tt.want)
   540  		}
   541  	}
   542  }
   543  
   544  func TestIPSetFuzz(t *testing.T) {
   545  	t.Parallel()
   546  	if testing.Short() {
   547  		doIPSetFuzz(t, 100)
   548  	} else {
   549  		doIPSetFuzz(t, 5000)
   550  	}
   551  }
   552  
   553  func BenchmarkIPSetFuzz(b *testing.B) {
   554  	b.ReportAllocs()
   555  	doIPSetFuzz(b, b.N)
   556  }
   557  
   558  func doIPSetFuzz(t testing.TB, iters int) {
   559  	var buf bytes.Buffer
   560  	logger := log.New(&buf, "", 0)
   561  	debugf = logger.Printf
   562  	defer func() { debugf = discardf }()
   563  	for i := 0; i < iters; i++ {
   564  		buf.Reset()
   565  		steps, set, wantContains := newRandomIPSet()
   566  		for b, want := range wantContains {
   567  			ip := IPv4(0, 0, 0, uint8(b))
   568  			got := set.Contains(ip)
   569  			if got != want {
   570  				t.Fatalf("for steps %q, contains(%v) = %v; want %v\n%s", steps, ip, got, want, buf.Bytes())
   571  			}
   572  		}
   573  	}
   574  }
   575  
   576  func newRandomIPSet() (steps []string, s *IPSet, wantContains [256]bool) {
   577  	b := new(IPSetBuilder)
   578  	nstep := 2 + rand.Intn(10)
   579  	for i := 0; i < nstep; i++ {
   580  		op := rand.Intn(2)
   581  		ip1 := uint8(rand.Intn(256))
   582  		ip2 := uint8(rand.Intn(256))
   583  		if ip2 < ip1 {
   584  			ip1, ip2 = ip2, ip1
   585  		}
   586  		var v bool
   587  		switch op {
   588  		case 0:
   589  			steps = append(steps, fmt.Sprintf("add 0.0.0.%d-0.0.0.%d", ip1, ip2))
   590  			b.AddRange(IPRangeFrom(IPv4(0, 0, 0, ip1), IPv4(0, 0, 0, ip2)))
   591  			v = true
   592  		case 1:
   593  			steps = append(steps, fmt.Sprintf("remove 0.0.0.%d-0.0.0.%d", ip1, ip2))
   594  			b.RemoveRange(IPRangeFrom(IPv4(0, 0, 0, ip1), IPv4(0, 0, 0, ip2)))
   595  		}
   596  		for i := ip1; i <= ip2; i++ {
   597  			wantContains[i] = v
   598  			if i == ip2 {
   599  				break
   600  			}
   601  		}
   602  	}
   603  	s = buildIPSet(b)
   604  	return
   605  }
   606  
   607  // TestIPSetRanges tests IPSet.Ranges against 64k
   608  // patterns of sets of ranges, checking the real implementation
   609  // against the test's separate implementation.
   610  //
   611  // For each of uint16 pattern, each set bit is treated as an IP that
   612  // should be in the set's resultant Ranges.
   613  func TestIPSetRanges(t *testing.T) {
   614  	t.Parallel()
   615  	upper := 0x0fff
   616  	if *long {
   617  		upper = 0xffff
   618  	}
   619  	for pat := 0; pat <= upper; pat++ {
   620  		var build IPSetBuilder
   621  		var from, to IP
   622  		ranges := make([]IPRange, 0)
   623  		flush := func() {
   624  			r := IPRangeFrom(from, to)
   625  			build.AddRange(r)
   626  			ranges = append(ranges, r)
   627  			from, to = IP{}, IP{}
   628  		}
   629  		for b := uint16(0); b < 16; b++ {
   630  			if uint16(pat)&(1<<b) != 0 {
   631  				ip := IPv4(1, 0, 0, uint8(b))
   632  				to = ip
   633  				if !from.IsValid() {
   634  					from = ip
   635  				}
   636  				continue
   637  			}
   638  			if from.IsValid() {
   639  				flush()
   640  			}
   641  		}
   642  		if from.IsValid() {
   643  			flush()
   644  		}
   645  		got := buildIPSet(&build).Ranges()
   646  		if !reflect.DeepEqual(got, ranges) {
   647  			t.Errorf("for %016b: got %v; want %v", pat, got, ranges)
   648  		}
   649  	}
   650  }
   651  
   652  func TestIPSetRangesStress(t *testing.T) {
   653  	t.Parallel()
   654  	n := 50
   655  	if testing.Short() {
   656  		n /= 10
   657  	} else if *long {
   658  		n = 500
   659  	}
   660  	const numIPs = 1 << 16 // we test lower 16 bits only
   661  	randRange := func() (a, b int, r IPRange) {
   662  		a, b = rand.Intn(numIPs), rand.Intn(numIPs)
   663  		if a > b {
   664  			a, b = b, a
   665  		}
   666  		from := IPv4(0, 0, uint8(a>>8), uint8(a))
   667  		to := IPv4(0, 0, uint8(b>>8), uint8(b))
   668  		return a, b, IPRangeFrom(from, to)
   669  	}
   670  	for i := 0; i < n; i++ {
   671  		var build IPSetBuilder
   672  		var want [numIPs]bool
   673  		// Add some ranges
   674  		for i := 0; i < 1+rand.Intn(2); i++ {
   675  			a, b, r := randRange()
   676  			for i := a; i <= b; i++ {
   677  				want[i] = true
   678  			}
   679  			build.AddRange(r)
   680  		}
   681  		// Remove some ranges
   682  		for i := 0; i < rand.Intn(3); i++ {
   683  			a, b, r := randRange()
   684  			for i := a; i <= b; i++ {
   685  				want[i] = false
   686  			}
   687  			build.RemoveRange(r)
   688  		}
   689  		ranges := buildIPSet(&build).Ranges()
   690  
   691  		// Make sure no ranges are adjacent or overlapping
   692  		for i, r := range ranges {
   693  			if i == 0 {
   694  				continue
   695  			}
   696  			if ranges[i-1].To().Compare(r.From()) != -1 {
   697  				t.Fatalf("overlapping ranges: %v", ranges)
   698  			}
   699  		}
   700  
   701  		// Copy the ranges back to a new set before using
   702  		// ContainsFunc, in case the ContainsFunc implementation
   703  		// changes in the future to not use Ranges itself:
   704  		var build2 IPSetBuilder
   705  		for _, r := range ranges {
   706  			build2.AddRange(r)
   707  		}
   708  		s2 := buildIPSet(&build2)
   709  		for i, want := range want {
   710  			if got := s2.Contains(IPv4(0, 0, uint8(i>>8), uint8(i))); got != want {
   711  				t.Fatal("failed")
   712  			}
   713  		}
   714  	}
   715  }
   716  
   717  func TestIPSetEqual(t *testing.T) {
   718  	a := new(IPSetBuilder)
   719  	b := new(IPSetBuilder)
   720  	MustParseIP := netip.MustParseAddr
   721  	MustParseIPPrefix := netip.MustParsePrefix
   722  
   723  	assertEqual := func(want bool) {
   724  		t.Helper()
   725  		if got := buildIPSet(a).Equal(buildIPSet(b)); got != want {
   726  			t.Errorf("%v.Equal(%v) = %v want %v", a, b, got, want)
   727  		}
   728  	}
   729  
   730  	a.Add(MustParseIP("1.1.1.0"))
   731  	a.Add(MustParseIP("1.1.1.1"))
   732  	a.Add(MustParseIP("1.1.1.2"))
   733  	b.AddPrefix(MustParseIPPrefix("1.1.1.0/31"))
   734  	b.Add(MustParseIP("1.1.1.2"))
   735  	assertEqual(true)
   736  
   737  	a.RemoveSet(buildIPSet(a))
   738  	assertEqual(false)
   739  	b.RemoveSet(buildIPSet(b))
   740  	assertEqual(true)
   741  
   742  	a.Add(MustParseIP("1.1.1.0"))
   743  	a.Add(MustParseIP("1.1.1.1"))
   744  	a.Add(MustParseIP("1.1.1.2"))
   745  
   746  	b.AddPrefix(MustParseIPPrefix("1.1.1.0/30"))
   747  	b.Remove(MustParseIP("1.1.1.3"))
   748  	assertEqual(true)
   749  }
   750  

View as plain text