fix(declarative): allow bundled custom components without AIRBYTE_ENABLE_UNSAFE_CODE (#1083)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Daryna Ishchenko <darina.ishchenko17@gmail.com>

Author: 3 people

airbyte_cdk/cli/source_declarative_manifest/_run.py

@@ -43,6 +43,10 @@
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
     ConcurrentDeclarativeSource,
 )
+from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    AirbyteCustomCodeNotPermittedError,
+    custom_code_execution_permitted,
+)
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
 from airbyte_cdk.sources.source import TState
 from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
@@ -212,6 +216,10 @@ def create_declarative_source(
             state=state,
             source_config=cast(dict[str, Any], config["__injected_declarative_manifest"]),
             config_path=parsed_args.config if hasattr(parsed_args, "config") else None,
+            # The manifest is supplied by the caller (via config or --manifest-path) rather than
+            # bundled in a connector image, so any custom components it references are untrusted
+            # code. Bundled manifest-only connectors take the local-manifest path instead.
+            custom_components_trusted=False,
         )
     except Exception as error:
         print(
@@ -275,10 +283,17 @@ def _parse_manifest_from_file(filepath: str) -> dict[str, Any] | None:
 
 
 def _register_components_from_file(filepath: str) -> None:
-    """Load and register components from a Python file specified in the args."""
+    """Load and register components from a Python file specified in the args.
+
+    The file is caller-supplied Python that is executed via `exec_module`, so it is
+    untrusted code and must be gated behind `AIRBYTE_ENABLE_UNSAFE_CODE`.
+    """
     import importlib.util
     import sys
 
+    if not custom_code_execution_permitted():
+        raise AirbyteCustomCodeNotPermittedError
+
     components_path = Path(filepath)
 
     module_name = "components"

airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -84,6 +84,9 @@ def create_source(
         migrate_manifest=should_migrate_manifest(config),
         normalize_manifest=should_normalize_manifest(config),
         limits=limits,
+        # The manifest is supplied by the Connector Builder caller rather than bundled in a
+        # connector image, so any custom components it references are untrusted code.
+        custom_components_trusted=False,
     )
 
 

airbyte_cdk/manifest_server/command_processor/utils.py

@@ -79,6 +79,9 @@ def build_source(
         state=state,
         source_config=definition,
         config=config,
+        # The manifest is supplied by the API caller rather than bundled in a connector image, so
+        # any custom components it references are untrusted code.
+        custom_components_trusted=False,
         normalize_manifest=should_normalize,
         migrate_manifest=should_migrate,
         emit_connector_builder_messages=True,

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -150,6 +150,7 @@ def __init__(
         normalize_manifest: bool = False,
         limits: Optional[TestLimits] = None,
         config_path: Optional[str] = None,
+        custom_components_trusted: bool = True,
         **kwargs: Any,
     ) -> None:
         self.logger = logging.getLogger(f"airbyte.{self.name}")
@@ -174,6 +175,7 @@ def __init__(
         # the declarative models into runtime components. Concurrent sources will continue to checkpoint
         # incremental streams running in full refresh.
         component_factory = ModelToComponentFactory(
+            custom_components_trusted=custom_components_trusted,
             emit_connector_builder_messages=emit_connector_builder_messages,
             message_repository=ConcurrentMessageRepository(queue, message_repository),
             configured_catalog=catalog,
@@ -197,6 +199,7 @@ def __init__(
             component_factory
             if component_factory
             else ModelToComponentFactory(
+                custom_components_trusted=custom_components_trusted,
                 emit_connector_builder_messages=emit_connector_builder_messages,
                 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"),
             )

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -478,6 +478,7 @@
     ZipfileDecoder as ZipfileDecoderModel,
 )
 from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    INJECTED_MANIFEST,
     AirbyteCustomCodeNotPermittedError,
     custom_code_execution_permitted,
 )
@@ -704,8 +705,10 @@ def __init__(
         max_concurrent_async_job_count: Optional[int] = None,
         configured_catalog: Optional[ConfiguredAirbyteCatalog] = None,
         api_budget: Optional[APIBudget] = None,
+        custom_components_trusted: bool = True,
     ):
         self._init_mappings()
+        self._custom_components_trusted = custom_components_trusted
         self._limit_pages_fetched_per_slice = limit_pages_fetched_per_slice
         self._limit_slices_fetched = limit_slices_fetched
         self._emit_connector_builder_messages = emit_connector_builder_messages
@@ -1814,11 +1817,15 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) ->
         :return: The declarative component built from the Pydantic model to be used at runtime
         """
         # Instantiating a custom component means importing and executing arbitrary code referenced
-        # by `class_name`. This is only permitted when custom code execution is explicitly enabled,
-        # mirroring the gate applied to injected `components.py` code. Without this check, a manifest
-        # could point `class_name` at any importable callable and have it invoked, bypassing the
-        # `AIRBYTE_ENABLE_UNSAFE_CODE` protection.
-        if not custom_code_execution_permitted():
+        # by `class_name`. Manifests supplied by a caller, whether through the config or directly to
+        # the manifest server, are untrusted input and could point `class_name` at any importable
+        # callable, so they honor the same `AIRBYTE_ENABLE_UNSAFE_CODE` gate as injected
+        # `components.py` code. Manifests bundled in a published connector image are trusted and may
+        # always use their bundled custom components.
+        manifest_is_untrusted = not self._custom_components_trusted or bool(
+            config.get(INJECTED_MANIFEST)
+        )
+        if manifest_is_untrusted and not custom_code_execution_permitted():
             raise AirbyteCustomCodeNotPermittedError
 
         custom_component_class = self._get_class_from_fully_qualified_class_name(model.class_name)
@@ -4140,6 +4147,7 @@ def create_parent_stream_config_with_substream_wrapper(
         )
 
         substream_factory = ModelToComponentFactory(
+            custom_components_trusted=self._custom_components_trusted,
             connector_state_manager=connector_state_manager,
             limit_pages_fetched_per_slice=self._limit_pages_fetched_per_slice,
             limit_slices_fetched=self._limit_slices_fetched,

unit_tests/connector_builder/test_connector_builder_handler.py

@@ -56,6 +56,10 @@
     ConcurrentDeclarativeSource,
     TestLimits,
 )
+from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    ENV_VAR_ALLOW_CUSTOM_CODE,
+    AirbyteCustomCodeNotPermittedError,
+)
 from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever
 from airbyte_cdk.sources.declarative.stream_slicers import StreamSlicerTestReadDecorator
 from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
@@ -1031,6 +1035,117 @@ def test_create_source():
     assert source._constructor._disable_cache
 
 
+def test_create_source_marks_manifest_untrusted():
+    """A Connector Builder manifest is caller-supplied, so its factory must be untrusted.
+
+    This is the invariant that gates custom-component execution independently of the config
+    contents, so it cannot be defeated by a manifest that manipulates its own config.
+    """
+    source = create_source(
+        config={"__injected_declarative_manifest": MANIFEST}, limits=None, catalog=None, state=None
+    )
+
+    assert source._constructor._custom_components_trusted is False
+
+
+_GATE_BYPASS_STREAM = {
+    "type": "DeclarativeStream",
+    "name": "s",
+    "retriever": {
+        "type": "SimpleRetriever",
+        "requester": {
+            "type": "HttpRequester",
+            "url_base": "https://example.com",
+            "path": "/",
+        },
+        "record_selector": {
+            "type": "RecordSelector",
+            "extractor": {
+                "type": "CustomRecordExtractor",
+                "class_name": "not_a_real_module.NotAThing",
+            },
+        },
+    },
+    "schema_loader": {
+        "type": "InlineSchemaLoader",
+        "schema": {"type": "object", "properties": {}},
+    },
+}
+
+
+def test_spec_level_custom_component_is_gated(monkeypatch):
+    """A `Custom*` component in `spec.config_normalization_rules` must honor the gate.
+
+    The spec component is built with an empty config, so the config-key provenance signal is
+    absent there; the untrusted-factory flag is what keeps the gate active. A non-existent
+    `class_name` proves the gate fires *before* the class is resolved (otherwise the failure
+    would be a class-resolution `ValueError`, not `AirbyteCustomCodeNotPermittedError`).
+    """
+    monkeypatch.delenv(ENV_VAR_ALLOW_CUSTOM_CODE, raising=False)
+    manifest = {
+        "type": "DeclarativeSource",
+        "version": "6.0.0",
+        "check": {"type": "CheckStream", "stream_names": ["s"]},
+        "spec": {
+            "type": "Spec",
+            "connection_specification": {"type": "object", "properties": {}},
+            "config_normalization_rules": {
+                "type": "ConfigNormalizationRules",
+                "transformations": [
+                    {
+                        "type": "CustomConfigTransformation",
+                        "class_name": "not_a_real_module.NotAThing",
+                    }
+                ],
+            },
+        },
+        "streams": [_GATE_BYPASS_STREAM],
+    }
+
+    with pytest.raises(AirbyteCustomCodeNotPermittedError):
+        create_source(
+            config={"__injected_declarative_manifest": manifest},
+            limits=None,
+            catalog=None,
+            state=None,
+        )
+
+
+def test_config_strip_does_not_bypass_custom_code_gate(monkeypatch):
+    """A manifest that strips its own provenance key must not escape the gate.
+
+    A plain `ConfigRemoveFields` removing `__injected_declarative_manifest` runs before streams
+    are built, defeating the config-key signal. The untrusted-factory flag does not live in the
+    config, so a stream-level `Custom*` component stays gated.
+    """
+    monkeypatch.delenv(ENV_VAR_ALLOW_CUSTOM_CODE, raising=False)
+    manifest = {
+        "type": "DeclarativeSource",
+        "version": "6.0.0",
+        "check": {"type": "CheckStream", "stream_names": ["s"]},
+        "spec": {
+            "type": "Spec",
+            "connection_specification": {"type": "object", "properties": {}},
+            "config_normalization_rules": {
+                "type": "ConfigNormalizationRules",
+                "transformations": [
+                    {
+                        "type": "ConfigRemoveFields",
+                        "field_pointers": [["__injected_declarative_manifest"]],
+                    }
+                ],
+            },
+        },
+        "streams": [_GATE_BYPASS_STREAM],
+    }
+    config = {"__injected_declarative_manifest": manifest}
+
+    source = create_source(config=config, limits=None, catalog=None, state=None)
+
+    with pytest.raises(AirbyteCustomCodeNotPermittedError):
+        source.streams(config)
+
+
 def request_log_message(request: dict) -> AirbyteMessage:
     return AirbyteMessage(
         type=Type.LOG,

unit_tests/legacy/sources/declarative/conftest.py

@@ -288,10 +288,34 @@ def _read_fn(*args, **kwargs):
             _read_fn()
 
 
-def test_register_components_from_file(components_file: str) -> None:
-    """Test that components can be properly loaded from a file."""
+def test_register_components_from_file(
+    components_file: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test that components can be properly loaded from a file when custom code is permitted."""
+    monkeypatch.setenv(ENV_VAR_ALLOW_CUSTOM_CODE, "true")
+
     # Register the components
     _register_components_from_file(components_file)
 
     # Verify the components were loaded correctly
     verify_components_loaded()
+
+
+def test_register_components_from_file_is_gated(
+    components_file: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A `--components-path` file is caller-supplied Python executed via `exec_module`.
+
+    It must be gated behind `AIRBYTE_ENABLE_UNSAFE_CODE`: with the env var unset the file must
+    not be executed and the loader must raise `AirbyteCustomCodeNotPermittedError` before any
+    module-level code runs.
+    """
+    monkeypatch.delenv(ENV_VAR_ALLOW_CUSTOM_CODE, raising=False)
+
+    with pytest.raises(AirbyteCustomCodeNotPermittedError):
+        _register_components_from_file(components_file)
+
+    assert "components" not in sys.modules
+    assert "source_declarative_manifest.components" not in sys.modules