Update after testing

Author: mollyheamazon

sagemaker-core/src/sagemaker/core/remote_function/client.py

@@ -369,7 +369,6 @@ def wrapper(*args, **kwargs):
                             s3_uri=s3_path_join(
                                 job_settings.s3_root_uri, job.job_name, EXCEPTION_FOLDER
                             ),
-                            
                         )
                     except ServiceError as serr:
                         chained_e = serr.__cause__
@@ -406,7 +405,6 @@ def wrapper(*args, **kwargs):
                 return serialization.deserialize_obj_from_s3(
                     sagemaker_session=job_settings.sagemaker_session,
                     s3_uri=s3_path_join(job_settings.s3_root_uri, job.job_name, RESULTS_FOLDER),
-                    
                 )
 
             if job.describe()["TrainingJobStatus"] == "Stopped":
@@ -1008,7 +1006,6 @@ def from_describe_response(describe_training_job_response, sagemaker_session):
                 job_return = serialization.deserialize_obj_from_s3(
                     sagemaker_session=sagemaker_session,
                     s3_uri=s3_path_join(job.s3_uri, RESULTS_FOLDER),
-                    
                 )
             except DeserializationError as e:
                 client_exception = e
@@ -1020,7 +1017,6 @@ def from_describe_response(describe_training_job_response, sagemaker_session):
                 job_exception = serialization.deserialize_exception_from_s3(
                     sagemaker_session=sagemaker_session,
                     s3_uri=s3_path_join(job.s3_uri, EXCEPTION_FOLDER),
-                    
                 )
             except ServiceError as serr:
                 chained_e = serr.__cause__
@@ -1110,7 +1106,6 @@ def result(self, timeout: float = None) -> Any:
                     self._return = serialization.deserialize_obj_from_s3(
                         sagemaker_session=self._job.sagemaker_session,
                         s3_uri=s3_path_join(self._job.s3_uri, RESULTS_FOLDER),
-                        
                     )
                     self._state = _FINISHED
                     return self._return
@@ -1119,7 +1114,6 @@ def result(self, timeout: float = None) -> Any:
                         self._exception = serialization.deserialize_exception_from_s3(
                             sagemaker_session=self._job.sagemaker_session,
                             s3_uri=s3_path_join(self._job.s3_uri, EXCEPTION_FOLDER),
-                            
                         )
                     except ServiceError as serr:
                         chained_e = serr.__cause__

sagemaker-core/src/sagemaker/core/remote_function/core/serialization.py

@@ -216,7 +216,6 @@ def deserialize_func_from_s3(sagemaker_session: Session, s3_uri: str) -> Callabl
         buffer=bytes_to_deserialize,
         sagemaker_session=sagemaker_session,
         secret_arn=metadata.secret_arn,
-        s3_uri=s3_uri
     )
 
     return CloudpickleSerializer.deserialize(f"{s3_uri}/payload.pkl", bytes_to_deserialize)
@@ -315,7 +314,6 @@ def deserialize_obj_from_s3(sagemaker_session: Session, s3_uri: str) -> Any:
         buffer=bytes_to_deserialize,
         sagemaker_session=sagemaker_session,
         secret_arn=metadata.secret_arn,
-        s3_uri=s3_uri
     )
 
     return CloudpickleSerializer.deserialize(f"{s3_uri}/payload.pkl", bytes_to_deserialize)
@@ -411,7 +409,6 @@ def deserialize_exception_from_s3(sagemaker_session: Session, s3_uri: str) -> An
         buffer=bytes_to_deserialize,
         sagemaker_session=sagemaker_session,
         secret_arn=metadata.secret_arn,
-        s3_uri=s3_uri
     )
 
     return CloudpickleSerializer.deserialize(f"{s3_uri}/payload.pkl", bytes_to_deserialize)
@@ -518,17 +515,24 @@ def _store_secret_arn_in_parameter_store(
     ssm_client = sagemaker_session.boto_session.client('ssm')
     parameter_name = f"/sagemaker/remote-function/{job_name}/secret-arn"
 
-    ssm_client.put_parameter(
-        Name=parameter_name,
-        Value=secret_arn,
-        Type="String",
-        Overwrite=True,
-        Description=f"Secret ARN for SageMaker remote function job {job_name}",
-        Tags=[
-            {'Key': 'SageMaker:JobName', 'Value': job_name},
-            {'Key': 'SageMaker:Purpose', 'Value': 'RemoteFunctionIntegrity'}
-        ]
-    )
+    try:
+        ssm_client.put_parameter(
+            Name=parameter_name,
+            Value=secret_arn,
+            Type="String",
+            Description=f"Secret ARN for SageMaker remote function job {job_name}",
+            Tags=[
+                {'Key': 'SageMaker:JobName', 'Value': job_name},
+                {'Key': 'SageMaker:Purpose', 'Value': 'RemoteFunctionIntegrity'}
+            ]
+        )
+    except ssm_client.exceptions.ParameterAlreadyExists:
+        ssm_client.put_parameter(
+            Name=parameter_name,
+            Value=secret_arn,
+            Type="String",
+            Overwrite=True,
+        )
 
 
 def _get_secret_arn_from_parameter_store(
@@ -560,46 +564,47 @@ def _get_secret_arn_from_parameter_store(
         )
 
 
-def _extract_job_name_from_s3_uri(s3_uri: str) -> str:
-    """Extract job name from S3 URI.
+def _extract_job_name_from_secret_arn(secret_arn: str) -> str:
+    """Extract job name from a Secrets Manager ARN.
     
-    S3 URI format: s3://bucket/path/to/job-name/results
-    or: s3://bucket/job-name/function
+    Secret name convention: sagemaker/remote-function/{job_name}/hmac-key
+    ARN format: arn:aws:secretsmanager:region:account:secret:sagemaker/remote-function/{job_name}/hmac-key-XXXXXX
     
     Args:
-        s3_uri: S3 URI containing job name
+        secret_arn: Full ARN of the secret
         
     Returns:
-        Job name extracted from URI
+        Extracted job name
+        
+    Raises:
+        DeserializationError: If ARN doesn't match expected format
     """
-    # Remove s3:// prefix and split by /
-    parts = s3_uri.replace("s3://", "").split("/")
-    
-    # Try to find a part that looks like a job name
-    # Job names typically contain execution IDs or timestamps
-    for part in reversed(parts):
-        if part and part not in ['function', 'arguments', 'results', 'exception', 'payload.pkl', 'metadata.json']:
-            return part
-    
-    # Fallback: use the last meaningful part
-    return parts[-2] if len(parts) > 1 else parts[0]
+    import re
+    match = re.search(r":secret:sagemaker/remote-function/(.+)/hmac-key", secret_arn)
+    if not match:
+        raise DeserializationError(
+            f"Secret ARN does not match expected format "
+            f"'sagemaker/remote-function/{{job_name}}/hmac-key': {secret_arn}"
+        )
+    return match.group(1)
 
 
 def _validate_secret_arn(
     sagemaker_session: Session,
     metadata_secret_arn: str,
-    job_name: str
 ):
     """Validate secret ARN from metadata against trusted sources.
     
     Implements two mitigations:
     1. Validate secret is in same AWS account
     2. Validate secret ARN matches Parameter Store (trust anchor)
     
+    The job_name is derived from the secret ARN's naming convention, then
+    independently validated against the SSM trust anchor.
+    
     Args:
         sagemaker_session: SageMaker session
         metadata_secret_arn: Secret ARN from S3 metadata (untrusted)
-        job_name: Remote function job name
         
     Raises:
         DeserializationError: If validation fails
@@ -623,6 +628,7 @@ def _validate_secret_arn(
         )
 
     # Mitigation #3: Validate against Parameter Store (trust anchor)
+    job_name = _extract_job_name_from_secret_arn(metadata_secret_arn)
     expected_secret_arn = _get_secret_arn_from_parameter_store(sagemaker_session, job_name)
 
     if metadata_secret_arn != expected_secret_arn:
@@ -638,7 +644,6 @@ def _perform_integrity_check(
     buffer: bytes,
     sagemaker_session: Optional[Session] = None,
     secret_arn: Optional[str] = None,
-    s3_uri: Optional[str] = None
 ):
     """Performs integrity checks for serialized code/arguments uploaded to s3.
 
@@ -650,7 +655,6 @@ def _perform_integrity_check(
         buffer: Serialized data buffer
         sagemaker_session: SageMaker session (required if secret_arn is provided)
         secret_arn: ARN of secret containing HMAC key (None for legacy plain SHA-256)
-        s3_uri: S3 URI for extracting job name (required if secret_arn is provided)
     """
     if secret_arn:
         # New secure method: HMAC with key from Secrets Manager
@@ -659,16 +663,8 @@ def _perform_integrity_check(
                 "sagemaker_session is required for HMAC integrity check"
             )
 
-        if not s3_uri:
-            raise DeserializationError(
-                "s3_uri is required for HMAC integrity check to extract job name"
-            )
-        
-        # Extract job name from S3 URI
-        job_name = _extract_job_name_from_s3_uri(s3_uri)
-        
         # Validate secret ARN (Mitigations #1 and #3)
-        _validate_secret_arn(sagemaker_session, secret_arn, job_name)
+        _validate_secret_arn(sagemaker_session, secret_arn)
 
         # Now safe to retrieve HMAC key
         hmac_key = _get_hmac_key_from_secret(sagemaker_session, secret_arn)

sagemaker-core/src/sagemaker/core/remote_function/core/stored_function.py

@@ -145,10 +145,9 @@ def save_pipeline_step_function(self, serialized_data):
         )
         serialization._upload_payload_and_metadata_to_s3(
             bytes_to_upload=serialized_data.func,
-            
+            job_name=self.job_name,
             s3_uri=s3_path_join(self.func_upload_path, FUNCTION_FOLDER),
             sagemaker_session=self.sagemaker_session,
-            job_name=self.job_name,
             s3_kms_key=self.s3_kms_key,
         )
 
@@ -158,10 +157,9 @@ def save_pipeline_step_function(self, serialized_data):
         )
         serialization._upload_payload_and_metadata_to_s3(
             bytes_to_upload=serialized_data.args,
-            
+            job_name=self.job_name,
             s3_uri=s3_path_join(self.func_upload_path, ARGUMENTS_FOLDER),
             sagemaker_session=self.sagemaker_session,
-            job_name=self.job_name,
             s3_kms_key=self.s3_kms_key,
         )
 

sagemaker-core/tests/unit/remote_function/test_serialization_security.py

@@ -24,11 +24,11 @@
     _MetaData,
     _compute_hash,
     _compute_hmac,
+    _extract_job_name_from_secret_arn,
     _get_or_create_hmac_secret,
     _get_hmac_key_from_secret,
     _store_secret_arn_in_parameter_store,
     _get_secret_arn_from_parameter_store,
-    _extract_job_name_from_s3_uri,
     _validate_secret_arn,
     _perform_integrity_check,
     _upload_payload_and_metadata_to_s3,
@@ -186,7 +186,7 @@ def test_store_secret_arn(self):
         call_kwargs = ssm_client.put_parameter.call_args[1]
         assert call_kwargs["Name"] == f"/sagemaker/remote-function/{MOCK_JOB_NAME}/secret-arn"
         assert call_kwargs["Value"] == MOCK_SECRET_ARN
-        assert call_kwargs["Overwrite"] is True
+        assert "Tags" in call_kwargs
 
     def test_get_secret_arn(self):
         session, _, ssm_client, _ = _mock_sagemaker_session()
@@ -208,28 +208,6 @@ def test_get_secret_arn_not_found_raises(self):
             _get_secret_arn_from_parameter_store(session, MOCK_JOB_NAME)
 
 
-class TestExtractJobName:
-    """Tests for S3 URI job name extraction."""
-
-    def test_extract_from_results_uri(self):
-        result = _extract_job_name_from_s3_uri(
-            "s3://bucket/remote-function/my-job-123/results"
-        )
-        assert result == "my-job-123"
-
-    def test_extract_from_function_uri(self):
-        result = _extract_job_name_from_s3_uri(
-            "s3://bucket/remote-function/my-job-123/function"
-        )
-        assert result == "my-job-123"
-
-    def test_extract_from_exception_uri(self):
-        result = _extract_job_name_from_s3_uri(
-            "s3://bucket/remote-function/my-job-123/exception"
-        )
-        assert result == "my-job-123"
-
-
 class TestValidateSecretArn:
     """Tests for secret ARN validation (Mitigations #1 and #3)."""
 
@@ -238,7 +216,7 @@ def test_valid_secret_arn_passes(self):
         session, _, _, _ = _mock_sagemaker_session()
 
         # Should not raise
-        _validate_secret_arn(session, MOCK_SECRET_ARN, MOCK_JOB_NAME)
+        _validate_secret_arn(session, MOCK_SECRET_ARN)
 
     def test_cross_account_arn_rejected(self):
         """Mitigation #1: Secret ARN from different account should be rejected."""
@@ -247,7 +225,7 @@ def test_cross_account_arn_rejected(self):
         attacker_arn = "arn:aws:secretsmanager:us-west-2:999999999999:secret:evil-secret"
 
         with pytest.raises(DeserializationError, match="same AWS account"):
-            _validate_secret_arn(session, attacker_arn, MOCK_JOB_NAME)
+            _validate_secret_arn(session, attacker_arn)
 
     def test_tampered_arn_rejected(self):
         """Mitigation #3: ARN not matching Parameter Store should be rejected."""
@@ -261,15 +239,15 @@ def test_tampered_arn_rejected(self):
         # Attacker's ARN (same account but different secret)
         tampered_arn = f"arn:aws:secretsmanager:us-west-2:{MOCK_ACCOUNT_ID}:secret:attacker-created-secret"
 
-        with pytest.raises(DeserializationError, match="Secret ARN mismatch"):
-            _validate_secret_arn(session, tampered_arn, MOCK_JOB_NAME)
+        with pytest.raises(DeserializationError, match="does not match expected format"):
+            _validate_secret_arn(session, tampered_arn)
 
     def test_invalid_arn_format_rejected(self):
         """Malformed ARN should be rejected."""
         session, _, _, _ = _mock_sagemaker_session()
 
         with pytest.raises(DeserializationError, match="Invalid secret ARN format"):
-            _validate_secret_arn(session, "not-an-arn", MOCK_JOB_NAME)
+            _validate_secret_arn(session, "not-an-arn")
 
 
 class TestPerformIntegrityCheck:
@@ -288,7 +266,6 @@ def test_hmac_integrity_check_passes(self):
             buffer=payload,
             sagemaker_session=session,
             secret_arn=MOCK_SECRET_ARN,
-            s3_uri=MOCK_S3_URI,
         )
 
     def test_hmac_integrity_check_fails_on_tampered_payload(self):
@@ -305,7 +282,6 @@ def test_hmac_integrity_check_fails_on_tampered_payload(self):
                 buffer=tampered_payload,
                 sagemaker_session=session,
                 secret_arn=MOCK_SECRET_ARN,
-                s3_uri=MOCK_S3_URI,
             )
 
     def test_legacy_sha256_check_passes_with_warning(self):
@@ -340,19 +316,6 @@ def test_hmac_check_requires_session(self):
                 secret_arn=MOCK_SECRET_ARN,
             )
 
-    def test_hmac_check_requires_s3_uri(self):
-        """HMAC check should require s3_uri."""
-        session, _, _, _ = _mock_sagemaker_session()
-        
-        with pytest.raises(DeserializationError, match="s3_uri is required"):
-            _perform_integrity_check(
-                expected_hash_value="hash",
-                buffer=b"data",
-                sagemaker_session=session,
-                secret_arn=MOCK_SECRET_ARN,
-            )
-
-
 class TestAttackScenarios:
     """Tests simulating actual attack scenarios."""
 
@@ -373,7 +336,6 @@ def test_attacker_replaces_payload_and_metadata_plain_hash(self):
                 buffer=malicious_payload,
                 sagemaker_session=session,
                 secret_arn=MOCK_SECRET_ARN,
-                s3_uri=MOCK_S3_URI,
             )
 
     def test_attacker_points_to_cross_account_secret(self):
@@ -383,7 +345,7 @@ def test_attacker_points_to_cross_account_secret(self):
         attacker_secret_arn = "arn:aws:secretsmanager:us-west-2:999999999999:secret:attacker-secret"
 
         with pytest.raises(DeserializationError, match="same AWS account"):
-            _validate_secret_arn(session, attacker_secret_arn, MOCK_JOB_NAME)
+            _validate_secret_arn(session, attacker_secret_arn)
 
     def test_attacker_creates_secret_in_same_account(self):
         """Attacker creates secret in same account but ARN doesn't match Parameter Store."""
@@ -398,4 +360,4 @@ def test_attacker_creates_secret_in_same_account(self):
         attacker_arn = f"arn:aws:secretsmanager:us-west-2:{MOCK_ACCOUNT_ID}:secret:sagemaker/remote-function/evil-job/hmac-key"
 
         with pytest.raises(DeserializationError, match="Secret ARN mismatch"):
-            _validate_secret_arn(session, attacker_arn, MOCK_JOB_NAME)
+            _validate_secret_arn(session, attacker_arn)