fix(declarative): require AIRBYTE_ENABLE_UNSAFE_CODE to instantiate Custom* components from YAML (#1076)

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

Author: aaronsteers

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -477,6 +477,10 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ZipfileDecoder as ZipfileDecoderModel,
 )
+from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    AirbyteCustomCodeNotPermittedError,
+    custom_code_execution_permitted,
+)
 from airbyte_cdk.sources.declarative.partition_routers import (
     CartesianProductStreamSlicer,
     GroupingPartitionRouter,
@@ -1809,6 +1813,14 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) ->
         :param config: The custom defined connector config
         :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():
+            raise AirbyteCustomCodeNotPermittedError
+
         custom_component_class = self._get_class_from_fully_qualified_class_name(model.class_name)
         component_fields = get_type_hints(custom_component_class)
         model_args = model.dict()

unit_tests/legacy/sources/declarative/conftest.py

@@ -0,0 +1,33 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import os
+
+import pytest
+
+from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    ENV_VAR_ALLOW_CUSTOM_CODE,
+)
+
+
+@pytest.fixture(autouse=True)
+def _allow_custom_code():
+    """Enable custom-component code execution for the declarative unit tests.
+
+    `ModelToComponentFactory.create_custom_component` requires
+    `AIRBYTE_ENABLE_UNSAFE_CODE` to be set before it will resolve and instantiate a
+    component's `class_name`, matching the gate already applied to injected
+    `components.py` code. These tests exercise custom components extensively, so enable
+    the flag by default here. Tests that verify the gate itself opt out by deleting the
+    env var (e.g. through their own `monkeypatch` fixture).
+    """
+    previous = os.environ.get(ENV_VAR_ALLOW_CUSTOM_CODE)
+    os.environ[ENV_VAR_ALLOW_CUSTOM_CODE] = "true"
+    try:
+        yield
+    finally:
+        if previous is None:
+            os.environ.pop(ENV_VAR_ALLOW_CUSTOM_CODE, None)
+        else:
+            os.environ[ENV_VAR_ALLOW_CUSTOM_CODE] = previous

unit_tests/sources/declarative/conftest.py

@@ -0,0 +1,33 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import os
+
+import pytest
+
+from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    ENV_VAR_ALLOW_CUSTOM_CODE,
+)
+
+
+@pytest.fixture(autouse=True)
+def _allow_custom_code():
+    """Enable custom-component code execution for the declarative unit tests.
+
+    `ModelToComponentFactory.create_custom_component` requires
+    `AIRBYTE_ENABLE_UNSAFE_CODE` to be set before it will resolve and instantiate a
+    component's `class_name`, matching the gate already applied to injected
+    `components.py` code. These tests exercise custom components extensively, so enable
+    the flag by default here. Tests that verify the gate itself opt out by deleting the
+    env var (e.g. through their own `monkeypatch` fixture).
+    """
+    previous = os.environ.get(ENV_VAR_ALLOW_CUSTOM_CODE)
+    os.environ[ENV_VAR_ALLOW_CUSTOM_CODE] = "true"
+    try:
+        yield
+    finally:
+        if previous is None:
+            os.environ.pop(ENV_VAR_ALLOW_CUSTOM_CODE, None)
+        else:
+            os.environ[ENV_VAR_ALLOW_CUSTOM_CODE] = previous

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -107,6 +107,10 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     SelectiveAuthenticator,
 )
+from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    ENV_VAR_ALLOW_CUSTOM_CODE,
+    AirbyteCustomCodeNotPermittedError,
+)
 from airbyte_cdk.sources.declarative.parsers.manifest_component_transformer import (
     ManifestComponentTransformer,
 )
@@ -2509,6 +2513,47 @@ def test_create_custom_components(manifest, field_name, expected_value, expected
         assert getattr(custom_component, field_name) == expected_value
 
 
+@pytest.mark.parametrize(
+    "class_name",
+    [
+        pytest.param(
+            "unit_tests.sources.declarative.parsers.testing_components.TestingSomeComponent",
+            id="bundled_custom_component_class",
+        ),
+        pytest.param(
+            "os.getcwd",
+            id="arbitrary_importable_callable",
+        ),
+    ],
+)
+def test_create_custom_component_requires_custom_code_enabled(class_name, monkeypatch):
+    """A `Custom*` component must not be instantiated unless custom code execution is
+    explicitly enabled via `AIRBYTE_ENABLE_UNSAFE_CODE`.
+
+    Resolving and instantiating a component's `class_name` executes arbitrary
+    importable code, so it must honor the same gate as injected `components.py` code.
+    The gate must fire regardless of whether `class_name` points at a bundled custom
+    component or at an arbitrary importable callable, and it must fire before the
+    referenced module is imported.
+    """
+    monkeypatch.delenv(ENV_VAR_ALLOW_CUSTOM_CODE, raising=False)
+
+    def _fail_if_resolved(*args, **kwargs):
+        raise AssertionError(
+            "`class_name` must not be resolved or imported when custom code is disabled"
+        )
+
+    monkeypatch.setattr(factory, "_get_class_from_fully_qualified_class_name", _fail_if_resolved)
+
+    manifest = {
+        "type": "CustomErrorHandler",
+        "class_name": class_name,
+    }
+
+    with pytest.raises(AirbyteCustomCodeNotPermittedError):
+        factory.create_component(CustomErrorHandlerModel, manifest, input_config)
+
+
 def test_custom_components_do_not_contain_extra_fields():
     custom_substream_partition_router_manifest = {
         "type": "CustomPartitionRouter",