Merge pull request #23 from arbortech/offset-lookups

Add Reader.LookupOffset & Decode

Author: oschwald

decoder.go

@@ -35,7 +35,13 @@ const (
 
 func (d *decoder) decode(offset uint, result reflect.Value) (uint, error) {
 	typeNum, size, newOffset := d.decodeCtrlData(offset)
-	return d.decodeFromType(typeNum, size, newOffset, result)
+
+	if typeNum != _Pointer && result.Kind() == reflect.Uintptr {
+		result.Set(reflect.ValueOf(uintptr(offset)))
+		return d.nextValueOffset(offset, 1), nil
+	} else {
+		return d.decodeFromType(typeNum, size, newOffset, result)
+	}
 }
 
 func (d *decoder) decodeCtrlData(offset uint) (dataType, uint, uint) {

reader.go

@@ -8,7 +8,12 @@ import (
 	"reflect"
 )
 
-const dataSectionSeparatorSize = 16
+const (
+	// Returned by LookupOffset when a matched root record offset cannot be found.
+	NotFound = ^uintptr(0)
+
+	dataSectionSeparatorSize = 16
+)
 
 var metadataStartMarker = []byte("\xAB\xCD\xEFMaxMind.com")
 
@@ -99,39 +104,72 @@ func (r *Reader) startNode() (uint, error) {
 }
 
 // Lookup takes an IP address as a net.IP structure and a pointer to the
-// result value to decode into. 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 data or an
-// empty interface{} value.
+// result value to Decode into.
+func (r *Reader) Lookup(ipAddress net.IP, result interface{}) error {
+	if pointer, err := r.lookupPointer(ipAddress); pointer == 0 {
+		return err
+	} else {
+		return 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
+// is an advanced API, which exists to provide clients with a means to cache
+// previously-decoded records.
+func (r *Reader) LookupOffset(ipAddress net.IP) (uintptr, error) {
+	if pointer, err := r.lookupPointer(ipAddress); pointer == 0 {
+		return NotFound, err
+	} else {
+		return r.resolveDataPointer(pointer)
+	}
+}
+
+// Decodes 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
+// data or an empty interface{} value.
 //
 // If result is a pointer to a struct, the struct need not include a field
 // for every value that may be in the database. If a field is not present in
 // the structure, the decoder will not decode that field, reducing the time
 // required to decode the record.
 //
+// As a special case, a struct field of type uintptr will be used to capture
+// the offset of the value. Decode may later be used to extract the stored
+// value from the offset. MaxmindDBs are highly normalized: for example in the
+// Cities database, all records of the same country will reference a single
+// representative record for that country. This uintptr behavior allows clients
+// to leverage this normalization in their own sub-record caching.
+//
 // Currently the decoder expect most data types to correspond exactly (e.g.,
-// a uint64 database type must be decoded into a uint64 Go type). In the
-// future, this may be made more flexible.
-func (r *Reader) Lookup(ipAddress net.IP, result interface{}) error {
+// a uint64 database type must be decoded into a uint64 Go type), with the
+// exception of uintptr. In the future, this may be made more flexible.
+func (r *Reader) Decode(offset uintptr, result interface{}) error {
+	rv := reflect.ValueOf(result)
+	if rv.Kind() != reflect.Ptr || rv.IsNil() {
+		return errors.New("result param must be a pointer")
+	}
+
+	_, err := r.decoder.decode(uint(offset), reflect.ValueOf(result))
+	return err
+}
+
+func (r *Reader) lookupPointer(ipAddress net.IP) (uint, error) {
 	if ipAddress == nil {
-		return errors.New("ipAddress passed to Lookup cannot be nil")
+		return 0, 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 fmt.Errorf("error looking up '%s': you attempted to look up an IPv6 address in an IPv4-only database", ipAddress.String())
+		return 0, fmt.Errorf("error looking up '%s': you attempted to look up an IPv6 address in an IPv4-only database", ipAddress.String())
 	}
 
-	pointer, err := r.findAddressInTree(ipAddress)
-
-	if pointer == 0 {
-		return err
-	}
-
-	return r.retrieveData(pointer, result)
+	return r.findAddressInTree(ipAddress)
 }
 
 func (r *Reader) findAddressInTree(ipAddress net.IP) (uint, error) {
@@ -194,28 +232,18 @@ func (r *Reader) readNode(nodeNumber uint, index uint) (uint, error) {
 }
 
 func (r *Reader) retrieveData(pointer uint, result interface{}) error {
-	rv := reflect.ValueOf(result)
-	if rv.Kind() != reflect.Ptr || rv.IsNil() {
-		return errors.New("result param must be a pointer")
-	}
-
-	offset, err := r.resolveDataPointer(pointer)
-	if err != nil {
+	if offset, err := r.resolveDataPointer(pointer); err != nil {
 		return err
+	} else {
+		return r.Decode(offset, result)
 	}
-
-	_, err = r.decoder.decode(offset, rv)
-	return err
 }
 
-func (r *Reader) resolveDataPointer(pointer uint) (uint, error) {
-	nodeCount := r.Metadata.NodeCount
-
-	resolved := pointer - nodeCount - dataSectionSeparatorSize
+func (r *Reader) resolveDataPointer(pointer uint) (uintptr, error) {
+	var resolved = uintptr(pointer - r.Metadata.NodeCount - dataSectionSeparatorSize)
 
-	if resolved > uint(len(r.buffer)) {
+	if resolved > uintptr(len(r.buffer)) {
 		return 0, newInvalidDatabaseError("the MaxMind DB file's search tree is corrupt")
 	}
-
 	return resolved, nil
 }

reader_test.go

@@ -116,43 +116,89 @@ type TestType struct {
 
 func (s *MySuite) TestDecoder(c *C) {
 	reader, err := Open("test-data/test-data/MaxMind-DB-test-decoder.mmdb")
-	if err != nil {
-		c.Logf("unexpected error while opening database: %v", err)
-		c.Fail()
-	}
+	c.Assert(err, IsNil)
+
+	verify := func(result TestType) {
+		c.Assert(result.Array, DeepEquals, []uint{uint(1), uint(2), uint(3)})
+		c.Assert(result.Boolean, Equals, true)
+		c.Assert(result.Bytes, DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
+		c.Assert(result.Double, Equals, 42.123456)
+		c.Assert(result.Float, Equals, float32(1.1))
+		c.Assert(result.Int32, Equals, int32(-268435456))
+
+		c.Assert(result.Map, DeepEquals,
+			map[string]interface{}{
+				"mapX": map[string]interface{}{
+					"arrayX":       []interface{}{uint64(7), uint64(8), uint64(9)},
+					"utf8_stringX": "hello",
+				}})
+
+		c.Assert(result.Uint16, Equals, uint16(100))
+		c.Assert(result.Uint32, Equals, uint32(268435456))
+		c.Assert(result.Uint64, Equals, uint64(1152921504606846976))
+		c.Assert(result.Utf8String, Equals, "unicode! ☯ - ♫")
+		bigInt := new(big.Int)
+		bigInt.SetString("1329227995784915872903807060280344576", 10)
+		c.Assert(&result.Uint128, DeepEquals, bigInt)
+	}
+
+	{
+		// Directly lookup and decode.
+		var result TestType
+		c.Assert(reader.Lookup(net.ParseIP("::1.1.1.0"), &result), IsNil)
+		verify(result)
+	}
+	{
+		// Lookup record offset, then Decode.
+		var result TestType
+		offset, err := reader.LookupOffset(net.ParseIP("::1.1.1.0"))
+		c.Assert(err, IsNil)
+		c.Assert(offset, Not(Equals), NotFound)
+
+		c.Assert(reader.Decode(offset, &result), IsNil)
+		verify(result)
+	}
+
+	c.Assert(reader.Close(), IsNil)
+}
 
-	var result TestType
-	err = reader.Lookup(net.ParseIP("::1.1.1.0"), &result)
-	if err != nil {
-		c.Log(err)
-		c.Fail()
-	}
+func (s *MySuite) TestNestedOffsetDecode(c *C) {
+	db, err := Open("test-data/test-data/GeoIP2-City-Test.mmdb")
+	c.Assert(err, IsNil)
 
-	c.Assert(result.Array, DeepEquals, []uint{uint(1), uint(2), uint(3)})
-	c.Assert(result.Boolean, Equals, true)
-	c.Assert(result.Bytes, DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
-	c.Assert(result.Double, Equals, 42.123456)
-	c.Assert(result.Float, Equals, float32(1.1))
-	c.Assert(result.Int32, Equals, int32(-268435456))
+	off, err := db.LookupOffset(net.ParseIP("81.2.69.142"))
+	c.Assert(off, Not(Equals), NotFound)
+	c.Check(err, IsNil)
 
-	c.Assert(result.Map, DeepEquals,
-		map[string]interface{}{
-			"mapX": map[string]interface{}{
-				"arrayX":       []interface{}{uint64(7), uint64(8), uint64(9)},
-				"utf8_stringX": "hello",
-			}})
+	var root struct {
+		CountryOffset uintptr `maxminddb:"country"`
 
-	c.Assert(result.Uint16, Equals, uint16(100))
-	c.Assert(result.Uint32, Equals, uint32(268435456))
-	c.Assert(result.Uint64, Equals, uint64(1152921504606846976))
-	c.Assert(result.Utf8String, Equals, "unicode! ☯ - ♫")
-	bigInt := new(big.Int)
-	bigInt.SetString("1329227995784915872903807060280344576", 10)
-	c.Assert(&result.Uint128, DeepEquals, bigInt)
+		Location struct {
+			Latitude float64 `maxminddb:"latitude"`
+			// Longitude is directly nested within the parent map.
+			LongitudeOffset uintptr `maxminddb:"longitude"`
+			// TimeZone is indirected via a pointer.
+			TimeZoneOffset uintptr `maxminddb:"time_zone"`
+		} `maxminddb:"location"`
+	}
+	c.Check(db.Decode(off, &root), IsNil)
+	c.Check(root.Location.Latitude, Equals, 51.5142)
 
-	if err = reader.Close(); err != nil {
-		c.Assert(err, nil, "no error on close")
+	var longitude float64
+	c.Check(db.Decode(root.Location.LongitudeOffset, &longitude), IsNil)
+	c.Check(longitude, Equals, -0.0931)
+
+	var timeZone string
+	c.Check(db.Decode(root.Location.TimeZoneOffset, &timeZone), IsNil)
+	c.Check(timeZone, Equals, "Europe/London")
+
+	var country struct {
+		IsoCode string `maxminddb:"iso_code"`
 	}
+	c.Check(db.Decode(root.CountryOffset, &country), IsNil)
+	c.Check(country.IsoCode, Equals, "GB")
+
+	c.Check(db.Close(), IsNil)
 }
 
 func (s *MySuite) TestDecodingUint16IntoInt(c *C) {

verifier.go

@@ -103,7 +103,7 @@ func (v *verifier) verifySearchTree() (map[uint]bool, error) {
 		if err != nil {
 			return nil, err
 		}
-		offsets[offset] = true
+		offsets[uint(offset)] = true
 	}
 	if err := it.Err(); err != nil {
 		return nil, err