Unix domain socket support

ddtestpy/internal/constants.py

@@ -4,6 +4,7 @@
 
 DEFAULT_AGENT_HOSTNAME = "localhost"
 DEFAULT_AGENT_PORT = 8126
+DEFAULT_AGENT_SOCKET_FILE = "/var/run/datadog/apm.socket"
 
 TAG_TRUE = "true"
 TAG_FALSE = "false"

ddtestpy/internal/http.py

@@ -8,6 +8,7 @@
 import json
 import logging
 import os
+import socket
 import threading
 import time
 import typing as t
@@ -17,6 +18,7 @@
 
 from ddtestpy.internal.constants import DEFAULT_AGENT_HOSTNAME
 from ddtestpy.internal.constants import DEFAULT_AGENT_PORT
+from ddtestpy.internal.constants import DEFAULT_AGENT_SOCKET_FILE
 from ddtestpy.internal.constants import DEFAULT_SITE
 from ddtestpy.internal.errors import SetupError
 from ddtestpy.internal.utils import asbool
@@ -74,15 +76,19 @@ def _detect_evp_proxy_setup(cls) -> BackendConnectorSetup:
         """
         agent_url = os.environ.get("DD_TRACE_AGENT_URL")
         if not agent_url:
-            agent_host = (
-                os.environ.get("DD_TRACE_AGENT_HOSTNAME") or os.environ.get("DD_AGENT_HOST") or DEFAULT_AGENT_HOSTNAME
-            )
-            agent_port = (
-                os.environ.get("DD_TRACE_AGENT_PORT") or os.environ.get("DD_AGENT_PORT") or str(DEFAULT_AGENT_PORT)
-            )
-            agent_url = f"http://{agent_host}:{agent_port}"
+            user_provided_host = os.environ.get("DD_TRACE_AGENT_HOSTNAME") or os.environ.get("DD_AGENT_HOST")
+            user_provided_port = os.environ.get("DD_TRACE_AGENT_PORT") or os.environ.get("DD_AGENT_PORT")
+
+            if user_provided_host or user_provided_port:
+                host = user_provided_host or DEFAULT_AGENT_HOSTNAME
+                port = user_provided_port or DEFAULT_AGENT_PORT
+                agent_url = f"http://{host}:{port}"
 
-        agent_url = agent_url.rstrip("/")  # Avoid an extra / when concatenating with the base path
+            elif os.path.exists(DEFAULT_AGENT_SOCKET_FILE):
+                agent_url = f"unix://{DEFAULT_AGENT_SOCKET_FILE}"
+
+            else:
+                agent_url = f"http://{DEFAULT_AGENT_HOSTNAME}:{DEFAULT_AGENT_PORT}"
 
         # Get info from agent to check if the agent is there, and which EVP proxy version it supports.
         try:
@@ -100,10 +106,10 @@ def _detect_evp_proxy_setup(cls) -> BackendConnectorSetup:
             )
 
         if "/evp_proxy/v4/" in endpoints:
-            return BackendConnectorEVPProxySetup(url=f"{agent_url}/evp_proxy/v4", use_gzip=True)
+            return BackendConnectorEVPProxySetup(url=agent_url, base_path="/evp_proxy/v4", use_gzip=True)
 
         if "/evp_proxy/v2/" in endpoints:
-            return BackendConnectorEVPProxySetup(url=f"{agent_url}/evp_proxy/v2", use_gzip=False)
+            return BackendConnectorEVPProxySetup(url=agent_url, base_path="/evp_proxy/v2", use_gzip=False)
 
         raise SetupError(f"Datadog agent at {agent_url} does not support EVP proxy mode")
 
@@ -122,13 +128,15 @@ def get_connector_for_subdomain(self, subdomain: str) -> BackendConnector:
 
 
 class BackendConnectorEVPProxySetup(BackendConnectorSetup):
-    def __init__(self, url: str, use_gzip: bool) -> None:
+    def __init__(self, url: str, base_path: str, use_gzip: bool) -> None:
         self.url = url
+        self.base_path = base_path
         self.use_gzip = use_gzip
 
     def get_connector_for_subdomain(self, subdomain: str) -> BackendConnector:
         return BackendConnector(
             url=self.url,
+            base_path=self.base_path,
             default_headers={"X-Datadog-EVP-Subdomain": subdomain},
             use_gzip=self.use_gzip,
         )
@@ -140,12 +148,13 @@ def __init__(
         url: str,
         default_headers: t.Optional[t.Dict[str, str]] = None,
         timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
+        base_path: t.Optional[str] = None,
         use_gzip: bool = False,
     ):
         parsed_url = urlparse(url)
         self.conn = self._make_connection(parsed_url, timeout_seconds)
         self.default_headers = default_headers or {}
-        self.base_path = parsed_url.path.rstrip("/")
+        self.base_path = base_path if base_path is not None else parsed_url.path.rstrip("/")
         self.use_gzip = use_gzip
         if self.use_gzip:
             self.default_headers["Accept-Encoding"] = "gzip"
@@ -170,7 +179,17 @@ def _make_connection(self, parsed_url: ParseResult, timeout_seconds: float) -> h
                 host=parsed_url.hostname, port=parsed_url.port or 443, timeout=timeout_seconds
             )
 
-        # TODO: Unix domain socket support.
+        if parsed_url.scheme == "unix":
+            # These URLs usually have an empty hostname component, e.g., unix:///var/run/datadog/apm.socket, that is,
+            # unix:// + empty hostname + /var/run/datadog/apm.socket (pathname of the socket). We need to ensure some
+            # hostname is used so the `Host` header will be passed correctly in requests to the agent, so we use the
+            # default hostname (localhost) if none is provided.
+            return UnixDomainSocketHTTPConnection(
+                host=parsed_url.hostname or DEFAULT_AGENT_HOSTNAME,
+                port=parsed_url.port or 80,
+                timeout=timeout_seconds,
+                path=parsed_url.path,
+            )
 
         raise SetupError(f"Unknown scheme {parsed_url.scheme} in {parsed_url.geturl()}")
 
@@ -256,3 +275,18 @@ class FileAttachment:
     filename: t.Optional[str]
     content_type: str
     data: bytes
+
+
+class UnixDomainSocketHTTPConnection(http.client.HTTPConnection):
+    """An HTTP connection established over a Unix Domain Socket."""
+
+    # It's important to keep the hostname and port arguments here; while there are not used by the connection
+    # mechanism, they are actually used as HTTP headers such as `Host`.
+    def __init__(self, path: str, *args: t.Any, **kwargs: t.Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.path = path
+
+    def connect(self) -> None:
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        sock.connect(self.path)
+        self.sock = sock

tests/internal/test_http.py

@@ -14,6 +14,7 @@
 from ddtestpy.internal.http import BackendConnectorEVPProxySetup
 from ddtestpy.internal.http import BackendConnectorSetup
 from ddtestpy.internal.http import FileAttachment
+from ddtestpy.internal.http import UnixDomainSocketHTTPConnection
 from tests.mocks import mock_backend_connector
 
 
@@ -35,13 +36,24 @@ def test_init_default_parameters(self, mock_https_connection: Mock) -> None:
 
     @patch("http.client.HTTPSConnection")
     def test_init_custom_parameters(self, mock_https_connection: Mock) -> None:
-        """Test BackendConnector initialization with default parameters."""
+        """Test BackendConnector initialization with custom parameters."""
         connector = BackendConnector(url="https://api.example.com/some-path", use_gzip=False)
 
         mock_https_connection.assert_called_once_with(host="api.example.com", port=443, timeout=DEFAULT_TIMEOUT_SECONDS)
         assert connector.default_headers == {}
         assert connector.base_path == "/some-path"
 
+    @patch("ddtestpy.internal.http.UnixDomainSocketHTTPConnection")
+    def test_init_unix_domain_socket(self, mock_unix_connection: Mock) -> None:
+        """Test BackendConnector initialization with Unix domain socket URL."""
+        connector = BackendConnector(url="unix:///some/path/name", base_path="/evp_proxy/over9000", use_gzip=False)
+
+        mock_unix_connection.assert_called_once_with(
+            host="localhost", port=80, timeout=DEFAULT_TIMEOUT_SECONDS, path="/some/path/name"
+        )
+        assert connector.default_headers == {}
+        assert connector.base_path == "/evp_proxy/over9000"
+
     @patch("http.client.HTTPSConnection")
     @patch("uuid.uuid4")
     def test_post_files_multiple_files(self, mock_uuid: Mock, mock_https_connection: Mock) -> None:
@@ -121,38 +133,75 @@ def test_detect_agentless_setup_no_api_key(self, monkeypatch: pytest.MonkeyPatch
         with pytest.raises(SetupError, match="DD_API_KEY environment variable is not set"):
             BackendConnectorSetup.detect_setup()
 
-    def test_detect_evp_proxy_mode_v4(self, monkeypatch: pytest.MonkeyPatch) -> None:
+    def test_detect_evp_proxy_mode_v4_via_unix_domain_socket(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.setattr(os, "environ", {})
 
         backend_connector_mock = (
             mock_backend_connector().with_get_json_response("/info", {"endpoints": ["/evp_proxy/v4/"]}).build()
         )
         with patch("ddtestpy.internal.http.BackendConnector", return_value=backend_connector_mock):
-            connector_setup = BackendConnectorSetup.detect_setup()
+            # Ensure Unix domain socket WILL be detected.
+            with patch("os.path.exists", return_value=True) as mock_path_exists:
+                connector_setup = BackendConnectorSetup.detect_setup()
 
         assert isinstance(connector_setup, BackendConnectorEVPProxySetup)
 
+        path_exists_args, _ = mock_path_exists.call_args
+        assert path_exists_args == ("/var/run/datadog/apm.socket",)
+
+        connector = connector_setup.get_connector_for_subdomain("api")
+        assert isinstance(connector.conn, UnixDomainSocketHTTPConnection)
+        assert connector.conn.host == "localhost"
+        assert connector.conn.port == 80
+        assert connector.conn.path == "/var/run/datadog/apm.socket"
+        assert connector.base_path == "/evp_proxy/v4"
+        assert connector.use_gzip is True
+        assert connector.default_headers["X-Datadog-EVP-Subdomain"] == "api"
+
+    def test_detect_evp_proxy_mode_v4_via_http(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setattr(os, "environ", {})
+
+        backend_connector_mock = (
+            mock_backend_connector().with_get_json_response("/info", {"endpoints": ["/evp_proxy/v4/"]}).build()
+        )
+        with patch("ddtestpy.internal.http.BackendConnector", return_value=backend_connector_mock):
+            # Ensure Unix domain socket WILL NOT be detected.
+            with patch("os.path.exists", return_value=False) as mock_path_exists:
+                connector_setup = BackendConnectorSetup.detect_setup()
+
+        assert isinstance(connector_setup, BackendConnectorEVPProxySetup)
+
+        path_exists_args, _ = mock_path_exists.call_args
+        assert path_exists_args == ("/var/run/datadog/apm.socket",)
+
         connector = connector_setup.get_connector_for_subdomain("api")
         assert isinstance(connector.conn, http.client.HTTPConnection)
+        assert not isinstance(connector.conn, UnixDomainSocketHTTPConnection)
         assert connector.conn.host == "localhost"
         assert connector.conn.port == 8126
         assert connector.base_path == "/evp_proxy/v4"
         assert connector.use_gzip is True
         assert connector.default_headers["X-Datadog-EVP-Subdomain"] == "api"
 
-    def test_detect_evp_proxy_mode_v2(self, monkeypatch: pytest.MonkeyPatch) -> None:
+    def test_detect_evp_proxy_mode_v2_via_http(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.setattr(os, "environ", {})
 
         backend_connector_mock = (
             mock_backend_connector().with_get_json_response("/info", {"endpoints": ["/evp_proxy/v2/"]}).build()
         )
         with patch("ddtestpy.internal.http.BackendConnector", return_value=backend_connector_mock):
-            connector_setup = BackendConnectorSetup.detect_setup()
+            # Ensure Unix domain socket WILL NOT be detected.
+            with patch("os.path.exists", return_value=False) as mock_path_exists:
+                connector_setup = BackendConnectorSetup.detect_setup()
 
         assert isinstance(connector_setup, BackendConnectorEVPProxySetup)
 
+        path_exists_args, _ = mock_path_exists.call_args
+        assert path_exists_args == ("/var/run/datadog/apm.socket",)
+
         connector = connector_setup.get_connector_for_subdomain("api")
         assert isinstance(connector.conn, http.client.HTTPConnection)
+        assert not isinstance(connector.conn, UnixDomainSocketHTTPConnection)
         assert connector.conn.host == "localhost"
         assert connector.conn.port == 8126
         assert connector.base_path == "/evp_proxy/v2"
@@ -230,3 +279,24 @@ def test_detect_evp_proxy_mode_v4_custom_dd_agent_host(self, monkeypatch: pytest
         assert connector.base_path == "/evp_proxy/v4"
         assert connector.use_gzip is True
         assert connector.default_headers["X-Datadog-EVP-Subdomain"] == "api"
+
+    def test_detect_evp_proxy_mode_v4_custom_dd_trace_agent_unix_url(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setattr(os, "environ", {"DD_TRACE_AGENT_URL": "unix:///some/file/name.socket"})
+
+        backend_connector_mock = (
+            mock_backend_connector().with_get_json_response("/info", {"endpoints": ["/evp_proxy/v4/"]}).build()
+        )
+        with patch("ddtestpy.internal.http.BackendConnector", return_value=backend_connector_mock):
+            with patch("os.path.exists", return_value=True):
+                connector_setup = BackendConnectorSetup.detect_setup()
+
+        assert isinstance(connector_setup, BackendConnectorEVPProxySetup)
+
+        connector = connector_setup.get_connector_for_subdomain("api")
+        assert isinstance(connector.conn, UnixDomainSocketHTTPConnection)
+        assert connector.conn.host == "localhost"
+        assert connector.conn.port == 80
+        assert connector.conn.path == "/some/file/name.socket"
+        assert connector.base_path == "/evp_proxy/v4"
+        assert connector.use_gzip is True
+        assert connector.default_headers["X-Datadog-EVP-Subdomain"] == "api"