diff --git a/airbyte_cdk/entrypoint.py b/airbyte_cdk/entrypoint.py index 57820f005..645c11117 100644 --- a/airbyte_cdk/entrypoint.py +++ b/airbyte_cdk/entrypoint.py @@ -49,6 +49,7 @@ VALID_URL_SCHEMES = ["https"] CLOUD_DEPLOYMENT_MODE = "cloud" _HAS_LOGGED_FOR_SERIALIZATION_ERROR = False +_RESOLVED_HOSTNAMES: set[str] = set() class AirbyteEntrypoint(object): @@ -430,6 +431,13 @@ def filtered_send(self: Any, request: PreparedRequest, **kwargs: Any) -> Respons message="Invalid URL endpoint: The endpoint that data is being requested from belongs to a private network. Source connectors only support requesting data from public API endpoints.", ) except socket.gaierror as exception: + if parsed_url.hostname in _RESOLVED_HOSTNAMES: + raise AirbyteTracedException( + internal_message=f"DNS resolution failed for hostname {parsed_url.hostname!r}: {exception}", + failure_type=FailureType.transient_error, + message=f"DNS resolution temporarily failed for hostname {parsed_url.hostname!r}.", + exception=exception, + ) # This is a special case where the developer specifies an IP address string that is not formatted correctly like trailing # whitespace which will fail the socket IP lookup. This only happens when using IP addresses and not text hostnames. # Knowing that this is a request using the requests library, we will mock the exception without calling the lib @@ -445,6 +453,7 @@ def _is_private_url(hostname: str, port: int) -> bool: Helper method that checks if any of the IP addresses associated with a hostname belong to a private network. """ address_info_entries = socket.getaddrinfo(hostname, port) + _RESOLVED_HOSTNAMES.add(hostname) for entry in address_info_entries: # getaddrinfo() returns entries in the form of a 5-tuple where the IP is stored as the sockaddr. For IPv4 this # is a 2-tuple and for IPv6 it is a 4-tuple, but the address is always the first value of the tuple at 0. diff --git a/unit_tests/test_entrypoint.py b/unit_tests/test_entrypoint.py index fcfb44915..4262490e7 100644 --- a/unit_tests/test_entrypoint.py +++ b/unit_tests/test_entrypoint.py @@ -3,6 +3,7 @@ # import os +import socket from argparse import Namespace from collections import defaultdict from copy import deepcopy @@ -117,6 +118,95 @@ def test_airbyte_entrypoint_init(mocker): ) +@pytest.fixture +def internal_request_filter(): + original_send = requests.Session.send + original_resolved_hostnames = entrypoint_module._RESOLVED_HOSTNAMES + send_mock = MagicMock() + requests.Session.send = send_mock + entrypoint_module._RESOLVED_HOSTNAMES = set() + entrypoint_module._init_internal_request_filter() + try: + yield send_mock + finally: + requests.Session.send = original_send + entrypoint_module._RESOLVED_HOSTNAMES = original_resolved_hostnames + + +def test_internal_request_filter_transient_dns_failure(internal_request_filter, mocker): + hostname = "graph.facebook.com" + token = "secret-token" + request = requests.Request("GET", f"https://{hostname}/endpoint?access_token={token}").prepare() + mocker.patch.object( + socket, + "getaddrinfo", + side_effect=[ + [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("157.240.241.17", 443))], + socket.gaierror( + getattr(socket, "EAI_AGAIN", -3), "Temporary failure in name resolution" + ), + ], + ) + + requests.Session().send(request) + assert hostname in entrypoint_module._RESOLVED_HOSTNAMES + + with pytest.raises(AirbyteTracedException) as exc_info: + requests.Session().send(request) + + exception = exc_info.value + assert exception.failure_type == FailureType.transient_error + assert hostname in exception.message + assert "DNS resolution temporarily failed" in exception.message + assert token not in (exception.message or "") + assert token not in (exception.internal_message or "") + + +def test_internal_request_filter_unresolved_hostname_keeps_invalid_url( + internal_request_filter, mocker +): + hostname = "domainwithoutextension" + request = requests.Request("GET", f"https://{hostname}/endpoint").prepare() + mocker.patch.object( + socket, + "getaddrinfo", + side_effect=socket.gaierror( + getattr(socket, "EAI_AGAIN", -3), "Temporary failure in name resolution" + ), + ) + + with pytest.raises(requests.exceptions.InvalidURL): + requests.Session().send(request) + + +def test_internal_request_filter_private_ip(internal_request_filter, mocker): + request = requests.Request("GET", "https://127.0.0.1/endpoint").prepare() + mocker.patch.object( + socket, + "getaddrinfo", + return_value=[(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 443))], + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + requests.Session().send(request) + + assert exc_info.value.failure_type == FailureType.config_error + + +def test_internal_request_filter_public_ip_passes_through(internal_request_filter, mocker): + request = requests.Request("GET", "https://graph.facebook.com/endpoint").prepare() + mocker.patch.object( + socket, + "getaddrinfo", + return_value=[(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("157.240.241.17", 443))], + ) + + requests.Session().send(request) + + assert internal_request_filter.call_count == 1 + assert internal_request_filter.call_args.args[1] == request + + @pytest.mark.parametrize( ["cmd", "args", "expected_args"], [