NLClassifier C API: Fix memory leakage by releasing the copied c_str …

…from `strdup`. `strdup` is called in nl_classifier_c_api.cc:76 PiperOrigin-RevId: 374313022
tensorflow · May 18, 2021 · 0f5fa0f · 0f5fa0f
1 parent de407f1
commit 0f5fa0f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier_c_api_common.cc b/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier_c_api_common.cc
@@ -15,12 +15,18 @@ limitations under the License.
 
 #include "tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier_c_api_common.h"
 
+#include <memory>
 
 #ifdef __cplusplus
 extern "C" {
 #endif  // __cplusplus
 
 void NLClassifierCategoriesDelete(Categories* categories) {
+  for (int i = 0; i < categories->size; i++) {
+    // `strdup` obtains memory using `malloc` and the memory needs to be
+    // released using `free`.
+    free(categories->categories[i].text);
+  }
   delete[] categories->categories;
   delete categories;
 }

diff --git a/tensorflow_lite_support/cc/task/text/qa/bert_qa_c_api.cc b/tensorflow_lite_support/cc/task/text/qa/bert_qa_c_api.cc
@@ -70,6 +70,11 @@ void BertQuestionAnswererDelete(BertQuestionAnswerer* bert_question_answerer) {
 }
 
 void BertQuestionAnswererQaAnswersDelete(QaAnswers* qa_answers) {
+  for (int i = 0; i < qa_answers->size; i++) {
+    // `strdup` obtains memory using `malloc` and the memory needs to be
+    // released using `free`.
+    free(qa_answers->answers[i].text);
+  }
   delete[] qa_answers->answers;
   delete qa_answers;
 }