[Common] check mem size of flex tensor

Check memory size before parsing tensor-meta.

Signed-off-by: Jaeyun Jung <jy1210.jung@samsung.com>

diff --git a/gst/nnstreamer/nnstreamer_plugin_api_impl.c b/gst/nnstreamer/nnstreamer_plugin_api_impl.c
index 5da68d1..0b065cc 100644 (file)
--- 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 (file)
--- 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)