fix: remove unsafe pickle deserialization in txt_to_obj (CVE CWE-502)

Signed-off-by: yashasvi <yashasvi@ibm.com>

Author: YashasviChaurasia

build/utils.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 # Standard
-import os
-import logging
-import pickle
 import base64
+import json
+import logging
+import os
 
 # Third Party
 import torch
@@ -67,14 +67,21 @@ def get_highest_checkpoint(dir_path):
     return checkpoint_dir
 
 
+def _json_default(obj):
+    """Fallback serializer for objects not natively JSON-serializable."""
+    if hasattr(obj, "__dict__"):
+        return obj.__dict__
+    return str(obj)
+
+
 def serialize_args(args_json):
     """Given dict, converts to base64 byte representation.
 
     Args:
         args_json: dict
     Returns: str
     """
-    message_bytes = pickle.dumps(args_json)
+    message_bytes = json.dumps(args_json, default=_json_default).encode("utf-8")
     base64_bytes = base64.b64encode(message_bytes)
     return base64_bytes.decode("ascii")
 

tests/utils/test_config_utils.py

@@ -17,7 +17,7 @@
 
 # Standard
 import base64
-import pickle
+import json
 
 # Third Party
 from datasets import Dataset, Features, Value
@@ -224,7 +224,7 @@ def test_get_json_config_can_load_from_envvar(monkeypatch):
     the json path from env var SFT_TRAINER_CONFIG_JSON_ENV_VAR
     """
     config_json = {"model_name_or_path": "foobar"}
-    message_bytes = pickle.dumps(config_json)
+    message_bytes = json.dumps(config_json).encode("utf-8")
     base64_bytes = base64.b64encode(message_bytes)
     encoded_json = base64_bytes.decode("ascii")
     monkeypatch.delenv("SFT_TRAINER_CONFIG_JSON_PATH", raising=False)

tuning/utils/config_utils.py

@@ -17,7 +17,6 @@
 import base64
 import json
 import os
-import pickle
 
 # Third Party
 from peft import PromptTuningConfig as HFPromptTuningConfig
@@ -159,9 +158,4 @@ def txt_to_obj(txt):
     """
     base64_bytes = txt.encode("ascii")
     message_bytes = base64.b64decode(base64_bytes)
-    try:
-        # If the bytes represent JSON string
-        return json.loads(message_bytes)
-    except UnicodeDecodeError:
-        # Otherwise the bytes are a pickled python dictionary
-        return pickle.loads(message_bytes)
+    return json.loads(message_bytes)