Redact attributes that are added after span creation

attribute.go

@@ -56,29 +56,21 @@ func NewAttributeCensor(replacements map[attribute.Key]attribute.Value) Attribut
 	return a
 }
 
-// OnStart censors the attributes of s matching the Replacements keys of c.
-func (c AttributeCensor) OnStart(_ context.Context, s trace.ReadWriteSpan) {
-	// The SetAttributes method only overwrites attributes that already exists,
-	// it does not set the attributes to only the values passed. Therefore,
-	// determine if there are any attributes that need to be redacted and set
-	// overrides for only those values.
-	redacted := c.args[:0]
-	for _, a := range s.Attributes() {
-		if v, ok := c.replacements[a.Key]; ok {
-			redacted = append(redacted, attribute.KeyValue{
-				Key:   a.Key,
-				Value: v,
-			})
+// OnStart does nothing.
+func (c AttributeCensor) OnStart(_ context.Context, _ trace.ReadWriteSpan) {
+}
+
+// OnEnd censors the attributes of s matching the Replacements keys of c.
+func (c AttributeCensor) OnEnd(s trace.ReadOnlySpan) {
+	// We can't change the span snapshot in OnEnd, but we can change attribute value
+	attributes := s.Attributes()
+	for i := range attributes {
+		if v, ok := c.replacements[attributes[i].Key]; ok {
+			attributes[i].Value = v
 		}
 	}
-	if len(redacted) > 0 {
-		s.SetAttributes(redacted...)
-	}
 }
 
-// OnEnd does nothing.
-func (AttributeCensor) OnEnd(trace.ReadOnlySpan) {}
-
 // Shutdown does nothing.
 func (AttributeCensor) Shutdown(context.Context) error { return nil }
 

attribute_test.go

@@ -29,12 +29,12 @@ type attrRecorder struct {
 	attrs []attribute.KeyValue
 }
 
-func (*attrRecorder) OnEnd(trace.ReadOnlySpan)         {}
-func (*attrRecorder) Shutdown(context.Context) error   { return nil }
-func (*attrRecorder) ForceFlush(context.Context) error { return nil }
-func (r *attrRecorder) OnStart(_ context.Context, s trace.ReadWriteSpan) {
+func (r *attrRecorder) OnEnd(s trace.ReadOnlySpan) {
 	r.attrs = s.Attributes()
 }
+func (*attrRecorder) Shutdown(context.Context) error                   { return nil }
+func (*attrRecorder) ForceFlush(context.Context) error                 { return nil }
+func (*attrRecorder) OnStart(_ context.Context, _ trace.ReadWriteSpan) {}
 
 func TestAttributes(t *testing.T) {
 	const key = "password"
@@ -57,21 +57,37 @@ func TestAttributes(t *testing.T) {
 		got := testAttributes(Attributes(), name, passStr, eID)
 		contains(t, got, name, eID, passStr)
 	})
+	t.Run("EmptyAfterCreation", func(t *testing.T) {
+		got := testAttributesAfterCreation(Attributes(), name, passStr, eID)
+		contains(t, got, name, eID, passStr)
+	})
 
 	t.Run("SingleStringAttribute", func(t *testing.T) {
 		got := testAttributes(Attributes(key), name, passStr, eID)
 		contains(t, got, name, eID, replaced)
 	})
+	t.Run("SingleStringAttributeAfterCreation", func(t *testing.T) {
+		got := testAttributesAfterCreation(Attributes(key), name, passStr, eID)
+		contains(t, got, name, eID, replaced)
+	})
 
 	t.Run("NoMatchingKey", func(t *testing.T) {
 		got := testAttributes(Attributes("secret"), name, passStr, eID)
 		contains(t, got, name, eID, passStr)
 	})
+	t.Run("NoMatchingKeyAfterCreation", func(t *testing.T) {
+		got := testAttributesAfterCreation(Attributes("secret"), name, passStr, eID)
+		contains(t, got, name, eID, passStr)
+	})
 
 	t.Run("DifferentValueTypes", func(t *testing.T) {
 		got := testAttributes(Attributes(key), name, passBool, eID)
 		contains(t, got, name, eID, replaced)
 	})
+	t.Run("DifferentValueTypesAfterCreation", func(t *testing.T) {
+		got := testAttributesAfterCreation(Attributes(key), name, passBool, eID)
+		contains(t, got, name, eID, replaced)
+	})
 }
 
 func testAttributes(opt trace.TracerProviderOption, attrs ...attribute.KeyValue) []attribute.KeyValue {
@@ -86,6 +102,19 @@ func testAttributes(opt trace.TracerProviderOption, attrs ...attribute.KeyValue)
 	return r.attrs
 }
 
+func testAttributesAfterCreation(opt trace.TracerProviderOption, attrs ...attribute.KeyValue) []attribute.KeyValue {
+	r := &attrRecorder{}
+	tp := trace.NewTracerProvider(opt, trace.WithSpanProcessor(r))
+	defer func() { _ = tp.Shutdown(context.Background()) }()
+
+	ctx := context.Background()
+	tracer := tp.Tracer("testAttributes")
+	_, s := tracer.Start(ctx, "span name")
+	s.SetAttributes(attrs...)
+	s.End()
+	return r.attrs
+}
+
 func BenchmarkAttributeCensorOnStart(b *testing.B) {
 	b.Run("0/16", benchAttributeCensorOnStart(0, 16))
 	b.Run("1/16", benchAttributeCensorOnStart(1, 16))
@@ -133,6 +162,7 @@ func benchAttributeCensorOnStart(redacted, total int) func(*testing.B) {
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			ac.OnStart(ctx, s)
+			ac.OnEnd(s)
 		}
 	}
 }