fix: classify transient DNS failures in the SSRF filter as transient errors

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.

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"],
     [