Views
No views yet
orbax.checkpoint.CheckpointManager.restore method, when used with a crafted step argument, is susceptible to a path traversal vulnerability. This allows an attacker to manipulate the path resolution logic, causing Orbax to attempt to load checkpoint data from arbitrary locations on the filesystem outside the intended base directory.step argument containing path traversal sequences (e.g., ../../), an attacker can force the CheckpointManager to read sensitive files (e.g., /etc/passwd, private keys) from the system. If the content of these files can be interpreted as a valid checkpoint item (e.g., JSON for JsonCheckpointer or a PyTree for PyTreeCheckpointer), the data can be exfiltrated or used to further compromise the system. In scenarios where a malicious checkpoint is loaded, this could also lead to Arbitrary Code Execution (ACE) if the loaded data contains executable code that is subsequently processed by the application.orbax_disclosure_v2.py):flax, jax, orbax-checkpoint, msgpack, and numpy are installed./tmp/secret_item/data with content like SECRET_DATA_CONTENT.
1import os
2secret_dir = "/tmp/secret_item"
3if not os.path.exists(secret_dir):
4 os.makedirs(secret_dir)
5with open(os.path.join(secret_dir, "data"), "w") as f:
6 f.write("PWNED_DATA")orbax.checkpoint.CheckpointManager with a PyTreeCheckpointer and a base directory (e.g., /tmp/orbax_base).
1import orbax.checkpoint
2base_dir = "/tmp/orbax_base"
3if not os.path.exists(base_dir):
4 os.makedirs(base_dir)
5mngr = orbax.checkpoint.CheckpointManager(
6 base_dir,
7 item_names=(\'data\',),
8 item_handlers={\'data\': orbax.checkpoint.PyTreeCheckpointer()}
9)restore method with a crafted step argument that includes path traversal sequences to point to the sensitive file.
1traversal_path = "../../tmp/secret_item"
2res = mngr.restore(traversal_path)
3print(f"[+] Successfully read data from outside base_dir: {res}")res variable will contain the content of /tmp/secret_item/data, demonstrating successful information disclosure.CheckpointManager.restore. The step argument should be strictly validated to be a simple integer or a sanitized string that does not contain directory traversal characters.flax.serialization.from_state_dict function, responsible for reconstructing Python objects from a serialized state dictionary, exhibits a type confusion vulnerability. When a leaf node in the target object (the template for reconstruction) is of a type not explicitly registered in Flax's _STATE_DICT_REGISTRY, the function directly returns the corresponding value from the untrusted state_dict without type validation or conversion. This bypasses expected type safety mechanisms.state_dict to inject arbitrary Python objects (e.g., strings, integers, or even custom malicious classes if they can be instantiated) into a PyTree structure where the application expects a specific, registered type (like a JAX array or a Flax module). While not directly leading to ACE, this type confusion can be chained with other vulnerabilities or application logic flaws that make assumptions about the type of data being processed. For instance, if a downstream function expects a numerical array but receives a string, it could lead to unexpected behavior, crashes, or further exploitation.exploit_partial_v2.py):flax, jax, orbax-checkpoint, msgpack, and numpy are installed.class Secret: passPyTree (e.g., a dictionary) where one of the leaves is an instance of the unregistered type.
target = {'key': Secret()}state_dict where the value corresponding to the unregistered type's key is an arbitrary, attacker-controlled value (e.g., a string).
payload = {'key': 'INJECTED_MALICIOUS_STRING'}flax.serialization.from_state_dict with the target template and the malicious payload.
1restored = flax.serialization.from_state_dict(target, payload)
2print(f"[+] Restored with injection: {restored}")restored['key'] will now contain 'INJECTED_MALICIOUS_STRING' instead of an instance of Secret, demonstrating successful type injection.flax.serialization.from_state_dict for all leaf nodes, regardless of whether their type is explicitly registered. If an unregistered type is encountered, it should either raise an error or be handled with a default safe deserialization mechanism, preventing arbitrary type injection.