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

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

Author: YashasviChaurasia

.github/workflows/test.yaml

@@ -26,4 +26,4 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install tox
       - name: Run unit tests
-        run: tox -e py
+        run: tox -e py -- --ignore=tests/trackers

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/test_sft_trainer.py

@@ -37,7 +37,6 @@
 import yaml
 
 # First Party
-from build.utils import serialize_args
 from scripts.run_inference import TunedCausalLM
 from tests.artifacts.language_models import MAYKEYE_TINY_LLAMA_CACHED, TINYMIXTRAL_MOE
 from tests.artifacts.predefined_data_configs import (
@@ -875,14 +874,36 @@ def test_successful_lora_target_modules_default_from_main(monkeypatch):
     """
     with tempfile.TemporaryDirectory() as tempdir:
         TRAIN_KWARGS = {
-            **MODEL_ARGS.__dict__,
-            **TRAIN_ARGS.__dict__,
-            **DATA_ARGS.__dict__,
-            **PEFT_LORA_ARGS.__dict__,
-            **{"peft_method": "lora", "output_dir": tempdir},
+            "model_name_or_path": MODEL_ARGS.model_name_or_path,
+            "use_flash_attn": MODEL_ARGS.use_flash_attn,
+            "torch_dtype": MODEL_ARGS.torch_dtype,
+            "training_data_path": DATA_ARGS.training_data_path,
+            "response_template": DATA_ARGS.response_template,
+            "dataset_text_field": DATA_ARGS.dataset_text_field,
+            "num_train_epochs": TRAIN_ARGS.num_train_epochs,
+            "per_device_train_batch_size": TRAIN_ARGS.per_device_train_batch_size,
+            "per_device_eval_batch_size": TRAIN_ARGS.per_device_eval_batch_size,
+            "gradient_accumulation_steps": TRAIN_ARGS.gradient_accumulation_steps,
+            "learning_rate": TRAIN_ARGS.learning_rate,
+            "weight_decay": TRAIN_ARGS.weight_decay,
+            "warmup_ratio": TRAIN_ARGS.warmup_ratio,
+            "lr_scheduler_type": TRAIN_ARGS.lr_scheduler_type,
+            "logging_steps": TRAIN_ARGS.logging_steps,
+            "include_tokens_per_second": TRAIN_ARGS.include_tokens_per_second,
+            "packing": TRAIN_ARGS.packing,
+            "max_seq_length": TRAIN_ARGS.max_seq_length,
+            "save_strategy": TRAIN_ARGS.save_strategy,
+            "r": PEFT_LORA_ARGS.r,
+            "lora_alpha": PEFT_LORA_ARGS.lora_alpha,
+            "lora_dropout": PEFT_LORA_ARGS.lora_dropout,
+            "peft_method": "lora",
+            "output_dir": tempdir,
         }
-        serialized_args = serialize_args(TRAIN_KWARGS)
-        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
+        config_path = os.path.join(tempdir, "config.json")
+        with open(config_path, "w", encoding="utf-8") as f:
+            json.dump(TRAIN_KWARGS, f)
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_PATH", config_path)
+        monkeypatch.delenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", raising=False)
 
         sft_trainer.main()
 
@@ -902,8 +923,6 @@ def test_successful_lora_target_modules_default_from_main(monkeypatch):
             "v_proj",
         }, "target_modules are not set to the default values."
 
-        os.environ.pop("SFT_TRAINER_CONFIG_JSON_ENV_VAR", None)
-
 
 def test_run_causallm_lora_add_special_tokens():
     """Check if embed layer is added as modules_to_save when special tokens are added"""

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)