From a72839ef990205898ea613e261d64525b9d25f79 Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <mrschmidt@google.com>
Date: Tue, 3 Oct 2023 09:30:12 -0700
Subject: [PATCH] See memory of freed result to nullptr

PiperOrigin-RevId: 570410751
---
 .../components/containers/category_converter.cc  |  2 ++
 .../classification_result_converter.cc           |  5 ++++-
 .../classification_result_converter_test.cc      | 16 ++++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/mediapipe/tasks/c/components/containers/category_converter.cc b/mediapipe/tasks/c/components/containers/category_converter.cc
index a38ef10b0..b04f86c9a 100644
--- 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
index 64e62c309..b7190c932 100644
--- 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
index 59eb53aaf..cecef3ccd 100644
--- 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