From 0f2cf2331820ec8645b3457a6ab86332f2048ab7 Mon Sep 17 00:00:00 2001
From: Michael Steenbeek <m.o.steenbeek@gmail.com>
Date: Thu, 10 Mar 2022 01:01:59 +0100
Subject: [PATCH] Safely look up widget for drawing

---
 src/openrct2-ui/interface/Widget.cpp | 38 +++++++++++++++++++++-------
 src/openrct2-ui/interface/Widget.h   |  1 +
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/src/openrct2-ui/interface/Widget.cpp b/src/openrct2-ui/interface/Widget.cpp
index 07532565bb..7893317c63 100644
--- a/src/openrct2-ui/interface/Widget.cpp
+++ b/src/openrct2-ui/interface/Widget.cpp
@@ -48,8 +48,14 @@ static void WidgetDrawImage(rct_drawpixelinfo* dpi, rct_window* w, rct_widgetind
  */
 void WidgetDraw(rct_drawpixelinfo* dpi, rct_window* w, rct_widgetindex widgetIndex)
 {
-    const auto& widget = w->widgets[widgetIndex];
-    switch (widget.type)
+    const auto* widget = GetWidgetByIndex(*w, widgetIndex);
+    if (widget == nullptr)
+    {
+        log_error("Tried drawing an out-of-bounds widget index!");
+        return;
+    }
+
+    switch (widget->type)
     {
         case WindowWidgetType::Frame:
             WidgetFrameDraw(dpi, w, widgetIndex);
@@ -1030,22 +1036,36 @@ void WidgetScrollGetPart(
     }
 }
 
-static void SafeSetWidgetFlag(rct_window* w, rct_widgetindex widgetIndex, WidgetFlags mask, bool value)
+rct_widget* GetWidgetByIndex(const rct_window& w, rct_widgetindex widgetIndex)
 {
     // Make sure we don't go out of bounds if we are given a bad widget index
     rct_widgetindex index = 0;
-    for (auto* widget = w->widgets; widget->type != WindowWidgetType::Last; widget++)
+    for (auto* widget = w.widgets; widget->type != WindowWidgetType::Last; widget++)
     {
         if (index == widgetIndex)
         {
-            if (value)
-                widget->flags |= mask;
-            else
-                widget->flags &= ~mask;
-            break;
+            return widget;
         }
         index++;
     }
+
+    log_error("Widget index %i out of bounds for window class %u", widgetIndex, w.classification);
+
+    return nullptr;
+}
+
+static void SafeSetWidgetFlag(rct_window* w, rct_widgetindex widgetIndex, WidgetFlags mask, bool value)
+{
+    rct_widget* widget = GetWidgetByIndex(*w, widgetIndex);
+    if (widget == nullptr)
+    {
+        return;
+    }
+
+    if (value)
+        widget->flags |= mask;
+    else
+        widget->flags &= ~mask;
 }
 
 void WidgetSetEnabled(rct_window* w, rct_widgetindex widgetIndex, bool enabled)
diff --git a/src/openrct2-ui/interface/Widget.h b/src/openrct2-ui/interface/Widget.h
index 6c6de4eac2..cd4b6d8bab 100644
--- a/src/openrct2-ui/interface/Widget.h
+++ b/src/openrct2-ui/interface/Widget.h
@@ -12,3 +12,4 @@
 #include <openrct2/interface/Widget.h>
 
 uint32_t GetColourButtonImage(colour_t colour);
+rct_widget* GetWidgetByIndex(const rct_window& w, rct_widgetindex widgetIndex);