From d2b0af69156598744b527ff640e9088b8d0fd553 Mon Sep 17 00:00:00 2001
From: Pyarelal Knowles <pknowles@nvidia.com>
Date: Tue, 27 Dec 2022 16:38:48 -0800
Subject: [PATCH] propagate image writing failures

Modifies UpdateImageObject() so that Returning false from the
WriteImageDataFunction callback results in the WriteGltfScene*() call
returning false.

Does not call WriteImageDataFunction if there is no image data.

Adds test case serialize-image-failure to verify the callback return
code is able to cause an overall failure to save the gltf.
---
 tests/tester.cc | 27 ++++++++++++++++++++++++++-
 tiny_gltf.h     | 29 ++++++++++++++++++++---------
 2 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/tests/tester.cc b/tests/tester.cc
index b43cf4c..00197d6 100644
--- a/tests/tester.cc
+++ b/tests/tester.cc
@@ -593,13 +593,38 @@ TEST_CASE("serialize-image-callback", "[issue-394]") {
 
   tinygltf::TinyGLTF ctx;
   ctx.SetImageWriter(writer, (void *)0xba5e1e55);
-  ctx.WriteGltfSceneToStream(const_cast<const tinygltf::Model *>(&m), os, false,
+  bool result = ctx.WriteGltfSceneToStream(const_cast<const tinygltf::Model *>(&m), os, false,
                              false);
 
   // use nlohmann json
   nlohmann::json j = nlohmann::json::parse(os.str());
 
+  REQUIRE(true == result);
   REQUIRE(1 == j["images"].size());
   REQUIRE(j["images"][0].is_object());
   REQUIRE(j["images"][0]["uri"].get<std::string>() == "bar");
 }
+
+TEST_CASE("serialize-image-failure", "[issue-394]") {
+  tinygltf::Model m;
+  tinygltf::Image i;
+  // Set some data so the ImageWriter callback will be called
+  i.image = {255, 255, 255, 255};
+  m.images.push_back(i);
+
+  std::stringstream os;
+
+  auto writer = [](const std::string *basepath, const std::string *filename,
+                   const tinygltf::Image *image, bool embedImages,
+                   std::string *out_uri, void *user_pointer) -> bool {
+    return false;
+  };
+
+  tinygltf::TinyGLTF ctx;
+  ctx.SetImageWriter(writer, (void *)0xba5e1e55);
+  bool result = ctx.WriteGltfSceneToStream(const_cast<const tinygltf::Model *>(&m), os, false,
+                             false);
+
+  REQUIRE(false == result);
+  REQUIRE(os.str().size() == 0);
+}
diff --git a/tiny_gltf.h b/tiny_gltf.h
index b83daa1..9680b00 100644
--- a/tiny_gltf.h
+++ b/tiny_gltf.h
@@ -2834,7 +2834,7 @@ static std::string MimeToExt(const std::string &mimeType) {
   return "";
 }
 
-static void UpdateImageObject(const Image &image, std::string &baseDir,
+static bool UpdateImageObject(const Image &image, std::string &baseDir,
                               int index, bool embedImages,
                               WriteImageDataFunction *WriteImageData,
                               std::string *out_uri, void *user_data) {
@@ -2857,17 +2857,24 @@ static void UpdateImageObject(const Image &image, std::string &baseDir,
     filename = std::to_string(index) + "." + ext;
   }
 
-  // If callback is set, modify image data object
+  // If callback is set and image data exists, modify image data object. If
+  // image data does not exist, this is not considered a failure and the
+  // original uri should be maintained.
   bool imageWritten = false;
-  if (*WriteImageData != nullptr && !filename.empty()) {
+  if (*WriteImageData != nullptr && !filename.empty() && !image.image.empty()) {
     imageWritten = (*WriteImageData)(&baseDir, &filename, &image, embedImages,
                                      out_uri, user_data);
+    if (!imageWritten) {
+      return false;
+    }
   }
 
   // Use the original uri if the image was not written.
   if (!imageWritten) {
     *out_uri = image.uri;
   }
+
+  return true;
 }
 
 bool IsDataURI(const std::string &in) {
@@ -7817,9 +7824,11 @@ bool TinyGLTF::WriteGltfSceneToStream(const Model *model, std::ostream &stream,
       // enabled, since we won't write separate images when writing to a stream
       // we
       std::string uri;
-      UpdateImageObject(model->images[i], dummystring, int(i), true,
-                        &this->WriteImageData, &uri,
-                        this->write_image_user_data_);
+      if (!UpdateImageObject(model->images[i], dummystring, int(i), true,
+                             &this->WriteImageData, &uri,
+                             this->write_image_user_data_)) {
+        return false;
+      }
       SerializeGltfImage(model->images[i], uri, image);
       JsonPushBack(images, std::move(image));
     }
@@ -7908,9 +7917,11 @@ bool TinyGLTF::WriteGltfSceneToFile(const Model *model,
       json image;
 
       std::string uri;
-      UpdateImageObject(model->images[i], baseDir, int(i), embedImages,
-                        &this->WriteImageData, &uri,
-                        this->write_image_user_data_);
+      if (!UpdateImageObject(model->images[i], baseDir, int(i), embedImages,
+                             &this->WriteImageData, &uri,
+                             this->write_image_user_data_)) {
+        return false;
+      }
       SerializeGltfImage(model->images[i], uri, image);
       JsonPushBack(images, std::move(image));
     }