Skip to content

fix(api): replace pickle with safetensors for inference request bodies#62

Open
vlobstein-vc wants to merge 1 commit into
nv-tlabs:mainfrom
vlobstein-vc:fix/api-pickle-rce
Open

fix(api): replace pickle with safetensors for inference request bodies#62
vlobstein-vc wants to merge 1 commit into
nv-tlabs:mainfrom
vlobstein-vc:fix/api-pickle-rce

Conversation

@vlobstein-vc

Copy link
Copy Markdown

What

The inference API server (gui/api/server.py) deserializes raw HTTP request bodies with pickle.loads() on two endpoints, /request-inference and /seed-model. pickle.loads() runs arbitrary code carried in the stream (a __reduce__ gadget executes during deserialization), and these endpoints have no authentication, so any client that can reach the API port gets code execution as the inference process (CWE-502).

@app.post("/request-inference")
async def request_inference(request: Request):
    req: bytes = await request.body()
    req = pickle.loads(req)   # untrusted body -> arbitrary code

Fix

This replaces pickle on the request path with a small safetensors-based codec (gui/api/safe_serialization.py):

  • numpy arrays (cameras, focal lengths, images, depths, ...) are stored as safetensors tensors;
  • scalar / enum / compressed-buffer fields are carried in a JSON header in the safetensors metadata;
  • loads() only reconstructs an allow-listed set of request dataclasses (InferenceRequest, SeedingRequest, CompressedSeedingRequest), so a body can neither carry an executable gadget nor instantiate an arbitrary class.

safetensors is already a dependency (requirements.txt), so nothing new is added. The matching client serialization in gui/api/client.py is updated in lockstep.

Scope

This covers the request path, the unauthenticated input the server reads off the network. The response path (server to client) still uses pickle; that is a separate trust boundary (the client decoding its own server's reply) and can be migrated the same way as a follow-up.

Testing

Round-trips of InferenceRequest and CompressedSeedingRequest preserve every field (arrays, dtypes, enums, compressed buffers). A pickle os.system gadget posted to the new path is inert: it decodes to nothing and raises instead of executing. The codec test needs no GPU.

Context

Reported to the NVIDIA PSIRT through coordinated disclosure. Opening the fix here so it is available to anyone running the GUI inference server.

Signed-off-by: Valentin Lobstein <281638514+vlobstein-vc@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant