T09 · Insecure Skill Coding Practices
Warning
- Location
- scripts/lowrank_compress.py:99
- Finding
- Unsafe NumPy Archive Deserialization with Pickle Enabled<![CDATA[ ## Vulnerability Details **File Location**: `scripts/lowrank_compress.py`, lines 99-104 **Vulnerability Type**: Unsafe deserialization **Risk Level**: Medium ### Vulnerable Code ```python def load_compressed(self, filepath: str): """从文件加载压缩数据""" data = np.load(filepath, allow_pickle=True) self.U = data["U"] self.S = data["S"] self.Vt = data["Vt"] self.original_shape = tuple(data["original_shape"]) print(f"已加载压缩数据，形状: U={self.U.shape}, S={self.S.shape}, Vt={self.Vt.shape}") ``` ### Technical Analysis The method loads a caller-supplied NumPy archive using `allow_pickle=True`. This setting permits object arrays in `.npy` or `.npz` files to be reconstructed through Python's pickle mechanism. Pickle is not a safe data format for untrusted input. A crafted object can define a reduction operation that invokes arbitrary Python functions during deserialization. In an `.npz` archive, loading is generally deferred until an archive member such as `data["U"]` is accessed. Therefore, the subsequent array accesses can trigger the malicious pickle payload. The archive format used by `save_compressed()` contains only numeric arrays and does not require pickle support. Enabling it unnecessarily expands the trust boundary from numeric data parsing to arbitrary Python object reconstruction. ### Attack Path 1. An attacker creates a malicious `.npz` archive containing an object array with a crafted pickle reducer. 2. The attacker convinces a user or integrating application to treat that archive as compressed KV-cache data, or replaces an existing archive in a writable location. 3. The application calls `KVCacheLowRank.load_compressed()` with the attacker's file path. 4. `np.load()` accepts the archive because `allow_pickle=True` is enabled. 5. Accessing an object-backed member such as `data["U"]` causes the embedded pickle object to be deserialized. 6. The attacker's reducer executes arbitrary Python code under the identity of the proce ...[truncated 694 chars]
- Remediation
- <![CDATA[ ## Remediation Suggestions 1. Disable pickle support because the expected archive members are numeric arrays: ```python with np.load(filepath, allow_pickle=False) as data: required = {"U", "S", "Vt", "original_shape"} if not required.issubset(data.files): raise ValueError("Compressed archive is missing required fields") U = data["U"] S = data["S"] Vt = data["Vt"] original_shape = data["original_shape"] ``` 2. Reject object dtypes explicitly: ```python for name, array in { "U": U, "S": S, "Vt": Vt, "original_shape": original_shape, }.items(): if array.dtype.hasobject: raise ValueError(f"Object dtype is not permitted for {name}") ``` 3. Validate that arrays have the expected dimensions, compatible shapes, finite values, and reasonable element counts before retaining them. 4. Enforce a maximum input file size and maximum decompressed array dimensions to reduce memory-exhaustion risk. 5. If archives cross a trust boundary, distribute a cryptographic hash or signature and verify it before loading. 6. Use a context manager for `np.load()` so the underlying archive is closed reliably. ]]>