Skip to content

bugfix: add more stringent out-of-bounds checks for GGUF parser #2159

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 36 additions & 8 deletions engine/config/gguf_parser.cc
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,11 @@ void GGUFHandler::OpenFile(const std::string& file_path) {
#endif
}

void GGUFHandler::CheckOffset(int offset) const {
if (offset > file_size_)
throw std::runtime_error("Unexpected EOF");
}

void GGUFHandler::CloseFile() {
#ifdef _WIN32
if (data_ != nullptr) {
Expand All @@ -101,14 +106,11 @@ void GGUFHandler::CloseFile() {

std::pair<std::size_t, std::string> GGUFHandler::ReadString(
std::size_t offset) const {
CheckOffset(offset + 8);
uint64_t length;
std::memcpy(&length, data_ + offset, sizeof(uint64_t));

if (offset + 8 + length > file_size_) {
throw std::runtime_error(
"GGUF metadata string length exceeds file size.\n");
}

CheckOffset(offset + 8 + length);
std::string value(reinterpret_cast<const char*>(data_ + offset + 8), length);
return {8 + static_cast<std::size_t>(length), value};
}
Expand All @@ -117,29 +119,37 @@ size_t GGUFHandler::ReadMetadataValue(int type, std::size_t offset,
const std::string& key) {
switch (type) {
case 0: // UINT8
CheckOffset(offset + 1);
metadata_uint8_[key] = data_[offset];
return 1;
case 1: // INT8
CheckOffset(offset + 1);
metadata_int8_[key] = static_cast<int8_t>(data_[offset]);
return 1;
case 2: // UINT16
CheckOffset(offset + 2);
metadata_uint16_[key] =
*reinterpret_cast<const uint16_t*>(data_ + offset);
return 2;
case 3: // INT16
CheckOffset(offset + 2);
metadata_int16_[key] = *reinterpret_cast<const int16_t*>(data_ + offset);
return 2;
case 4: // UINT32
CheckOffset(offset + 4);
metadata_uint32_[key] =
*reinterpret_cast<const uint32_t*>(data_ + offset);
return 4;
case 5: // INT32
CheckOffset(offset + 4);
metadata_int32_[key] = *reinterpret_cast<const int32_t*>(data_ + offset);
return 4;
case 6: // FLOAT32
CheckOffset(offset + 4);
metadata_float_[key] = *reinterpret_cast<const float*>(data_ + offset);
return 4;
case 7: // BOOL
CheckOffset(offset + 1);
metadata_bool_[key] = data_[offset] != 0;
return 1;
case 8: // STRING
Expand All @@ -152,13 +162,16 @@ size_t GGUFHandler::ReadMetadataValue(int type, std::size_t offset,

return ReadArray(offset, key);
case 10: // UINT64
CheckOffset(offset + 8);
metadata_uint64_[key] =
*reinterpret_cast<const uint64_t*>(data_ + offset);
return 8;
case 11: // INT64
CheckOffset(offset + 8);
metadata_int64_[key] = *reinterpret_cast<const int64_t*>(data_ + offset);
return 8;
case 12: // FLOAT64
CheckOffset(offset + 8);
metadata_double_[key] = *reinterpret_cast<const double*>(data_ + offset);
return 8;
default:
Expand All @@ -168,9 +181,11 @@ size_t GGUFHandler::ReadMetadataValue(int type, std::size_t offset,
}

size_t GGUFHandler::ReadArray(std::size_t offset, const std::string& key) {
CheckOffset(offset + 4);
uint32_t array_type = *reinterpret_cast<const uint32_t*>(data_ + offset);
// std::memcpy(&array_type, data_ + offset, sizeof(uint32_t));

CheckOffset(offset + 4 + 8);
uint64_t array_length =
*reinterpret_cast<const uint64_t*>(data_ + offset + 4);
// std::memcpy(&array_length, data_ + offset + 4, sizeof(uint64_t));
Expand Down Expand Up @@ -199,53 +214,62 @@ size_t GGUFHandler::ReadArray(std::size_t offset, const std::string& key) {
// assume that array ony has 2 type string and int
switch (array_type) {
case 0:
CheckOffset(offset + array_offset + 1);
uint8_value = data_[offset + array_offset];
length = 1;
array_values_float.push_back(static_cast<float>(uint8_value));
break;
case 1: {
CheckOffset(offset + array_offset + 1);
int8_value = static_cast<int8_t>(data_[offset + array_offset]);
length = 1;
array_values_float.push_back(static_cast<float>(int8_value));
}

break;
case 2:
CheckOffset(offset + array_offset + 2);
uint16_value =
*reinterpret_cast<const uint16_t*>(data_ + offset + array_offset);
length = 2;
array_values_float.push_back(static_cast<float>(uint16_value));
break;
case 3:
CheckOffset(offset + array_offset + 2);
int16_value =
*reinterpret_cast<const int16_t*>(data_ + offset + array_offset);
length = 2;
array_values_float.push_back(static_cast<float>(int16_value));
break;
case 4:
CheckOffset(offset + array_offset + 4);
uint32_value =
*reinterpret_cast<const uint32_t*>(data_ + offset + array_offset);
length = 4;
array_values_float.push_back(static_cast<float>(uint32_value));
break;
case 5:
CheckOffset(offset + array_offset + 4);
int32_value =
*reinterpret_cast<const int32_t*>(data_ + offset + array_offset);
length = 4;
array_values_float.push_back(static_cast<float>(int32_value));
break;
case 6:
CheckOffset(offset + array_offset + 4);
float_value =
*reinterpret_cast<const float*>(data_ + offset + array_offset);
length = 4;
array_values_float.push_back(static_cast<float>(float_value));
break;
case 7:
CheckOffset(offset + array_offset + 1);
bool_value = data_[offset + array_offset] != 0;
length = 1;
array_values_float.push_back(static_cast<float>(bool_value));
break;
case 8: {
CheckOffset(offset + array_offset + 8);
uint64_t length_ =
*reinterpret_cast<const uint64_t*>(data_ + offset + array_offset);
std::string value(
Expand All @@ -255,18 +279,21 @@ size_t GGUFHandler::ReadArray(std::size_t offset, const std::string& key) {
array_values_string.push_back(value);
} break;
case 10:
CheckOffset(offset + array_offset + 8);
uint64_value =
*reinterpret_cast<const uint64_t*>(data_ + offset + array_offset);
length = 8;
array_values_float.push_back(static_cast<float>(uint64_value));
break;
case 11:
CheckOffset(offset + array_offset + 8);
int64_value =
*reinterpret_cast<const int64_t*>(data_ + offset + array_offset);
length = 8;
array_values_float.push_back(static_cast<float>(int64_value));
break;
case 12:
CheckOffset(offset + array_offset + 8);
double_value =
*reinterpret_cast<const double*>(data_ + offset + array_offset);
length = 8;
Expand All @@ -279,9 +306,6 @@ size_t GGUFHandler::ReadArray(std::size_t offset, const std::string& key) {
}

array_offset += length;
if (offset + array_offset > file_size_) {
throw std::runtime_error("GGUF Parser Array exceeded file size.\n");
}
}
if (array_values_string.size() > 0)
metadata_array_string_[key] = array_values_string;
Expand All @@ -290,8 +314,11 @@ size_t GGUFHandler::ReadArray(std::size_t offset, const std::string& key) {
return array_offset;
}

// https://github.com/ggml-org/ggml/blob/master/docs/gguf.md
void GGUFHandler::Parse(const std::string& file_path) {
OpenFile(file_path);
CheckOffset(4 + 4 + 8 + 8);

LOG_INFO << "GGUF magic number: " << *reinterpret_cast<const uint32_t*>(data_)
<< "\n";
if (*reinterpret_cast<const uint32_t*>(data_) != GGUF_MAGIC_NUMBER) {
Expand All @@ -311,6 +338,7 @@ void GGUFHandler::Parse(const std::string& file_path) {
auto [key_byte_length, key] = ReadString(offset);
offset += key_byte_length;
LOG_INFO << "key: " << key << "\n";
CheckOffset(offset + 4);
uint32_t value_type = *reinterpret_cast<const uint32_t*>(data_ + offset);
offset += 4;
LOG_INFO << "value type number: " << value_type << "\n";
Expand Down
1 change: 1 addition & 0 deletions engine/config/gguf_parser.h
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ class GGUFHandler {
size_t ReadArray(std::size_t offset, const std::string& key);
void ModelConfigFromMetadata();
void OpenFile(const std::string& file_path);
void CheckOffset(int offset) const;

uint8_t* data_;
size_t file_size_;
Expand Down
22 changes: 18 additions & 4 deletions engine/e2e-test/api/model/test_api_model_import.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
from pathlib import Path
import json

import pytest
import requests
from utils.test_runner import start_server, stop_server
from utils.test_runner import start_server, stop_server, run

class TestApiModelImport:
@pytest.fixture(autouse=True)
Expand All @@ -18,7 +21,7 @@ def setup_and_teardown(self):
def test_model_import_should_be_success(self):
body_json = {'model': 'tinyllama:1b',
'modelPath': '/path/to/local/gguf'}
response = requests.post("http://localhost:3928/v1/models/import", json=body_json)
response = requests.post("http://localhost:3928/v1/models/import", json=body_json)
assert response.status_code == 200

@pytest.mark.skipif(True, reason="Expensive test. Only test when you have local gguf file.")
Expand Down Expand Up @@ -53,5 +56,16 @@ def test_model_import_with_invalid_path_should_fail(self):
def test_model_import_with_missing_model_should_fail(self):
body_json = {'modelPath': '/path/to/local/gguf'}
response = requests.post("http://localhost:3928/v1/models/import", json=body_json)
print(response)
assert response.status_code == 409
assert response.status_code == 409

def test_model_import_with_invalid_gguf(self, tmp_path: Path):
run("Delete model", ["models", "delete", "model"])
gguf_path = tmp_path / "model.gguf"
with open(gguf_path, "wb") as f:
f.write(b"GGUF") # only GGUF magic number
body_json = {'model': 'model', 'modelPath': str(gguf_path.absolute())}
response = requests.post("http://localhost:3928/v1/models/import", json=body_json)
print(response.content.decode())
assert response.status_code == 400
assert json.loads(response.content)["message"].startswith("Error importing model")
run("Delete model", ["models", "delete", "model"])
Loading