Skip to content
Open
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
17 changes: 16 additions & 1 deletion airbyte_cdk/cli/source_declarative_manifest/_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Comment thread
coderabbitai[bot] marked this conversation as resolved.
)
except Exception as error:
print(
Expand Down Expand Up @@ -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"
Expand Down
3 changes: 3 additions & 0 deletions airbyte_cdk/connector_builder/connector_builder_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down
3 changes: 3 additions & 0 deletions airbyte_cdk/manifest_server/command_processor/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand All @@ -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,
Expand All @@ -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"),
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -478,6 +478,7 @@
ZipfileDecoder as ZipfileDecoderModel,
)
from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
INJECTED_MANIFEST,
AirbyteCustomCodeNotPermittedError,
custom_code_execution_permitted,
)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
115 changes: 115 additions & 0 deletions unit_tests/connector_builder/test_connector_builder_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
33 changes: 0 additions & 33 deletions unit_tests/legacy/sources/declarative/conftest.py

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -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
33 changes: 0 additions & 33 deletions unit_tests/sources/declarative/conftest.py

This file was deleted.

Loading
Loading