gguf : harden loader against malformed tensor dims and metadata types (#25596)

* gguf : harden loader against malformed tensor dims and metadata types

* gguf: address review on malformed-metadata hardening

- report the expected vs. actual type when general.alignment is not u32
- use ggml_nelements() > 0 for the zero-element guard and keep the
  representability checks visually aligned
- add test-gguf cases for a wrong-typed alignment key and a zero-dim
  tensor (both used to crash: assert-abort and SIGFPE respectively)

Ran tests/test-gguf: 164/164 pass. Used an AI assistant to help draft
these edits; reviewed and verified by me.

* cont : less comments

Co-authored-by: Georgi Gerganov <redacted>
---------

Co-authored-by: Georgi Gerganov <redacted>

diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp
index 9f9e4fe5d104fe753e9fbc646efb231c3d78ee76..6c7b5817812b3246ecb9172e20b3dcf83dc216eb 100644 (file)
--- a/ggml/src/gguf.cpp
+++ b/ggml/src/gguf.cpp
@@ -611,6 +611,13 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr
         GGML_ASSERT(int64_t(ctx->kv.size()) == n_kv);
 
         const int alignment_idx = gguf_find_key(ctx, GGUF_KEY_GENERAL_ALIGNMENT);
+        if (alignment_idx != -1 && gguf_get_kv_type(ctx, alignment_idx) != GGUF_TYPE_UINT32) {
+            GGML_LOG_ERROR("%s: key '%s' must be of type %s but is %s\n",
+                __func__, GGUF_KEY_GENERAL_ALIGNMENT, gguf_type_name(GGUF_TYPE_UINT32),
+                gguf_type_name(gguf_get_kv_type(ctx, alignment_idx)));
+            gguf_free(ctx);
+            return nullptr;
+        }
         ctx->alignment = alignment_idx == -1 ? GGUF_DEFAULT_ALIGNMENT : gguf_get_val_u32(ctx, alignment_idx);
 
         if (ctx->alignment == 0 || (ctx->alignment & (ctx->alignment - 1)) != 0) {
@@ -682,9 +689,11 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr
             }
 
             // check that the total number of elements is representable
-            if (ok && ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) ||
-                       (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) ||
-                       (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) {
+            // (a zero-element tensor is trivially representable; the guard also avoids a division by zero below)
+            if (ok && ggml_nelements(&info.t) > 0 &&
+                ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) ||
+                 (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) ||
+                 (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) {
 
                 GGML_LOG_ERROR("%s: total number of elements in tensor '%s' with shape "
                     "(%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") is >= %" PRIi64 "\n",

diff --git a/tests/test-gguf.cpp b/tests/test-gguf.cpp
index 2875dec806da769cf7bed849a9a1db9850689706..fc636186f4c55c448782bf63d0a084d96efc22a1 100644 (file)
--- a/tests/test-gguf.cpp
+++ b/tests/test-gguf.cpp
@@ -31,11 +31,13 @@ enum handcrafted_file_type {
     // HANDCRAFTED_KV_BAD_VALUE_SIZE          =  30 + offset_has_kv, // removed because it can result in allocations > 1 TB (default sanitizer limit)
     HANDCRAFTED_KV_DUPLICATE_KEY           =  40 + offset_has_kv,
     HANDCRAFTED_KV_BAD_ALIGN               =  50 + offset_has_kv,
+    HANDCRAFTED_KV_WRONG_TYPE_ALIGN        =  55 + offset_has_kv,
     HANDCRAFTED_KV_SUCCESS                 = 800 + offset_has_kv,
 
     HANDCRAFTED_TENSORS_BAD_NAME_SIZE      =  10 + offset_has_tensors,
     HANDCRAFTED_TENSORS_BAD_N_DIMS         =  20 + offset_has_tensors,
     HANDCRAFTED_TENSORS_BAD_SHAPE          =  30 + offset_has_tensors,
+    HANDCRAFTED_TENSORS_ZERO_DIM           =  35 + offset_has_tensors,
     HANDCRAFTED_TENSORS_NE_TOO_BIG         =  40 + offset_has_tensors,
     HANDCRAFTED_TENSORS_NBYTES_TOO_BIG     =  45 + offset_has_tensors,
     HANDCRAFTED_TENSORS_BAD_TYPE           =  50 + offset_has_tensors,
@@ -69,11 +71,13 @@ static std::string handcrafted_file_type_name(const enum handcrafted_file_type h
         case HANDCRAFTED_KV_BAD_TYPE:                return "KV_BAD_TYPE";
         case HANDCRAFTED_KV_DUPLICATE_KEY:           return "KV_DUPLICATE_KEY";
         case HANDCRAFTED_KV_BAD_ALIGN:               return "KV_BAD_ALIGN";
+        case HANDCRAFTED_KV_WRONG_TYPE_ALIGN:        return "KV_WRONG_TYPE_ALIGN";
         case HANDCRAFTED_KV_SUCCESS:                 return "KV_RANDOM_KV";
 
         case HANDCRAFTED_TENSORS_BAD_NAME_SIZE:      return "TENSORS_BAD_NAME_SIZE";
         case HANDCRAFTED_TENSORS_BAD_N_DIMS:         return "TENSORS_BAD_N_DIMS";
         case HANDCRAFTED_TENSORS_BAD_SHAPE:          return "TENSORS_BAD_SHAPE";
+        case HANDCRAFTED_TENSORS_ZERO_DIM:           return "TENSORS_ZERO_DIM";
         case HANDCRAFTED_TENSORS_NE_TOO_BIG:         return "TENSORS_NE_TOO_BIG";
         case HANDCRAFTED_TENSORS_NBYTES_TOO_BIG:     return "TENSORS_NBYTES_TOO_BIG";
         case HANDCRAFTED_TENSORS_BAD_TYPE:           return "TENSORS_BAD_TYPE";
@@ -95,6 +99,9 @@ static std::string handcrafted_file_type_name(const enum handcrafted_file_type h
 }
 
 static bool expect_context_not_null(const enum handcrafted_file_type hft) {
+    if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
+        return true;
+    }
     if (hft < offset_has_kv) {
         return hft >= HANDCRAFTED_HEADER_EMPTY;
     }
@@ -257,9 +264,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft
     }
     {
         uint64_t n_kv = kv_types.size();
-        if (hft == HANDCRAFTED_KV_BAD_ALIGN      ||
-            hft == HANDCRAFTED_TENSORS_BAD_ALIGN || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN ||
-            hft == HANDCRAFTED_DATA_BAD_ALIGN    || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) {
+        if (hft == HANDCRAFTED_KV_BAD_ALIGN        || hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN ||
+            hft == HANDCRAFTED_TENSORS_BAD_ALIGN   || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN ||
+            hft == HANDCRAFTED_DATA_BAD_ALIGN      || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) {
 
             n_kv += 1;
         } else if (hft == HANDCRAFTED_HEADER_BAD_N_KV) {
@@ -344,15 +351,17 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft
         helper_write(file, data, hft == HANDCRAFTED_KV_BAD_TYPE ? 1 : gguf_type_size(type));
     }
 
-    if (hft == HANDCRAFTED_KV_BAD_ALIGN      ||
-        hft == HANDCRAFTED_TENSORS_BAD_ALIGN || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN ||
-        hft == HANDCRAFTED_DATA_BAD_ALIGN    || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) {
+    if (hft == HANDCRAFTED_KV_BAD_ALIGN        || hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN ||
+        hft == HANDCRAFTED_TENSORS_BAD_ALIGN   || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN ||
+        hft == HANDCRAFTED_DATA_BAD_ALIGN      || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) {
 
         const uint64_t n = strlen(GGUF_KEY_GENERAL_ALIGNMENT);
         helper_write(file, n);
         helper_write(file, GGUF_KEY_GENERAL_ALIGNMENT, n);
 
-        const int32_t type = gguf_type(GGUF_TYPE_UINT32);
+        // HANDCRAFTED_KV_WRONG_TYPE_ALIGN declares general.alignment with a non-UINT32 type,
+        // which the loader must reject cleanly instead of aborting on an assertion
+        const int32_t type = hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN ? int32_t(GGUF_TYPE_INT32) : int32_t(GGUF_TYPE_UINT32);
         helper_write(file, type);
 
         alignment = expect_context_not_null(hft) ? 1 : 13;
@@ -403,6 +412,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft
                 break;
             }
         }
+        if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
+            n_dims = 2;
+        }
         if (hft == HANDCRAFTED_TENSORS_BAD_N_DIMS) {
             const uint32_t n_dims_bad = GGML_MAX_DIMS + 1;
             helper_write(file, n_dims_bad);
@@ -415,6 +427,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft
             for (uint32_t j = 0; j < n_dims; ++j) {
                 helper_write(file, bad_dim);
             }
+        } else if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
+            const int64_t zero_shape[2] = { shape[0], 0 };
+            helper_write(file, zero_shape, 2*sizeof(int64_t));
         } else if (hft == HANDCRAFTED_TENSORS_NE_TOO_BIG){
             const int64_t big_dim = 4*int64_t(INT32_MAX);
             for (uint32_t j = 0; j < n_dims; ++j) {
@@ -446,6 +461,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft
         for (uint32_t i = 1; i < n_dims; ++i) {
             ne *= shape[i];
         }
+        if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
+            ne = 0;
+        }
 
         offset += GGML_PAD(ggml_row_size(type, ne), (uint64_t) alignment);
     }
@@ -747,11 +765,13 @@ static std::pair<int, int> test_handcrafted_file(const unsigned int seed) {
         HANDCRAFTED_KV_BAD_TYPE,
         HANDCRAFTED_KV_DUPLICATE_KEY,
         HANDCRAFTED_KV_BAD_ALIGN,
+        HANDCRAFTED_KV_WRONG_TYPE_ALIGN,
         HANDCRAFTED_KV_SUCCESS,
 
         HANDCRAFTED_TENSORS_BAD_NAME_SIZE,
         HANDCRAFTED_TENSORS_BAD_N_DIMS,
         HANDCRAFTED_TENSORS_BAD_SHAPE,
+        HANDCRAFTED_TENSORS_ZERO_DIM,
         HANDCRAFTED_TENSORS_NE_TOO_BIG,
         HANDCRAFTED_TENSORS_NBYTES_TOO_BIG,
         HANDCRAFTED_TENSORS_BAD_TYPE,
@@ -840,7 +860,9 @@ static std::pair<int, int> test_handcrafted_file(const unsigned int seed) {
             ntest++;
         }
 
-        if (expect_context_not_null(hft) && hft >= offset_has_tensors) {
+        // HANDCRAFTED_TENSORS_ZERO_DIM deliberately mangles the tensor shapes to 0 elements,
+        // so only assert that it loads without crashing; skip the exact-geometry comparison.
+        if (expect_context_not_null(hft) && hft >= offset_has_tensors && hft != HANDCRAFTED_TENSORS_ZERO_DIM) {
             printf("%s:   - check_tensors: ", __func__);
             if (handcrafted_check_tensors(gguf_ctx, seed)) {
                 printf("\033[1;32mOK\033[0m\n");