diff --git a/src/sentry/identity/datadog/provider.py b/src/sentry/identity/datadog/provider.py index f13ef64c3e3c7c..1aaab48772b66e 100644 --- a/src/sentry/identity/datadog/provider.py +++ b/src/sentry/identity/datadog/provider.py @@ -28,6 +28,19 @@ from sentry.users.models.identity import Identity from sentry.utils.http import absolute_uri +DATADOG_VALID_SITES = frozenset( + { + "datadoghq.com", + "us3.datadoghq.com", + "us5.datadoghq.com", + "datadoghq.eu", + "ddog-gov.com", + "us2.ddog-gov.com", + "ap1.datadoghq.com", + "ap2.datadoghq.com", + } +) + MCP_REGISTER_PATH = "/api/unstable/mcp-server/register" MCP_AUTHORIZE_PATH = "/api/unstable/mcp-server/authorize" MCP_TOKEN_PATH = "/api/unstable/mcp-server/token" @@ -38,9 +51,9 @@ def _basic_auth_header(client_id: str, client_secret: str) -> str: return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode("ascii") -def get_user_info(access_token: str, site: str) -> dict[str, Any]: +def get_user_info(access_token: str, mcp_base_url: str) -> dict[str, Any]: """Fetch the current Datadog user via the MCP ``datadog://mcp/whoami`` resource.""" - url = f"https://mcp.{site}{MCP_ENDPOINT_PATH}" + url = f"{mcp_base_url}{MCP_ENDPOINT_PATH}" headers = {"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"} init_resp = safe_urlopen( @@ -264,10 +277,15 @@ def get_pipeline_config(self, data: dict[str, Any]) -> dict[str, str]: site = data.get("site") if not site: raise ValueError("Datadog requires a 'site' parameter (e.g. 'datadoghq.com').") + elif site not in DATADOG_VALID_SITES: + raise ValueError(f"Invalid Datadog site: {site}") return {"site": site} def _get_mcp_base_url(self) -> str: - return f"https://mcp.{self._get_oauth_parameter('site')}" + site = self._get_oauth_parameter("site") + if site not in DATADOG_VALID_SITES: + raise ValueError(f"Invalid Datadog site: {site}") + return f"https://mcp.{site}" def get_oauth_authorize_url(self) -> str: return self._get_mcp_base_url() + MCP_AUTHORIZE_PATH @@ -302,14 +320,13 @@ def build_identity(self, data: dict[str, Any]) -> dict[str, Any]: if not access_token: raise ValueError("Datadog token exchange did not return an access_token") - site = self._get_oauth_parameter("site") - user = get_user_info(access_token, site) - + user = get_user_info(access_token, self._get_mcp_base_url()) if "user_uuid" not in user or "org_uuid" not in user: raise IdentityNotValid( "User info response missing required fields (user_uuid, org_uuid)" ) + site = self._get_oauth_parameter("site") oauth_data = self.get_oauth_data(token_data) # Persist DCR credentials and site so refresh_identity can access them outside a pipeline context. @@ -356,6 +373,8 @@ def refresh_identity(self, identity: Identity | RpcIdentity, **kwargs: Any) -> N site = identity.data.get("site") if not site: raise IdentityNotValid("Missing Datadog site") + elif site not in DATADOG_VALID_SITES: + raise IdentityNotValid(f"Invalid Datadog site: {site}") self.config["site"] = site client_id = identity.data.get("client_id") diff --git a/tests/sentry/api/endpoints/test_organization_monitoring_provider_details.py b/tests/sentry/api/endpoints/test_organization_monitoring_provider_details.py index be5bc919ee8bea..970428495cb534 100644 --- a/tests/sentry/api/endpoints/test_organization_monitoring_provider_details.py +++ b/tests/sentry/api/endpoints/test_organization_monitoring_provider_details.py @@ -92,11 +92,19 @@ def test_connect_datadog_requires_site(self) -> None: assert response.status_code == 400 assert "site" in response.data["detail"] + def test_connect_datadog_invalid_site(self) -> None: + with self.feature("organizations:seer-infra-telemetry"): + response = self.get_response(self.organization.slug, "datadog", site="evil.example.com") + + assert response.status_code == 400 + assert "Invalid Datadog site" in response.data["detail"] + def test_connect_unknown_provider(self) -> None: with self.feature("organizations:seer-infra-telemetry"): response = self.get_response(self.organization.slug, "unknown") assert response.status_code == 400 + assert "Unknown monitoring provider" in response.data["detail"] @patch( "sentry.api.endpoints.organization_monitoring_provider_details.IdentityPipeline.current_step" @@ -200,12 +208,14 @@ def test_disconnect_unknown_provider(self) -> None: response = self.get_response(self.organization.slug, "unknown") assert response.status_code == 400 + assert "Unknown monitoring provider" in response.data["detail"] def test_disconnect_not_connected(self) -> None: with self.feature("organizations:seer-infra-telemetry"): response = self.get_response(self.organization.slug, "gcp") assert response.status_code == 404 + assert "Not connected to this provider" in response.data["detail"] def test_disconnect_allowed_for_org_read_member(self) -> None: member_user = self.create_user() diff --git a/tests/sentry/identity/datadog/test_provider.py b/tests/sentry/identity/datadog/test_provider.py index 46fdf603a90d27..8f3f6ea7077d13 100644 --- a/tests/sentry/identity/datadog/test_provider.py +++ b/tests/sentry/identity/datadog/test_provider.py @@ -493,7 +493,7 @@ def test_build_identity(self, mock_get_user_info: MagicMock) -> None: assert result["data"]["client_id"] == "dcr-client-id" assert result["data"]["client_secret"] == "dcr-client-secret" assert result["data"]["site"] == "datadoghq.com" - mock_get_user_info.assert_called_once_with("token-abc", "datadoghq.com") + mock_get_user_info.assert_called_once_with("token-abc", "https://mcp.datadoghq.com") @patch("sentry.identity.datadog.provider.get_user_info") def test_build_identity_missing_access_token(self, mock_get_user_info: MagicMock) -> None: @@ -636,6 +636,12 @@ def test_refresh_identity_missing_site(self) -> None: with pytest.raises(IdentityNotValid, match="Missing Datadog site"): self.provider.refresh_identity(identity) + def test_refresh_identity_invalid_site(self) -> None: + identity = self._make_identity(site="evil.example.com") + + with pytest.raises(IdentityNotValid, match="Invalid Datadog site"): + self.provider.refresh_identity(identity) + def test_refresh_identity_missing_dcr_credentials(self) -> None: identity = self._make_identity(client_id=None, client_secret=None) @@ -647,3 +653,9 @@ def test_refresh_identity_missing_refresh_token(self) -> None: with pytest.raises(IdentityNotValid, match="Missing refresh token"): self.provider.refresh_identity(identity) + + def test_invalid_site_rejected(self) -> None: + self.provider.config = {"site": "evil.example.com"} + + with pytest.raises(ValueError, match="Invalid Datadog site"): + self.provider._get_mcp_base_url()