Skip to content

Commit a0de8c7

Browse files
authored
wasi-nn: fix buffer overruns in load_by_name/load_by_name_with_config (#4491)
also, ensure NUL terminations here to simplify backend implementations. (note that, although names with '\0' in the middle are usually valid in wasm, wamr doesn't support them in general. we might revisit this later.) also, add a missing address validation in load_by_name_with_config. cf. #4487
1 parent c6afa13 commit a0de8c7

File tree

1 file changed

+64
-29
lines changed

1 file changed

+64
-29
lines changed

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

Lines changed: 64 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -507,36 +507,57 @@ wasi_nn_load(wasm_exec_env_t exec_env, graph_builder_array_wasm *builder,
507507
return res;
508508
}
509509

510+
static wasi_nn_error
511+
copyin_and_nul_terminate(wasm_module_inst_t inst, char *name, uint32_t name_len,
512+
char **resultp)
513+
{
514+
char *nul_terminated_name;
515+
if (!wasm_runtime_validate_native_addr(inst, name, name_len)) {
516+
return invalid_argument;
517+
}
518+
nul_terminated_name = wasm_runtime_malloc(name_len + 1);
519+
if (nul_terminated_name == NULL) {
520+
return runtime_error;
521+
}
522+
bh_memcpy_s(nul_terminated_name, name_len + 1, name, name_len);
523+
nul_terminated_name[name_len] = '\0'; /* ensure NUL termination */
524+
if (strlen(nul_terminated_name) != name_len) {
525+
/* reject names containing '\0' for now */
526+
wasm_runtime_free(nul_terminated_name);
527+
return invalid_argument;
528+
}
529+
*resultp = nul_terminated_name;
530+
return success;
531+
}
532+
510533
wasi_nn_error
511534
wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
512535
graph *g)
513536
{
537+
WASINNContext *wasi_nn_ctx = NULL;
538+
char *nul_terminated_name = NULL;
514539
wasi_nn_error res;
515540

516541
wasm_module_inst_t instance = wasm_runtime_get_module_inst(exec_env);
517542
if (!instance) {
518543
return runtime_error;
519544
}
520545

521-
if (!wasm_runtime_validate_native_addr(instance, name, name_len)) {
522-
NN_ERR_PRINTF("name is invalid");
523-
return invalid_argument;
524-
}
525-
526546
if (!wasm_runtime_validate_native_addr(instance, g,
527547
(uint64)sizeof(graph))) {
528548
NN_ERR_PRINTF("graph is invalid");
529549
return invalid_argument;
530550
}
531551

532-
if (name_len == 0 || name[name_len] != '\0') {
533-
NN_ERR_PRINTF("Invalid filename");
534-
return invalid_argument;
552+
res = copyin_and_nul_terminate(instance, name, name_len,
553+
&nul_terminated_name);
554+
if (res != success) {
555+
goto fail;
535556
}
536557

537-
NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME %s...", name);
558+
NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME %s...", nul_terminated_name);
538559

539-
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
560+
wasi_nn_ctx = lock_ctx(instance);
540561
if (wasi_nn_ctx == NULL) {
541562
res = busy;
542563
goto fail;
@@ -547,14 +568,20 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
547568
goto fail;
548569

549570
call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name, res,
550-
wasi_nn_ctx->backend_ctx, name, name_len, g);
571+
wasi_nn_ctx->backend_ctx, nul_terminated_name, name_len,
572+
g);
551573
if (res != success)
552574
goto fail;
553575

554576
wasi_nn_ctx->is_model_loaded = true;
555577
res = success;
556578
fail:
557-
unlock_ctx(wasi_nn_ctx);
579+
if (nul_terminated_name != NULL) {
580+
wasm_runtime_free(nul_terminated_name);
581+
}
582+
if (wasi_nn_ctx != NULL) {
583+
unlock_ctx(wasi_nn_ctx);
584+
}
558585
return res;
559586
}
560587

@@ -563,37 +590,37 @@ wasi_nn_load_by_name_with_config(wasm_exec_env_t exec_env, char *name,
563590
int32_t name_len, char *config,
564591
int32_t config_len, graph *g)
565592
{
593+
WASINNContext *wasi_nn_ctx = NULL;
594+
char *nul_terminated_name = NULL;
595+
char *nul_terminated_config = NULL;
566596
wasi_nn_error res;
567597

568598
wasm_module_inst_t instance = wasm_runtime_get_module_inst(exec_env);
569599
if (!instance) {
570600
return runtime_error;
571601
}
572602

573-
if (!wasm_runtime_validate_native_addr(instance, name, name_len)) {
574-
NN_ERR_PRINTF("name is invalid");
575-
return invalid_argument;
576-
}
577-
578603
if (!wasm_runtime_validate_native_addr(instance, g,
579604
(uint64)sizeof(graph))) {
580605
NN_ERR_PRINTF("graph is invalid");
581606
return invalid_argument;
582607
}
583608

584-
if (name_len == 0 || name[name_len] != '\0') {
585-
NN_ERR_PRINTF("Invalid filename");
586-
return invalid_argument;
609+
res = copyin_and_nul_terminate(instance, name, name_len,
610+
&nul_terminated_name);
611+
if (res != success) {
612+
goto fail;
587613
}
588-
589-
if (!config || config_len == 0 || config[config_len] != '\0') {
590-
NN_ERR_PRINTF("Invalid config");
591-
return invalid_argument;
614+
res = copyin_and_nul_terminate(instance, config, config_len,
615+
&nul_terminated_config);
616+
if (res != success) {
617+
goto fail;
592618
}
593619

594-
NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME_WITH_CONFIG %s %s...", name, config);
620+
NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME_WITH_CONFIG %s %s...",
621+
nul_terminated_name, nul_terminated_config);
595622

596-
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
623+
wasi_nn_ctx = lock_ctx(instance);
597624
if (wasi_nn_ctx == NULL) {
598625
res = busy;
599626
goto fail;
@@ -605,15 +632,23 @@ wasi_nn_load_by_name_with_config(wasm_exec_env_t exec_env, char *name,
605632
;
606633

607634
call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name_with_config, res,
608-
wasi_nn_ctx->backend_ctx, name, name_len, config,
609-
config_len, g);
635+
wasi_nn_ctx->backend_ctx, nul_terminated_name, name_len,
636+
nul_terminated_config, config_len, g);
610637
if (res != success)
611638
goto fail;
612639

613640
wasi_nn_ctx->is_model_loaded = true;
614641
res = success;
615642
fail:
616-
unlock_ctx(wasi_nn_ctx);
643+
if (nul_terminated_name != NULL) {
644+
wasm_runtime_free(nul_terminated_name);
645+
}
646+
if (nul_terminated_config != NULL) {
647+
wasm_runtime_free(nul_terminated_config);
648+
}
649+
if (wasi_nn_ctx != NULL) {
650+
unlock_ctx(wasi_nn_ctx);
651+
}
617652
return res;
618653
}
619654

0 commit comments

Comments
 (0)