From 9c916cebaa7c56b7f53f3978af11d6ac5593da4d Mon Sep 17 00:00:00 2001
From: pintaoz <pintaoz@amazon.com>
Date: Thu, 18 Dec 2025 14:59:32 -0800
Subject: [PATCH] Add validation in conda env name

---
 .../sagemaker/core/remote_function/client.py  | 25 +++++++++++++++++++
 .../tests/unit/remote_function/test_client.py | 23 ++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/sagemaker-core/src/sagemaker/core/remote_function/client.py b/sagemaker-core/src/sagemaker/core/remote_function/client.py
index a38b57662a..3cfa5e3b23 100644
--- a/sagemaker-core/src/sagemaker/core/remote_function/client.py
+++ b/sagemaker-core/src/sagemaker/core/remote_function/client.py
@@ -303,6 +303,9 @@ def remote(
     """
 
     def _remote(func):
+        
+        if job_conda_env:
+            RemoteExecutor._validate_env_name(job_conda_env)
 
         job_settings = _JobSettings(
             dependencies=dependencies,
@@ -774,6 +777,9 @@ def __init__(
                 + "without spark_config or use_torchrun or use_mpirun. "
                 + "Please provide instance_count = 1"
             )
+        
+        if job_conda_env:
+            self._validate_env_name(job_conda_env)
 
         self.job_settings = _JobSettings(
             dependencies=dependencies,
@@ -951,6 +957,25 @@ def _validate_submit_args(func, *args, **kwargs):
                 + f"{'arguments' if len(missing_kwargs) > 1 else 'argument'}: "
                 + f"{missing_kwargs_string}"
             )
+    
+    @staticmethod
+    def _validate_env_name(env_name: str) -> None:
+        """Validate conda environment name to prevent command injection.
+        
+        Args:
+            env_name (str): The environment name to validate
+            
+        Raises:
+            ValueError: If the environment name contains invalid characters
+        """
+        
+        # Allow only alphanumeric, underscore, and hyphen
+        import re
+        if not re.match(r'^[a-zA-Z0-9_-]+$', env_name):
+            raise ValueError(
+                f"Invalid environment name '{env_name}'. "
+                "Only alphanumeric characters, underscores, and hyphens are allowed."
+            )
 
 
 class Future(object):
diff --git a/sagemaker-core/tests/unit/remote_function/test_client.py b/sagemaker-core/tests/unit/remote_function/test_client.py
index 83e1a2db80..8621b98063 100644
--- a/sagemaker-core/tests/unit/remote_function/test_client.py
+++ b/sagemaker-core/tests/unit/remote_function/test_client.py
@@ -64,7 +64,28 @@ def my_function(x):
 
         with pytest.raises(TypeError):
             RemoteExecutor._validate_submit_args(my_function, 1, 2)
-
+        
+    def test_validate_env_names_valid(self):
+        """Test valid conda environment names"""
+        valid_names = [
+            "myenv",
+            "base",
+            "py39",
+            "env123",
+        ]
+        for name in valid_names:
+            RemoteExecutor._validate_env_name(name)
+
+    def test_validate_env_names_invalid(self):
+        """Test invalid conda environment names"""
+        invalid_names = [
+            "env && echo PWNED",
+            "env > /tmp/output.txt",
+            "sagemaker-rce-env; echo PWNED_FROM_CONDA_ENV > /tmp/conda_rce.txt #",
+        ]
+        for name in invalid_names:
+            with pytest.raises(ValueError):
+                RemoteExecutor._validate_env_name(name)
 
 class TestWorkerFunctions:
     """Test worker thread functions"""