Unsafe Deserialization in JsonPickleSerializer Enables Remote Code Execution in run-llama/llama_index

@@ -64,7 +64,7 @@ def deserialize(self, value: str) -> Any:
         return self._deserialize_value(data)
 
 
-class JsonPickleSerializer(JsonSerializer):
+class PickleSerializer(JsonSerializer):
     def serialize(self, value: Any) -> str:
         """Serialize while prioritizing JSON, falling back to Pickle."""
         try:
@@ -73,8 +73,15 @@ def serialize(self, value: Any) -> str:
             return base64.b64encode(pickle.dumps(value)).decode("utf-8")
 
     def deserialize(self, value: str) -> Any:
-        """Deserialize while prioritizing Pickle, falling back to JSON."""
+        """
+        Deserialize while prioritizing Pickle, falling back to JSON.
+        To avoid malicious exploits of the deserialization, deserialize objects
+        only when you deem it safe to do so.
+        """
         try:
             return pickle.loads(base64.b64decode(value))
         except Exception:
             return super().deserialize(value)
+
+
+JsonPickleSerializer = PickleSerializer

@@ -8,7 +8,7 @@
 import pytest
 from llama_index.core import MockEmbedding
 from llama_index.core.llms.mock import MockLLM
-from llama_index.core.workflow.context_serializers import JsonPickleSerializer
+from llama_index.core.workflow.context_serializers import PickleSerializer
 from llama_index.core.workflow.decorators import step
 from llama_index.core.workflow.events import (
     Event,
@@ -534,9 +534,9 @@ async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
         state_dict = handler.ctx.to_dict()
 
     # if we allow pickle, then we can pickle the LLM/embedding object
-    state_dict = handler.ctx.to_dict(serializer=JsonPickleSerializer())
+    state_dict = handler.ctx.to_dict(serializer=PickleSerializer())
     new_handler = WorkflowHandler(
-        ctx=Context.from_dict(wf, state_dict, serializer=JsonPickleSerializer())
+        ctx=Context.from_dict(wf, state_dict, serializer=PickleSerializer())
     )
     assert new_handler.ctx