Add benchmark cases

cidranger_test.go

+package cidranger
+
+import (
+	"encoding/json"
+	"io/ioutil"
+	"math/rand"
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/yl2chen/cidranger/util/ip"
+)
+
+type AWSRanges struct {
+	Prefixes []Prefix `json:"prefixes"`
+}
+
+type Prefix struct {
+	IPPrefix string `json:"ip_prefix"`
+	Region   string `json:"region"`
+	Service  string `json:"service"`
+}
+
+func TestContains(t *testing.T) {
+	rangers := []Ranger{NewLPCTrieRanger()}
+	groundRanger := NewBruteRanger()
+	for _, ranger := range rangers {
+		configureRangerWithAWSRanges(t, ranger)
+	}
+	configureRangerWithAWSRanges(t, groundRanger)
+
+	for i := 0; i < 100000; i++ {
+		nn := ip.Uint32ToIPv4(rand.Uint32())
+		expected, err := groundRanger.Contains(nn)
+		for _, ranger := range rangers {
+			assert.NoError(t, err)
+			actual, err := ranger.Contains(nn)
+			assert.NoError(t, err)
+			assert.Equal(t, expected, actual)
+		}
+	}
+}
+
+func TestContainingNetworks(t *testing.T) {
+	rangers := []Ranger{NewLPCTrieRanger()}
+	groundRanger := NewBruteRanger()
+	for _, ranger := range rangers {
+		configureRangerWithAWSRanges(t, ranger)
+	}
+	configureRangerWithAWSRanges(t, groundRanger)
+
+	for i := 0; i < 100000; i++ {
+		nn := ip.Uint32ToIPv4(rand.Uint32())
+		expected, err := groundRanger.ContainingNetworks(nn)
+		for _, ranger := range rangers {
+			assert.NoError(t, err)
+			actual, err := ranger.ContainingNetworks(nn)
+			assert.NoError(t, err)
+			assert.Equal(t, len(expected), len(actual))
+			for _, network := range actual {
+				assert.Contains(t, expected, network)
+			}
+		}
+	}
+}
+
+func BenchmarkLPCTrieUsingAWSRanges(b *testing.B) {
+	benchmarkContainsUsingAWSRanges(b, NewLPCTrieRanger())
+}
+
+func BenchmarkBruteRangerUsingAWSRanges(b *testing.B) {
+	benchmarkContainsUsingAWSRanges(b, NewBruteRanger())
+}
+
+func configureRangerWithAWSRanges(tb testing.TB, ranger Ranger) {
+	ranges := loadAWSRanges(tb)
+	for _, prefix := range ranges.Prefixes {
+		_, network, err := net.ParseCIDR(prefix.IPPrefix)
+		assert.NoError(tb, err)
+		ranger.Insert(*network)
+	}
+}
+
+func benchmarkContainsUsingAWSRanges(tb testing.TB, ranger Ranger) {
+	configureRangerWithAWSRanges(tb, ranger)
+	ip := net.ParseIP("52.95.110.1")
+	for n := 0; n < tb.(*testing.B).N; n++ {
+		ranger.Contains(ip)
+	}
+}
+
+func loadAWSRanges(tb testing.TB) *AWSRanges {
+	file, err := ioutil.ReadFile("./testdata/aws_ip_ranges.json")
+	assert.NoError(tb, err)
+
+	var ranges AWSRanges
+	err = json.Unmarshal(file, &ranges)
+	assert.NoError(tb, err)
+	return &ranges
+}

ranger/trie/trie.go

@@ -65,7 +65,7 @@ func newPathPrefixTrie(network *net.IPNet, numBitsSkipped uint8) (*PrefixTrie, e
 	path := NewPrefixTree()
 	path.numBitsSkipped = numBitsSkipped
 	path.network = cidr.MaskNetwork(network, int(numBitsSkipped))
-	networkNumber, err := iputil.IPv4ToBigEndianUint32(path.network.IP)
+	networkNumber, err := iputil.IPv4ToUint32(path.network.IP)
 	if err != nil {
 		return nil, err
 	}
@@ -86,7 +86,7 @@ func newEntryTrie(network *net.IPNet) (*PrefixTrie, error) {
 
 // Insert inserts the given cidr range into prefix trie.
 func (p *PrefixTrie) Insert(network net.IPNet) error {
-	networkNumber, err := iputil.IPv4ToBigEndianUint32(network.IP)
+	networkNumber, err := iputil.IPv4ToUint32(network.IP)
 	if err != nil {
 		return err
 	}
@@ -95,7 +95,7 @@ func (p *PrefixTrie) Insert(network net.IPNet) error {
 
 // Remove removes network from trie.
 func (p *PrefixTrie) Remove(network net.IPNet) (*net.IPNet, error) {
-	networkNumber, err := iputil.IPv4ToBigEndianUint32(network.IP)
+	networkNumber, err := iputil.IPv4ToUint32(network.IP)
 	if err != nil {
 		return nil, err
 	}
@@ -147,7 +147,7 @@ func (p *PrefixTrie) childrenCount() int {
 // Contains returns boolean indicating whether given ip is contained in any
 // of the inserted networks.
 func (p *PrefixTrie) Contains(ip net.IP) (bool, error) {
-	ipUint32, err := iputil.IPv4ToBigEndianUint32(ip)
+	ipUint32, err := iputil.IPv4ToUint32(ip)
 	if err != nil {
 		return false, err
 	}
@@ -161,7 +161,7 @@ func (p *PrefixTrie) Contains(ip net.IP) (bool, error) {
 // ContainingNetworks returns the list of networks given ip is a part of in
 // ascending prefix order.
 func (p *PrefixTrie) ContainingNetworks(ip net.IP) ([]net.IPNet, error) {
-	ipUint32, err := iputil.IPv4ToBigEndianUint32(ip)
+	ipUint32, err := iputil.IPv4ToUint32(ip)
 	if err != nil {
 		return nil, err
 	}

util/cidr/cidr.go

@@ -17,11 +17,11 @@ var ErrNoGreatestCommonBit = fmt.Errorf("No greatest common bit")
 // GreatestCommonBitPosition returns the greatest common bit position of
 // given cidr blocks.
 func GreatestCommonBitPosition(network1 *net.IPNet, network2 *net.IPNet) (uint8, error) {
-	ip1, err := ip.IPv4ToBigEndianUint32(network1.IP)
+	ip1, err := ip.IPv4ToUint32(network1.IP)
 	if err != nil {
 		return 0, err
 	}
-	ip2, err := ip.IPv4ToBigEndianUint32(network2.IP)
+	ip2, err := ip.IPv4ToUint32(network2.IP)
 	if err != nil {
 		return 0, err
 	}

util/ip/ip.go

@@ -17,8 +17,8 @@ var ErrBitsNotValid = fmt.Errorf("bits requested not valid")
 
 const ipv4BitLength = 32
 
-// IPv4ToBigEndianUint32 converts IPv4 to uint32.
-func IPv4ToBigEndianUint32(ip net.IP) (uint32, error) {
+// IPv4ToUint32 converts ipV4 to uint32.
+func IPv4ToUint32(ip net.IP) (uint32, error) {
 	ip = ip.To4()
 	if ip == nil {
 		return 0, ErrNotIPv4Error
@@ -26,6 +26,13 @@ func IPv4ToBigEndianUint32(ip net.IP) (uint32, error) {
 	return binary.BigEndian.Uint32(ip), nil
 }
 
+// Uint32ToIPv4 converts uint32 to ipV4 net.IP.
+func Uint32ToIPv4(nn uint32) net.IP {
+	ip := make(net.IP, 4)
+	binary.BigEndian.PutUint32(ip, nn)
+	return ip
+}
+
 // IPv4BitsAsUint returns uint32 representing bits at position of length
 // numberOfBits, position is a number in [0, 31] representing the starting
 // position in ip, with 31 being the most significant bit.

util/ip/ip_test.go

@@ -7,7 +7,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestIPv4ToLittleEndianUint32(t *testing.T) {
+func TestIPv4ToUint32(t *testing.T) {
 	cases := []struct {
 		ip          string
 		ipUint32    uint32
@@ -20,13 +20,31 @@ func TestIPv4ToLittleEndianUint32(t *testing.T) {
 
 	for _, c := range cases {
 		t.Run(c.ip, func(t *testing.T) {
-			ret, err := IPv4ToBigEndianUint32(net.ParseIP(c.ip))
+			ret, err := IPv4ToUint32(net.ParseIP(c.ip))
 			assert.Equal(t, c.expectedErr, err)
 			assert.Equal(t, c.ipUint32, ret)
 		})
 	}
 }
 
+func TestUint32ToIPv4(t *testing.T) {
+	cases := []struct {
+		ip       uint32
+		expected net.IP
+		name     string
+	}{
+		{2147483648, net.ParseIP("128.0.0.0").To4(), "128.0.0.0"},
+		{2147483649, net.ParseIP("128.0.0.1").To4(), "128.0.0.1"},
+		{4294967295, net.ParseIP("255.255.255.255").To4(), "255.255.255.255"},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.expected, Uint32ToIPv4(tc.ip))
+		})
+	}
+}
+
 func TestIPv4BitsAsUint(t *testing.T) {
 	cases := []struct {
 		ip           string