See memory of freed result to nullptr

PiperOrigin-RevId: 570410751
google-ai-edge · Oct 3, 2023 · a72839e · a72839e
1 parent d0183b2
commit a72839e
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/mediapipe/tasks/c/components/containers/category_converter.cc b/mediapipe/tasks/c/components/containers/category_converter.cc
@@ -36,7 +36,9 @@ void CppConvertToCategory(
 
 void CppCloseCategory(Category* in) {
   free(in->category_name);
+  in->category_name = nullptr;
   free(in->display_name);
+  in->display_name = 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
@@ -60,17 +60,20 @@ void CppConvertToClassificationResult(
 
 void CppCloseClassificationResult(ClassificationResult* in) {
   for (uint32_t i = 0; i < in->classifications_count; ++i) {
-    auto classification_in = in->classifications[i];
+    auto& classification_in = in->classifications[i];
 
     for (uint32_t j = 0; j < classification_in.categories_count; ++j) {
       CppCloseCategory(&classification_in.categories[j]);
     }
     delete[] classification_in.categories;
+    classification_in.categories = nullptr;
 
     free(classification_in.head_name);
+    classification_in.head_name = nullptr;
   }
 
   delete[] in->classifications;
+  in->classifications = nullptr;
 }
 
 }  // namespace mediapipe::tasks::c::components::containers
diff --git a/mediapipe/tasks/c/components/containers/classification_result_converter_test.cc b/mediapipe/tasks/c/components/containers/classification_result_converter_test.cc
@@ -99,4 +99,20 @@ TEST(ClassificationResultConverterTest,
   CppCloseClassificationResult(&c_classification_result);
 }
 
+TEST(ClassificationResultConverterTest, FreesMemory) {
+  mediapipe::tasks::components::containers::ClassificationResult
+      cpp_classification_result = {
+          /* classifications= */ {{{}, 0, "foo"}},
+          /* timestamp_ms= */ 42,
+      };
+
+  ClassificationResult c_classification_result;
+  CppConvertToClassificationResult(cpp_classification_result,
+                                   &c_classification_result);
+  EXPECT_NE(c_classification_result.classifications, nullptr);
+
+  CppCloseClassificationResult(&c_classification_result);
+  EXPECT_EQ(c_classification_result.classifications, nullptr);
+}
+
 }  // namespace mediapipe::tasks::c::components::containers