Skip to content

fix: reject invalid n_dims in tensor header to prevent stack-buffer-overflow - #220

Open
shafiuzzaman-md wants to merge 1 commit into
PABannier:mainfrom
shafiuzzaman-md:patch-1
Open

fix: reject invalid n_dims in tensor header to prevent stack-buffer-overflow#220
shafiuzzaman-md wants to merge 1 commit into
PABannier:mainfrom
shafiuzzaman-md:patch-1

Conversation

@shafiuzzaman-md

Copy link
Copy Markdown

Problem

bark_model_load() (bark.cpp:1019) reads a tensor header's n_dims field as a raw int32_t directly from the model file and uses it, unbounded, as the loop count writing into a fixed-size 2-element stack array int32_t ne[2]:

int32_t nelements = 1;
int32_t ne[2] = {1, 1};
for (int i = 0; i < n_dims; ++i) {
    read_safe(fin, ne[i]);   // n_dims is fully attacker-controlled, no upper bound
    nelements *= ne[i];
}

A model file with n_dims > 2 writes past the ne[2] array on the stack. Reproduced with a crafted 144-byte model file (n_dims = 20) through the stock main -m <file> -p "text" command line, under an AddressSanitizer build:

==...==ERROR: AddressSanitizer: stack-buffer-overflow
WRITE of size 4 at ... thread T0
    #0 memcpy
    #1-3 std::istream::read / filebuf / streambuf
    #4 read_safe<int>          bark.cpp:168
    #5 bark_model_load         bark.cpp:1019
    #6 bark_load_model_from_file  bark.cpp:1124
    #7 bark_load_model         bark.cpp:1174
    #8 main                    examples/main/main.cpp:55

The identical copy-pasted loop also exists in ggml_quantize_weights (the offline quantize tool, bark.cpp:357, int32_t ne[4]), fixed here too.

Fix

Add an explicit n_dims range check immediately after it is read, before the loop that populates ne[], at both sites sharing this pattern:

if (n_dims < 1 || n_dims > 2) {   // 4 for ggml_quantize_weights
    fprintf(stderr, "%s: invalid n_dims %d in model file (expected 1 <= n_dims <= 2)\n", __func__, n_dims);
    return false;
}

For normal, well-formed model files this is a no-op; a malformed file is now rejected with a clear error message instead of corrupting the stack.

Testing

  • Before: the crafted model file (n_dims=20) crashes with an ASan stack-buffer-overflow WRITE at bark_model_load (bark.cpp:1019).
  • After: the same file is rejected cleanly ("invalid n_dims 20 in model file"), no crash.
  • Regression check: a file with a valid n_dims (2) passes the new check unaffected (fails later only because the synthetic test file has no real tensor data, unrelated to this fix).
  • Both main and quantize targets build cleanly with the change.

@shafiuzzaman-md

Copy link
Copy Markdown
Author

The 3 failing checks (ios, emscripten, ubuntu-latest-run) are unrelated to this change. Each fails before reaching the modified code:

  • ios: fatal error: 'ggml-alloc.h' file not found in the Xcode SwiftUI example project's include paths.
  • emscripten: error: "embind requires -std=c++17 or newer", the emscripten build isn't passing -std=c++17.
  • ubuntu-latest-run: fails at the apt-get dependencies step (Unable to locate package python3.10-venv), before build/run even starts.

The two checks that actually compile and run the changed code (windows-msys2 CLANG64 and UCRT64) both pass. Happy to help fix the CI separately if useful, but wanted to flag that these failures predate and are unrelated to this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant