From 3a3c4f60bf53d05566177759ee8992889afac3b8 Mon Sep 17 00:00:00 2001
From: jdgleaver <james@leaver.myzen.co.uk>
Date: Thu, 30 Apr 2020 11:25:53 +0100
Subject: [PATCH] (task_screenshot) Fix heap-use-after-free error when widgets
 are disabled

---
 tasks/task_screenshot.c | 82 +++++++++++++++++++++++++++++------------
 1 file changed, 59 insertions(+), 23 deletions(-)

diff --git a/tasks/task_screenshot.c b/tasks/task_screenshot.c
index f1a34deb3a..9a0bf713d6 100644
--- a/tasks/task_screenshot.c
+++ b/tasks/task_screenshot.c
@@ -146,31 +146,27 @@ static bool screenshot_dump_direct(screenshot_task_state_t *state)
  **/
 static void task_screenshot_handler(retro_task_t *task)
 {
-   screenshot_task_state_t *state = (screenshot_task_state_t*)task->state;
+   screenshot_task_state_t *state = NULL;
    bool ret                       = false;
 
-   if (task_get_progress(task) == 100)
-   {
-      task_set_finished(task, true);
-
-      if (task->title)
-         task_free_title(task);
-
-      if (state->userbuf)
-         free(state->userbuf);
-
-#if defined(HAVE_GFX_WIDGETS)
-      /* If display widgets are enabled, state is freed
-         in the callback after the notification
-         is displayed */
-      if (!state->widgets_ready)
-#endif
-         free(state);
+   if (!task)
       return;
-   }
 
+   state = (screenshot_task_state_t*)task->state;
+
+   if (!state)
+      goto task_finished;
+
+   if (task_get_cancelled(task))
+      goto task_finished;
+
+   if (task_get_progress(task) == 100)
+      goto task_finished;
+
+   /* Take screenshot */
    ret = screenshot_dump_direct(state);
 
+   /* Push screenshot to image history playlist */
 #ifdef HAVE_IMAGEVIEWER
    if (  ret                        &&
          !state->silence            &&
@@ -193,6 +189,7 @@ static void task_screenshot_handler(retro_task_t *task)
 
    task_set_progress(task, 100);
 
+   /* Report any errors */
    if (!ret)
    {
       char *msg = strdup(msg_hash_to_str(MSG_FAILED_TO_TAKE_SCREENSHOT));
@@ -202,6 +199,32 @@ static void task_screenshot_handler(retro_task_t *task)
 
    if (task->title)
       task_free_title(task);
+
+   return;
+
+task_finished:
+
+   task_set_finished(task, true);
+
+   if (task->title)
+      task_free_title(task);
+
+   if (state && state->userbuf)
+      free(state->userbuf);
+
+#if defined(HAVE_GFX_WIDGETS)
+   /* If display widgets are enabled, state is freed
+      in the callback after the notification
+      is displayed */
+   if (state && !state->widgets_ready)
+#endif
+   {
+      free(state);
+      /* Must explicitly set task->state to NULL here,
+       * to avoid potential heap-use-after-free errors */
+      state       = NULL;
+      task->state = NULL;
+   }
 }
 
 #if defined(HAVE_GFX_WIDGETS)
@@ -209,15 +232,24 @@ static void task_screenshot_callback(retro_task_t *task,
       void *task_data,
       void *user_data, const char *error)
 {
-   screenshot_task_state_t *state = (screenshot_task_state_t*)task->state;
+   screenshot_task_state_t *state = NULL;
 
-   if (!state->widgets_ready)
+   if (!task)
       return;
 
-   if (state && !state->silence && state->widgets_ready)
+   state = (screenshot_task_state_t*)task->state;
+
+   if (!state)
+      return;
+
+   if (!state->silence && state->widgets_ready)
       gfx_widget_screenshot_taken(state->shotname, state->filename);
 
    free(state);
+   /* Must explicitly set task->state to NULL here,
+    * to avoid potential heap-use-after-free errors */
+   state       = NULL;
+   task->state = NULL;
 }
 #endif
 
@@ -340,8 +372,12 @@ static bool screenshot_dump(
       task->type        = TASK_TYPE_BLOCKING;
       task->state       = state;
       task->handler     = task_screenshot_handler;
+      task->mute        = savestate;
 #if defined(HAVE_GFX_WIDGETS)
-      task->callback    = task_screenshot_callback;
+      /* This callback is only required when
+       * widgets are enabled */
+      task->callback    = state->widgets_ready ?
+            task_screenshot_callback : NULL;
       if (state->widgets_ready && !savestate)
          task_free_title(task);
       else