Skip to content

Commit 7ecf573

Browse files
QiuYuan HanQiuYuan Han
authored andcommitted
Add a new error check for wasi_nn_load_by_name
1 parent 5aef8e5 commit 7ecf573

File tree

1 file changed

+11
-2
lines changed

1 file changed

+11
-2
lines changed

core/iwasm/libraries/wasi-nn/src/wasi_nn.c

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -622,9 +622,10 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
622622
bool is_loaded = false;
623623
uint32 model_idx = 0;
624624
char *global_model_path_i;
625+
uint32_t global_n_graphs = wasm_runtime_get_wasi_nn_global_ctx_ngraphs(wasi_nn_global_ctx);
625626
// Assume filename got from user wasm app : max; sum; average; ...
626627
// Assume file path got from user cmd opt: /your/path1/max.tflite; /your/path2/sum.tflite; ......
627-
for (model_idx = 0; model_idx < wasm_runtime_get_wasi_nn_global_ctx_ngraphs(wasi_nn_global_ctx); model_idx++)
628+
for (model_idx = 0; model_idx < global_n_graphs; model_idx++)
628629
{
629630
// Extract filename from file path
630631
global_model_path_i = wasm_runtime_get_wasi_nn_global_ctx_graph_paths_i(wasi_nn_global_ctx, model_idx);
@@ -654,7 +655,9 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
654655
}
655656
}
656657

657-
if (!is_loaded && (model_idx < MAX_GLOBAL_GRAPHS_PER_INST))
658+
if (!is_loaded && \
659+
(model_idx < MAX_GLOBAL_GRAPHS_PER_INST) && \
660+
(model_idx < global_n_graphs))
658661
{
659662
NN_DBG_PRINTF("Model is not yet loaded, will add to global context");
660663
call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name, res,
@@ -679,6 +682,12 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
679682
NN_ERR_PRINTF("No enough space for new model");
680683
res = too_large;
681684
}
685+
else if (model_idx >= global_n_graphs)
686+
{
687+
NN_ERR_PRINTF("Cannot find model %s, you should pass its path through --wasi-nn-graph",
688+
nul_terminated_name);
689+
res = not_found;
690+
}
682691
goto fail;
683692
}
684693
fail:

0 commit comments

Comments
 (0)