From 9a502184007d841e3a6d738b8bce925f41e17627 Mon Sep 17 00:00:00 2001
From: KamFretoZ <14798312+kamfretoz@users.noreply.github.com>
Date: Thu, 22 Aug 2024 22:52:50 +0700
Subject: [PATCH] Image: Fix crash loading corrupted/invalid JPEG files

Backport from: https://github.com/stenzek/duckstation/commit/afea18f65ecd4d1abf06ab834deccd2d6c9223bd
---
 common/Image.cpp | 64 ++++++++++++++++++++++++++----------------------
 1 file changed, 35 insertions(+), 29 deletions(-)

diff --git a/common/Image.cpp b/common/Image.cpp
index 1807cda3e4..b84643c9c6 100644
--- a/common/Image.cpp
+++ b/common/Image.cpp
@@ -8,6 +8,7 @@
 #include "ScopedGuard.h"
 #include "StringUtil.h"
 
+#include <common/FastJmp.h>
 #include <jpeglib.h>
 #include <png.h>
 #include <webp/decode.h>
@@ -383,39 +384,40 @@ namespace
 	struct JPEGErrorHandler
 	{
 		jpeg_error_mgr err;
-		jmp_buf jbuf;
+		fastjmp_buf jbuf;
+
+		JPEGErrorHandler()
+		{
+			jpeg_std_error(&err);
+			err.error_exit = &ErrorExit;
+		}
+
+		static void ErrorExit(j_common_ptr cinfo)
+		{
+			JPEGErrorHandler* eh = (JPEGErrorHandler*)cinfo->err;
+			char msg[JMSG_LENGTH_MAX];
+			eh->err.format_message(cinfo, msg);
+			Console.ErrorFmt("libjpeg fatal error: {}", msg);
+			fastjmp_jmp(&eh->jbuf, 1);
+		}
 	};
 } // namespace
 
-static bool HandleJPEGError(JPEGErrorHandler* eh)
-{
-	jpeg_std_error(&eh->err);
-
-	eh->err.error_exit = [](j_common_ptr cinfo) {
-		JPEGErrorHandler* eh = (JPEGErrorHandler*)cinfo->err;
-		char msg[JMSG_LENGTH_MAX];
-		eh->err.format_message(cinfo, msg);
-		Console.ErrorFmt("libjpeg fatal error: {}", msg);
-		longjmp(eh->jbuf, 1);
-	};
-
-	if (setjmp(eh->jbuf) == 0)
-		return true;
-
-	return false;
-}
-
 template <typename T>
 static bool WrapJPEGDecompress(RGBA8Image* image, T setup_func)
 {
 	std::vector<u8> scanline;
+	jpeg_decompress_struct info = {};
 
-	JPEGErrorHandler err;
-	if (!HandleJPEGError(&err))
+	// NOTE: Be **very** careful not to allocate memory after calling this function.
+	// It won't get freed, because fastjmp does not unwind the stack.
+	JPEGErrorHandler errhandler;
+	if (fastjmp_set(&errhandler.jbuf) != 0)
+	{
+		jpeg_destroy_decompress(&info);
 		return false;
-
-	jpeg_decompress_struct info;
-	info.err = &err.err;
+	}
+	info.err = &errhandler.err;
 	jpeg_create_decompress(&info);
 	setup_func(info);
 
@@ -541,13 +543,17 @@ template <typename T>
 static bool WrapJPEGCompress(const RGBA8Image& image, u8 quality, T setup_func)
 {
 	std::vector<u8> scanline;
+	jpeg_compress_struct info = {};
 
-	JPEGErrorHandler err;
-	if (!HandleJPEGError(&err))
+	// NOTE: Be **very** careful not to allocate memory after calling this function.
+	// It won't get freed, because fastjmp does not unwind the stack.
+	JPEGErrorHandler errhandler;
+	if (fastjmp_set(&errhandler.jbuf) != 0)
+	{
+		jpeg_destroy_compress(&info);
 		return false;
-
-	jpeg_compress_struct info;
-	info.err = &err.err;
+	}
+	info.err = &errhandler.err;
 	jpeg_create_compress(&info);
 	setup_func(info);