Skip to content

Commit b9e02e8

Browse files
authored
ggml : fix memory leaks when loading invalid gguf files (#10094)
* ggml : fix gguf string leak when reading kv pairs fails * ggml : avoid crashing with GGML_ABORT when the KV has an invalid type * ggml : avoid crashing on failed memory allocations when loading a gguf file
1 parent 6763f71 commit b9e02e8

File tree

1 file changed

+54
-13
lines changed

1 file changed

+54
-13
lines changed

ggml/src/ggml.c

Lines changed: 54 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22136,7 +22136,11 @@ static bool gguf_fread_str(FILE * file, struct gguf_str * p, size_t * offset) {
2213622136
return false;
2213722137
}
2213822138

22139-
p->data = GGML_CALLOC(p->n + 1, 1);
22139+
p->data = calloc(p->n + 1, 1);
22140+
if (!p->data) {
22141+
fprintf(stderr, "%s: failed to allocate memory for string of length %" PRIu64 "\n", __func__, p->n);
22142+
return false;
22143+
}
2214022144

2214122145
ok = ok && gguf_fread_el(file, p->data, p->n, offset);
2214222146

@@ -22170,7 +22174,11 @@ static void gguf_free_kv(struct gguf_kv * kv) {
2217022174
}
2217122175

2217222176
struct gguf_context * gguf_init_empty(void) {
22173-
struct gguf_context * ctx = GGML_CALLOC(1, sizeof(struct gguf_context));
22177+
struct gguf_context * ctx = calloc(1, sizeof(struct gguf_context));
22178+
if (!ctx) {
22179+
fprintf(stderr, "%s: failed to allocate memory for context\n", __func__);
22180+
return NULL;
22181+
}
2217422182

2217522183
memcpy(ctx->header.magic, GGUF_MAGIC, sizeof(ctx->header.magic));
2217622184
ctx->header.version = GGUF_VERSION;
@@ -22216,7 +22224,12 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
2221622224

2221722225
bool ok = true;
2221822226

22219-
struct gguf_context * ctx = GGML_CALLOC(1, sizeof(struct gguf_context));
22227+
struct gguf_context * ctx = calloc(1, sizeof(struct gguf_context));
22228+
if (!ctx) {
22229+
fprintf(stderr, "%s: failed to allocate memory for context\n", __func__);
22230+
fclose(file);
22231+
return NULL;
22232+
}
2222022233

2222122234
// read the header
2222222235
{
@@ -22255,9 +22268,13 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
2225522268
{
2225622269
const uint64_t n_kv = ctx->header.n_kv;
2225722270

22258-
// header.n_kv will hold the actual value of pairs that were successfully read in the loop below
22259-
ctx->header.n_kv = 0;
22260-
ctx->kv = GGML_CALLOC(n_kv, sizeof(struct gguf_kv));
22271+
ctx->kv = calloc(n_kv, sizeof(struct gguf_kv));
22272+
if (!ctx->kv) {
22273+
fprintf(stderr, "%s: failed to allocate memory for kv pairs\n", __func__);
22274+
fclose(file);
22275+
gguf_free(ctx);
22276+
return NULL;
22277+
}
2226122278

2226222279
for (uint64_t i = 0; i < n_kv; ++i) {
2226322280
struct gguf_kv * kv = &ctx->kv[i];
@@ -22308,7 +22325,13 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
2230822325
return NULL;
2230922326
}
2231022327

22311-
kv->value.arr.data = GGML_CALLOC(kv->value.arr.n, gguf_type_size(kv->value.arr.type));
22328+
kv->value.arr.data = calloc(kv->value.arr.n, gguf_type_size(kv->value.arr.type));
22329+
if (!kv->value.arr.data) {
22330+
fprintf(stderr, "%s: failed to allocate memory for array\n", __func__);
22331+
fclose(file);
22332+
gguf_free(ctx);
22333+
return NULL;
22334+
}
2231222335

2231322336
ok = ok && gguf_fread_el(file, kv->value.arr.data, kv->value.arr.n * gguf_type_size(kv->value.arr.type), &offset);
2231422337
} break;
@@ -22322,24 +22345,36 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
2232222345
return NULL;
2232322346
}
2232422347

22325-
kv->value.arr.data = GGML_CALLOC(kv->value.arr.n, sizeof(struct gguf_str));
22348+
kv->value.arr.data = calloc(kv->value.arr.n, sizeof(struct gguf_str));
22349+
if (!kv->value.arr.data) {
22350+
fprintf(stderr, "%s: failed to allocate memory for array\n", __func__);
22351+
fclose(file);
22352+
gguf_free(ctx);
22353+
return NULL;
22354+
}
2232622355

2232722356
for (uint64_t j = 0; j < kv->value.arr.n; ++j) {
2232822357
ok = ok && gguf_fread_str(file, &((struct gguf_str *) kv->value.arr.data)[j], &offset);
2232922358
}
2233022359
} break;
2233122360
case GGUF_TYPE_ARRAY:
22332-
default: GGML_ABORT("invalid type");
22361+
default:
22362+
{
22363+
fprintf(stderr, "%s: invalid array type %d\n", __func__, kv->value.arr.type);
22364+
ok = false;
22365+
} break;
2233322366
}
2233422367
} break;
22335-
default: GGML_ABORT("invalid type");
22368+
default:
22369+
{
22370+
fprintf(stderr, "%s: invalid type %d\n", __func__, kv->type);
22371+
ok = false;
22372+
} break;
2233622373
}
2233722374

2233822375
if (!ok) {
2233922376
break;
2234022377
}
22341-
22342-
ctx->header.n_kv++;
2234322378
}
2234422379

2234522380
if (!ok) {
@@ -22352,7 +22387,13 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
2235222387

2235322388
// read the tensor infos
2235422389
if (ctx->header.n_tensors > 0) {
22355-
ctx->infos = GGML_CALLOC(ctx->header.n_tensors, sizeof(struct gguf_tensor_info));
22390+
ctx->infos = calloc(ctx->header.n_tensors, sizeof(struct gguf_tensor_info));
22391+
if (!ctx->infos) {
22392+
fprintf(stderr, "%s: failed to allocate memory for tensor infos\n", __func__);
22393+
fclose(file);
22394+
gguf_free(ctx);
22395+
return NULL;
22396+
}
2235622397

2235722398
for (uint64_t i = 0; i < ctx->header.n_tensors; ++i) {
2235822399
struct gguf_tensor_info * info = &ctx->infos[i];

0 commit comments

Comments
 (0)