Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down
33 changes: 33 additions & 0 deletions unit_tests/legacy/sources/declarative/conftest.py
Original file line number Diff line number Diff line change
@@ -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
33 changes: 33 additions & 0 deletions unit_tests/sources/declarative/conftest.py
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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",
Expand Down
Loading