From e47bb165442a61a93af2d83d0c081bc00c6b215a Mon Sep 17 00:00:00 2001
From: Prianka Liz Kariat <prianka.kariat@ymedialabs.com>
Date: Thu, 4 May 2023 16:52:58 +0530
Subject: [PATCH] Added validation of C++ image classification result packet in
 MPPImageClassifierResult+Helpers.mm

---
 .../sources/MPPImageClassifierResult+Helpers.h      |  4 +++-
 .../sources/MPPImageClassifierResult+Helpers.mm     | 13 +++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.h b/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.h
index 0375ac2a5..68d939f45 100644
--- a/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.h
+++ b/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.h
@@ -18,6 +18,8 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+static const int kMicroSecondsPerMilliSecond = 1000;
+
 @interface MPPImageClassifierResult (Helpers)
 
 /**
@@ -28,7 +30,7 @@ NS_ASSUME_NONNULL_BEGIN
  *
  * @return  An `MPPImageClassifierResult` object that contains a list of image classifications.
  */
-+ (MPPImageClassifierResult *)imageClassifierResultWithClassificationsPacket:
++ (nullable MPPImageClassifierResult *)imageClassifierResultWithClassificationsPacket:
     (const mediapipe::Packet &)packet;
 
 @end
diff --git a/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.mm b/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.mm
index f5199765d..f43ad0d7a 100644
--- a/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.mm
+++ b/mediapipe/tasks/ios/vision/image_classifier/utils/sources/MPPImageClassifierResult+Helpers.mm
@@ -17,8 +17,6 @@
 
 #include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h"
 
-static const int kMicroSecondsPerMilliSecond = 1000;
-
 namespace {
 using ClassificationResultProto =
     ::mediapipe::tasks::components::containers::proto::ClassificationResult;
@@ -27,9 +25,16 @@ using ::mediapipe::Packet;
 
 @implementation MPPImageClassifierResult (Helpers)
 
-+ (MPPImageClassifierResult *)imageClassifierResultWithClassificationsPacket:
++ (nullable MPPImageClassifierResult *)imageClassifierResultWithClassificationsPacket:
     (const Packet &)packet {
-  MPPClassificationResult *classificationResult = [MPPClassificationResult
+  MPPClassificationResult *classificationResult;
+  MPPImageClassifierResult *imageClassifierResult;
+
+  if (!packet.ValidateAsType<ClassificationResultProto>().ok()) {
+    return nil;
+  }
+
+  classificationResult = [MPPClassificationResult
       classificationResultWithProto:packet.Get<ClassificationResultProto>()];
 
   return [[MPPImageClassifierResult alloc]