depr: Disable pickle deserialization by default for security (#1245)

Adds `allow_pickle=False` parameter to `deserialize_values()` and all calling functions (`job.result()`, `load_job_result()`, `load_job_checkpoint()`). When data is in `PICKLED_V4` format and `allow_pickle` is not explicitly `True`, a `RuntimeError` is raised with clear instructions.

Co-authored-by: Cody Wang <speller26@gmail.com>

Author: AjinkyaGawali

src/braket/aws/aws_quantum_job.py

@@ -532,6 +532,7 @@ def result(
         self,
         poll_timeout_seconds: float = QuantumJob.DEFAULT_RESULTS_POLL_TIMEOUT,
         poll_interval_seconds: float = QuantumJob.DEFAULT_RESULTS_POLL_INTERVAL,
+        allow_pickle: bool = False,
     ) -> dict[str, Any]:
         """Retrieves the hybrid job result persisted using the `save_job_result` function.
 
@@ -540,12 +541,16 @@ def result(
                 Default: 10 days.
             poll_interval_seconds (float): The polling interval, in seconds, for `result()`.
                 Default: 5 seconds.
+            allow_pickle (bool): Whether to allow deserialization of pickled data. Pickle
+                deserialization can execute arbitrary code and is unsafe on untrusted data.
+                Default: False.
 
         Returns:
             dict[str, Any]: Dict specifying the job results.
 
         Raises:
-            RuntimeError: if hybrid job is in a FAILED or CANCELLED state.
+            RuntimeError: if hybrid job is in a FAILED or CANCELLED state, or if data is in
+                PICKLED_V4 format and allow_pickle is False.
             TimeoutError: if hybrid job execution exceeds the polling timeout period.
         """
         with tempfile.TemporaryDirectory() as temp_dir:
@@ -557,11 +562,16 @@ def result(
                 if e.response["Error"]["Code"] == "404":
                     return {}
                 raise
-            return AwsQuantumJob._read_and_deserialize_results(temp_dir, job_name)
+            return AwsQuantumJob._read_and_deserialize_results(temp_dir, job_name, allow_pickle)
 
     @staticmethod
-    def _read_and_deserialize_results(temp_dir: str, job_name: str) -> dict[str, Any]:
-        return load_job_result(Path(temp_dir, job_name, AwsQuantumJob.RESULTS_FILENAME))
+    def _read_and_deserialize_results(
+        temp_dir: str, job_name: str, allow_pickle: bool = False
+    ) -> dict[str, Any]:
+        return load_job_result(
+            Path(temp_dir, job_name, AwsQuantumJob.RESULTS_FILENAME),
+            allow_pickle=allow_pickle,
+        )
 
     def download_result(
         self,

src/braket/jobs/_entry_point_template.py

@@ -19,7 +19,7 @@ def {function_name}():
     finally:
         clean_links(links)
     if result is not None:
-        save_job_result(result, data_format=PersistedJobDataFormat.PICKLED_V4)
+        save_job_result(result, data_format=PersistedJobDataFormat.{data_format})
     return result
 """
 

src/braket/jobs/data_persistence.py

@@ -66,7 +66,9 @@ def save_job_checkpoint(
 
 
 def load_job_checkpoint(
-    job_name: str | None = None, checkpoint_file_suffix: str = ""
+    job_name: str | None = None,
+    checkpoint_file_suffix: str = "",
+    allow_pickle: bool = False,
 ) -> dict[str, Any]:
     """Loads the job checkpoint data stored for the job named 'job_name', with the checkpoint
     file that ends with the `checkpoint_file_suffix`. The `job_name` can refer to any job whose
@@ -87,6 +89,10 @@ def load_job_checkpoint(
             `f"{job_name}(_{checkpoint_file_suffix}).json"` is used to locate the
             checkpoint file. Default: ""
 
+        allow_pickle (bool): Whether to allow deserialization of pickled data. Pickle
+            deserialization can execute arbitrary code and is unsafe on untrusted data.
+            Default: False.
+
     Returns:
         dict[str, Any]: Dict that contains the checkpoint data persisted in the checkpoint file.
 
@@ -95,6 +101,7 @@ def load_job_checkpoint(
             in the directory specified by the container environment variable `CHECKPOINT_DIR`.
         ValueError: If the data stored in the checkpoint file can't be deserialized (possibly due to
             corruption).
+        RuntimeError: If data is in PICKLED_V4 format and allow_pickle is False.
     """
     job_name = job_name or get_job_name()
     checkpoint_directory = get_checkpoint_dir()
@@ -105,7 +112,9 @@ def load_job_checkpoint(
     )
     with open(checkpoint_file_path, encoding="utf-8") as f:
         persisted_data = PersistedJobData.parse_raw(f.read())
-        return deserialize_values(persisted_data.dataDictionary, persisted_data.dataFormat)
+        return deserialize_values(
+            persisted_data.dataDictionary, persisted_data.dataFormat, allow_pickle
+        )
 
 
 def _load_persisted_data(filename: str | Path | None = None) -> PersistedJobData:
@@ -120,19 +129,29 @@ def _load_persisted_data(filename: str | Path | None = None) -> PersistedJobData
         )
 
 
-def load_job_result(filename: str | Path | None = None) -> dict[str, Any]:
+def load_job_result(
+    filename: str | Path | None = None, allow_pickle: bool = False
+) -> dict[str, Any]:
     """Loads job result of currently running job.
 
     Args:
         filename (str | Path | None): Location of job results. Default `results.json` in job
             results directory in a job instance or in working directory locally. This file
             must be in the format used by `save_job_result`.
+        allow_pickle (bool): Whether to allow deserialization of pickled data. Pickle
+            deserialization can execute arbitrary code and is unsafe on untrusted data.
+            Default: False.
 
     Returns:
         dict[str, Any]: Job result data of current job
+
+    Raises:
+        RuntimeError: If data is in PICKLED_V4 format and allow_pickle is False.
     """
     persisted_data = _load_persisted_data(filename)
-    return deserialize_values(persisted_data.dataDictionary, persisted_data.dataFormat)
+    return deserialize_values(
+        persisted_data.dataDictionary, persisted_data.dataFormat, allow_pickle
+    )
 
 
 def save_job_result(
@@ -180,6 +199,7 @@ def save_job_result(
     current_results = deserialize_values(
         current_persisted_data.dataDictionary,
         current_persisted_data.dataFormat,
+        allow_pickle=True,  # safe: data was written by this SDK inside the job container
     )
     updated_results = current_results | result_data
 

src/braket/jobs/hybrid_job.py

@@ -30,6 +30,7 @@
 from typing import Any
 
 import cloudpickle
+from braket.jobs_data import PersistedJobDataFormat
 
 from braket.aws.aws_session import AwsSession
 from braket.jobs._entry_point_template import run_entry_point, symlink_input_data
@@ -72,6 +73,7 @@ def hybrid_job(
     logger: Logger = getLogger(__name__),
     quiet: bool | None = None,
     reservation_arn: str | None = None,
+    data_format: PersistedJobDataFormat | str = PersistedJobDataFormat.PLAINTEXT,
 ) -> Callable:
     """Defines a hybrid job by decorating the entry point function. The job will be created
     when the decorated function is called.
@@ -207,7 +209,7 @@ def job_wrapper(*args: Any, **kwargs: Any) -> Callable:
                 ) as entry_point_file:
                     template = "\n".join([
                         _process_input_data(input_data),
-                        _serialize_entry_point(entry_point, args, kwargs),
+                        _serialize_entry_point(entry_point, args, kwargs, data_format),
                     ])
                     entry_point_file.write(template)
 
@@ -417,8 +419,32 @@ def __exit__(self, exc_type, exc_val, exc_tb):  # noqa: ANN001
             cloudpickle.unregister_pickle_by_value(module)
 
 
-def _serialize_entry_point(entry_point: Callable, args: tuple, kwargs: dict) -> str:
+_DATA_FORMAT_ALIASES = {
+    "pickle": PersistedJobDataFormat.PICKLED_V4,
+    "plaintext": PersistedJobDataFormat.PLAINTEXT,
+}
+
+
+def _resolve_data_format(data_format: PersistedJobDataFormat | str) -> PersistedJobDataFormat:
+    if isinstance(data_format, PersistedJobDataFormat):
+        return data_format
+    resolved = _DATA_FORMAT_ALIASES.get(data_format.lower())
+    if resolved is None:
+        raise ValueError(
+            f"Unknown data_format '{data_format}'. "
+            f"Use {list(_DATA_FORMAT_ALIASES.keys())} or a PersistedJobDataFormat enum."
+        )
+    return resolved
+
+
+def _serialize_entry_point(
+    entry_point: Callable,
+    args: tuple,
+    kwargs: dict,
+    data_format: PersistedJobDataFormat | str = PersistedJobDataFormat.PLAINTEXT,
+) -> str:
     """Create an entry point from a function"""
+    data_format = _resolve_data_format(data_format)
     wrapped_entry_point = functools.partial(entry_point, *args, **kwargs)
 
     try:
@@ -434,6 +460,7 @@ def _serialize_entry_point(entry_point: Callable, args: tuple, kwargs: dict) ->
     return run_entry_point.format(
         serialized=serialized,
         function_name=entry_point.__name__,
+        data_format=data_format.name,
     )
 
 

src/braket/jobs/local/local_job.py

@@ -274,6 +274,7 @@ def result(
         self,
         poll_timeout_seconds: float = QuantumJob.DEFAULT_RESULTS_POLL_TIMEOUT,
         poll_interval_seconds: float = QuantumJob.DEFAULT_RESULTS_POLL_INTERVAL,
+        allow_pickle: bool = False,
     ) -> dict[str, Any]:
         """Retrieves the `LocalQuantumJob` result persisted using `save_job_result` function.
 
@@ -282,17 +283,23 @@ def result(
                 Default: 10 days.
             poll_interval_seconds (float): The polling interval, in seconds, for `result()`.
                 Default: 5 seconds.
+            allow_pickle (bool): Whether to allow deserialization of pickled data. Pickle
+                deserialization can execute arbitrary code and is unsafe on untrusted data.
+                Default: False.
 
         Raises:
             ValueError: The local job directory does not exist.
+            RuntimeError: If data is in PICKLED_V4 format and allow_pickle is False.
 
         Returns:
             dict[str, Any]: Dict specifying the hybrid job results.
         """
         try:
             with open(os.path.join(self.name, "results.json"), encoding="utf-8") as f:
                 persisted_data = PersistedJobData.parse_raw(f.read())
-                return deserialize_values(persisted_data.dataDictionary, persisted_data.dataFormat)
+                return deserialize_values(
+                    persisted_data.dataDictionary, persisted_data.dataFormat, allow_pickle
+                )
         except FileNotFoundError as e:
             raise ValueError(
                 f"Unable to find results in the local job directory {self.name}."

src/braket/jobs/quantum_job.py

@@ -155,6 +155,7 @@ def result(
         self,
         poll_timeout_seconds: float = DEFAULT_RESULTS_POLL_TIMEOUT,
         poll_interval_seconds: float = DEFAULT_RESULTS_POLL_INTERVAL,
+        allow_pickle: bool = False,
     ) -> dict[str, Any]:
         """Retrieves the hybrid job result persisted using save_job_result() function.
 
@@ -165,12 +166,16 @@ def result(
             poll_interval_seconds (float): The polling interval, in seconds, for `result()`.
                 Default: 5 seconds.
 
+            allow_pickle (bool): Whether to allow deserialization of pickled data. Pickle
+                deserialization can execute arbitrary code and is unsafe on untrusted data.
+                Default: False.
 
         Returns:
             dict[str, Any]: Dict specifying the hybrid job results.
 
         Raises:
-            RuntimeError: if hybrid job is in a FAILED or CANCELLED state.
+            RuntimeError: if hybrid job is in a FAILED or CANCELLED state, or if data is in
+                PICKLED_V4 format and allow_pickle is False.
             TimeoutError: if hybrid job execution exceeds the polling timeout period.
         """
 

src/braket/jobs/serialization.py

@@ -45,21 +45,39 @@ def serialize_values(
 
 
 def deserialize_values(
-    data_dictionary: dict[str, Any], data_format: PersistedJobDataFormat
+    data_dictionary: dict[str, Any],
+    data_format: PersistedJobDataFormat,
+    allow_pickle: bool = False,
 ) -> dict[str, Any]:
     """Deserializes the `data_dictionary` values from the format specified by `data_format`.
 
     Args:
         data_dictionary (dict[str, Any]): Dict whose values are to be deserialized.
         data_format (PersistedJobDataFormat): The data format that the `data_dictionary` values
             are currently serialized with.
+        allow_pickle (bool): Whether to allow deserialization of pickled data. Pickle
+            deserialization can execute arbitrary code and is unsafe on untrusted data.
+            Default: False.
 
     Returns:
         dict[str, Any]: Dict with same keys as `data_dictionary` and values deserialized from
         the specified `data_format` to plaintext.
+
+    Raises:
+        RuntimeError: If data format is PICKLED_V4 and allow_pickle is False.
     """
-    return (
-        {k: pickle.loads(codecs.decode(v.encode(), "base64")) for k, v in data_dictionary.items()}  # noqa: S301
-        if data_format == PersistedJobDataFormat.PICKLED_V4
-        else data_dictionary
-    )
+    if data_format == PersistedJobDataFormat.PICKLED_V4:
+        if not allow_pickle:
+            raise RuntimeError(
+                "Data is in PICKLED_V4 format, but pickle deserialization is disabled by "
+                "default due to security concerns. Pickle deserialization can execute arbitrary "
+                "code and is unsafe on untrusted data. To enable pickle deserialization, pass "
+                "allow_pickle=True to the calling function (e.g. job.result(allow_pickle=True), "
+                "load_job_result(allow_pickle=True), or load_job_checkpoint(allow_pickle=True)). "
+                "Only do this if you trust the source of the data."
+            )
+        return {
+            k: pickle.loads(codecs.decode(v.encode(), "base64"))  # noqa: S301
+            for k, v in data_dictionary.items()
+        }
+    return data_dictionary

test/unit_tests/braket/jobs/test_data_persistence.py

@@ -142,10 +142,11 @@ def test_load_job_checkpoint(
         with open(file_path, "w") as f:
             f.write(saved_data)
 
+        allow_pickle = data_format == PersistedJobDataFormat.PICKLED_V4
         with patch.dict(
             os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
         ):
-            loaded_data = load_job_checkpoint(job_name, file_suffix)
+            loaded_data = load_job_checkpoint(job_name, file_suffix, allow_pickle=allow_pickle)
             assert loaded_data == expected_checkpoint_data
 
 
@@ -188,7 +189,7 @@ def test_load_job_checkpoint_raises_error_corrupted_data():
         with patch.dict(
             os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
         ):
-            load_job_checkpoint(job_name, file_suffix)
+            load_job_checkpoint(job_name, file_suffix, allow_pickle=True)
 
 
 @dataclass
@@ -211,7 +212,7 @@ def test_save_and_load_job_checkpoint():
             os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
         ):
             save_job_checkpoint(data, data_format=PersistedJobDataFormat.PICKLED_V4)
-            retrieved = load_job_checkpoint(job_name)
+            retrieved = load_job_checkpoint(job_name, allow_pickle=True)
             assert retrieved == data
 
 
@@ -308,7 +309,10 @@ def test_update_result_data(
             save_job_result(first_result_data, first_data_format)
             save_job_result(second_result_data, second_data_format)
 
-            assert load_job_result() == expected_result_data
+            allow_pickle = first_data_format == PersistedJobDataFormat.PICKLED_V4 or (
+                second_data_format == PersistedJobDataFormat.PICKLED_V4
+            )
+            assert load_job_result(allow_pickle=allow_pickle) == expected_result_data
 
 
 def test_update_pickled_results_as_plaintext_error():
@@ -322,3 +326,22 @@ def test_update_pickled_results_as_plaintext_error():
             )
             with pytest.raises(TypeError, match=cannot_convert_pickled_to_plaintext):
                 save_job_result("hello", PersistedJobDataFormat.PLAINTEXT)
+
+
+def test_load_job_result_raises_without_allow_pickle():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        with patch.dict(os.environ, {"AMZN_BRAKET_JOB_RESULTS_DIR": tmp_dir}):
+            save_job_result({"key": "value"}, PersistedJobDataFormat.PICKLED_V4)
+            with pytest.raises(RuntimeError, match="pickle deserialization is disabled by default"):
+                load_job_result()
+
+
+def test_load_job_checkpoint_raises_without_allow_pickle():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        job_name = "test_job"
+        with patch.dict(
+            os.environ, {"AMZN_BRAKET_CHECKPOINT_DIR": tmp_dir, "AMZN_BRAKET_JOB_NAME": job_name}
+        ):
+            save_job_checkpoint({"key": "value"}, data_format=PersistedJobDataFormat.PICKLED_V4)
+            with pytest.raises(RuntimeError, match="pickle deserialization is disabled by default"):
+                load_job_checkpoint(job_name)