From e2baeea5834b9266e00b8452f970686e901a8c2c Mon Sep 17 00:00:00 2001 From: Gregory Oschwald Date: Tue, 12 Nov 2024 13:57:17 -0800 Subject: [PATCH] Add IncludeEmptyNetworks as Networks option This can be useful when doing things like comparing two databases. --- traverse.go | 37 ++++++++++++++++++++++++++++++++++--- traverse_test.go | 45 ++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/traverse.go b/traverse.go index 13e1e52..9317baf 100644 --- a/traverse.go +++ b/traverse.go @@ -16,6 +16,7 @@ type netNode struct { type networkOptions struct { includeAliasedNetworks bool + includeEmptyNetworks bool } var ( @@ -33,13 +34,22 @@ func IncludeAliasedNetworks(networks *networkOptions) { networks.includeAliasedNetworks = true } -// Networks returns an iterator that can be used to traverse all networks in +// IncludeEmptyNetworks is an option for Networks and NetworksWithin +// that makes them include networks without any data in the iteration. +func IncludeEmptyNetworks(networks *networkOptions) { + networks.includeEmptyNetworks = true +} + +// Networks returns an iterator that can be used to traverse the networks in // the database. // // Please note that a MaxMind DB may map IPv4 networks into several locations // in an IPv6 database. This iterator will only iterate over these once by // default. To iterate over all the IPv4 network locations, use the // [IncludeAliasedNetworks] option. +// +// Networks without data are excluded by default. To include them, use +// [IncludeEmptyNetworks]. func (r *Reader) Networks(options ...NetworksOption) iter.Seq[Result] { if r.Metadata.IPVersion == 6 { return r.NetworksWithin(allIPv6, options...) @@ -47,7 +57,7 @@ func (r *Reader) Networks(options ...NetworksOption) iter.Seq[Result] { return r.NetworksWithin(allIPv4, options...) } -// NetworksWithin returns an iterator that can be used to traverse all networks +// NetworksWithin returns an iterator that can be used to traverse the networks // in the database which are contained in a given prefix. // // Please note that a MaxMind DB may map IPv4 networks into several locations @@ -57,6 +67,9 @@ func (r *Reader) Networks(options ...NetworksOption) iter.Seq[Result] { // // If the provided prefix is contained within a network in the database, the // iterator will iterate over exactly one network, the containing network. +// +// Networks without data are excluded by default. To include them, use +// [IncludeEmptyNetworks]. func (r *Reader) NetworksWithin(prefix netip.Prefix, options ...NetworksOption) iter.Seq[Result] { return func(yield func(Result) bool) { if r.Metadata.IPVersion == 4 && prefix.Addr().Is6() { @@ -106,7 +119,25 @@ func (r *Reader) NetworksWithin(prefix netip.Prefix, options ...NetworksOption) node := nodes[len(nodes)-1] nodes = nodes[:len(nodes)-1] - for node.pointer != r.Metadata.NodeCount { + for { + if node.pointer == r.Metadata.NodeCount { + if n.includeEmptyNetworks { + ip := node.ip + if isInIPv4Subtree(ip) { + ip = v6ToV4(ip) + } + + ok := yield(Result{ + ip: ip, + offset: notFound, + prefixLen: uint8(node.bit), + }) + if !ok { + return + } + } + break + } // This skips IPv4 aliases without hardcoding the networks that the writer // currently aliases. if !n.includeAliasedNetworks && r.ipv4Start != 0 && diff --git a/traverse_test.go b/traverse_test.go index 963d710..9a16116 100644 --- a/traverse_test.go +++ b/traverse_test.go @@ -3,6 +3,8 @@ package maxminddb import ( "fmt" "net/netip" + "reflect" + "runtime" "strconv" "strings" "testing" @@ -244,6 +246,43 @@ var tests = []networkTest{ "::2:0:58/127", }, }, + { + Network: "1.0.0.0/8", + Database: "mixed", + Expected: []string{ + "1.0.0.0/16", + "1.1.0.0/24", + "1.1.1.0/32", + "1.1.1.1/32", + "1.1.1.2/31", + "1.1.1.4/30", + "1.1.1.8/29", + "1.1.1.16/28", + "1.1.1.32/32", + "1.1.1.33/32", + "1.1.1.34/31", + "1.1.1.36/30", + "1.1.1.40/29", + "1.1.1.48/28", + "1.1.1.64/26", + "1.1.1.128/25", + "1.1.2.0/23", + "1.1.4.0/22", + "1.1.8.0/21", + "1.1.16.0/20", + "1.1.32.0/19", + "1.1.64.0/18", + "1.1.128.0/17", + "1.2.0.0/15", + "1.4.0.0/14", + "1.8.0.0/13", + "1.16.0.0/12", + "1.32.0.0/11", + "1.64.0.0/10", + "1.128.0.0/9", + }, + Options: []NetworksOption{IncludeEmptyNetworks}, + }, { Network: "1.1.1.16/28", Database: "mixed", @@ -263,12 +302,16 @@ var tests = []networkTest{ func TestNetworksWithin(t *testing.T) { for _, v := range tests { for _, recordSize := range []uint{24, 28, 32} { + var opts []string + for _, o := range v.Options { + opts = append(opts, runtime.FuncForPC(reflect.ValueOf(o).Pointer()).Name()) + } name := fmt.Sprintf( "%s-%d: %s, options: %v", v.Database, recordSize, v.Network, - len(v.Options) != 0, + opts, ) t.Run(name, func(t *testing.T) { fileName := testFile(fmt.Sprintf("MaxMind-DB-test-%s-%d.mmdb", v.Database, recordSize))