From 06a7052c703df6dfd5d71b304bc70ad9a7b1d1b5 Mon Sep 17 00:00:00 2001
From: Andreas Bergmeier <andreas.bergmeier@otto.de>
Date: Fri, 5 Apr 2024 16:00:18 +0200
Subject: [PATCH] Ensure that editing fetched Attributes is race-free in any
 case.

Signed-off-by: Andreas Bergmeier <andreas.bergmeier@otto.de>
---
 protocol/pubsub/v2/attributes.go | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/protocol/pubsub/v2/attributes.go b/protocol/pubsub/v2/attributes.go
index 714f9c040..88b900de7 100644
--- a/protocol/pubsub/v2/attributes.go
+++ b/protocol/pubsub/v2/attributes.go
@@ -14,12 +14,34 @@ import (
 type withCustomAttributes struct{}
 
 func AttributesFrom(ctx context.Context) map[string]string {
-	return binding.GetOrDefaultFromCtx(ctx, withCustomAttributes{}, make(map[string]string)).(map[string]string)
+	ctxVal := binding.GetOrDefaultFromCtx(ctx, withCustomAttributes{}, nil)
+	if ctxVal == nil {
+		return make(map[string]string, 0)
+	}
+
+	m := ctxVal.(map[string]string)
+
+	// Since it is possible that we get the same map from one ctx multiple times
+	// we need to make sure, that it is race-free to modify returned map.
+	cp := make(map[string]string, len(m))
+	for k, v := range m {
+		cp[k] = v
+	}
+	return cp
 }
 
 // WithCustomAttributes sets Message Attributes without any CloudEvent logic.
 // Note that this function is not intended for CloudEvent Extensions or any `ce-`-prefixed Attributes.
 // For these please see `Event` and `Event.SetExtension`.
 func WithCustomAttributes(ctx context.Context, attrs map[string]string) context.Context {
+	if attrs != nil {
+		// Since it is likely that the map gets used in another goroutine
+		// ensure that modifying passed in map is race-free.
+		cp := make(map[string]string, len(attrs))
+		for k, v := range attrs {
+			cp[k] = v
+		}
+		attrs = cp
+	}
 	return context.WithValue(ctx, withCustomAttributes{}, attrs)
 }