Skip to content

Commit 4e38809

Browse files
authored
Fix memory bug in grammar parser (#7194)
The llama.cpp grammar parser had a bug where forgetting to add a closing quotation mark to strings would cause parsing to crash. Anyone running a server on a public endpoint is advised to upgrade. To reproduce this bug ./llamafile -m foo.gguf -p bar --grammar 'root::="' Credit for discovering and reporting this issue goes to Eclypsium Security Researcher Richard Johnson <[email protected]>.
1 parent f89fe27 commit 4e38809

File tree

4 files changed

+21
-5
lines changed

4 files changed

+21
-5
lines changed

common/common.cpp

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1371,14 +1371,12 @@ bool gpt_params_parse_ex(int argc, char ** argv, gpt_params & params) {
13711371
if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) {
13721372
std::replace(arg.begin(), arg.end(), '_', '-');
13731373
}
1374-
13751374
if (!gpt_params_find_arg(argc, argv, arg, params, i, invalid_param)) {
13761375
throw std::invalid_argument("error: unknown argument: " + arg);
13771376
}
1378-
}
1379-
1380-
if (invalid_param) {
1381-
throw std::invalid_argument("error: invalid parameter for argument: " + arg);
1377+
if (invalid_param) {
1378+
throw std::invalid_argument("error: invalid parameter for argument: " + arg);
1379+
}
13821380
}
13831381

13841382
if (params.prompt_cache_all &&

common/grammar-parser.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,9 @@ namespace grammar_parser {
142142
pos++;
143143
last_sym_start = out_elements.size();
144144
while (*pos != '"') {
145+
if (!*pos) {
146+
throw std::runtime_error("unexpected end of input");
147+
}
145148
auto char_pair = parse_char(pos);
146149
pos = char_pair.second;
147150
out_elements.push_back({LLAMA_GRETYPE_CHAR, char_pair.first});
@@ -156,6 +159,9 @@ namespace grammar_parser {
156159
}
157160
last_sym_start = out_elements.size();
158161
while (*pos != ']') {
162+
if (!*pos) {
163+
throw std::runtime_error("unexpected end of input");
164+
}
159165
auto char_pair = parse_char(pos);
160166
pos = char_pair.second;
161167
enum llama_gretype type = last_sym_start < out_elements.size()
@@ -164,6 +170,9 @@ namespace grammar_parser {
164170

165171
out_elements.push_back({type, char_pair.first});
166172
if (pos[0] == '-' && pos[1] != ']') {
173+
if (!pos[1]) {
174+
throw std::runtime_error("unexpected end of input");
175+
}
167176
auto endchar_pair = parse_char(pos + 1);
168177
pos = endchar_pair.second;
169178
out_elements.push_back({LLAMA_GRETYPE_CHAR_RNG_UPPER, endchar_pair.first});

examples/llava/llava-cli.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,11 @@ static void process_prompt(struct llava_context * ctx_llava, struct llava_image_
189189
LOG_TEE("\n");
190190

191191
struct llama_sampling_context * ctx_sampling = llama_sampling_init(params->sparams);
192+
if (!ctx_sampling) {
193+
fprintf(stderr, "%s: failed to initialize sampling subsystem\n", __func__);
194+
exit(1);
195+
}
196+
192197
std::string response = "";
193198
for (int i = 0; i < max_tgt_len; i++) {
194199
const char * tmp = sample(ctx_sampling, ctx_llava->ctx_llama, &n_past);

examples/main/main.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,6 +523,10 @@ int main(int argc, char ** argv) {
523523
}
524524

525525
struct llama_sampling_context * ctx_sampling = llama_sampling_init(sparams);
526+
if (!ctx_sampling) {
527+
fprintf(stderr, "%s: failed to initialize sampling subsystem\n", __func__);
528+
exit(1);
529+
}
526530

527531
while ((n_remain != 0 && !is_antiprompt) || params.interactive) {
528532
// predict

0 commit comments

Comments
 (0)