From b8416703e48ee53290d04876fff5e68ff7999101 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Fri, 21 Mar 2025 09:38:07 +0800 Subject: [PATCH] add more stringent oob checks --- engine/config/gguf_parser.cc | 44 +++++++++++++++---- engine/config/gguf_parser.h | 1 + .../api/model/test_api_model_import.py | 22 ++++++++-- 3 files changed, 55 insertions(+), 12 deletions(-) diff --git a/engine/config/gguf_parser.cc b/engine/config/gguf_parser.cc index e07ddecc0..81424ba1f 100644 --- a/engine/config/gguf_parser.cc +++ b/engine/config/gguf_parser.cc @@ -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) { @@ -101,14 +106,11 @@ void GGUFHandler::CloseFile() { std::pair 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(data_ + offset + 8), length); return {8 + static_cast(length), value}; } @@ -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(data_[offset]); return 1; case 2: // UINT16 + CheckOffset(offset + 2); metadata_uint16_[key] = *reinterpret_cast(data_ + offset); return 2; case 3: // INT16 + CheckOffset(offset + 2); metadata_int16_[key] = *reinterpret_cast(data_ + offset); return 2; case 4: // UINT32 + CheckOffset(offset + 4); metadata_uint32_[key] = *reinterpret_cast(data_ + offset); return 4; case 5: // INT32 + CheckOffset(offset + 4); metadata_int32_[key] = *reinterpret_cast(data_ + offset); return 4; case 6: // FLOAT32 + CheckOffset(offset + 4); metadata_float_[key] = *reinterpret_cast(data_ + offset); return 4; case 7: // BOOL + CheckOffset(offset + 1); metadata_bool_[key] = data_[offset] != 0; return 1; case 8: // STRING @@ -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(data_ + offset); return 8; case 11: // INT64 + CheckOffset(offset + 8); metadata_int64_[key] = *reinterpret_cast(data_ + offset); return 8; case 12: // FLOAT64 + CheckOffset(offset + 8); metadata_double_[key] = *reinterpret_cast(data_ + offset); return 8; default: @@ -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(data_ + offset); // std::memcpy(&array_type, data_ + offset, sizeof(uint32_t)); + CheckOffset(offset + 4 + 8); uint64_t array_length = *reinterpret_cast(data_ + offset + 4); // std::memcpy(&array_length, data_ + offset + 4, sizeof(uint64_t)); @@ -199,11 +214,13 @@ 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(uint8_value)); break; case 1: { + CheckOffset(offset + array_offset + 1); int8_value = static_cast(data_[offset + array_offset]); length = 1; array_values_float.push_back(static_cast(int8_value)); @@ -211,41 +228,48 @@ size_t GGUFHandler::ReadArray(std::size_t offset, const std::string& key) { break; case 2: + CheckOffset(offset + array_offset + 2); uint16_value = *reinterpret_cast(data_ + offset + array_offset); length = 2; array_values_float.push_back(static_cast(uint16_value)); break; case 3: + CheckOffset(offset + array_offset + 2); int16_value = *reinterpret_cast(data_ + offset + array_offset); length = 2; array_values_float.push_back(static_cast(int16_value)); break; case 4: + CheckOffset(offset + array_offset + 4); uint32_value = *reinterpret_cast(data_ + offset + array_offset); length = 4; array_values_float.push_back(static_cast(uint32_value)); break; case 5: + CheckOffset(offset + array_offset + 4); int32_value = *reinterpret_cast(data_ + offset + array_offset); length = 4; array_values_float.push_back(static_cast(int32_value)); break; case 6: + CheckOffset(offset + array_offset + 4); float_value = *reinterpret_cast(data_ + offset + array_offset); length = 4; array_values_float.push_back(static_cast(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(bool_value)); break; case 8: { + CheckOffset(offset + array_offset + 8); uint64_t length_ = *reinterpret_cast(data_ + offset + array_offset); std::string value( @@ -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(data_ + offset + array_offset); length = 8; array_values_float.push_back(static_cast(uint64_value)); break; case 11: + CheckOffset(offset + array_offset + 8); int64_value = *reinterpret_cast(data_ + offset + array_offset); length = 8; array_values_float.push_back(static_cast(int64_value)); break; case 12: + CheckOffset(offset + array_offset + 8); double_value = *reinterpret_cast(data_ + offset + array_offset); length = 8; @@ -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; @@ -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(data_) << "\n"; if (*reinterpret_cast(data_) != GGUF_MAGIC_NUMBER) { @@ -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(data_ + offset); offset += 4; LOG_INFO << "value type number: " << value_type << "\n"; diff --git a/engine/config/gguf_parser.h b/engine/config/gguf_parser.h index d9997c797..f7e28f4b5 100644 --- a/engine/config/gguf_parser.h +++ b/engine/config/gguf_parser.h @@ -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_; diff --git a/engine/e2e-test/api/model/test_api_model_import.py b/engine/e2e-test/api/model/test_api_model_import.py index 34746dbe9..4bf02605d 100644 --- a/engine/e2e-test/api/model/test_api_model_import.py +++ b/engine/e2e-test/api/model/test_api_model_import.py @@ -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) @@ -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.") @@ -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 \ No newline at end of file + 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"])