Merge pull request #72 from agruzdev/save-via-tmp

FreeImage_Save uses tmp file to avoid corrupted target file states

Author: agruzdev

Source/FreeImage/Plugin.cpp

@@ -35,6 +35,8 @@
 #include <ctype.h>
 #endif // _WIN32
 
+#include <filesystem>
+
 #include "FreeImage.h"
 #include "Utilities.h"
 #include "FreeImageIO.h"
@@ -769,41 +771,107 @@ FreeImage_SaveToHandle(FREE_IMAGE_FORMAT fif, FIBITMAP *dib, FreeImageIO *io, fi
 }
 
 
-FIBOOL DLL_CALLCONV
-FreeImage_Save(FREE_IMAGE_FORMAT fif, FIBITMAP *dib, const char *filename, int flags) {
-	FreeImageIO io;
-	SetDefaultIO(&io);
+namespace {
 
-	FIBOOL success{FALSE};
-	if (auto *handle = fopen(filename, "w+b")) {
-		success = FreeImage_SaveToHandle(fif, dib, &io, (fi_handle)handle, flags);
 
-		fclose(handle);
-	} else {
-		FreeImage_OutputMessageProc((int)fif, "FreeImage_Save: failed to open file %s", filename);
+	std::filesystem::path MakeRandomSuffix()
+	{
+		std::stringstream strs{};
+		strs << ".fitmp" << std::hex << static_cast<uint32_t>(std::rand());
+		return strs.str();
 	}
 
-	return success;
-}
+
+	std::filesystem::path MakeTmpName(const std::filesystem::path& target)
+	{
+		std::filesystem::path res{};
+
+		int attempts = 16;
+		do {
+			res = target.stem();
+			res += MakeRandomSuffix();
+			res += target.extension();
+
+			if (std::filesystem::status(res).type() != std::filesystem::file_type::not_found) {
+				continue;
+			}
+
+		} while (--attempts);
+
+		return res;
+	}
+
+
+	bool MoveOrCopy(const std::filesystem::path& oldp, const std::filesystem::path& newp) noexcept
+	{
+		std::error_code err{};
+		std::filesystem::rename(oldp, newp, err);
+		if (!err) {
+			return true;
+		}
+		const auto copied = std::filesystem::copy_file(oldp, newp, std::filesystem::copy_options::overwrite_existing, err);
+		if (copied) {
+			std::filesystem::remove(oldp, err);
+		}
+		return copied;
+	}
+
+
+	FIBOOL SaveImpl(FREE_IMAGE_FORMAT fif, FIBITMAP* dib, const std::filesystem::path& filename, int flags) 
+	try
+	{
+		FreeImageIO io;
+		SetDefaultIO(&io);
+
+		FIBOOL success{ FALSE };
+
+		const auto tmpname = MakeTmpName(filename);
+		if (auto* handle = FreeImage_FOpen(tmpname, "w+b")) {
+			success = FreeImage_SaveToHandle(fif, dib, &io, (fi_handle)handle, flags);
+			fclose(handle);
+
+			if (success) {
+				success = MoveOrCopy(tmpname, filename);
+				if (!success) {
+					FreeImage_OutputMessageProc((int)fif, "FreeImage_Save: failed to rename output file %s", filename.u8string().c_str());
+				}
+			}
+			else {
+				std::error_code err{};
+				std::filesystem::remove(tmpname, err);
+			}
+		}
+		else {
+			FreeImage_OutputMessageProc((int)fif, "FreeImage_Save: failed to open file %s", filename.u8string().c_str());
+		}
+
+		return success;
+	}
+	catch (...) {
+		return FALSE;
+	}
+
+} // namespace
+
 
 FIBOOL DLL_CALLCONV
-FreeImage_SaveU(FREE_IMAGE_FORMAT fif, FIBITMAP *dib, const wchar_t *filename, int flags) {
-	FreeImageIO io;
-	SetDefaultIO(&io);
+FreeImage_Save(FREE_IMAGE_FORMAT fif, FIBITMAP *dib, const char *filename, int flags) {
+	if (!dib || !filename) {
+		return FALSE;
+	}
+	return SaveImpl(fif, dib, filename, flags);
+}
 
-	FIBOOL success{FALSE};
-#ifdef _WIN32	
-	if (auto *handle = _wfopen(filename, L"w+b")) {
-		success = FreeImage_SaveToHandle(fif, dib, &io, (fi_handle)handle, flags);
 
-		fclose(handle);
-	} else {
-		FreeImage_OutputMessageProc((int)fif, "FreeImage_SaveU: failed to open output file");
+FIBOOL DLL_CALLCONV
+FreeImage_SaveU(FREE_IMAGE_FORMAT fif, FIBITMAP *dib, const wchar_t *filename, int flags) {
+	if (!dib || !filename) {
+		return FALSE;
 	}
-#endif
-	return success;
+	return SaveImpl(fif, dib, filename, flags);
 }
 
+
 // =====================================================================
 // Plugin construction + enable/disable functions
 // =====================================================================

TestAPI/MainTestSuite.cpp

@@ -89,6 +89,12 @@ int main(int argc, char *argv[]) {
 	// test internal image types
 	testImageType(width, height);
 
+
+	auto bmp = FreeImage_AllocateT(FIT_COMPLEX, 128, 128, 128);
+	FreeImage_Save(FIF_JPEG, bmp, "failed_to_save.jpg");
+	FreeImage_Unload(bmp);
+
+
 #if FREEIMAGE_WITH_LIBJPEG
 	// test the clone function
 	testAllocateCloneUnload("exif.jpg");

TestAPI/testHeaderOnly.cpp

@@ -174,7 +174,9 @@ testExifRawFile(const char *lpszPathName, int load_flags, int save_flags) {
 			uint8_t *value = (uint8_t*)FreeImage_GetTagValue(tag);
 
 			// save as jpeg : Exif data should be preserved
-			FreeImage_Save(fif, dib, lpszDstPathName, save_flags); 
+			if (!FreeImage_Save(fif, dib, lpszDstPathName, save_flags)) {
+				throw (1);
+			}
 
 			// load and check Exif raw data
 			fif = FreeImage_GetFileType(lpszDstPathName);