Enable nested UnmarshalMaxMindDB support

Extend the UnmarshalMaxMindDB interface to work recursively with nested
types, matching the behavior of encoding/json's UnmarshalJSON. Custom
unmarshalers are now called for:

- Struct fields that implement Unmarshaler
- Pointer fields (creates value if nil, then checks for Unmarshaler)
- Slice elements that implement Unmarshaler
- Map values that implement Unmarshaler

This enhancement allows for more flexible custom decoding strategies in
complex data structures, improving performance optimization opportunities
for nested types.

Author: oschwald

CHANGELOG.md

@@ -13,6 +13,10 @@
   `json.Unmarshaler`.
 - Added public `Decoder` type with methods for manual decoding including
   `DecodeMap()`, `DecodeSlice()`, `DecodeString()`, `DecodeUInt32()`, etc.
+- Enhanced `UnmarshalMaxMindDB` to work with nested struct fields, slice
+  elements, and map values. The custom unmarshaler is now called recursively
+  for any type that implements the `Unmarshaler` interface, similar to
+  `encoding/json`.
 
 ## 2.0.0-beta.3 - 2025-02-16
 

internal/decoder/decoder.go

@@ -1,6 +1,7 @@
 package decoder
 
 import (
+	"errors"
 	"fmt"
 	"iter"
 
@@ -362,6 +363,13 @@ func (d *Decoder) setNextOffset(offset uint) {
 	}
 }
 
+func (d *Decoder) getNextOffset() (uint, error) {
+	if !d.hasNextOffset {
+		return 0, errors.New("no next offset available")
+	}
+	return d.nextOffset, nil
+}
+
 func unexpectedTypeErr(expectedType, actualType Type) error {
 	return fmt.Errorf("unexpected type %d, expected %d", actualType, expectedType)
 }

internal/decoder/decoder_test.go

@@ -382,16 +382,22 @@ func TestBoundsChecking(t *testing.T) {
 	require.Contains(t, err.Error(), "exceeds buffer length")
 
 	// Test DecodeBytes bounds checking with a separate buffer
-	bytesBuffer := []byte{0x84, 0x41} // Type bytes (4 << 5 = 0x80), size 4 (0x04), but only 2 bytes total
+	bytesBuffer := []byte{
+		0x84,
+		0x41,
+	} // Type bytes (4 << 5 = 0x80), size 4 (0x04), but only 2 bytes total
 	dd3 := NewDataDecoder(bytesBuffer)
 	decoder3 := &Decoder{d: dd3, offset: 0}
-	
+
 	_, err = decoder3.DecodeBytes()
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "exceeds buffer length")
 
-	// Test DecodeUInt128 bounds checking  
-	uint128Buffer := []byte{0x0B, 0x03} // Extended type (0x0), size 11, TypeUint128-7=3, but only 2 bytes total
+	// Test DecodeUInt128 bounds checking
+	uint128Buffer := []byte{
+		0x0B,
+		0x03,
+	} // Extended type (0x0), size 11, TypeUint128-7=3, but only 2 bytes total
 	dd2 := NewDataDecoder(uint128Buffer)
 	decoder2 := &Decoder{d: dd2, offset: 0}
 

internal/decoder/nested_unmarshaler_test.go

@@ -0,0 +1,222 @@
+package decoder
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// Inner type with UnmarshalMaxMindDB.
+type testInnerNested struct {
+	Value  string
+	custom bool // track if custom unmarshaler was called
+}
+
+func (i *testInnerNested) UnmarshalMaxMindDB(d *Decoder) error {
+	i.custom = true
+	str, err := d.DecodeString()
+	if err != nil {
+		return err
+	}
+	i.Value = "custom:" + str
+	return nil
+}
+
+// TestNestedUnmarshaler tests that UnmarshalMaxMindDB is called for nested struct fields.
+func TestNestedUnmarshaler(t *testing.T) {
+	// Outer type without UnmarshalMaxMindDB
+	type Outer struct {
+		Field testInnerNested
+		Name  string
+	}
+
+	// Create test data: a map with "Field" -> "test" and "Name" -> "example"
+	data := []byte{
+		// Map with 2 items
+		0xe2,
+		// Key "Field"
+		0x45, 'F', 'i', 'e', 'l', 'd',
+		// Value "test" (string)
+		0x44, 't', 'e', 's', 't',
+		// Key "Name"
+		0x44, 'N', 'a', 'm', 'e',
+		// Value "example" (string)
+		0x47, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
+	}
+
+	t.Run("nested field with UnmarshalMaxMindDB", func(t *testing.T) {
+		d := New(data)
+		var result Outer
+
+		err := d.Decode(0, &result)
+		require.NoError(t, err)
+
+		// Check that custom unmarshaler WAS called for nested field
+		require.True(
+			t,
+			result.Field.custom,
+			"Custom unmarshaler should be called for nested fields",
+		)
+		require.Equal(t, "custom:test", result.Field.Value)
+		require.Equal(t, "example", result.Name)
+	})
+}
+
+// testInnerPointer with UnmarshalMaxMindDB for pointer test.
+type testInnerPointer struct {
+	Value  string
+	custom bool
+}
+
+func (i *testInnerPointer) UnmarshalMaxMindDB(d *Decoder) error {
+	i.custom = true
+	str, err := d.DecodeString()
+	if err != nil {
+		return err
+	}
+	i.Value = "ptr:" + str
+	return nil
+}
+
+// TestNestedUnmarshalerPointer tests UnmarshalMaxMindDB with pointer fields.
+func TestNestedUnmarshalerPointer(t *testing.T) {
+	type Outer struct {
+		Field *testInnerPointer
+		Name  string
+	}
+
+	// Test data
+	data := []byte{
+		// Map with 2 items
+		0xe2,
+		// Key "Field"
+		0x45, 'F', 'i', 'e', 'l', 'd',
+		// Value "test"
+		0x44, 't', 'e', 's', 't',
+		// Key "Name"
+		0x44, 'N', 'a', 'm', 'e',
+		// Value "example"
+		0x47, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
+	}
+
+	t.Run("pointer field with UnmarshalMaxMindDB", func(t *testing.T) {
+		d := New(data)
+		var result Outer
+		err := d.Decode(0, &result)
+		require.NoError(t, err)
+
+		// The pointer should be created and unmarshaled with custom unmarshaler
+		require.NotNil(t, result.Field)
+		require.True(
+			t,
+			result.Field.custom,
+			"Custom unmarshaler should be called for pointer fields",
+		)
+		require.Equal(t, "ptr:test", result.Field.Value)
+		require.Equal(t, "example", result.Name)
+	})
+}
+
+// testItem with UnmarshalMaxMindDB for slice test.
+type testItem struct {
+	ID     int
+	custom bool
+}
+
+func (item *testItem) UnmarshalMaxMindDB(d *Decoder) error {
+	item.custom = true
+	id, err := d.DecodeUInt32()
+	if err != nil {
+		return err
+	}
+	item.ID = int(id) * 2
+	return nil
+}
+
+// TestNestedUnmarshalerInSlice tests UnmarshalMaxMindDB for slice elements.
+func TestNestedUnmarshalerInSlice(t *testing.T) {
+	type Container struct {
+		Items []testItem
+	}
+
+	// Test data: a map with "Items" -> [1, 2, 3]
+	data := []byte{
+		// Map with 1 item (TypeMap=7 << 5 | size=1)
+		0xe1,
+		// Key "Items" (TypeString=2 << 5 | size=5)
+		0x45, 'I', 't', 'e', 'm', 's',
+		// Slice with 3 items - TypeSlice=11, which is > 7, so we need extended type
+		// Extended type: ctrl_byte = (TypeExtended << 5) | size = (0 << 5) | 3 = 0x03
+		// Next byte: TypeSlice - 7 = 11 - 7 = 4
+		0x03, 0x04,
+		// Value 1 (TypeUint32=6 << 5 | size=1)
+		0xc1, 0x01,
+		// Value 2 (TypeUint32=6 << 5 | size=1)
+		0xc1, 0x02,
+		// Value 3 (TypeUint32=6 << 5 | size=1)
+		0xc1, 0x03,
+	}
+
+	t.Run("slice elements with UnmarshalMaxMindDB", func(t *testing.T) {
+		d := New(data)
+		var result Container
+		err := d.Decode(0, &result)
+		require.NoError(t, err)
+
+		require.Len(t, result.Items, 3)
+		// With custom unmarshaler, values should be doubled
+		require.True(
+			t,
+			result.Items[0].custom,
+			"Custom unmarshaler should be called for slice elements",
+		)
+		require.Equal(t, 2, result.Items[0].ID) // 1 * 2
+		require.Equal(t, 4, result.Items[1].ID) // 2 * 2
+		require.Equal(t, 6, result.Items[2].ID) // 3 * 2
+	})
+}
+
+// testValue with UnmarshalMaxMindDB for map test.
+type testValue struct {
+	Data   string
+	custom bool
+}
+
+func (v *testValue) UnmarshalMaxMindDB(d *Decoder) error {
+	v.custom = true
+	str, err := d.DecodeString()
+	if err != nil {
+		return err
+	}
+	v.Data = "map:" + str
+	return nil
+}
+
+// TestNestedUnmarshalerInMap tests UnmarshalMaxMindDB for map values.
+func TestNestedUnmarshalerInMap(t *testing.T) {
+	// Test data: {"key1": "value1", "key2": "value2"}
+	data := []byte{
+		// Map with 2 items
+		0xe2,
+		// Key "key1"
+		0x44, 'k', 'e', 'y', '1',
+		// Value "value1"
+		0x46, 'v', 'a', 'l', 'u', 'e', '1',
+		// Key "key2"
+		0x44, 'k', 'e', 'y', '2',
+		// Value "value2"
+		0x46, 'v', 'a', 'l', 'u', 'e', '2',
+	}
+
+	t.Run("map values with UnmarshalMaxMindDB", func(t *testing.T) {
+		d := New(data)
+		var result map[string]testValue
+		err := d.Decode(0, &result)
+		require.NoError(t, err)
+
+		require.Len(t, result, 2)
+		require.True(t, result["key1"].custom, "Custom unmarshaler should be called for map values")
+		require.Equal(t, "map:value1", result["key1"].Data)
+		require.Equal(t, "map:value2", result["key2"].Data)
+	})
+}

internal/decoder/reflection.go

@@ -143,6 +143,37 @@ func (d *ReflectionDecoder) decode(offset uint, result reflect.Value, depth int)
 			"exceeded maximum data structure depth; database is likely corrupt",
 		)
 	}
+
+	// First handle pointers by creating the value if needed, similar to indirect()
+	// but we don't want to fully indirect yet as we need to check for Unmarshaler
+	if result.Kind() == reflect.Ptr {
+		if result.IsNil() {
+			result.Set(reflect.New(result.Type().Elem()))
+		}
+		// Now check if the pointed-to type implements Unmarshaler
+		if unmarshaler, ok := result.Interface().(Unmarshaler); ok {
+			decoder := &Decoder{d: d.DataDecoder, offset: offset}
+			if err := unmarshaler.UnmarshalMaxMindDB(decoder); err != nil {
+				return 0, err
+			}
+			return decoder.getNextOffset()
+		}
+		// Continue with the pointed-to value
+		return d.decode(offset, result.Elem(), depth)
+	}
+
+	// Check if the value implements Unmarshaler interface
+	// We need to check if result can be addressed and if the pointer type implements Unmarshaler
+	if result.CanAddr() {
+		if unmarshaler, ok := result.Addr().Interface().(Unmarshaler); ok {
+			decoder := &Decoder{d: d.DataDecoder, offset: offset}
+			if err := unmarshaler.UnmarshalMaxMindDB(decoder); err != nil {
+				return 0, err
+			}
+			return decoder.getNextOffset()
+		}
+	}
+
 	typeNum, size, newOffset, err := d.decodeCtrlData(offset)
 	if err != nil {
 		return 0, err