Add support for getting record network

Author: oschwald

reader.go

@@ -110,13 +110,36 @@ func (r *Reader) Lookup(ipAddress net.IP, result interface{}) error {
 	if r.buffer == nil {
 		return errors.New("cannot call Lookup on a closed database")
 	}
-	pointer, err := r.lookupPointer(ipAddress)
+	pointer, _, _, err := r.lookupPointer(ipAddress)
 	if pointer == 0 || err != nil {
 		return err
 	}
 	return r.retrieveData(pointer, result)
 }
 
+// LookupNetwork retrieves the database record for ipAddress and stores it in
+// the value pointed to be result. The network returned is the network
+// associated with the data record in the database. The ok return value
+// indicates whether the database contained a record for the ipAddress.
+//
+// If result is nil or not a pointer, an error is returned. If the data in the
+// database record cannot be stored in result because of type differences, an
+// UnmarshalTypeError is returned. If the database is invalid or otherwise
+// cannot be read, an InvalidDatabaseError is returned.
+func (r *Reader) LookupNetwork(ipAddress net.IP, result interface{}) (network *net.IPNet, ok bool, err error) {
+	if r.buffer == nil {
+		return nil, false, errors.New("cannot call Lookup on a closed database")
+	}
+	pointer, prefixLength, ipAddress, err := r.lookupPointer(ipAddress)
+
+	network = r.cidr(ipAddress, prefixLength)
+	if pointer == 0 || err != nil {
+		return network, false, err
+	}
+
+	return network, true, r.retrieveData(pointer, result)
+}
+
 // LookupOffset maps an argument net.IP to a corresponding record offset in the
 // database. NotFound is returned if no such record is found, and a record may
 // otherwise be extracted by passing the returned offset to Decode. LookupOffset
@@ -126,13 +149,20 @@ func (r *Reader) LookupOffset(ipAddress net.IP) (uintptr, error) {
 	if r.buffer == nil {
 		return 0, errors.New("cannot call LookupOffset on a closed database")
 	}
-	pointer, err := r.lookupPointer(ipAddress)
+	pointer, _, _, err := r.lookupPointer(ipAddress)
 	if pointer == 0 || err != nil {
 		return NotFound, err
 	}
 	return r.resolveDataPointer(pointer)
 }
 
+func (r *Reader) cidr(ipAddress net.IP, prefixLength int) *net.IPNet {
+	ipBitLength := len(ipAddress) * 8
+	mask := net.CIDRMask(prefixLength, ipBitLength)
+
+	return &net.IPNet{IP: ipAddress.Mask(mask), Mask: mask}
+}
+
 // Decode the record at |offset| into |result|. The result value pointed to
 // must be a data value that corresponds to a record in the database. This may
 // include a struct representation of the data, a map capable of holding the
@@ -166,24 +196,19 @@ func (r *Reader) decode(offset uintptr, result interface{}) error {
 	return err
 }
 
-func (r *Reader) lookupPointer(ipAddress net.IP) (uint, error) {
+func (r *Reader) lookupPointer(ipAddress net.IP) (uint, int, net.IP, error) {
 	if ipAddress == nil {
-		return 0, errors.New("ipAddress passed to Lookup cannot be nil")
+		return 0, 0, ipAddress, errors.New("ipAddress passed to Lookup cannot be nil")
 	}
 
 	ipV4Address := ipAddress.To4()
 	if ipV4Address != nil {
 		ipAddress = ipV4Address
 	}
 	if len(ipAddress) == 16 && r.Metadata.IPVersion == 4 {
-		return 0, fmt.Errorf("error looking up '%s': you attempted to look up an IPv6 address in an IPv4-only database", ipAddress.String())
+		return 0, 0, ipAddress, fmt.Errorf("error looking up '%s': you attempted to look up an IPv6 address in an IPv4-only database", ipAddress.String())
 	}
 
-	return r.findAddressInTree(ipAddress)
-}
-
-func (r *Reader) findAddressInTree(ipAddress net.IP) (uint, error) {
-
 	bitCount := uint(len(ipAddress) * 8)
 
 	var node uint
@@ -193,23 +218,24 @@ func (r *Reader) findAddressInTree(ipAddress net.IP) (uint, error) {
 
 	nodeCount := r.Metadata.NodeCount
 
-	for i := uint(0); i < bitCount && node < nodeCount; i++ {
+	i := uint(0)
+	for ; i < bitCount && node < nodeCount; i++ {
 		bit := uint(1) & (uint(ipAddress[i>>3]) >> (7 - (i % 8)))
 
 		var err error
 		node, err = r.readNode(node, bit)
 		if err != nil {
-			return 0, err
+			return 0, int(i), ipAddress, err
 		}
 	}
 	if node == nodeCount {
 		// Record is empty
-		return 0, nil
+		return 0, int(i), ipAddress, nil
 	} else if node > nodeCount {
-		return node, nil
+		return node, int(i), ipAddress, nil
 	}
 
-	return 0, newInvalidDatabaseError("invalid node in search tree")
+	return 0, int(i), ipAddress, newInvalidDatabaseError("invalid node in search tree")
 }
 
 func (r *Reader) readNode(nodeNumber uint, index uint) (uint, error) {

reader_test.go

@@ -51,6 +51,119 @@ func TestReaderBytes(t *testing.T) {
 	}
 }
 
+func TestLookupNetwork(t *testing.T) {
+	bigInt := new(big.Int)
+	bigInt.SetString("1329227995784915872903807060280344576", 10)
+	decoderRecord := map[string]interface{}{"array": []interface{}{uint64(1),
+		uint64(2),
+		uint64(3)},
+		"boolean": true,
+		"bytes": []uint8{
+			0x0,
+			0x0,
+			0x0,
+			0x2a,
+		},
+		"double": 42.123456,
+		"float":  float32(1.1),
+		"int32":  -268435456,
+		"map": map[string]interface{}{
+			"mapX": map[string]interface{}{
+				"arrayX": []interface{}{
+					uint64(0x7),
+					uint64(0x8),
+					uint64(0x9)},
+				"utf8_stringX": "hello",
+			},
+		},
+		"uint128":     bigInt,
+		"uint16":      uint64(0x64),
+		"uint32":      uint64(0x10000000),
+		"uint64":      uint64(0x1000000000000000),
+		"utf8_string": "unicode! ☯ - ♫",
+	}
+
+	tests := []struct {
+		IP             net.IP
+		DBFile         string
+		ExpectedCIDR   string
+		ExpectedRecord interface{}
+		ExpectedOK     bool
+	}{
+		// XXX - add test of IPv4 lookup in IPv6 database with no IPv4 subtree
+		{
+			IP:             net.ParseIP("1.1.1.1"),
+			DBFile:         "MaxMind-DB-test-ipv6-32.mmdb",
+			ExpectedCIDR:   "1.0.0.0/8",
+			ExpectedRecord: nil,
+			ExpectedOK:     false,
+		},
+		{
+			IP:             net.ParseIP("::1:ffff:ffff"),
+			DBFile:         "MaxMind-DB-test-ipv6-24.mmdb",
+			ExpectedCIDR:   "::1:ffff:ffff/128",
+			ExpectedRecord: map[string]interface{}{"ip": "::1:ffff:ffff"},
+			ExpectedOK:     true,
+		},
+		{
+			IP:             net.ParseIP("::2:0:1"),
+			DBFile:         "MaxMind-DB-test-ipv6-24.mmdb",
+			ExpectedCIDR:   "::2:0:0/122",
+			ExpectedRecord: map[string]interface{}{"ip": "::2:0:0"},
+			ExpectedOK:     true,
+		},
+		{
+			IP:             net.ParseIP("1.1.1.1"),
+			DBFile:         "MaxMind-DB-test-ipv4-24.mmdb",
+			ExpectedCIDR:   "1.1.1.1/32",
+			ExpectedRecord: map[string]interface{}{"ip": "1.1.1.1"},
+			ExpectedOK:     true,
+		},
+		{
+			IP:             net.ParseIP("1.1.1.3"),
+			DBFile:         "MaxMind-DB-test-ipv4-24.mmdb",
+			ExpectedCIDR:   "1.1.1.2/31",
+			ExpectedRecord: map[string]interface{}{"ip": "1.1.1.2"},
+			ExpectedOK:     true,
+		},
+		{
+			IP:             net.ParseIP("1.1.1.3"),
+			DBFile:         "MaxMind-DB-test-decoder.mmdb",
+			ExpectedCIDR:   "1.1.1.0/24",
+			ExpectedRecord: decoderRecord,
+			ExpectedOK:     true,
+		},
+		{
+			IP:             net.ParseIP("::ffff:1.1.1.128"),
+			DBFile:         "MaxMind-DB-test-decoder.mmdb",
+			ExpectedCIDR:   "1.1.1.0/24",
+			ExpectedRecord: decoderRecord,
+			ExpectedOK:     true,
+		},
+		{
+			IP:             net.ParseIP("::1.1.1.128"),
+			DBFile:         "MaxMind-DB-test-decoder.mmdb",
+			ExpectedCIDR:   "::101:100/120",
+			ExpectedRecord: decoderRecord,
+			ExpectedOK:     true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%s - %s", test.DBFile, test.IP), func(t *testing.T) {
+			var record interface{}
+			reader, err := Open(testFile(test.DBFile))
+			require.NoError(t, err)
+
+			network, ok, err := reader.LookupNetwork(test.IP, &record)
+			require.NoError(t, err)
+			assert.Equal(t, test.ExpectedOK, ok)
+			assert.Equal(t, test.ExpectedCIDR, network.String())
+			assert.Equal(t, test.ExpectedRecord, record)
+		})
+	}
+}
+
 func TestDecodingToInterface(t *testing.T) {
 	reader, err := Open(testFile("MaxMind-DB-test-decoder.mmdb"))
 	require.NoError(t, err, "unexpected error while opening database: %v", err)