More aggressively mask user code errors when masking enabled (#27183)

## Summary

Follow-on to the functionality introduced in #18688. The original PR
ensures that if the `DAGSTER_REDACT_USER_CODE_ERRORS` env var is set to
`1`, instances of `DagsterUserCodeExecutionError` which are serialized
are masked so that their potentially sensitive contents are not leaked.
However, there is a class of errors that can still expose elements of
user code: system/framework errors or interrupts, for example, will
include their stack trace and snippets of user code.

This PR attempts to handle this second class of errors, in a few ways:
- Any error which is raised in a user code context has its stacktrace
logged and then cleared, so if the error makes it to serialization, it
will only include the error message.
- User code exceptions have their `original_exc_info` purged.
- A global mapping stores the error, so that a unique error ID can be
raised in the UI, which can be used to locate the original error in
logs.

## Test Plan

New unit tests which specifically handle the intterrupt/system error
cases.

Author: benpankow

python_modules/dagster/dagster/_core/errors.py

@@ -280,8 +280,9 @@ def user_code_error_boundary(
     """
     check.callable_param(msg_fn, "msg_fn")
     check.class_param(error_cls, "error_cls", superclass=DagsterUserCodeExecutionError)
+    from dagster._utils.error import redact_user_stacktrace_if_enabled
 
-    with raise_execution_interrupts():
+    with redact_user_stacktrace_if_enabled(), raise_execution_interrupts():
         if log_manager:
             log_manager.begin_python_log_capture()
         try:

python_modules/dagster/dagster/_core/execution/plan/utils.py

@@ -42,12 +42,13 @@ def op_execution_error_boundary(
     as respecting the RetryPolicy if present.
     """
     from dagster._core.execution.context.system import StepExecutionContext
+    from dagster._utils.error import redact_user_stacktrace_if_enabled
 
     check.callable_param(msg_fn, "msg_fn")
     check.class_param(error_cls, "error_cls", superclass=DagsterUserCodeExecutionError)
     check.inst_param(step_context, "step_context", StepExecutionContext)
 
-    with raise_execution_interrupts():
+    with redact_user_stacktrace_if_enabled(), raise_execution_interrupts():
         step_context.log.begin_python_log_capture()
         retry_policy = step_context.op_retry_policy
 

python_modules/dagster/dagster/_utils/error.py

@@ -1,5 +1,7 @@
+import contextlib
 import logging
 import os
+import sys
 import traceback
 import uuid
 from collections.abc import Sequence
@@ -9,6 +11,7 @@
 from typing_extensions import TypeAlias
 
 import dagster._check as check
+from dagster._core.errors import DagsterUserCodeExecutionError
 from dagster._serdes import whitelist_for_serdes
 
 
@@ -94,6 +97,104 @@ def _should_redact_user_code_error() -> bool:
 )
 
 
+USER_FACING_ERROR_ID_ATTR_NAME = "_redacted_error_uuid"
+
+
+class DagsterRedactedUserCodeError(DagsterUserCodeExecutionError):
+    pass
+
+
+@contextlib.contextmanager
+def redact_user_stacktrace_if_enabled():
+    """Context manager which, if a user has enabled redacting user code errors, logs exceptions raised from within,
+    and clears the stacktrace from the exception. It also marks the exception to be redacted if it was to be persisted
+    or otherwise serialized to be sent to Dagster Plus. This is useful for preventing sensitive information from
+    being leaked in error messages.
+    """
+    if not _should_redact_user_code_error():
+        yield
+    else:
+        try:
+            yield
+        except BaseException as e:
+            exc_info = sys.exc_info()
+
+            # Generate a unique error ID for this error, or re-use an existing one
+            # if this error has already been seen
+            existing_error_id = getattr(e, USER_FACING_ERROR_ID_ATTR_NAME, None)
+
+            if not existing_error_id:
+                error_id = str(uuid.uuid4())
+
+                # Track the error ID for this exception so we can redact it later
+                setattr(e, USER_FACING_ERROR_ID_ATTR_NAME, error_id)
+                masked_logger = logging.getLogger(_REDACTED_ERROR_LOGGER_NAME)
+
+                masked_logger.error(
+                    f"Error occurred during user code execution, error ID {error_id}",
+                    exc_info=exc_info,
+                )
+            else:
+                error_id = existing_error_id
+
+            if isinstance(e, DagsterUserCodeExecutionError):
+                # To be especially sure that user code error information doesn't leak from
+                # outside the context, we raise a new exception with a cleared original_exc_info
+                # The only remnant that remains is user_exception, which we only use to allow the user
+                # to retrieve exceptions in hooks
+                try:
+                    raise Exception("Masked").with_traceback(None) from None
+                except Exception as dummy_exception:
+                    redacted_exception = DagsterRedactedUserCodeError(
+                        f"Error occurred during user code execution, error ID {error_id}. "
+                        "The error has been masked to prevent leaking sensitive information. "
+                        "Search in logs for this error ID for more details.",
+                        user_exception=e.user_exception,
+                        original_exc_info=sys.exc_info(),
+                    ).with_traceback(None)
+                    setattr(dummy_exception, USER_FACING_ERROR_ID_ATTR_NAME, error_id)
+                    setattr(redacted_exception, USER_FACING_ERROR_ID_ATTR_NAME, error_id)
+                    raise redacted_exception from None
+
+            # Redact the stacktrace to ensure it will not be passed to Dagster Plus
+            raise e.with_traceback(None) from None
+
+
+def _generate_redacted_user_code_error_message(err_id: str) -> SerializableErrorInfo:
+    return SerializableErrorInfo(
+        message=(
+            f"Error occurred during user code execution, error ID {err_id}. "
+            "The error has been masked to prevent leaking sensitive information. "
+            "Search in logs for this error ID for more details."
+        ),
+        stack=[],
+        cls_name="DagsterRedactedUserCodeError",
+        cause=None,
+        context=None,
+    )
+
+
+def _generate_partly_redacted_framework_error_message(
+    exc_info: ExceptionInfo, err_id: str
+) -> SerializableErrorInfo:
+    exc_type, e, tb = exc_info
+    tb_exc = traceback.TracebackException(check.not_none(exc_type), check.not_none(e), tb)
+    error_info = _serializable_error_info_from_tb(tb_exc)
+
+    return SerializableErrorInfo(
+        message=error_info.message
+        + (
+            f"Error ID {err_id}. "
+            "The error has been masked to prevent leaking sensitive information. "
+            "Search in logs for this error ID for more details."
+        ),
+        stack=[],
+        cls_name=error_info.cls_name,
+        cause=None,
+        context=None,
+    )
+
+
 def serializable_error_info_from_exc_info(
     exc_info: ExceptionInfo,
     # Whether to forward serialized errors thrown from subprocesses
@@ -116,27 +217,18 @@ def serializable_error_info_from_exc_info(
     e = check.not_none(e, additional_message=additional_message)
     tb = check.not_none(tb, additional_message=additional_message)
 
-    from dagster._core.errors import DagsterUserCodeExecutionError, DagsterUserCodeProcessError
-
-    if isinstance(e, DagsterUserCodeExecutionError) and _should_redact_user_code_error():
-        error_id = str(uuid.uuid4())
-        masked_logger = logging.getLogger(_REDACTED_ERROR_LOGGER_NAME)
-
-        masked_logger.error(
-            f"Error occurred during user code execution, error ID {error_id}",
-            exc_info=exc_info,
-        )
-        return SerializableErrorInfo(
-            message=(
-                f"Error occurred during user code execution, error ID {error_id}. "
-                "The error has been masked to prevent leaking sensitive information. "
-                "Search in logs for this error ID for more details."
-            ),
-            stack=[],
-            cls_name="DagsterRedactedUserCodeError",
-            cause=None,
-            context=None,
-        )
+    from dagster._core.errors import DagsterUserCodeProcessError
+
+    err_id = getattr(e, USER_FACING_ERROR_ID_ATTR_NAME, None)
+    if err_id:
+        if isinstance(e, DagsterUserCodeExecutionError):
+            # For user code, we want to completely mask the error message, since
+            # both the stacktrace and the message could contain sensitive information
+            return _generate_redacted_user_code_error_message(err_id)
+        else:
+            # For all other errors (framework errors, interrupts),
+            # we want to redact the error message, but keep the stacktrace
+            return _generate_partly_redacted_framework_error_message(exc_info, err_id)
 
     if (
         hoist_user_code_error

python_modules/dagster/dagster_tests/core_tests/test_mask_user_code_errors.py

@@ -1,13 +1,20 @@
+import logging
 import re
 import sys
 import time
 import traceback
-from typing import Any
+from typing import Any, Callable
 
 import pytest
 from dagster import Config, RunConfig, config_mapping, job, op
+from dagster._core.definitions.events import Failure
 from dagster._core.definitions.timestamp import TimestampWithTimezone
-from dagster._core.errors import DagsterUserCodeProcessError
+from dagster._core.errors import (
+    DagsterExecutionInterruptedError,
+    DagsterUserCodeExecutionError,
+    DagsterUserCodeProcessError,
+    user_code_error_boundary,
+)
 from dagster._core.test_utils import environ, instance_for_test
 from dagster._utils.error import (
     _serializable_error_info_from_tb,
@@ -30,32 +37,148 @@ def __init__(self):
         )
 
 
+class hunter2:
+    pass
+
+
 @pytest.fixture(scope="function")
 def enable_masking_user_code_errors() -> Any:
     with environ({"DAGSTER_REDACT_USER_CODE_ERRORS": "1"}):
         yield
 
 
-def test_masking_op_execution(enable_masking_user_code_errors) -> Any:
+def test_masking_basic(enable_masking_user_code_errors):
+    try:
+        with user_code_error_boundary(
+            error_cls=DagsterUserCodeExecutionError,
+            msg_fn=lambda: "hunter2",
+        ):
+
+            def hunter2():
+                raise UserError()
+
+            hunter2()
+    except Exception:
+        exc_info = sys.exc_info()
+        err_info = serializable_error_info_from_exc_info(exc_info)
+
+    assert "hunter2" not in str(err_info)
+
+
+def test_masking_nested_user_code_err_boundaries(enable_masking_user_code_errors):
+    try:
+        with user_code_error_boundary(
+            error_cls=DagsterUserCodeExecutionError,
+            msg_fn=lambda: "hunter2 as well",
+        ):
+            with user_code_error_boundary(
+                error_cls=DagsterUserCodeExecutionError,
+                msg_fn=lambda: "hunter2",
+            ):
+
+                def hunter2():
+                    raise UserError()
+
+                hunter2()
+    except Exception:
+        exc_info = sys.exc_info()
+        err_info = serializable_error_info_from_exc_info(exc_info)
+
+    assert "hunter2" not in str(err_info)
+
+
+def test_masking_nested_user_code_err_boundaries_reraise(enable_masking_user_code_errors):
+    try:
+        try:
+            with user_code_error_boundary(
+                error_cls=DagsterUserCodeExecutionError,
+                msg_fn=lambda: "hunter2",
+            ):
+
+                def hunter2():
+                    raise UserError()
+
+                hunter2()
+        except Exception as e:
+            # Mimics behavior of resource teardown, which runs in a
+            # user_code_error_boundary after the user code raises an error
+            with user_code_error_boundary(
+                error_cls=DagsterUserCodeExecutionError,
+                msg_fn=lambda: "teardown after we raised hunter2 error",
+            ):
+                # do teardown stuff
+                raise e
+
+    except Exception:
+        exc_info = sys.exc_info()
+        err_info = serializable_error_info_from_exc_info(exc_info)
+
+    assert "hunter2" not in str(err_info)
+
+
+ERROR_ID_REGEX = r"[Ee]rror ID ([a-z0-9\-]+)"
+
+
+@pytest.mark.parametrize(
+    "exc_name, expect_exc_name_in_error, build_exc",
+    [
+        ("UserError", False, lambda: UserError()),
+        ("TypeError", False, lambda: TypeError("hunter2")),
+        ("KeyboardInterrupt", True, lambda: KeyboardInterrupt()),
+        ("DagsterExecutionInterruptedError", True, lambda: DagsterExecutionInterruptedError()),
+        ("Failure", True, lambda: Failure("asdf")),
+    ],
+)
+def test_masking_op_execution(
+    enable_masking_user_code_errors,
+    exc_name: str,
+    expect_exc_name_in_error: bool,
+    build_exc: Callable[[], BaseException],
+    caplog,
+) -> Any:
     @op
     def throws_user_error(_):
-        raise UserError()
+        def hunter2():
+            raise build_exc()
+
+        hunter2()
 
     @job
     def job_def():
         throws_user_error()
 
-    result = job_def.execute_in_process(raise_on_error=False)
+    with caplog.at_level(logging.ERROR):
+        result = job_def.execute_in_process(raise_on_error=False)
     assert not result.success
-    assert not any("hunter2" in str(event) for event in result.all_events)
+
+    # Ensure error message and contents of user code don't leak (e.g. hunter2 text or function name)
+    assert not any("hunter2" in str(event).lower() for event in result.all_events), [
+        str(event) for event in result.all_events if "hunter2" in str(event)
+    ]
+
     step_error = next(event for event in result.all_events if event.is_step_failure)
-    assert (
-        step_error.step_failure_data.error
-        and step_error.step_failure_data.error.cls_name == "DagsterRedactedUserCodeError"
-    )
 
+    # Certain exceptions will not be fully redacted, just the stack trace
+    # For example, system errors and interrupts may contain useful information
+    # or information that the framework itself relies on
+    if expect_exc_name_in_error:
+        assert (
+            step_error.step_failure_data.error
+            and step_error.step_failure_data.error.cls_name == exc_name
+        )
+    else:
+        assert (
+            step_error.step_failure_data.error
+            and step_error.step_failure_data.error.cls_name == "DagsterRedactedUserCodeError"
+        )
+
+    # Ensures we can match the error ID in the Dagster+ UI surfaced message to the rich error message
+    # in logs which includes the redacted error message
+    assert "Search in logs for this error ID for more details" in str(step_error)
+    error_id = re.search(ERROR_ID_REGEX, str(step_error)).group(1)  # type: ignore
 
-ERROR_ID_REGEX = r"Error occurred during user code execution, error ID ([a-z0-9\-]+)"
+    assert f"Error occurred during user code execution, error ID {error_id}" in caplog.text
+    assert "hunter2" in caplog.text
 
 
 def test_masking_sensor_execution(instance, enable_masking_user_code_errors, capsys) -> None: