Enforce :safe in Nx.deserialize and EXLA disk cache deserialization

[blasphemetheus]

been experimenting with sandboxing an agent and telling it it's in a ctf find vulnerabilities, make a report. It found one on nx that seemed worth upstreaming

---

Summary

Nx.deserialize/2 passes opts directly to :erlang.binary_to_term/2, defaulting to []. This means callers who write Nx.deserialize(data) — the natural API — get no atom safety. A crafted .nx file can exhaust the atom table and crash the VM.

Similarly, EXLA.Defn.Disk calls binary_to_term(blob) with no flags when loading cached executables.

This PR adds automatic :safe enforcement in both paths. The flag only restricts creation of new atoms — all atoms used by legitimate tensors (:f32, :s64, axis names, etc.) already exist in any running Nx application, so this is a no-op for valid data.

Changes

• nx/lib/nx.ex: deserialize/2 now ensures :safe is always in the opts list passed to binary_to_term/2
• exla/lib/exla/defn/disk.ex: Disk cache loading now uses binary_to_term(blob, [:safe])
• nx/test/nx_test.exs: Updated test that calls binary_to_term directly to also pass [:safe]

I have a runnable PoC demonstrating atom table exhaustion via a crafted payload — happy to share privately if useful for review.

Test plan

• Existing Nx.deserialize round-trip tests pass (:safe is a no-op for valid data since all relevant atoms already exist)
• EXLA disk cache tests unaffected
• Maintainers: optionally verify that Nx.deserialize(malicious_payload) now raises ArgumentError instead of crashing the VM

🤖 Generated with Claude Code

[josevalim]

Thank you but I don't think this moves the needle. You have to trust the source of the data, otherwise you will have bigger trouble. For example, EXLA is serializing a whole executable, that can do anything at all, preventing the atom ingestion is the last of your concerns.

Furthermore, safe introduces the complication that any atom that you load must have been seen before hand, which is hard to guarantee when deserialisation.

Perhaps we should add disclaimers to the docs that say you must trust the source of the data, if we don't have those yet, but overall, :safe is the least of our concerns here.

[blasphemetheus]

hmm as for a warning. could add on deserialize/2 something along the lines of:

    @doc """                                                                               
    Deserializes a serialized representation of a tensor or a container                    
    with the given options.                                                                
                  
    It is the opposite of `Nx.serialize/2`.                                                
   
  + > ### Warning {: .warning}                                                             
  + >             
  + > This function uses `:erlang.binary_to_term/2` under the hood.                        
  + > Only deserialize data from trusted sources. A malicious payload                      
  + > could crash the VM or cause unexpected behavior.                                     
                                                                                           
    Note: This function cannot be used in `defn`. 

and serialize/1 something like:

    @doc """                                                                               
    Serializes the given tensor or container of tensors to iodata.

  + > ### Warning {: .warning}
  + >
  + > The data produced by this function is intended to be deserialized
  + > by `Nx.deserialize/2`. Only deserialize data from trusted sources.                   
   
    You may pass any tensor or `Nx.Container` to serialization.

I'll close this one tho

[josevalim]

:+:1 for the warnings! ❤️