drpcmetadata: ensure metadata immutability by returning copy

drpcmetadata/metadata.go

@@ -11,10 +11,15 @@ import (
 
 // AddPairs attaches metadata onto a context and return the context.
 func AddPairs(ctx context.Context, metadata map[string]string) context.Context {
-	for key, val := range metadata {
-		ctx = Add(ctx, key, val)
+	// Get returns a copy of metadata
+	newMetadata, ok := Get(ctx)
+	if !ok {
+		newMetadata = make(map[string]string)
+	}
+	for k, v := range metadata {
+		newMetadata[k] = v
 	}
-	return ctx
+	return context.WithValue(ctx, metadataKey{}, newMetadata)
 }
 
 // Encode generates byte form of the metadata and appends it onto the passed in buffer.
@@ -74,19 +79,27 @@ func ClearContextExcept(ctx context.Context, key string) context.Context {
 
 // Add associates a key/value pair on the context.
 func Add(ctx context.Context, key, value string) context.Context {
+	// Get returns a copy of metadata
 	metadata, ok := Get(ctx)
 	if !ok {
 		metadata = make(map[string]string)
-		ctx = context.WithValue(ctx, metadataKey{}, metadata)
 	}
 	metadata[key] = value
-	return ctx
+	return context.WithValue(ctx, metadataKey{}, metadata)
 }
 
 // Get returns all key/value pairs on the given context.
 func Get(ctx context.Context) (map[string]string, bool) {
 	metadata, ok := ctx.Value(metadataKey{}).(map[string]string)
-	return metadata, ok
+	if !ok {
+		return nil, false
+	}
+	// Return a copy to prevent mutation of the original map
+	copy := make(map[string]string)
+	for k, v := range metadata {
+		copy[k] = v
+	}
+	return copy, true
 }
 
 // GetValue retrieves a specific value by key from the context's metadata.

drpcmetadata/metadata_test.go

@@ -77,3 +77,62 @@ func TestDecode(t *testing.T) {
 		assert.DeepEqual(t, metadata, map[string]string{"test": "a"})
 	})
 }
+
+func TestMetadataImmutability(t *testing.T) {
+	ctx := context.Background()
+	ctx = Add(ctx, "foo", "bar")
+
+	metadata1, ok := Get(ctx)
+	assert.That(t, ok)
+	assert.Equal(t, metadata1["foo"], "bar")
+
+	metadata1["foo"] = "modified"
+	metadata1["new"] = "value"
+
+	metadata2, ok := Get(ctx)
+	assert.That(t, ok)
+	assert.Equal(t, metadata2["foo"], "bar")
+	assert.Equal(t, len(metadata2), 1)
+}
+
+func TestAddImmutability(t *testing.T) {
+	ctx := context.Background()
+	ctx = Add(ctx, "original", "value")
+
+	originalCtx := ctx
+	newCtx := Add(ctx, "new", "key")
+
+	originalMd, ok := Get(originalCtx)
+	assert.That(t, ok)
+	assert.Equal(t, len(originalMd), 1)
+	assert.Equal(t, originalMd["original"], "value")
+
+	newMd, ok := Get(newCtx)
+	assert.That(t, ok)
+	assert.Equal(t, len(newMd), 2)
+	assert.Equal(t, newMd["original"], "value")
+	assert.Equal(t, newMd["new"], "key")
+}
+
+func TestAddPairsImmutability(t *testing.T) {
+	ctx := context.Background()
+	ctx = Add(ctx, "existing", "value")
+
+	originalCtx := ctx
+	newCtx := AddPairs(ctx, map[string]string{
+		"key1": "val1",
+		"key2": "val2",
+	})
+
+	originalMd, ok := Get(originalCtx)
+	assert.That(t, ok)
+	assert.Equal(t, len(originalMd), 1)
+	assert.Equal(t, originalMd["existing"], "value")
+
+	newMd, ok := Get(newCtx)
+	assert.That(t, ok)
+	assert.Equal(t, len(newMd), 3)
+	assert.Equal(t, newMd["existing"], "value")
+	assert.Equal(t, newMd["key1"], "val1")
+	assert.Equal(t, newMd["key2"], "val2")
+}