Skip to content
This repository was archived by the owner on Jul 6, 2026. It is now read-only.

Commit 91d1cd4

Browse files
liu-shaojunclaude
liu-shaojun
and
claude
authored
Bug Fix: Fix unsafe pickle deserialization and CI workflow unauthorized execution (#5236)
* Fix unsafe pickle deserialization in distributed training workers Replace raw cloudpickle.load() with HMAC-verified SafePickle across all four distributed training backends (MPI, Horovod, PyTorch DDP, TensorFlow). The vulnerability allowed arbitrary code execution via crafted pickle files passed through CLI-controlled paths. Worker scripts deserialized these files without integrity verification, enabling cluster-wide RCE in distributed deployments. Changes: - Upgrade SafePickle to use per-session random HMAC-SHA256 keys instead of the hardcoded b'shared-key' with SHA1 - SafePickle.dump() now auto-generates .sig sidecar files for integrity - SafePickle.load() mandates signature verification, rejecting unsigned or tampered pickle files - Use timing-safe hmac.compare_digest() to prevent timing attacks - Propagate HMAC key to MPI workers via mpiexec -genv and to subprocesses via inherited environment - SCP .sig files alongside .pkl files for remote MPI nodes Affected files: - mpi_train.py, mpi_estimator.py, mpi_runner.py (MPI training) - horovod_worker.py, multiprocs_backend.py (Horovod) - worker.py, ddp_subprocess.py (PyTorch DDP) - subprocess_worker.py, backend.py (TensorFlow) - safepickle.py in nano, ppml, and friesian packages Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix CI workflow allowing unauthorized code execution on self-hosted runner The performance-regression-test workflow had two vulnerabilities that allowed any GitHub user to execute arbitrary code on Intel's self-hosted CI runner (Rohan): 1. The issue_comment trigger had no author_association check, so anyone could comment 'APRT' on any PR to trigger benchmark execution of fork PR code on the self-hosted runner. 2. The pull_request trigger caused the workflow to run the PR author's version of the workflow file on the self-hosted runner, giving full control over the execution environment. Fix: - Remove the pull_request trigger entirely — it allowed fork PRs to override the workflow definition itself - Add author_association guard requiring MEMBER, OWNER, or COLLABORATOR status to trigger via issue_comment - Simplify the job-level condition to only handle the comment-based flow Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 240024c commit 91d1cd4

12 files changed

Lines changed: 209 additions & 108 deletions

File tree

‎.github/workflows/performance-regression-test.yml‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,24 +8,25 @@ on:
88
# - cron: '0 16 * * *'
99
issue_comment:
1010
types: [created]
11-
pull_request:
12-
branches: [ main ]
13-
paths:
14-
- '.github/workflows/performance-regression-test.yml'
1511

1612
jobs:
1713
performace-regression-test:
1814
runs-on: [Rohan, ubuntu-20.04-lts]
19-
if: github.event.pull_request || github.event.schedule || github.event.issue.pull_request && github.event.comment.body=='APRT'
15+
if: >-
16+
github.event.issue.pull_request
17+
&& github.event.comment.body == 'APRT'
18+
&& (
19+
github.event.comment.author_association == 'MEMBER'
20+
|| github.event.comment.author_association == 'OWNER'
21+
|| github.event.comment.author_association == 'COLLABORATOR'
22+
)
2023
steps:
21-
# trigged by opening a PR which modifies this file or scheduling or commenting 'APRT' in a closed PR
22-
- name: checkout (open pr or schedule or closed pr comments)
23-
if: github.event.pull_request || github.event.schedule || github.event.issue.pull_request.merged_at
24+
- name: checkout (merged pr comments)
25+
if: github.event.issue.pull_request.merged_at
2426
uses: actions/checkout@f43a0e5ff2bd294095638e18286ca9a3d1956744 # actions/checkout@v3
2527

26-
# trigged by commenting 'APRT' in an opening PR
27-
- name: checkout (opening pr comments)
28-
if: github.event.issue.pull_request && !github.event.issue.pull_request.merged_at
28+
- name: checkout (open pr comments)
29+
if: "!github.event.issue.pull_request.merged_at"
2930
uses: actions/checkout@f43a0e5ff2bd294095638e18286ca9a3d1956744 # actions/checkout@v3
3031
with:
3132
ref: ${{ format('refs/pull/{0}/merge', github.event.issue.number) }}
@@ -62,5 +63,5 @@ jobs:
6263
COMMENT_URL: ${{ github.event.comment.html_url }}
6364
JOB_URL: https://github.com/intel-analytics/BigDL/actions/runs/${{ github.run_id }}
6465
ANALYTICS_ZOO_ROOT: ${{ github.workspace }}
65-
IS_PR: ${{ github.event.pull_request != null || github.event.issue.pull_request != null }}
66+
IS_PR: ${{ github.event.issue.pull_request != null }}
6667
IS_COMMENTS: ${{ github.event.issue.pull_request != null }}

‎python/friesian/src/bigdl/friesian/utils/safepickle.py‎

Lines changed: 53 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
# limitations under the License.
1515
#
1616

17+
import os
1718
import pickle
1819
import hmac
1920
import hashlib
@@ -24,31 +25,61 @@
2425

2526

2627
class SafePickle:
27-
key = b'shared-key'
28-
"""
29-
Example:
30-
>>> from bigdl.friesian.utils import SafePickle
31-
>>> with open(file_path, 'wb') as file:
32-
>>> signature = SafePickle.dump(data, file, return_digest=True)
33-
>>> with open(file_path, 'rb') as file:
34-
>>> data = SafePickle.load(file, signature)
35-
"""
28+
_key = None
29+
30+
@classmethod
31+
def _get_key(cls):
32+
if cls._key is None:
33+
env_key = os.environ.get('BIGDL_SAFE_PICKLE_KEY')
34+
if env_key:
35+
cls._key = bytes.fromhex(env_key)
36+
else:
37+
cls._key = os.urandom(32)
38+
os.environ['BIGDL_SAFE_PICKLE_KEY'] = cls._key.hex()
39+
return cls._key
40+
3641
@classmethod
37-
def dump(self, obj, file, return_digest=False, *args, **kwargs):
42+
def dump(cls, obj, file, return_digest=False, *args, **kwargs):
43+
"""
44+
Example:
45+
>>> from bigdl.friesian.utils import SafePickle
46+
>>> with open(file_path, 'wb') as file:
47+
>>> SafePickle.dump(data, file)
48+
"""
49+
pickled_data = pickle.dumps(obj)
50+
file.write(pickled_data)
51+
digest = hmac.new(cls._get_key(), pickled_data, hashlib.sha256).hexdigest()
3852
if return_digest:
39-
pickled_data = pickle.dumps(obj)
40-
file.write(pickled_data)
41-
digest = hmac.new(self.key, pickled_data, hashlib.sha1).hexdigest()
4253
return digest
43-
else:
44-
pickle.dump(obj, file, *args, **kwargs)
54+
sig_path = file.name + '.sig'
55+
with open(sig_path, 'w') as sig_file:
56+
sig_file.write(digest)
4557

4658
@classmethod
47-
def load(self, file, digest=None, *args, **kwargs):
59+
def load(cls, file, digest=None, *args, **kwargs):
60+
"""
61+
Example:
62+
>>> from bigdl.friesian.utils import SafePickle
63+
>>> with open(file_path, 'rb') as file:
64+
>>> data = SafePickle.load(file)
65+
"""
66+
content = file.read()
67+
68+
if digest is None:
69+
sig_path = file.name + '.sig'
70+
if os.path.exists(sig_path):
71+
with open(sig_path, 'r') as sig_file:
72+
digest = sig_file.read().strip()
73+
4874
if digest:
49-
content = file.read()
50-
new_digest = hmac.new(self.key, content, hashlib.sha1).hexdigest()
51-
if digest != new_digest:
52-
invalidInputError(False, 'Pickle safe check failed')
53-
file.seek(0)
54-
return pickle.load(file, *args, **kwargs)
75+
new_digest = hmac.new(cls._get_key(), content, hashlib.sha256).hexdigest()
76+
if not hmac.compare_digest(digest, new_digest):
77+
invalidInputError(False,
78+
'Pickle integrity check failed: '
79+
'file may have been tampered with')
80+
else:
81+
invalidInputError(False,
82+
'No HMAC signature found for pickle file: '
83+
'cannot verify integrity')
84+
85+
return pickle.loads(content, *args, **kwargs)

‎python/nano/src/bigdl/nano/deps/horovod/horovod_worker.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,18 @@
1515
#
1616

1717
import os
18-
import cloudpickle
1918
import sys
19+
from bigdl.nano.utils.common import SafePickle
2020

2121

2222
if __name__ == '__main__':
2323
temp_dir = sys.argv[1]
2424

2525
with open(os.path.join(temp_dir, "args.pkl"), 'rb') as f:
26-
args = cloudpickle.load(f)
26+
args = SafePickle.load(f)
2727

2828
with open(os.path.join(temp_dir, "target.pkl"), 'rb') as f:
29-
target = cloudpickle.load(f)
29+
target = SafePickle.load(f)
3030

3131
import horovod.tensorflow.keras as hvd
3232
hvd.init()
@@ -36,4 +36,4 @@
3636

3737
with open(os.path.join(temp_dir,
3838
f"history_{idx}"), "wb") as f:
39-
cloudpickle.dump(history, f)
39+
SafePickle.dump(history, f)

‎python/nano/src/bigdl/nano/pytorch/strategies/ddp_subprocess.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
# See the License for the specific language governing permissions and
3232
# limitations under the License.
3333

34-
import cloudpickle
3534
import os
3635
import multiprocessing
3736
import subprocess
@@ -46,6 +45,7 @@
4645
from bigdl.nano.pytorch.strategies.ddp_spawn import DDPSpawnStrategy, _DDPSpawnLauncher
4746
from bigdl.nano.utils.common import schedule_processors
4847
from bigdl.nano.utils.common import invalidInputError
48+
from bigdl.nano.utils.common import SafePickle
4949
from bigdl.nano.pytorch.dispatcher import _get_patch_status
5050

5151
import logging
@@ -100,9 +100,9 @@ def launch(self, function: Callable, *args: Any,
100100
# `self._wrapping_function`.
101101
with open(os.path.join(temp_dir, "args.pkl"), "wb") as f:
102102
if trainer is not None:
103-
cloudpickle.dump((None, args, error_queue), f)
103+
SafePickle.dump((None, args, error_queue), f)
104104
else:
105-
cloudpickle.dump((self._wrapping_function, args, error_queue), f)
105+
SafePickle.dump((self._wrapping_function, args, error_queue), f)
106106

107107
# we also need to pass sys.path to subprocess
108108
with open(os.path.join(temp_dir, "sys_path.json"), "w") as f:

‎python/nano/src/bigdl/nano/pytorch/strategies/worker.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,10 @@
1818
import sys
1919
import json
2020

21-
import cloudpickle
2221
import multiprocessing
2322
from torch.multiprocessing.spawn import _wrap
2423
from bigdl.nano.pytorch.dispatcher import patch_torch
24+
from bigdl.nano.utils.common import SafePickle
2525

2626

2727
if __name__ == '__main__':
@@ -44,7 +44,7 @@
4444
patch_torch(cuda_to_cpu=patch_status['patch_cuda'])
4545

4646
with open(os.path.join(temp_dir, "args.pkl"), "rb") as f:
47-
(fn, args, error_queue) = cloudpickle.load(f)
47+
(fn, args, error_queue) = SafePickle.load(f)
4848

4949
# args[0] is `trainer`, when it is None, it means when are using LightningLite,
5050
# otherwise we are using trainer, for the details here, see `ddp_subprocess.py`

‎python/nano/src/bigdl/nano/utils/common/safepickle.py‎

Lines changed: 53 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
# limitations under the License.
1515
#
1616

17+
import os
1718
import pickle
1819
import hmac
1920
import hashlib
@@ -24,31 +25,61 @@
2425

2526

2627
class SafePickle:
27-
key = b'shared-key'
28-
"""
29-
Example:
30-
>>> from bigdl.nano.utils.common import SafePickle
31-
>>> with open(file_path, 'wb') as file:
32-
>>> signature = SafePickle.dump(data, file, return_digest=True)
33-
>>> with open(file_path, 'rb') as file:
34-
>>> data = SafePickle.load(file, signature)
35-
"""
28+
_key = None
29+
30+
@classmethod
31+
def _get_key(cls):
32+
if cls._key is None:
33+
env_key = os.environ.get('BIGDL_SAFE_PICKLE_KEY')
34+
if env_key:
35+
cls._key = bytes.fromhex(env_key)
36+
else:
37+
cls._key = os.urandom(32)
38+
os.environ['BIGDL_SAFE_PICKLE_KEY'] = cls._key.hex()
39+
return cls._key
40+
3641
@classmethod
37-
def dump(self, obj, file, return_digest=False, *args, **kwargs):
42+
def dump(cls, obj, file, return_digest=False, *args, **kwargs):
43+
"""
44+
Example:
45+
>>> from bigdl.nano.utils.common import SafePickle
46+
>>> with open(file_path, 'wb') as file:
47+
>>> SafePickle.dump(data, file)
48+
"""
49+
pickled_data = pickle.dumps(obj)
50+
file.write(pickled_data)
51+
digest = hmac.new(cls._get_key(), pickled_data, hashlib.sha256).hexdigest()
3852
if return_digest:
39-
pickled_data = pickle.dumps(obj)
40-
file.write(pickled_data)
41-
digest = hmac.new(self.key, pickled_data, hashlib.sha1).hexdigest()
4253
return digest
43-
else:
44-
pickle.dump(obj, file, *args, **kwargs)
54+
sig_path = file.name + '.sig'
55+
with open(sig_path, 'w') as sig_file:
56+
sig_file.write(digest)
4557

4658
@classmethod
47-
def load(self, file, digest=None, *args, **kwargs):
59+
def load(cls, file, digest=None, *args, **kwargs):
60+
"""
61+
Example:
62+
>>> from bigdl.nano.utils.common import SafePickle
63+
>>> with open(file_path, 'rb') as file:
64+
>>> data = SafePickle.load(file)
65+
"""
66+
content = file.read()
67+
68+
if digest is None:
69+
sig_path = file.name + '.sig'
70+
if os.path.exists(sig_path):
71+
with open(sig_path, 'r') as sig_file:
72+
digest = sig_file.read().strip()
73+
4874
if digest:
49-
content = file.read()
50-
new_digest = hmac.new(self.key, content, hashlib.sha1).hexdigest()
51-
if digest != new_digest:
52-
invalidInputError(False, 'Pickle safe check failed')
53-
file.seek(0)
54-
return pickle.load(file, *args, **kwargs)
75+
new_digest = hmac.new(cls._get_key(), content, hashlib.sha256).hexdigest()
76+
if not hmac.compare_digest(digest, new_digest):
77+
invalidInputError(False,
78+
'Pickle integrity check failed: '
79+
'file may have been tampered with')
80+
else:
81+
invalidInputError(False,
82+
'No HMAC signature found for pickle file: '
83+
'cannot verify integrity')
84+
85+
return pickle.loads(content, *args, **kwargs)

‎python/nano/src/bigdl/nano/utils/tf/backend.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,15 +43,15 @@ def run(self, target, args=..., nprocs=1, envs=None) -> Any:
4343
return self.run_subprocess(target, args=args, nprocs=nprocs, envs=envs)
4444

4545
def run_subprocess(self, target, args=..., nprocs=1, envs=None) -> Any:
46-
import cloudpickle
46+
from bigdl.nano.utils.common import SafePickle
4747
import subprocess
4848
import sys
4949

5050
with TemporaryDirectory() as temp_dir:
5151
with open(os.path.join(temp_dir, "args.pkl"), 'wb') as f:
52-
cloudpickle.dump(args, f)
52+
SafePickle.dump(args, f)
5353
with open(os.path.join(temp_dir, "target.pkl"), 'wb') as f:
54-
cloudpickle.dump(target, f)
54+
SafePickle.dump(target, f)
5555

5656
ex_list = []
5757
cwd_path = os.path.dirname(__file__)
@@ -72,5 +72,5 @@ def run_subprocess(self, target, args=..., nprocs=1, envs=None) -> Any:
7272
results = []
7373
for i in range(nprocs):
7474
with open(os.path.join(temp_dir, f"history_{i}"), "rb") as f:
75-
results.append(cloudpickle.load(f))
75+
results.append(SafePickle.load(f))
7676
return results

‎python/nano/src/bigdl/nano/utils/tf/subprocess_worker.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616

1717
import json
1818
import os
19-
import cloudpickle
2019
import sys
2120
import tensorflow as tf
21+
from bigdl.nano.utils.common import SafePickle
2222

2323
if __name__ == '__main__':
2424
# Set number of threads in subprocess
@@ -27,14 +27,14 @@
2727
temp_dir = sys.argv[1]
2828

2929
with open(os.path.join(temp_dir, "args.pkl"), 'rb') as f:
30-
args = cloudpickle.load(f)
30+
args = SafePickle.load(f)
3131

3232
with open(os.path.join(temp_dir, "target.pkl"), 'rb') as f:
33-
target = cloudpickle.load(f)
33+
target = SafePickle.load(f)
3434

3535
history = target(*args)
3636
tf_config = json.loads(os.environ["TF_CONFIG"])
3737

3838
with open(os.path.join(temp_dir,
3939
f"history_{tf_config['task']['index']}"), "wb") as f:
40-
cloudpickle.dump(history, f)
40+
SafePickle.dump(history, f)

0 commit comments

Comments
 (0)