From c90bc64e78cf180a333e93ae0ff587ce4231fe82 Mon Sep 17 00:00:00 2001
From: Eric Yang <yijieyang@google.com>
Date: Wed, 30 Oct 2024 14:28:16 -0700
Subject: [PATCH] Check for flatbuffer loading failures

Will report "Failed to build model from buffer" if it's invalid flatbuffer.

PiperOrigin-RevId: 691554844
---
 src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc b/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc
index d914da13..8f203e0e 100644
--- a/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc
+++ b/src/builtin-adapter/direct_flatbuffer_to_json_graph_convert.cc
@@ -837,6 +837,9 @@ absl::StatusOr<std::string> ConvertFlatbufferDirectlyToJson(
   std::unique_ptr<FlatBufferModelAbslError> model_ptr =
       FlatBufferModelAbslError::VerifyAndBuildFromBuffer(
           model_content.data(), model_content.length());
+  if (model_ptr == nullptr) {
+    return absl::InvalidArgumentError("Failed to build model from buffer.");
+  }
 
   mlir::MLIRContext mlir_context;
   mlir::DialectRegistry registry;