Fix SSM endpoint for Gov regions (#693)

* only override ssm endpoint url in commercial fips

* lint

* check ssm-fips supported regions

* fix

Author: shreyamalpani

datadog_lambda/api.py

@@ -5,6 +5,14 @@
 
 logger = logging.getLogger(__name__)
 KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName"
+SSM_FIPS_SUPPORTED_REGIONS = {
+    "us-east-1",
+    "us-east-2",
+    "us-west-1",
+    "us-west-2",
+    "ca-central-1",
+    "ca-west-1",
+}
 api_key = None
 
 
@@ -92,11 +100,18 @@ def get_api_key() -> str:
         )["SecretString"]
     elif DD_API_KEY_SSM_NAME:
         # SSM endpoints: https://docs.aws.amazon.com/general/latest/gr/ssm.html
-        fips_endpoint = (
-            f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com"
-            if config.fips_mode_enabled
-            else None
-        )
+        fips_endpoint = None
+        if config.fips_mode_enabled:
+            if LAMBDA_REGION in SSM_FIPS_SUPPORTED_REGIONS:
+                fips_endpoint = f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com"
+            else:
+                # Log warning if SSM FIPS endpoint is not supported for commercial region
+                if not config.is_gov_region:
+                    logger.warning(
+                        "FIPS mode is enabled, but '%s' does not support SSM FIPS endpoints. "
+                        "Using standard SSM endpoint.",
+                        LAMBDA_REGION,
+                    )
         ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint)
         api_key = ssm_client.get_parameter(
             Name=DD_API_KEY_SSM_NAME, WithDecryption=True

tests/test_api.py

@@ -89,23 +89,65 @@ def test_secrets_manager_different_region_but_still_fips(self, mock_boto3_client
 
     @patch("datadog_lambda.config.Config.fips_mode_enabled", True)
     @patch("botocore.session.Session.create_client")
-    def test_ssm_fips_endpoint(self, mock_boto3_client):
+    def test_ssm_fips_endpoint_supported_region(self, mock_boto3_client):
         mock_client = MagicMock()
         mock_client.get_parameter.return_value = {
             "Parameter": {"Value": "test-api-key"}
         }
         mock_boto3_client.return_value = mock_client
 
-        os.environ["AWS_REGION"] = "us-gov-west-1"
+        os.environ["AWS_REGION"] = "us-east-1"
         os.environ["DD_API_KEY_SSM_NAME"] = "test-ssm-param"
 
         api_key = api.get_api_key()
 
         mock_boto3_client.assert_called_with(
-            "ssm", endpoint_url="https://ssm-fips.us-gov-west-1.amazonaws.com"
+            "ssm", endpoint_url="https://ssm-fips.us-east-1.amazonaws.com"
         )
         self.assertEqual(api_key, "test-api-key")
 
+    @patch("datadog_lambda.config.Config.fips_mode_enabled", True)
+    @patch("datadog_lambda.config.Config.is_gov_region", True)
+    @patch("botocore.session.Session.create_client")
+    def test_ssm_gov_endpoint(self, mock_boto3_client):
+        mock_client = MagicMock()
+        mock_client.get_parameter.return_value = {
+            "Parameter": {"Value": "test-api-key"}
+        }
+        mock_boto3_client.return_value = mock_client
+
+        os.environ["AWS_REGION"] = "us-gov-west-1"
+        os.environ["DD_API_KEY_SSM_NAME"] = "test-ssm-param"
+
+        api_key = api.get_api_key()
+
+        mock_boto3_client.assert_called_with("ssm", endpoint_url=None)
+        self.assertEqual(api_key, "test-api-key")
+
+    @patch("datadog_lambda.config.Config.fips_mode_enabled", True)
+    @patch("botocore.session.Session.create_client")
+    def test_ssm_fips_endpoint_unsupported_region(self, mock_boto3_client):
+        mock_client = MagicMock()
+        mock_client.get_parameter.return_value = {
+            "Parameter": {"Value": "test-api-key"}
+        }
+        mock_boto3_client.return_value = mock_client
+
+        os.environ["AWS_REGION"] = "eu-west-1"
+        os.environ["DD_API_KEY_SSM_NAME"] = "test-ssm-param"
+
+        with self.assertLogs("datadog_lambda.api", level="WARNING") as log_context:
+            api_key = api.get_api_key()
+
+        mock_boto3_client.assert_called_with("ssm", endpoint_url=None)
+        self.assertEqual(api_key, "test-api-key")
+        self.assertTrue(
+            any(
+                "does not support SSM FIPS endpoints" in log_msg
+                for log_msg in log_context.output
+            )
+        )
+
     @patch("datadog_lambda.config.Config.fips_mode_enabled", True)
     @patch("botocore.session.Session.create_client")
     @patch("datadog_lambda.api.decrypt_kms_api_key")