Fix memory access issue in C layer

PiperOrigin-RevId: 569279959
google-ai-edge · Sep 28, 2023 · 636cf99 · 636cf99
1 parent e169849
commit 636cf99
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/mediapipe/tasks/c/components/containers/category_converter.cc b/mediapipe/tasks/c/components/containers/category_converter.cc
@@ -25,10 +25,11 @@ void CppConvertToCategory(
     Category* out) {
   out->index = in.index;
   out->score = in.score;
-  out->category_name =
-      in.category_name.has_value() ? in.category_name->c_str() : nullptr;
+  out->category_name = in.category_name.has_value()
+                           ? strdup(in.category_name->c_str())
+                           : nullptr;
   out->display_name =
-      in.display_name.has_value() ? in.display_name->c_str() : nullptr;
+      in.display_name.has_value() ? strdup(in.display_name->c_str()) : nullptr;
 }
 
 }  // namespace mediapipe::tasks::c::components::containers
diff --git a/mediapipe/tasks/c/components/containers/classification_result_converter.cc b/mediapipe/tasks/c/components/containers/classification_result_converter.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "mediapipe/tasks/c/components/containers/classification_result_converter.h"
 
+#include <cstdint>
+
 #include "mediapipe/tasks/c/components/containers/category.h"
 #include "mediapipe/tasks/c/components/containers/category_converter.h"
 #include "mediapipe/tasks/cc/components/containers/classification_result.h"
@@ -32,22 +34,22 @@ void CppConvertToClassificationResult(
   out->classifications_count = in.classifications.size();
   out->classifications = new Classifications[out->classifications_count];
 
-  for (uint32_t i = 0; i <= out->classifications_count; ++i) {
+  for (uint32_t i = 0; i < out->classifications_count; ++i) {
     auto classification_in = in.classifications[i];
     auto classification_out = out->classifications[i];
 
     classification_out.categories_count = classification_in.categories.size();
     classification_out.categories =
         new Category[classification_out.categories_count];
-    for (uint32_t j = 0; j <= classification_out.categories_count; ++j) {
+    for (uint32_t j = 0; j < classification_out.categories_count; ++j) {
       CppConvertToCategory(classification_in.categories[j],
                            &(classification_out.categories[j]));
     }
 
     classification_out.head_index = classification_in.head_index;
     classification_out.head_name =
         classification_in.head_name.has_value()
-            ? classification_in.head_name.value().c_str()
+            ? strdup(classification_in.head_name->c_str())
             : nullptr;
   }
 }