github.com/tailscale/wireguard-go@v0.0.20201119-0.20210522003738-46b531feb08a/device/allowedips_rand_test.go (about)

     1  /* SPDX-License-Identifier: MIT
     2   *
     3   * Copyright (C) 2017-2021 WireGuard LLC. All Rights Reserved.
     4   */
     5  
     6  package device
     7  
     8  import (
     9  	"math/rand"
    10  	"sort"
    11  	"testing"
    12  )
    13  
    14  const (
    15  	NumberOfPeers     = 100
    16  	NumberOfAddresses = 250
    17  	NumberOfTests     = 10000
    18  )
    19  
    20  type SlowNode struct {
    21  	peer *Peer
    22  	cidr uint
    23  	bits []byte
    24  }
    25  
    26  type SlowRouter []*SlowNode
    27  
    28  func (r SlowRouter) Len() int {
    29  	return len(r)
    30  }
    31  
    32  func (r SlowRouter) Less(i, j int) bool {
    33  	return r[i].cidr > r[j].cidr
    34  }
    35  
    36  func (r SlowRouter) Swap(i, j int) {
    37  	r[i], r[j] = r[j], r[i]
    38  }
    39  
    40  func (r SlowRouter) Insert(addr []byte, cidr uint, peer *Peer) SlowRouter {
    41  	for _, t := range r {
    42  		if t.cidr == cidr && commonBits(t.bits, addr) >= cidr {
    43  			t.peer = peer
    44  			t.bits = addr
    45  			return r
    46  		}
    47  	}
    48  	r = append(r, &SlowNode{
    49  		cidr: cidr,
    50  		bits: addr,
    51  		peer: peer,
    52  	})
    53  	sort.Sort(r)
    54  	return r
    55  }
    56  
    57  func (r SlowRouter) Lookup(addr []byte) *Peer {
    58  	for _, t := range r {
    59  		common := commonBits(t.bits, addr)
    60  		if common >= t.cidr {
    61  			return t.peer
    62  		}
    63  	}
    64  	return nil
    65  }
    66  
    67  func TestTrieRandomIPv4(t *testing.T) {
    68  	var trie *trieEntry
    69  	var slow SlowRouter
    70  	var peers []*Peer
    71  
    72  	rand.Seed(1)
    73  
    74  	const AddressLength = 4
    75  
    76  	for n := 0; n < NumberOfPeers; n++ {
    77  		peers = append(peers, &Peer{})
    78  	}
    79  
    80  	for n := 0; n < NumberOfAddresses; n++ {
    81  		var addr [AddressLength]byte
    82  		rand.Read(addr[:])
    83  		cidr := uint(rand.Uint32() % (AddressLength * 8))
    84  		index := rand.Int() % NumberOfPeers
    85  		trie = trie.insert(addr[:], cidr, peers[index])
    86  		slow = slow.Insert(addr[:], cidr, peers[index])
    87  	}
    88  
    89  	for n := 0; n < NumberOfTests; n++ {
    90  		var addr [AddressLength]byte
    91  		rand.Read(addr[:])
    92  		peer1 := slow.Lookup(addr[:])
    93  		peer2 := trie.lookup(addr[:])
    94  		if peer1 != peer2 {
    95  			t.Error("Trie did not match naive implementation, for:", addr)
    96  		}
    97  	}
    98  }
    99  
   100  func TestTrieRandomIPv6(t *testing.T) {
   101  	var trie *trieEntry
   102  	var slow SlowRouter
   103  	var peers []*Peer
   104  
   105  	rand.Seed(1)
   106  
   107  	const AddressLength = 16
   108  
   109  	for n := 0; n < NumberOfPeers; n++ {
   110  		peers = append(peers, &Peer{})
   111  	}
   112  
   113  	for n := 0; n < NumberOfAddresses; n++ {
   114  		var addr [AddressLength]byte
   115  		rand.Read(addr[:])
   116  		cidr := uint(rand.Uint32() % (AddressLength * 8))
   117  		index := rand.Int() % NumberOfPeers
   118  		trie = trie.insert(addr[:], cidr, peers[index])
   119  		slow = slow.Insert(addr[:], cidr, peers[index])
   120  	}
   121  
   122  	for n := 0; n < NumberOfTests; n++ {
   123  		var addr [AddressLength]byte
   124  		rand.Read(addr[:])
   125  		peer1 := slow.Lookup(addr[:])
   126  		peer2 := trie.lookup(addr[:])
   127  		if peer1 != peer2 {
   128  			t.Error("Trie did not match naive implementation, for:", addr)
   129  		}
   130  	}
   131  }