ddtrace: support tracer.SetUser on a mockspan (#1480)

tracer.SetUser() was specific to the *tracer.span type and therefore was not working
with the mockup span.
To be able to use tracer.SetUser() on the mockup span type, we propose to implement
its behaviour through the now exported `SetUser()` method of the span struct types.
tracer.SetUser() does a type-assertion to check this method is present and delegates
its call to it. We can now provide a proper implementation for *mockspan.

Co-authored-by: Andrew Glaude <andrew.glaude@datadoghq.com>

ddtrace/mocktracer/mockspan.go

@@ -12,6 +12,7 @@ import (
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 )
 
 var _ ddtrace.Span = (*mockspan)(nil)
@@ -231,3 +232,41 @@ baggage: %#v
 
 // Context returns the SpanContext of this Span.
 func (s *mockspan) Context() ddtrace.SpanContext { return s.context }
+
+// SetUser associates user information to the current trace which the
+// provided span belongs to. The options can be used to tune which user
+// bit of information gets monitored. This mockup only sets the user
+// information as span tags of the root span of the current trace.
+func (s *mockspan) SetUser(id string, opts ...tracer.UserMonitoringOption) {
+	// Walk the span up to the root parent span
+	openSpans := s.tracer.openSpans
+	var current Span = s
+	for {
+		pid := current.ParentID()
+		if pid == 0 {
+			break
+		}
+		parent, ok := openSpans[pid]
+		if !ok {
+			break
+		}
+		current = parent
+	}
+
+	root, ok := current.(*mockspan)
+	if !ok {
+		return
+	}
+
+	var cfg tracer.UserMonitoringConfig
+	for _, fn := range opts {
+		fn(&cfg)
+	}
+
+	root.SetTag("usr.id", id)
+	root.SetTag("usr.email", cfg.Email)
+	root.SetTag("usr.name", cfg.Name)
+	root.SetTag("usr.role", cfg.Role)
+	root.SetTag("usr.scope", cfg.Scope)
+	root.SetTag("usr.session_id", cfg.SessionID)
+}

ddtrace/mocktracer/mockspan_test.go

@@ -15,6 +15,7 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // basicSpan returns a span with no configuration, having the set operation name.
@@ -198,3 +199,62 @@ func TestSpanWithID(t *testing.T) {
 	assert := assert.New(t)
 	assert.Equal(spanID, span.Context().SpanID())
 }
+
+func TestSetUser(t *testing.T) {
+	const (
+		id        = "john.doe#12345"
+		name      = "John Doe"
+		email     = "john.doe@hostname.com"
+		scope     = "read:message, write:files"
+		role      = "admin"
+		sessionID = "session#12345"
+	)
+	expected := []struct{ key, value string }{
+		{key: "usr.id", value: id},
+		{key: "usr.name", value: name},
+		{key: "usr.email", value: email},
+		{key: "usr.scope", value: scope},
+		{key: "usr.role", value: role},
+		{key: "usr.session_id", value: sessionID},
+	}
+
+	t.Run("root", func(t *testing.T) {
+		s := basicSpan("root operation")
+		tracer.SetUser(s,
+			id,
+			tracer.WithUserEmail(email),
+			tracer.WithUserName(name),
+			tracer.WithUserScope(scope),
+			tracer.WithUserRole(role),
+			tracer.WithUserSessionID(sessionID))
+		s.Finish()
+		for _, pair := range expected {
+			assert.Equal(t, pair.value, s.Tag(pair.key))
+		}
+	})
+
+	t.Run("nested", func(t *testing.T) {
+		tr := newMockTracer()
+		s0 := tr.StartSpan("root operation")
+		s1 := tr.StartSpan("nested operation", tracer.ChildOf(s0.Context()))
+		s2 := tr.StartSpan("nested nested operation", tracer.ChildOf(s1.Context()))
+		tracer.SetUser(s2,
+			id,
+			tracer.WithUserEmail(email),
+			tracer.WithUserName(name),
+			tracer.WithUserScope(scope),
+			tracer.WithUserRole(role),
+			tracer.WithUserSessionID(sessionID))
+		s2.Finish()
+		s1.Finish()
+		s0.Finish()
+		finished := tr.FinishedSpans()
+		require.Len(t, finished, 3)
+		for _, pair := range expected {
+			assert.Equal(t, pair.value, finished[2].Tag(pair.key))
+			assert.Nil(t, finished[1].Tag(pair.key))
+			assert.Nil(t, finished[0].Tag(pair.key))
+		}
+	})
+
+}

ddtrace/tracer/option.go

@@ -878,12 +878,12 @@ func StackFrames(n, skip uint) FinishOption {
 // UserMonitoringConfig is used to configure what is used to identify a user.
 // This configuration can be set by combining one or several UserMonitoringOption with a call to SetUser().
 type UserMonitoringConfig struct {
-	propagateID bool
-	email       string
-	name        string
-	role        string
-	sessionID   string
-	scope       string
+	PropagateID bool
+	Email       string
+	Name        string
+	Role        string
+	SessionID   string
+	Scope       string
 }
 
 // UserMonitoringOption represents a function that can be provided as a parameter to SetUser.
@@ -892,35 +892,35 @@ type UserMonitoringOption func(*UserMonitoringConfig)
 // WithUserEmail returns the option setting the email of the authenticated user.
 func WithUserEmail(email string) UserMonitoringOption {
 	return func(cfg *UserMonitoringConfig) {
-		cfg.email = email
+		cfg.Email = email
 	}
 }
 
 // WithUserName returns the option setting the name of the authenticated user.
 func WithUserName(name string) UserMonitoringOption {
 	return func(cfg *UserMonitoringConfig) {
-		cfg.name = name
+		cfg.Name = name
 	}
 }
 
 // WithUserSessionID returns the option setting the session ID of the authenticated user.
 func WithUserSessionID(sessionID string) UserMonitoringOption {
 	return func(cfg *UserMonitoringConfig) {
-		cfg.sessionID = sessionID
+		cfg.SessionID = sessionID
 	}
 }
 
 // WithUserRole returns the option setting the role of the authenticated user.
 func WithUserRole(role string) UserMonitoringOption {
 	return func(cfg *UserMonitoringConfig) {
-		cfg.role = role
+		cfg.Role = role
 	}
 }
 
 // WithUserScope returns the option setting the scope (authorizations) of the authenticated user.
 func WithUserScope(scope string) UserMonitoringOption {
 	return func(cfg *UserMonitoringConfig) {
-		cfg.scope = scope
+		cfg.Scope = scope
 	}
 }
 
@@ -930,6 +930,6 @@ func WithUserScope(scope string) UserMonitoringOption {
 // personal identifiable information or any kind of sensitive data, as it will be leaked to other services.
 func WithPropagation() UserMonitoringOption {
 	return func(cfg *UserMonitoringConfig) {
-		cfg.propagateID = true
+		cfg.PropagateID = true
 	}
 }

ddtrace/tracer/span.go

@@ -170,33 +170,41 @@ func (s *span) setSamplingPriority(priority int, sampler samplernames.SamplerNam
 	s.setSamplingPriorityLocked(priority, sampler)
 }
 
-// setUser sets the span user ID tag as well as some optional user monitoring tags depending on the configuration.
-// The function assumes that the span it is called on is the trace's root span.
-func (s *span) setUser(id string, cfg UserMonitoringConfig) {
-	trace := s.context.trace
-	s.Lock()
-	defer s.Unlock()
-	if cfg.propagateID {
+// SetUser associates user information to the current trace which the
+// provided span belongs to. The options can be used to tune which user
+// bit of information gets monitored. In case of distributed traces,
+// the user id can be propagated across traces using the WithPropagation() option.
+// See https://docs.datadoghq.com/security_platform/application_security/setup_and_configure/?tab=set_user#add-user-information-to-traces
+func (s *span) SetUser(id string, opts ...UserMonitoringOption) {
+	var cfg UserMonitoringConfig
+	for _, fn := range opts {
+		fn(&cfg)
+	}
+	root := s.context.trace.root
+	trace := root.context.trace
+	root.Lock()
+	defer root.Unlock()
+	if cfg.PropagateID {
 		// Delete usr.id from the tags since _dd.p.usr.id takes precedence
-		delete(s.Meta, keyUserID)
+		delete(root.Meta, keyUserID)
 		idenc := base64.StdEncoding.EncodeToString([]byte(id))
 		trace.setPropagatingTag(keyPropagatedUserID, idenc)
 	} else {
 		// Unset the propagated user ID so that a propagated user ID coming from upstream won't be propagated anymore.
 		trace.unsetPropagatingTag(keyPropagatedUserID)
-		delete(s.Meta, keyPropagatedUserID)
+		delete(root.Meta, keyPropagatedUserID)
 		// setMeta is used since the span is already locked
-		s.setMeta(keyUserID, id)
+		root.setMeta(keyUserID, id)
 	}
 	for k, v := range map[string]string{
-		keyUserEmail:     cfg.email,
-		keyUserName:      cfg.name,
-		keyUserScope:     cfg.scope,
-		keyUserRole:      cfg.role,
-		keyUserSessionID: cfg.sessionID,
+		keyUserEmail:     cfg.Email,
+		keyUserName:      cfg.Name,
+		keyUserScope:     cfg.Scope,
+		keyUserRole:      cfg.Role,
+		keyUserSessionID: cfg.SessionID,
 	} {
 		if v != "" {
-			s.setMeta(k, v)
+			root.setMeta(k, v)
 		}
 	}
 }
@@ -623,7 +631,7 @@ const (
 	keyPropagatedUserID = "_dd.p.usr.id"
 )
 
-// The following set of tags is used for user monitoring and set through calls to span.setUser().
+// The following set of tags is used for user monitoring and set through calls to span.SetUser().
 const (
 	keyUserID        = "usr.id"
 	keyUserEmail     = "usr.email"

ddtrace/tracer/tracer.go

@@ -168,16 +168,13 @@ func SetUser(s Span, id string, opts ...UserMonitoringOption) {
 	if s == nil {
 		return
 	}
-	sp, ok := s.(*span)
-	if !ok || sp.context == nil {
+	sp, ok := s.(interface {
+		SetUser(string, ...UserMonitoringOption)
+	})
+	if !ok {
 		return
 	}
-	sp = sp.context.trace.root
-	var cfg UserMonitoringConfig
-	for _, fn := range opts {
-		fn(&cfg)
-	}
-	sp.setUser(id, cfg)
+	sp.SetUser(id, opts...)
 }
 
 // payloadQueueSize is the buffer size of the trace channel.