GODRIVER-1751 Ensure codecs that cache are not shared across registries (#508)

Author: Divjot Arora

bson/bsoncodec/default_value_decoders.go

@@ -24,6 +24,16 @@ import (
 
 var defaultValueDecoders DefaultValueDecoders
 
+func newDefaultStructCodec() *StructCodec {
+	codec, err := NewStructCodec(DefaultStructTagParser)
+	if err != nil {
+		// This function is called from the codec registration path, so errors can't be propagated. If there's an error
+		// constructing the StructCodec, we panic to avoid losing it.
+		panic(fmt.Errorf("error creating default StructCodec: %v", err))
+	}
+	return codec
+}
+
 // DefaultValueDecoders is a namespace type for the default ValueDecoders used
 // when creating a registry.
 type DefaultValueDecoders struct{}
@@ -77,7 +87,7 @@ func (dvd DefaultValueDecoders) RegisterDefaultDecoders(rb *RegistryBuilder) {
 		RegisterDefaultDecoder(reflect.Map, defaultMapCodec).
 		RegisterDefaultDecoder(reflect.Slice, defaultSliceCodec).
 		RegisterDefaultDecoder(reflect.String, defaultStringCodec).
-		RegisterDefaultDecoder(reflect.Struct, defaultStructCodec).
+		RegisterDefaultDecoder(reflect.Struct, newDefaultStructCodec()).
 		RegisterDefaultDecoder(reflect.Ptr, NewPointerCodec()).
 		RegisterTypeMapEntry(bsontype.Double, tFloat64).
 		RegisterTypeMapEntry(bsontype.String, tString).

bson/bsoncodec/default_value_decoders_test.go

@@ -27,6 +27,10 @@ import (
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
+var (
+	defaultTestStructCodec = newDefaultStructCodec()
+)
+
 func TestDefaultValueDecoders(t *testing.T) {
 	var dvd DefaultValueDecoders
 	var wrong = func(string, string) string { return "wrong" }
@@ -2197,7 +2201,7 @@ func TestDefaultValueDecoders(t *testing.T) {
 		},
 		{
 			"StructCodec.DecodeValue",
-			defaultStructCodec,
+			defaultTestStructCodec,
 			[]subtest{
 				{
 					"Not struct",
@@ -3497,7 +3501,7 @@ func TestDefaultValueDecoders(t *testing.T) {
 				emptyInterfaceStruct{},
 				bsonrw.NewBSONDocumentReader(docBytes),
 				emptyInterfaceErrorRegistry,
-				defaultStructCodec,
+				defaultTestStructCodec,
 				emptyInterfaceStructErr,
 			},
 			{
@@ -3508,15 +3512,15 @@ func TestDefaultValueDecoders(t *testing.T) {
 				stringStruct{},
 				bsonrw.NewBSONDocumentReader(docBytes),
 				NewRegistryBuilder().Build(),
-				defaultStructCodec,
+				defaultTestStructCodec,
 				stringStructErr,
 			},
 			{
 				"deeply nested struct",
 				outer{},
 				bsonrw.NewBSONDocumentReader(outerDoc),
 				nestedRegistry,
-				defaultStructCodec,
+				defaultTestStructCodec,
 				nestedErr,
 			},
 		}
@@ -3546,7 +3550,7 @@ func TestDefaultValueDecoders(t *testing.T) {
 			dc := DecodeContext{Registry: buildDefaultRegistry()}
 			vr := bsonrw.NewBSONDocumentReader(outerBytes)
 			val := reflect.New(reflect.TypeOf(outer{})).Elem()
-			err := defaultStructCodec.DecodeValue(dc, vr, val)
+			err := defaultTestStructCodec.DecodeValue(dc, vr, val)
 
 			decodeErr, ok := err.(*DecodeError)
 			assert.True(t, ok, "expected DecodeError, got %v of type %T", err, err)

bson/bsoncodec/default_value_encoders.go

@@ -104,7 +104,7 @@ func (dve DefaultValueEncoders) RegisterDefaultEncoders(rb *RegistryBuilder) {
 		RegisterDefaultEncoder(reflect.Map, defaultMapCodec).
 		RegisterDefaultEncoder(reflect.Slice, defaultSliceCodec).
 		RegisterDefaultEncoder(reflect.String, defaultStringCodec).
-		RegisterDefaultEncoder(reflect.Struct, defaultStructCodec).
+		RegisterDefaultEncoder(reflect.Struct, newDefaultStructCodec()).
 		RegisterDefaultEncoder(reflect.Ptr, NewPointerCodec()).
 		RegisterHookEncoder(tValueMarshaler, ValueEncoderFunc(dve.ValueMarshalerEncodeValue)).
 		RegisterHookEncoder(tMarshaler, ValueEncoderFunc(dve.MarshalerEncodeValue)).

bson/bsoncodec/default_value_encoders_test.go

@@ -1084,7 +1084,7 @@ func TestDefaultValueEncoders(t *testing.T) {
 		},
 		{
 			"StructEncodeValue",
-			defaultStructCodec,
+			defaultTestStructCodec,
 			[]subtest{
 				{
 					"interface value",

bson/bsoncodec/pointer_codec.go

@@ -14,11 +14,6 @@ import (
 	"go.mongodb.org/mongo-driver/bson/bsontype"
 )
 
-var defaultPointerCodec = &PointerCodec{
-	ecache: make(map[reflect.Type]ValueEncoder),
-	dcache: make(map[reflect.Type]ValueDecoder),
-}
-
 var _ ValueEncoder = &PointerCodec{}
 var _ ValueDecoder = &PointerCodec{}
 

bson/bsoncodec/struct_codec.go

@@ -19,11 +19,6 @@ import (
 	"go.mongodb.org/mongo-driver/bson/bsontype"
 )
 
-var defaultStructCodec = &StructCodec{
-	cache:  make(map[reflect.Type]*structDescription),
-	parser: DefaultStructTagParser,
-}
-
 // DecodeError represents an error that occurs when unmarshalling BSON bytes into a native Go type.
 type DecodeError struct {
 	keys    []string

bson/marshal_test.go

@@ -8,12 +8,17 @@ package bson
 
 import (
 	"bytes"
+	"fmt"
+	"reflect"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
 	"go.mongodb.org/mongo-driver/bson/bsoncodec"
+	"go.mongodb.org/mongo-driver/bson/bsonrw"
 	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
+	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
 func TestMarshalAppendWithRegistry(t *testing.T) {
@@ -207,3 +212,58 @@ func TestMarshal_roundtripFromDoc(t *testing.T) {
 		t.Errorf("Documents to not match. got %v; want %v", after, before)
 	}
 }
+
+func TestCachingEncodersNotSharedAcrossRegistries(t *testing.T) {
+	// Encoders that have caches for recursive encoder lookup should not be shared across Registry instances. Otherwise,
+	// the first EncodeValue call would cache an encoder and a subsequent call would see that encoder even if a
+	// different Registry is used.
+
+	// Create a custom Registry that negates int32 values when encoding.
+	var encodeInt32 bsoncodec.ValueEncoderFunc = func(_ bsoncodec.EncodeContext, vw bsonrw.ValueWriter, val reflect.Value) error {
+		if val.Kind() != reflect.Int32 {
+			return fmt.Errorf("expected kind to be int32, got %v", val.Kind())
+		}
+
+		return vw.WriteInt32(int32(val.Int()) * -1)
+	}
+	customReg := NewRegistryBuilder().
+		RegisterTypeEncoder(tInt32, encodeInt32).
+		Build()
+
+	// Helper function to run the test and make assertions. The provided original value should result in the document
+	// {"x": {$numberInt: 1}} when marshalled with the default registry.
+	verifyResults := func(t *testing.T, original interface{}) {
+		// Marshal using the default and custom registries. Assert that the result is {x: 1} and {x: -1}, respectively.
+
+		first, err := Marshal(original)
+		assert.Nil(t, err, "Marshal error: %v", err)
+		expectedFirst := Raw(bsoncore.BuildDocumentFromElements(
+			nil,
+			bsoncore.AppendInt32Element(nil, "x", 1),
+		))
+		assert.Equal(t, expectedFirst, Raw(first), "expected document %v, got %v", expectedFirst, Raw(first))
+
+		second, err := MarshalWithRegistry(customReg, original)
+		assert.Nil(t, err, "Marshal error: %v", err)
+		expectedSecond := Raw(bsoncore.BuildDocumentFromElements(
+			nil,
+			bsoncore.AppendInt32Element(nil, "x", -1),
+		))
+		assert.Equal(t, expectedSecond, Raw(second), "expected document %v, got %v", expectedSecond, Raw(second))
+	}
+
+	t.Run("struct", func(t *testing.T) {
+		type Struct struct {
+			X int32
+		}
+		verifyResults(t, Struct{
+			X: 1,
+		})
+	})
+	t.Run("pointer", func(t *testing.T) {
+		i32 := int32(1)
+		verifyResults(t, M{
+			"x": &i32,
+		})
+	})
+}

bson/unmarshal_test.go

@@ -13,6 +13,8 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"go.mongodb.org/mongo-driver/bson/bsoncodec"
 	"go.mongodb.org/mongo-driver/bson/bsonrw"
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
+	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
 func TestUnmarshal(t *testing.T) {
@@ -106,3 +108,61 @@ func TestUnmarshalExtJSONWithContext(t *testing.T) {
 		}
 	})
 }
+
+func TestCachingDecodersNotSharedAcrossRegistries(t *testing.T) {
+	// Decoders that have caches for recursive decoder lookup should not be shared across Registry instances. Otherwise,
+	// the first DecodeValue call would cache an decoder and a subsequent call would see that decoder even if a
+	// different Registry is used.
+
+	// Create a custom Registry that negates BSON int32 values when decoding.
+	var decodeInt32 bsoncodec.ValueDecoderFunc = func(_ bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
+		i32, err := vr.ReadInt32()
+		if err != nil {
+			return err
+		}
+
+		val.SetInt(int64(-1 * i32))
+		return nil
+	}
+	customReg := NewRegistryBuilder().
+		RegisterTypeDecoder(tInt32, bsoncodec.ValueDecoderFunc(decodeInt32)).
+		Build()
+
+	docBytes := bsoncore.BuildDocumentFromElements(
+		nil,
+		bsoncore.AppendInt32Element(nil, "x", 1),
+	)
+
+	// For all sub-tests, unmarshal docBytes into a struct and assert that value for "x" is 1 when using the default
+	// registry and -1 when using the custom registry.
+	t.Run("struct", func(t *testing.T) {
+		type Struct struct {
+			X int32
+		}
+
+		var first Struct
+		err := Unmarshal(docBytes, &first)
+		assert.Nil(t, err, "Unmarshal error: %v", err)
+		assert.Equal(t, int32(1), first.X, "expected X value to be 1, got %v", first.X)
+
+		var second Struct
+		err = UnmarshalWithRegistry(customReg, docBytes, &second)
+		assert.Nil(t, err, "Unmarshal error: %v", err)
+		assert.Equal(t, int32(-1), second.X, "expected X value to be -1, got %v", second.X)
+	})
+	t.Run("pointer", func(t *testing.T) {
+		type Struct struct {
+			X *int32
+		}
+
+		var first Struct
+		err := Unmarshal(docBytes, &first)
+		assert.Nil(t, err, "Unmarshal error: %v", err)
+		assert.Equal(t, int32(1), *first.X, "expected X value to be 1, got %v", *first.X)
+
+		var second Struct
+		err = UnmarshalWithRegistry(customReg, docBytes, &second)
+		assert.Nil(t, err, "Unmarshal error: %v", err)
+		assert.Equal(t, int32(-1), *second.X, "expected X value to be -1, got %v", *second.X)
+	})
+}