[WIP] tag on service entry span instead of trace

lib/datadog/kit/appsec/events.rb

@@ -19,15 +19,15 @@ module Events
         #   Datadog::Kit::Identity.set_user. Must contain at least :id as key.
         # @param others [Hash<String || Symbol, String>] Additional free-form
         #   event information to attach to the trace.
-        def self.track_login_success(trace, user:, **others)
-          track(LOGIN_SUCCESS_EVENT, trace, **others)
+        def self.track_login_success(trace = nil, span = nil, user:, **others)
+          track(LOGIN_SUCCESS_EVENT, trace, span, **others)
 
           user_options = user.dup
           user_id = user_options.delete(:id)
 
           raise ArgumentError, 'missing required key: :user => { :id }' if user_id.nil?
 
-          Kit::Identity.set_user(trace, id: user_id, **user_options)
+          Kit::Identity.set_user(trace, span, id: user_id, **user_options)
         end
 
         # Attach login failure event information to the trace
@@ -39,13 +39,13 @@ def self.track_login_success(trace, user:, **others)
         # @param user_exists [bool] Whether the user id that did a login attempt exists.
         # @param others [Hash<String || Symbol, String>] Additional free-form
         #   event information to attach to the trace.
-        def self.track_login_failure(trace, user_id:, user_exists:, **others)
-          track(LOGIN_FAILURE_EVENT, trace, **others)
+        def self.track_login_failure(trace = nil, span = nil, user_id:, user_exists:, **others)
+          track(LOGIN_FAILURE_EVENT, trace, span, **others)
 
           raise ArgumentError, 'user_id cannot be nil' if user_id.nil?
 
-          trace.set_tag('appsec.events.users.login.failure.usr.id', user_id)
-          trace.set_tag('appsec.events.users.login.failure.usr.exists', user_exists)
+          span.set_tag('appsec.events.users.login.failure.usr.id', user_id)
+          span.set_tag('appsec.events.users.login.failure.usr.exists', user_exists)
         end
 
         # Attach custom event information to the trace
@@ -57,13 +57,13 @@ def self.track_login_failure(trace, user_id:, user_exists:, **others)
         # @param others [Hash<Symbol, String>] Additional free-form
         #   event information to attach to the trace. Key must not
         #   be :track.
-        def self.track(event, trace, **others)
-          trace.set_tag("appsec.events.#{event}.track", 'true')
+        def self.track(event, trace = nil, span = nil, **others)
+          span.set_tag("appsec.events.#{event}.track", 'true')
 
           others.each do |k, v|
             raise ArgumentError, 'key cannot be :track' if k.to_sym == :track
 
-            trace.set_tag("appsec.events.#{event}.#{k}", v) unless v.nil?
+            span.set_tag("appsec.events.#{event}.#{k}", v) unless v.nil?
           end
 
           trace.keep!

lib/datadog/kit/identity.rb

@@ -8,7 +8,12 @@ module Kit
     module Identity
       # Attach user information to the trace
       #
-      # @param trace [TraceOperation] Trace to attach data to.
+      # @param trace [TraceOperation] Trace to attach data to. Defaults to
+      #   active trace.
+      # @param span [SpanOperation] Span to attach data to. Defaults to
+      #   active span on trace. Note that this should be a service entry span.
+      #   When AppSec is enabled, the expected span and trace are automatically
+      #   used as defaults.
       # @param id [String] Mandatory. Username or client id extracted
       #   from the access token or Authorization header in the inbound request
       #   from outside the system.
@@ -29,7 +34,8 @@ module Identity
       #
       # rubocop:disable Metrics/CyclomaticComplexity
       # rubocop:disable Metrics/PerceivedComplexity
-      def self.set_user(trace, id:, email: nil, name: nil, session_id: nil, role: nil, scope: nil, **others)
+      # rubocop:disable Metrics/AbcSize
+      def self.set_user(trace = nil, span = nil, id:, email: nil, name: nil, session_id: nil, role: nil, scope: nil, **others)
         raise ArgumentError, 'missing required key: :id' if id.nil?
 
         # enforce types
@@ -45,24 +51,37 @@ def self.set_user(trace, id:, email: nil, name: nil, session_id: nil, role: nil,
           raise TypeError, "#{k.inspect} must be a String" unless v.nil? || v.is_a?(String)
         end
 
+        if (appsec_scope = Datadog::AppSec.active_scope)
+          trace = appsec_scope.trace
+          span = appsec_scope.span
+        end
+
+        trace ||= Datadog::Tracing.active_trace
+        span ||= trace.active_span || Datadog::Tracing.active_span
+
+        if trace.trace_id != span.trace_id
+          raise ArgumentError, "span #{span.span_id} does not belong to trace #{trace.trace_id}"
+        end
+
         # set tags once data is known consistent
 
-        trace.set_tag('usr.id', id)
-        trace.set_tag('usr.email', email)           unless email.nil?
-        trace.set_tag('usr.name', name)             unless name.nil?
-        trace.set_tag('usr.session_id', session_id) unless session_id.nil?
-        trace.set_tag('usr.role', role)             unless role.nil?
-        trace.set_tag('usr.scope', scope)           unless scope.nil?
+        span.set_tag('usr.id', id)
+        span.set_tag('usr.email', email)           unless email.nil?
+        span.set_tag('usr.name', name)             unless name.nil?
+        span.set_tag('usr.session_id', session_id) unless session_id.nil?
+        span.set_tag('usr.role', role)             unless role.nil?
+        span.set_tag('usr.scope', scope)           unless scope.nil?
 
         others.each do |k, v|
-          trace.set_tag("usr.#{k}", v) unless v.nil?
+          span.set_tag("usr.#{k}", v) unless v.nil?
         end
 
-        if Datadog.configuration.appsec.enabled
+        if appsec_scope
           user = ::Datadog::AppSec::Instrumentation::Gateway::User.new(id)
           ::Datadog::AppSec::Instrumentation.gateway.push('identity.set_user', user)
         end
       end
+      # rubocop:enable Metrics/AbcSize
       # rubocop:enable Metrics/PerceivedComplexity
       # rubocop:enable Metrics/CyclomaticComplexity
     end