From 0a50b74f8402278a43cc5340ec3ab19e586762fa Mon Sep 17 00:00:00 2001 From: Jaeyun Jung Date: Thu, 6 Jul 2023 12:31:12 +0900 Subject: [PATCH] [Common] check mem size of flex tensor Check memory size before parsing tensor-meta. Signed-off-by: Jaeyun Jung --- gst/nnstreamer/nnstreamer_plugin_api_impl.c | 9 +++++++++ tests/common/unittest_common.cc | 20 ++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/gst/nnstreamer/nnstreamer_plugin_api_impl.c b/gst/nnstreamer/nnstreamer_plugin_api_impl.c index 5da68d1..0b065cc 100644 --- a/gst/nnstreamer/nnstreamer_plugin_api_impl.c +++ b/gst/nnstreamer/nnstreamer_plugin_api_impl.c @@ -1473,6 +1473,7 @@ gboolean gst_tensor_meta_info_parse_memory (GstTensorMetaInfo * meta, GstMemory * mem) { GstMapInfo map; + gsize hsize, msize; gboolean ret; g_return_val_if_fail (mem != NULL, FALSE); @@ -1480,6 +1481,14 @@ gst_tensor_meta_info_parse_memory (GstTensorMetaInfo * meta, GstMemory * mem) gst_tensor_meta_info_init (meta); + /* Check header size of tensor-meta. */ + hsize = gst_tensor_meta_info_get_header_size (meta); + msize = gst_memory_get_sizes (mem, NULL, NULL); + if (msize < hsize) { + nns_loge ("Failed to get the meta, invalid memory size %zd.", msize); + return FALSE; + } + if (!gst_memory_map (mem, &map, GST_MAP_READ)) { nns_loge ("Failed to get the meta, cannot map the memory."); return FALSE; diff --git a/tests/common/unittest_common.cc b/tests/common/unittest_common.cc index 94f42eb..ef002e3 100644 --- a/tests/common/unittest_common.cc +++ b/tests/common/unittest_common.cc @@ -1552,6 +1552,26 @@ TEST (commonMetaInfo, parseMemInvalidParam_n) } /** + * @brief Test for tensor meta info (parse memory with invalid memory size). + */ +TEST (commonMetaInfo, parseMemInvalidSize_n) +{ + GstTensorMetaInfo meta; + GstMemory *mem; + gsize hsize; + gboolean ret; + + gst_tensor_meta_info_init (&meta); + hsize = gst_tensor_meta_info_get_header_size (&meta) / 2; + mem = gst_allocator_alloc (NULL, hsize, NULL); + + ret = gst_tensor_meta_info_parse_memory (&meta, mem); + EXPECT_FALSE (ret); + + gst_memory_unref (mem); +} + +/** * @brief Test for tensor meta info (append header to memory). */ TEST (commonMetaInfo, appendHeader) -- 2.7.4