pickle.loads fallback on dataset field executes arbitrary code in data processing

[EvolveAegis]

Summary

The LiveCodeBench dataset processing path in Agent0/executor_train/verl/recipe/r1/data_process.py unpickles an attacker-influenced field when it isn't valid JSON:

try:
    private_test_cases = json.loads(example["private_test_cases"])
except Exception as e:
    private_test_cases = json.loads(
        pickle.loads(zlib.decompress(base64.b64decode(example["private_test_cases"].encode("utf-8"))))
    )

private_test_cases comes from the HuggingFace dataset livecodebench/code_generation_lite (data_process.py:153-155) — outside the training pipeline's trust boundary, and load_dataset pins no revision, so the row content is whoever owns that dataset namespace/revision today. A row whose field fails json.loads falls into the pickle.loads branch, which executes arbitrary code during unpickling. The same fallback exists at tasks/livecodebench.py:63 for the test-case field, where the pickle branch is the working path (the harness itself writes that field via pickle.dumps at data_process.py:142).

I verified with the actual function: a row whose private_test_cases is base64(zlib(pickle(<<__reduce__ -> os.system>>))) executes the command in the data-processing worker (canary file created). A control row with valid JSON takes the non-pickle path and completes normally. (After the payload fires the worker crashes with a TypeError — irrelevant for the attacker, the command already ran.)

Details

• data_process.py:136-141 — build_livecodebench_dataset → process_livecodebench: JSON parse, on failure json.loads(pickle.loads(zlib.decompress(base64.b64decode(...)))).
• data_process.py:153-155 — load_dataset("livecodebench/code_generation_lite", split="test"), no revision pinned.
• tasks/livecodebench.py:63 — same fallback on the test-case field.
• The example_map_fn wrapper (data_process.py:101-115) runs with num_proc parallelism — the payload fires in whichever worker processes the row.
• Upstream note: verl-project/verl has no recipe/r1/data_process.py and no build_livecodebench_dataset (checked by code search; their LiveCodeBench eval code in verl/utils/reward_score/prime_code/ uses plain json.loads, no pickle). The same pickle fallback does exist in microsoft/LMOps oel/verl/recipe/r1/data_process.py:117 (this vendored code's apparent origin; no public issue reports it there, and Agent0's vendored copy is what runs here). Fixing in both places, or a shared comment, would cover the lineage.

How to reproduce

import base64, pickle, zlib
class P:
    def __reduce__(self):
        import os
        return (os.system, ("echo PWNED > /tmp/canary",))
blob = base64.b64encode(zlib.compress(pickle.dumps(P()))).decode()
row = {
    "question_content": "q", "starter_code": "",
    "public_test_cases": '[{"input":"1","output":"2"}]',
    "private_test_cases": blob,          # not JSON -> pickle branch
    "metadata": '{"func_name": null}',
}
# process_livecodebench(row) -> pickle.loads executes os.system -> /tmp/canary written

Impact

Code execution in the training-data pipeline, supply-chain class (a poisoned or replaced dataset/revision, or a compromised data contributor). The unsafe fallback turns any JSON-parse failure into a code-execution path; the producer side already writes pickle, so the fallback is load-bearing for legitimate data.

Suggested change

• Use a restricted unpickler in both data_process.py and tasks/livecodebench.py (allow only the built-in types the field legitimately carries — dict/list/str/int/float/bool) so the existing pickle-based flow keeps working without executing arbitrary code.
• If the format can change at the producer side, use plain compressed JSON instead of pickle.