chore(asm): improve user blocking for django (auth middleware)

ddtrace/appsec/_constants.py

@@ -90,6 +90,7 @@ class APPSEC(metaclass=Constant_Class):
     AUTO_LOGIN_EVENTS_FAILURE_MODE: Literal[
         "_dd.appsec.events.users.login.failure.auto.mode"
     ] = "_dd.appsec.events.users.login.failure.auto.mode"
+    AUTO_LOGIN_EVENTS_COLLECTION_MODE: Literal["_dd.appsec.user.collection_mode"] = "_dd.appsec.user.collection_mode"
     BLOCKED: Literal["appsec.blocked"] = "appsec.blocked"
     EVENT: Literal["appsec.event"] = "appsec.event"
     AUTO_USER_INSTRUMENTATION_MODE: Literal[

ddtrace/appsec/_trace_utils.py

@@ -133,10 +133,10 @@ def track_user_login_success_event(
         return
     if real_mode == LOGIN_EVENTS_MODE.ANON and isinstance(user_id, str):
         user_id = _hash_user_id(user_id)
-
+    span.set_tag_str(APPSEC.AUTO_LOGIN_EVENTS_COLLECTION_MODE, real_mode)
     if login_events_mode != LOGIN_EVENTS_MODE.SDK:
         span.set_tag_str(APPSEC.USER_LOGIN_USERID, str(user_id))
-    set_user(tracer, user_id, name, email, scope, role, session_id, propagate, span)
+    set_user(tracer, user_id, name, email, scope, role, session_id, propagate, span, may_block=False)
     if in_asm_context():
         res = call_waf_callback(
             custom_data={
@@ -185,6 +185,7 @@ def track_user_login_failure_event(
         if login_events_mode != LOGIN_EVENTS_MODE.SDK:
             span.set_tag_str(APPSEC.USER_LOGIN_USERID, str(user_id))
         span.set_tag_str("%s.failure.%s" % (APPSEC.USER_LOGIN_EVENT_PREFIX_PUBLIC, user.ID), str(user_id))
+        span.set_tag_str(APPSEC.AUTO_LOGIN_EVENTS_COLLECTION_MODE, real_mode)
     # if called from the SDK, set the login, email and name
     if login_events_mode in (LOGIN_EVENTS_MODE.SDK, LOGIN_EVENTS_MODE.AUTO):
         if login:
@@ -376,5 +377,57 @@ def _on_django_auth(result_user, mode, kwargs, pin, info_retriever, django_confi
     return False, None
 
 
+def _on_django_process(result_user, mode, kwargs, pin, info_retriever, django_config):
+    if (not asm_config._asm_enabled) or mode == LOGIN_EVENTS_MODE.DISABLED:
+        return
+    userid_list = info_retriever.possible_user_id_fields + info_retriever.possible_login_fields
+
+    for possible_key in userid_list:
+        if possible_key in kwargs:
+            user_id = kwargs[possible_key]
+            break
+    else:
+        user_id = None
+
+    user_id_found, user_extra = info_retriever.get_user_info(
+        login=django_config.include_user_login,
+        email=django_config.include_user_email,
+        name=django_config.include_user_realname,
+    )
+    if user_extra.get("login") is None:
+        user_extra["login"] = user_id
+    user_id = user_id_found or user_id
+    if result_user and result_user.is_authenticated:
+        span = pin.tracer.current_root_span()
+        if mode == LOGIN_EVENTS_MODE.ANON and isinstance(user_id, str):
+            hash_id = _hash_user_id(user_id)
+            span.set_tag_str(APPSEC.USER_LOGIN_USERID, hash_id)
+            span.set_tag_str(APPSEC.AUTO_LOGIN_EVENTS_COLLECTION_MODE, mode)
+            set_user(pin.tracer, hash_id, propagate=True, may_block=False, span=span)
+        elif mode == LOGIN_EVENTS_MODE.IDENT:
+            span.set_tag_str(APPSEC.USER_LOGIN_USERID, str(user_id))
+            span.set_tag_str(APPSEC.AUTO_LOGIN_EVENTS_COLLECTION_MODE, mode)
+            set_user(
+                pin.tracer,
+                str(user_id),
+                propagate=True,
+                email=user_extra.get("email"),
+                name=user_extra.get("name"),
+                may_block=False,
+                span=span,
+            )
+        if in_asm_context():
+            real_mode = mode if mode != LOGIN_EVENTS_MODE.AUTO else asm_config._user_event_mode
+            custom_data = {
+                "REQUEST_USER_ID": str(user_id) if user_id else None,
+                "REQUEST_USERNAME": user_extra.get("login"),
+                "LOGIN_SUCCESS": real_mode,
+            }
+            res = call_waf_callback(custom_data=custom_data, force_sent=True)
+            if res and any(action in [WAF_ACTIONS.BLOCK_ACTION, WAF_ACTIONS.REDIRECT_ACTION] for action in res.actions):
+                raise BlockingException(get_blocked())
+
+
 core.on("django.login", _on_django_login)
 core.on("django.auth", _on_django_auth, "user")
+core.on("django.process_request", _on_django_process)

ddtrace/appsec/trace_utils/__init__.py

@@ -1,3 +1,5 @@
+"""Public API for User events"""
+
 from ddtrace.appsec._trace_utils import block_request  # noqa: F401
 from ddtrace.appsec._trace_utils import block_request_if_user_blocked  # noqa: F401
 from ddtrace.appsec._trace_utils import should_block_user  # noqa: F401

ddtrace/contrib/internal/django/patch.py

@@ -832,6 +832,53 @@ def traced_authenticate(django, pin, func, instance, args, kwargs):
     return result_user
 
 
+@trace_utils.with_traced_module
+def traced_process_request(django, pin, func, instance, args, kwargs):
+    tags = {COMPONENT: config.django.integration_name}
+    with core.context_with_data(
+        "django.func.wrapped",
+        span_name="django.middleware",
+        resource="django.contrib.auth.middleware.AuthenticationMiddleware.process_request",
+        tags=tags,
+        pin=pin,
+    ) as ctx, ctx.span:
+        core.dispatch(
+            "django.func.wrapped",
+            (
+                args,
+                kwargs,
+                django.core.handlers.wsgi.WSGIRequest if hasattr(django.core.handlers, "wsgi") else object,
+                ctx,
+                None,
+            ),
+        )
+        func(*args, **kwargs)
+        mode = asm_config._user_event_mode
+        if mode == "disabled":
+            return
+        try:
+            request = get_argument_value(args, kwargs, 0, "request")
+            if request:
+                if hasattr(request, "user") and hasattr(request.user, "_setup"):
+                    request.user._setup()
+                    request_user = request.user._wrapped
+                else:
+                    request_user = request.user
+                core.dispatch(
+                    "django.process_request",
+                    (
+                        request_user,
+                        mode,
+                        kwargs,
+                        pin,
+                        _DjangoUserInfoRetriever(request_user, credentials=kwargs),
+                        config.django,
+                    ),
+                )
+        except Exception:
+            log.debug("Error while trying to trace Django AuthenticationMiddleware process_request", exc_info=True)
+
+
 def unwrap_views(func, instance, args, kwargs):
     """
     Django channels uses path() and re_path() to route asgi applications. This broke our initial
@@ -884,6 +931,10 @@ def _(m):
         trace_utils.wrap(m, "login", traced_login(django))
         trace_utils.wrap(m, "authenticate", traced_authenticate(django))
 
+    @when_imported("django.contrib.auth.middleware")
+    def _(m):
+        trace_utils.wrap(m, "AuthenticationMiddleware.process_request", traced_process_request(django))
+
     # Only wrap get_asgi_application if get_response_async exists. Otherwise we will effectively double-patch
     # because get_response and get_asgi_application will be used. We must rely on the version instead of coalescing
     # with the previous patching hook because of circular imports within `django.core.asgi`.

ddtrace/contrib/internal/trace_utils.py

@@ -639,6 +639,7 @@ def set_user(
     session_id=None,  # type: Optional[str]
     propagate=False,  # type bool
     span=None,  # type: Optional[Span]
+    may_block=True,  # type: bool
 ):
     # type: (...) -> None
     """Set user tags.
@@ -666,7 +667,7 @@ def set_user(
         if session_id:
             span.set_tag_str(user.SESSION_ID, session_id)
 
-        if asm_config._asm_enabled:
+        if may_block and asm_config._asm_enabled:
             exc = core.dispatch_with_results("set_user_for_asm", [tracer, user_id]).block_user.exception
             if exc:
                 raise exc