diff --git a/c/src/ml-api-common.c b/c/src/ml-api-common.c index 427bcaa2..00483289 100644 --- a/c/src/ml-api-common.c +++ b/c/src/ml-api-common.c @@ -479,7 +479,7 @@ ml_tensors_info_set_tensor_dimension (ml_tensors_info_h info, { ml_tensors_info_s *tensors_info; GstTensorInfo *_info; - guint i; + guint i, rank, max_rank; check_feature_state (ML_FEATURE); @@ -497,6 +497,27 @@ ml_tensors_info_set_tensor_dimension (ml_tensors_info_h info, tensors_info->info.num_tensors, index, index); } + /** + * Validate dimension. + * We cannot use util function to get the rank of tensor dimension here. + * The old rank limit is 4, and testcases or app may set old dimension. + */ + max_rank = tensors_info->is_extended ? + ML_TENSOR_RANK_LIMIT : ML_TENSOR_RANK_LIMIT_PREV; + rank = max_rank + 1; + for (i = 0; i < max_rank; i++) { + if (dimension[i] == 0) { + if (rank > max_rank) + rank = i; + } + + if (rank == 0 || (i > rank && dimension[i] > 0)) { + G_UNLOCK_UNLESS_NOLOCK (*tensors_info); + _ml_error_report_return (ML_ERROR_INVALID_PARAMETER, + "The parameter, dimension, is invalid. It should be a valid unsigned integer array."); + } + } + _info = gst_tensors_info_get_nth_info (&tensors_info->info, index); if (!_info) { G_UNLOCK_UNLESS_NOLOCK (*tensors_info); diff --git a/tests/capi/unittest_capi_inference.cc b/tests/capi/unittest_capi_inference.cc index 67bd1cdf..88959254 100644 --- a/tests/capi/unittest_capi_inference.cc +++ b/tests/capi/unittest_capi_inference.cc @@ -3305,6 +3305,54 @@ TEST (nnstreamer_capi_util, info_set_tdimension_02_n) ASSERT_EQ (status, ML_ERROR_NONE); } +/** + * @brief Test utility functions (public) + */ +TEST (nnstreamer_capi_util, info_set_tdimension_03_n) +{ + ml_tensors_info_h info; + ml_tensor_dimension dim = { 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0 }; + int status; + + /* max rank 4 (ML_TENSOR_RANK_LIMIT_PREV) */ + ml_tensors_info_create (&info); + ml_tensors_info_set_count (info, 1); + + dim[0] = 0; + status = ml_tensors_info_set_tensor_dimension (info, 0, dim); + EXPECT_NE (status, ML_ERROR_NONE); + + dim[0] = 2; + dim[2] = 0; + status = ml_tensors_info_set_tensor_dimension (info, 0, dim); + EXPECT_NE (status, ML_ERROR_NONE); + + dim[2] = 1; + status = ml_tensors_info_set_tensor_dimension (info, 0, dim); + EXPECT_EQ (status, ML_ERROR_NONE); + + ml_tensors_info_destroy (info); + + /* max rank 16 (ML_TENSOR_RANK_LIMIT) */ + ml_tensors_info_create_extended (&info); + ml_tensors_info_set_count (info, 1); + + dim[0] = 0; + status = ml_tensors_info_set_tensor_dimension (info, 0, dim); + EXPECT_NE (status, ML_ERROR_NONE); + + dim[0] = 2; + dim[7] = 0; + status = ml_tensors_info_set_tensor_dimension (info, 0, dim); + EXPECT_NE (status, ML_ERROR_NONE); + + dim[7] = 1; + status = ml_tensors_info_set_tensor_dimension (info, 0, dim); + EXPECT_EQ (status, ML_ERROR_NONE); + + ml_tensors_info_destroy (info); +} + /** * @brief Test utility functions (public) */