Diego Devesa commited on
Commit
f9baffc
·
1 Parent(s): 5a73b5f

ggml : fix memory leaks when loading invalid gguf files (llama/10094)

Browse files

* 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

Files changed (1) hide show
  1. ggml/src/ggml.c +54 -13
ggml/src/ggml.c CHANGED
@@ -22107,7 +22107,11 @@ static bool gguf_fread_str(FILE * file, struct gguf_str * p, size_t * offset) {
22107
  return false;
22108
  }
22109
 
22110
- p->data = GGML_CALLOC(p->n + 1, 1);
 
 
 
 
22111
 
22112
  ok = ok && gguf_fread_el(file, p->data, p->n, offset);
22113
 
@@ -22141,7 +22145,11 @@ static void gguf_free_kv(struct gguf_kv * kv) {
22141
  }
22142
 
22143
  struct gguf_context * gguf_init_empty(void) {
22144
- struct gguf_context * ctx = GGML_CALLOC(1, sizeof(struct gguf_context));
 
 
 
 
22145
 
22146
  memcpy(ctx->header.magic, GGUF_MAGIC, sizeof(ctx->header.magic));
22147
  ctx->header.version = GGUF_VERSION;
@@ -22187,7 +22195,12 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
22187
 
22188
  bool ok = true;
22189
 
22190
- struct gguf_context * ctx = GGML_CALLOC(1, sizeof(struct gguf_context));
 
 
 
 
 
22191
 
22192
  // read the header
22193
  {
@@ -22226,9 +22239,13 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
22226
  {
22227
  const uint64_t n_kv = ctx->header.n_kv;
22228
 
22229
- // header.n_kv will hold the actual value of pairs that were successfully read in the loop below
22230
- ctx->header.n_kv = 0;
22231
- ctx->kv = GGML_CALLOC(n_kv, sizeof(struct gguf_kv));
 
 
 
 
22232
 
22233
  for (uint64_t i = 0; i < n_kv; ++i) {
22234
  struct gguf_kv * kv = &ctx->kv[i];
@@ -22279,7 +22296,13 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
22279
  return NULL;
22280
  }
22281
 
22282
- kv->value.arr.data = GGML_CALLOC(kv->value.arr.n, gguf_type_size(kv->value.arr.type));
 
 
 
 
 
 
22283
 
22284
  ok = ok && gguf_fread_el(file, kv->value.arr.data, kv->value.arr.n * gguf_type_size(kv->value.arr.type), &offset);
22285
  } break;
@@ -22293,24 +22316,36 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
22293
  return NULL;
22294
  }
22295
 
22296
- kv->value.arr.data = GGML_CALLOC(kv->value.arr.n, sizeof(struct gguf_str));
 
 
 
 
 
 
22297
 
22298
  for (uint64_t j = 0; j < kv->value.arr.n; ++j) {
22299
  ok = ok && gguf_fread_str(file, &((struct gguf_str *) kv->value.arr.data)[j], &offset);
22300
  }
22301
  } break;
22302
  case GGUF_TYPE_ARRAY:
22303
- default: GGML_ABORT("invalid type");
 
 
 
 
22304
  }
22305
  } break;
22306
- default: GGML_ABORT("invalid type");
 
 
 
 
22307
  }
22308
 
22309
  if (!ok) {
22310
  break;
22311
  }
22312
-
22313
- ctx->header.n_kv++;
22314
  }
22315
 
22316
  if (!ok) {
@@ -22323,7 +22358,13 @@ struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_p
22323
 
22324
  // read the tensor infos
22325
  if (ctx->header.n_tensors > 0) {
22326
- ctx->infos = GGML_CALLOC(ctx->header.n_tensors, sizeof(struct gguf_tensor_info));
 
 
 
 
 
 
22327
 
22328
  for (uint64_t i = 0; i < ctx->header.n_tensors; ++i) {
22329
  struct gguf_tensor_info * info = &ctx->infos[i];
 
22107
  return false;
22108
  }
22109
 
22110
+ p->data = calloc(p->n + 1, 1);
22111
+ if (!p->data) {
22112
+ fprintf(stderr, "%s: failed to allocate memory for string of length %" PRIu64 "\n", __func__, p->n);
22113
+ return false;
22114
+ }
22115
 
22116
  ok = ok && gguf_fread_el(file, p->data, p->n, offset);
22117
 
 
22145
  }
22146
 
22147
  struct gguf_context * gguf_init_empty(void) {
22148
+ struct gguf_context * ctx = calloc(1, sizeof(struct gguf_context));
22149
+ if (!ctx) {
22150
+ fprintf(stderr, "%s: failed to allocate memory for context\n", __func__);
22151
+ return NULL;
22152
+ }
22153
 
22154
  memcpy(ctx->header.magic, GGUF_MAGIC, sizeof(ctx->header.magic));
22155
  ctx->header.version = GGUF_VERSION;
 
22195
 
22196
  bool ok = true;
22197
 
22198
+ struct gguf_context * ctx = calloc(1, sizeof(struct gguf_context));
22199
+ if (!ctx) {
22200
+ fprintf(stderr, "%s: failed to allocate memory for context\n", __func__);
22201
+ fclose(file);
22202
+ return NULL;
22203
+ }
22204
 
22205
  // read the header
22206
  {
 
22239
  {
22240
  const uint64_t n_kv = ctx->header.n_kv;
22241
 
22242
+ ctx->kv = calloc(n_kv, sizeof(struct gguf_kv));
22243
+ if (!ctx->kv) {
22244
+ fprintf(stderr, "%s: failed to allocate memory for kv pairs\n", __func__);
22245
+ fclose(file);
22246
+ gguf_free(ctx);
22247
+ return NULL;
22248
+ }
22249
 
22250
  for (uint64_t i = 0; i < n_kv; ++i) {
22251
  struct gguf_kv * kv = &ctx->kv[i];
 
22296
  return NULL;
22297
  }
22298
 
22299
+ kv->value.arr.data = calloc(kv->value.arr.n, gguf_type_size(kv->value.arr.type));
22300
+ if (!kv->value.arr.data) {
22301
+ fprintf(stderr, "%s: failed to allocate memory for array\n", __func__);
22302
+ fclose(file);
22303
+ gguf_free(ctx);
22304
+ return NULL;
22305
+ }
22306
 
22307
  ok = ok && gguf_fread_el(file, kv->value.arr.data, kv->value.arr.n * gguf_type_size(kv->value.arr.type), &offset);
22308
  } break;
 
22316
  return NULL;
22317
  }
22318
 
22319
+ kv->value.arr.data = calloc(kv->value.arr.n, sizeof(struct gguf_str));
22320
+ if (!kv->value.arr.data) {
22321
+ fprintf(stderr, "%s: failed to allocate memory for array\n", __func__);
22322
+ fclose(file);
22323
+ gguf_free(ctx);
22324
+ return NULL;
22325
+ }
22326
 
22327
  for (uint64_t j = 0; j < kv->value.arr.n; ++j) {
22328
  ok = ok && gguf_fread_str(file, &((struct gguf_str *) kv->value.arr.data)[j], &offset);
22329
  }
22330
  } break;
22331
  case GGUF_TYPE_ARRAY:
22332
+ default:
22333
+ {
22334
+ fprintf(stderr, "%s: invalid array type %d\n", __func__, kv->value.arr.type);
22335
+ ok = false;
22336
+ } break;
22337
  }
22338
  } break;
22339
+ default:
22340
+ {
22341
+ fprintf(stderr, "%s: invalid type %d\n", __func__, kv->type);
22342
+ ok = false;
22343
+ } break;
22344
  }
22345
 
22346
  if (!ok) {
22347
  break;
22348
  }
 
 
22349
  }
22350
 
22351
  if (!ok) {
 
22358
 
22359
  // read the tensor infos
22360
  if (ctx->header.n_tensors > 0) {
22361
+ ctx->infos = calloc(ctx->header.n_tensors, sizeof(struct gguf_tensor_info));
22362
+ if (!ctx->infos) {
22363
+ fprintf(stderr, "%s: failed to allocate memory for tensor infos\n", __func__);
22364
+ fclose(file);
22365
+ gguf_free(ctx);
22366
+ return NULL;
22367
+ }
22368
 
22369
  for (uint64_t i = 0; i < ctx->header.n_tensors; ++i) {
22370
  struct gguf_tensor_info * info = &ctx->infos[i];