Proof of concept of NetworksWithin

Author: oschwald

reader.go

@@ -249,7 +249,20 @@ func (r *Reader) lookupPointer(ip net.IP) (uint, int, net.IP, error) {
 	if bitCount == 32 {
 		node = r.ipv4Start
 	}
+	node, prefixLength := r.traverseTree(ip, node, bitCount)
 
+	nodeCount := r.Metadata.NodeCount
+	if node == nodeCount {
+		// Record is empty
+		return 0, prefixLength, ip, nil
+	} else if node > nodeCount {
+		return node, prefixLength, ip, nil
+	}
+
+	return 0, prefixLength, ip, newInvalidDatabaseError("invalid node in search tree")
+}
+
+func (r *Reader) traverseTree(ip net.IP, node uint, bitCount uint) (uint, int) {
 	nodeCount := r.Metadata.NodeCount
 
 	i := uint(0)
@@ -263,14 +276,8 @@ func (r *Reader) lookupPointer(ip net.IP) (uint, int, net.IP, error) {
 			node = r.nodeReader.readRight(offset)
 		}
 	}
-	if node == nodeCount {
-		// Record is empty
-		return 0, int(i), ip, nil
-	} else if node > nodeCount {
-		return node, int(i), ip, nil
-	}
 
-	return 0, int(i), ip, newInvalidDatabaseError("invalid node in search tree")
+	return node, int(i)
 }
 
 func (r *Reader) retrieveData(pointer uint, result interface{}) error {

traverse.go

@@ -17,22 +17,36 @@ type Networks struct {
 	err      error
 }
 
+var allIPv4 = &net.IPNet{IP: make(net.IP, 4), Mask: net.CIDRMask(0, 32)}
+var allIPv6 = &net.IPNet{IP: make(net.IP, 16), Mask: net.CIDRMask(0, 128)}
+
 // Networks returns an iterator that can be used to traverse all networks in
 // the database.
 //
 // Please note that a MaxMind DB may map IPv4 networks into several locations
 // in an IPv6 database. This iterator will iterate over all of these
 // locations separately.
 func (r *Reader) Networks() *Networks {
-	s := 4
+	var networks *Networks
 	if r.Metadata.IPVersion == 6 {
-		s = 16
+		networks = r.NetworksWithin(allIPv6)
+	} else {
+		networks = r.NetworksWithin(allIPv4)
 	}
+
+	return networks
+}
+
+func (r *Reader) NetworksWithin(network *net.IPNet) *Networks {
+	prefixLength, _ := network.Mask.Size()
+	pointer, bit := r.traverseTree(network.IP, 0, uint(prefixLength))
 	return &Networks{
 		reader: r,
 		nodes: []netNode{
 			{
-				ip: make(net.IP, s),
+				ip:      network.IP,
+				bit:     uint(bit),
+				pointer: pointer,
 			},
 		},
 	}

traverse_test.go

@@ -2,6 +2,7 @@ package maxminddb
 
 import (
 	"fmt"
+	"net"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -46,3 +47,75 @@ func TestNetworksWithInvalidSearchTree(t *testing.T) {
 	assert.NotNil(t, n.Err(), "no error received when traversing an broken search tree")
 	assert.Equal(t, n.Err().Error(), "invalid search tree at 128.128.128.128/32")
 }
+func TestNetworksWithin(t *testing.T) {
+	_, network, error := net.ParseCIDR("1.1.1.0/24")
+
+	assert.Nil(t, error)
+
+	for _, recordSize := range []uint{24, 28, 32} {
+		fileName := testFile(fmt.Sprintf("MaxMind-DB-test-ipv4-%d.mmdb", recordSize))
+		reader, err := Open(fileName)
+		require.Nil(t, err, "unexpected error while opening database: %v", err)
+		defer reader.Close()
+
+		n := reader.NetworksWithin(network)
+		var innerIPs []string
+
+		for n.Next() {
+			record := struct {
+				IP string `maxminddb:"ip"`
+			}{}
+			network, err := n.Network(&record)
+			assert.Nil(t, err)
+			assert.Equal(t, record.IP, network.IP.String(),
+				"expected %s got %s", record.IP, network.IP.String(),
+			)
+			innerIPs = append(innerIPs, record.IP)
+		}
+
+		expectedIPs := []string{
+			"1.1.1.1",
+			"1.1.1.2",
+			"1.1.1.4",
+			"1.1.1.8",
+			"1.1.1.16",
+			"1.1.1.32",
+		}
+
+		assert.Equal(t, expectedIPs, innerIPs)
+		assert.Nil(t, n.Err())
+	}
+}
+
+func TestNetworksWithinSlash32(t *testing.T) {
+	_, network, error := net.ParseCIDR("1.1.1.32/32")
+
+	assert.Nil(t, error)
+
+	for _, recordSize := range []uint{24, 28, 32} {
+		fileName := testFile(fmt.Sprintf("MaxMind-DB-test-ipv4-%d.mmdb", recordSize))
+		reader, err := Open(fileName)
+		require.Nil(t, err, "unexpected error while opening database: %v", err)
+		defer reader.Close()
+
+		n := reader.NetworksWithin(network)
+		var innerIPs []string
+
+		for n.Next() {
+			record := struct {
+				IP string `maxminddb:"ip"`
+			}{}
+			network, err := n.Network(&record)
+			assert.Nil(t, err)
+			assert.Equal(t, record.IP, network.IP.String(),
+				"expected %s got %s", record.IP, network.IP.String(),
+			)
+			innerIPs = append(innerIPs, record.IP)
+		}
+
+		expectedIPs := []string([]string{"1.1.1.32"})
+
+		assert.Equal(t, expectedIPs, innerIPs)
+		assert.Nil(t, n.Err())
+	}
+}