Skip to content

Commit

Permalink
Redact attributes that are added after span creation (#55)
Browse files Browse the repository at this point in the history
* adapt tests

* implement OnEnd function

* Update attribute.go

Signed-off-by: Tyler Yahn <[email protected]>

---------

Signed-off-by: Tyler Yahn <[email protected]>
Co-authored-by: Tyler Yahn <[email protected]>
  • Loading branch information
Squiry and MrAlias authored Sep 17, 2023
1 parent 80f5cda commit 8bd2a27
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 23 deletions.
31 changes: 12 additions & 19 deletions attribute.go
Original file line number Diff line number Diff line change
Expand Up @@ -56,29 +56,22 @@ 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 attribute slice of the span snapshot in OnEnd, but
// we can change the attribute value in the underlying array.
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 }

Expand Down
38 changes: 34 additions & 4 deletions attribute_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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 {
Expand All @@ -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))
Expand Down Expand Up @@ -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)
}
}
}

0 comments on commit 8bd2a27

Please sign in to comment.