diff --git a/ddtestpy/internal/constants.py b/ddtestpy/internal/constants.py index 1725551..e080be3 100644 --- a/ddtestpy/internal/constants.py +++ b/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" diff --git a/ddtestpy/internal/http.py b/ddtestpy/internal/http.py index 3de9cb4..d12b913 100644 --- a/ddtestpy/internal/http.py +++ b/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 diff --git a/tests/internal/test_http.py b/tests/internal/test_http.py index 92e40f4..1572c7b 100644 --- a/tests/internal/test_http.py +++ b/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"