Просмотр исходного кода

Add bounds checking for output tensor buffer in wasi-nn llama.cpp (#4847)

* Add bounds checking for output tensor buffer in wasi-nn llama.cpp

The get_output function copies LLM output into output_tensor->buf
without checking against output_tensor->size, allowing writes
past the buffer when the model generates output longer than the
caller-provided buffer. Add size checks for both the metadata
path and the token output loop.

Instead of silently truncating output when the buffer is too small,
return the too_large error with a diagnostic message. This makes the
behavior consistent with the OpenVINO backend's get_output and allows
callers to distinguish between successful completion and insufficient
buffer size.
Yi Liu 1 день назад
Родитель
Сommit
3d707f9173
1 измененных файлов с 18 добавлено и 4 удалено
  1. 18 4
      core/iwasm/libraries/wasi-nn/src/wasi_nn_llamacpp.c

+ 18 - 4
core/iwasm/libraries/wasi-nn/src/wasi_nn_llamacpp.c

@@ -623,8 +623,15 @@ get_output(void *ctx, graph_execution_context exec_ctx, uint32_t index,
             printf("%s\n", output_metadata);
         }
 
-        memcpy(output_tensor->buf, output_metadata, strlen(output_metadata));
-        *output_tensor_size = strlen(output_metadata);
+        size_t metadata_len = strlen(output_metadata);
+        if (metadata_len > output_tensor->size) {
+            NN_ERR_PRINTF("Output buffer too small for metadata: "
+                          "need %zu, got %zu",
+                          metadata_len, output_tensor->size);
+            return too_large;
+        }
+        memcpy(output_tensor->buf, output_metadata, metadata_len);
+        *output_tensor_size = metadata_len;
         return success;
     }
 
@@ -643,8 +650,15 @@ get_output(void *ctx, graph_execution_context exec_ctx, uint32_t index,
             printf("%s", buf);
         }
 
-        memcpy(output_tensor->buf + end_pos, buf, strlen(buf));
-        end_pos += strlen(buf);
+        size_t piece_len = strlen(buf);
+        if (end_pos + piece_len > output_tensor->size) {
+            NN_ERR_PRINTF("Output buffer too small: need at least %zu,"
+                          " got %zu",
+                          end_pos + piece_len, output_tensor->size);
+            return too_large;
+        }
+        memcpy(output_tensor->buf + end_pos, buf, piece_len);
+        end_pos += piece_len;
     }
 
     if (backend_ctx->config.stream_stdout) {