fix(api): replace pickle with safetensors for inference request bodies#62
Open
vlobstein-vc wants to merge 1 commit into
Open
fix(api): replace pickle with safetensors for inference request bodies#62vlobstein-vc wants to merge 1 commit into
vlobstein-vc wants to merge 1 commit into
Conversation
Signed-off-by: Valentin Lobstein <281638514+vlobstein-vc@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
The inference API server (
gui/api/server.py) deserializes raw HTTP request bodies withpickle.loads()on two endpoints,/request-inferenceand/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).Fix
This replaces
pickleon the request path with a small safetensors-based codec (gui/api/safe_serialization.py):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.safetensorsis already a dependency (requirements.txt), so nothing new is added. The matching client serialization ingui/api/client.pyis 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
InferenceRequestandCompressedSeedingRequestpreserve every field (arrays, dtypes, enums, compressed buffers). Apickleos.systemgadget 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.