From 9ee0e8ef2170c47f491eb2bf39593a6d0f7c2e3d Mon Sep 17 00:00:00 2001
From: glank <glankk@users.noreply.github.com>
Date: Wed, 21 Feb 2024 14:27:57 +0100
Subject: [PATCH] Add out-of-memory checks

---
 src/gz/explorer.c    |  6 ++++
 src/gz/files.c       | 25 ++++++++++-------
 src/gz/gz.c          | 67 +++++++++++++++++++++++++++++++-------------
 src/gz/gz.h          |  2 +-
 src/gz/gz_col_view.c | 59 ++++++++++++++++++++++++++++++++++++--
 src/gz/gz_command.c  | 10 +++++--
 src/gz/gz_scene.c    |  2 +-
 src/gz/sys.c         | 15 ++++++++++
 8 files changed, 149 insertions(+), 37 deletions(-)

diff --git a/src/gz/explorer.c b/src/gz/explorer.c
index 0fc68734..805642bf 100644
--- a/src/gz/explorer.c
+++ b/src/gz/explorer.c
@@ -241,6 +241,8 @@ static int draw_proc(struct menu_item *item,
       uint32_t scene_vrom_end = scene_entry->scene_vrom_end;
       uint32_t scene_vrom_size = scene_vrom_end - scene_vrom_start;
       data->scene_file = malloc(scene_vrom_size);
+      if (!data->scene_file)
+        return 0;
       zu_getfile(scene_vrom_start, data->scene_file, scene_vrom_size);
     }
     stab.seg[Z64_SEG_SCENE] = MIPS_KSEG0_TO_PHYS(data->scene_file);
@@ -260,6 +262,8 @@ static int draw_proc(struct menu_item *item,
       uint32_t room_vrom_end = room_files[data->room_index].vrom_end;
       uint32_t room_vrom_size = room_vrom_end - room_vrom_start;
       data->room_file = malloc(room_vrom_size);
+      if (!data->room_file)
+        return 0;
       zu_getfile(room_vrom_start, data->room_file, room_vrom_size);
       stab.seg[Z64_SEG_ROOM] = MIPS_KSEG0_TO_PHYS(data->room_file);
       /* populate mesh */
@@ -301,6 +305,8 @@ static int draw_proc(struct menu_item *item,
     static void *zbuf = NULL;
     if (!zbuf)
       zbuf = memalign(64, 2 * Z64_SCREEN_WIDTH * Z64_SCREEN_HEIGHT);
+    if (!zbuf)
+      return 0;
 
     static Gfx init_rcp[] = {
       /* rsp settings */
diff --git a/src/gz/files.c b/src/gz/files.c
index 4b3db70b..3f8de107 100644
--- a/src/gz/files.c
+++ b/src/gz/files.c
@@ -244,7 +244,7 @@ static int overwrite_prompt_proc(int option_index, void *data)
 static void return_path(const char *name)
 {
   char *path = malloc(PATH_MAX);
-  if (getcwd(path, PATH_MAX)) {
+  if (path && getcwd(path, PATH_MAX)) {
     int dl = strlen(path);
     int nl = strlen(name);
     if (dl + 1 + nl + gf_suffix_length < PATH_MAX) {
@@ -264,7 +264,8 @@ static void return_path(const char *name)
         menu_return(&gf_menu);
     }
   }
-  free(path);
+  if (path)
+    free(path);
 }
 
 static int file_enter_proc(struct menu_item *item,
@@ -343,15 +344,17 @@ static int file_activate_proc(struct menu_item *item)
     ds->index = index;
     int l = strlen(entry->name) - gf_suffix_length;
     char *name = malloc(l + 1);
-    memcpy(name, entry->name, l);
-    name[l] = 0;
-    if (gf_mode == GETFILE_SAVE || gf_mode == GETFILE_SAVE_PREFIX_INC) {
-      set_name(name, 1);
-      menu_select(&gf_menu, gf_accept);
+    if (name) {
+      memcpy(name, entry->name, l);
+      name[l] = 0;
+      if (gf_mode == GETFILE_SAVE || gf_mode == GETFILE_SAVE_PREFIX_INC) {
+        set_name(name, 1);
+        menu_select(&gf_menu, gf_accept);
+      }
+      else
+        return_path(name);
+      free(name);
     }
-    else
-      return_path(name);
-    free(name);
   }
   return 1;
 }
@@ -586,6 +589,8 @@ void menu_get_file(struct menu *menu, enum get_file_mode mode,
   if (gf_suffix)
     free(gf_suffix);
   gf_suffix = malloc(strlen(suffix) + 1);
+  if (!gf_suffix)
+    return;
   strcpy(gf_suffix, suffix);
   gf_suffix_length = strlen(gf_suffix);
   gf_callback_proc = callback_proc;
diff --git a/src/gz/gz.c b/src/gz/gz.c
index 3d1432a8..d1314c64 100644
--- a/src/gz/gz.c
+++ b/src/gz/gz.c
@@ -475,6 +475,13 @@ static void main_hook(void)
   gfx_flush();
 }
 
+static void macro_oom(void)
+{
+    gz_log("out of memory, stopping macro");
+    gz.frames_queued = 0;
+    gz.movie_state = MOVIE_IDLE;
+}
+
 HOOK int32_t room_load_sync_hook(OSMesgQueue *mq, OSMesg *msg, int32_t flag)
 {
   maybe_init_gp();
@@ -495,8 +502,12 @@ HOOK int32_t room_load_sync_hook(OSMesgQueue *mq, OSMesg *msg, int32_t flag)
         rl = vector_insert(&gz.movie_room_load,
                            gz.movie_room_load_pos, 1, NULL);
       }
-      rl->frame_idx = gz.movie_frame;
-      ++gz.movie_room_load_pos;
+      if (rl) {
+        rl->frame_idx = gz.movie_frame;
+        ++gz.movie_room_load_pos;
+      }
+      else
+        macro_oom();
       gz.room_load_flag = 1;
     }
     return result;
@@ -616,12 +627,16 @@ HOOK void input_hook(void)
           vector_reserve(&gz.movie_input, 128);
         vector_push_back(&gz.movie_input, 1, NULL);
       }
-      /* if the last recorded frame is not the previous frame,
-         increment the rerecord count */
-      if (gz.movie_last_recorded_frame >= gz.movie_frame)
-        ++gz.movie_rerecords;
-      gz.movie_last_recorded_frame = gz.movie_frame++;
-      z_to_movie(gz.movie_last_recorded_frame, &zi[0], gz.reset_flag);
+      if (gz.movie_frame < gz.movie_input.size) {
+        /* if the last recorded frame is not the previous frame,
+           increment the rerecord count */
+        if (gz.movie_last_recorded_frame >= gz.movie_frame)
+          ++gz.movie_rerecords;
+        gz.movie_last_recorded_frame = gz.movie_frame++;
+        z_to_movie(gz.movie_last_recorded_frame, &zi[0], gz.reset_flag);
+      }
+      else
+        macro_oom();
     }
     else if (gz.movie_state == MOVIE_PLAYING) {
       if (gz.movie_frame >= gz.movie_input.size) {
@@ -765,10 +780,14 @@ HOOK void srand_hook(uint32_t seed)
       /* insert a recorded seed */
       struct movie_seed *ms;
       ms = vector_insert(&gz.movie_seed, gz.movie_seed_pos, 1, NULL);
-      ms->frame_idx = gz.movie_frame;
-      ms->old_seed = z64_random;
-      ms->new_seed = seed;
-      ++gz.movie_seed_pos;
+      if (ms) {
+        ms->frame_idx = gz.movie_frame;
+        ms->old_seed = z64_random;
+        ms->new_seed = seed;
+        ++gz.movie_seed_pos;
+      }
+      else
+        macro_oom();
     }
     else if (gz.movie_state == MOVIE_PLAYING) {
       /* restore a recorded seed, if conditions match */
@@ -842,9 +861,13 @@ HOOK void ocarina_update_hook(void)
           os = vector_insert(&gz.movie_oca_sync,
                              gz.movie_oca_sync_pos, 1, NULL);
         }
-        os->frame_idx = gz.movie_frame;
-        os->audio_frames = audio_frames;
-        ++gz.movie_oca_sync_pos;
+        if (os) {
+          os->frame_idx = gz.movie_frame;
+          os->audio_frames = audio_frames;
+          ++gz.movie_oca_sync_pos;
+        }
+        else
+          macro_oom();
         gz.oca_sync_flag = 1;
       }
     }
@@ -922,11 +945,15 @@ HOOK void ocarina_input_hook(void *a0, z64_input_t *input, int a2)
           oi = vector_insert(&gz.movie_oca_input,
                              gz.movie_oca_input_pos, 1, NULL);
         }
-        oi->frame_idx = gz.movie_frame;
-        oi->pad = input->raw.pad;
-        oi->adjusted_x = input->adjusted_x;
-        oi->adjusted_y = input->adjusted_y;
-        ++gz.movie_oca_input_pos;
+        if (oi) {
+          oi->frame_idx = gz.movie_frame;
+          oi->pad = input->raw.pad;
+          oi->adjusted_x = input->adjusted_x;
+          oi->adjusted_y = input->adjusted_y;
+          ++gz.movie_oca_input_pos;
+        }
+        else
+          macro_oom();
         gz.oca_input_flag = 1;
       }
     }
diff --git a/src/gz/gz.h b/src/gz/gz.h
index 4369d7b5..0c930648 100644
--- a/src/gz/gz.h
+++ b/src/gz/gz.h
@@ -55,8 +55,8 @@ enum holl_view_state
   HOLLVIEW_INACTIVE,
   HOLLVIEW_START,
   HOLLVIEW_ACTIVE,
-  HOLLVIEW_BEGIN_STOP,
   HOLLVIEW_STOP,
+  HOLLVIEW_STOPPING,
 };
 
 enum path_view_state
diff --git a/src/gz/gz_col_view.c b/src/gz/gz_col_view.c
index 452a2efe..5873e032 100644
--- a/src/gz/gz_col_view.c
+++ b/src/gz/gz_col_view.c
@@ -937,6 +937,11 @@ void gz_col_view(void)
 
     /* allocate static polygon display list */
     stc_poly = malloc(sizeof(*stc_poly) * stc_poly_cap);
+    if (!stc_poly) {
+      gz.col_view_state = COLVIEW_STOPPING;
+      gz_log("out of memory, stopping colview");
+      return;
+    }
     Gfx *stc_poly_p = stc_poly;
     Gfx *stc_poly_d = stc_poly + stc_poly_cap;
     poly_writer_t *p_poly_writer = &poly_writer;
@@ -949,6 +954,11 @@ void gz_col_view(void)
     struct vector line_set;
     if (col_view_line) {
       stc_line = malloc(sizeof(*stc_line) * stc_line_cap);
+      if (!stc_line) {
+        gz.col_view_state = COLVIEW_STOPPING;
+        gz_log("out of memory, stopping colview");
+        return;
+      }
       stc_line_p = stc_line;
       stc_line_d = stc_line + stc_line_cap;
       p_line_writer = &line_writer;
@@ -960,9 +970,21 @@ void gz_col_view(void)
     dyn_poly_buf[0] = malloc(sizeof(*dyn_poly_buf[0]) * dyn_poly_cap);
     dyn_poly_buf[1] = malloc(sizeof(*dyn_poly_buf[1]) * dyn_poly_cap);
 
+    if (!dyn_poly_buf[0] || !dyn_poly_buf[1]) {
+      gz.col_view_state = COLVIEW_STOPPING;
+      gz_log("out of memory, stopping colview");
+      return;
+    }
+
     if (col_view_line) {
       dyn_line_buf[0] = malloc(sizeof(*dyn_line_buf[0]) * dyn_line_cap);
       dyn_line_buf[1] = malloc(sizeof(*dyn_line_buf[1]) * dyn_line_cap);
+
+      if (!dyn_line_buf[0] || !dyn_line_buf[1]) {
+        gz.col_view_state = COLVIEW_STOPPING;
+        gz_log("out of memory, stopping colview");
+        return;
+      }
     }
 
     /* generate static display lists */
@@ -1179,6 +1201,12 @@ void gz_hit_view(void)
     hit_gfx_buf[0] = malloc(sizeof(*hit_gfx_buf[0]) * hit_gfx_cap);
     hit_gfx_buf[1] = malloc(sizeof(*hit_gfx_buf[1]) * hit_gfx_cap);
 
+    if (!hit_gfx_buf[0] || !hit_gfx_buf[1]) {
+      gz.hit_view_state = HITVIEW_STOPPING;
+      gz_log("out of memory, stopping hitview");
+      return;
+    }
+
     gz.hit_view_state = HITVIEW_ACTIVE;
   }
   if (enable && gz.hit_view_state == HITVIEW_ACTIVE) {
@@ -1291,11 +1319,23 @@ void gz_path_view(void)
     if (path_view_points) {
       poly_gfx_buf[0] = malloc(sizeof(*poly_gfx_buf[0]) * poly_gfx_cap);
       poly_gfx_buf[1] = malloc(sizeof(*poly_gfx_buf[1]) * poly_gfx_cap);
+
+      if (!poly_gfx_buf[0] || !poly_gfx_buf[1]) {
+        gz.path_view_state = PATHVIEW_STOPPING;
+        gz_log("out of memory, stopping pathview");
+        return;
+      }
     }
 
     if (path_view_lines) {
       line_gfx_buf[0] = malloc(sizeof(*line_gfx_buf[0]) * line_gfx_cap);
       line_gfx_buf[1] = malloc(sizeof(*line_gfx_buf[1]) * line_gfx_cap);
+
+      if (!line_gfx_buf[0] || !line_gfx_buf[1]) {
+        gz.path_view_state = PATHVIEW_STOPPING;
+        gz_log("out of memory, stopping pathview");
+        return;
+      }
     }
 
     gz.path_view_state = PATHVIEW_ACTIVE;
@@ -1477,12 +1517,19 @@ void gz_cull_view(void)
     cull_gfx_buf[0] = malloc(sizeof(*cull_gfx_buf[0]) * cull_gfx_cap);
     cull_gfx_buf[1] = malloc(sizeof(*cull_gfx_buf[1]) * cull_gfx_cap);
 
+    if (!cull_gfx_buf[0] || !cull_gfx_buf[1]) {
+      gz.cull_view_state = CULLVIEW_STOPPING;
+      gz_log("out of memory, stopping cullview");
+      return;
+    }
+
     gz.cull_view_state = CULLVIEW_ACTIVE;
   }
   if (enable && gz.cull_view_state == CULLVIEW_ACTIVE) {
     /* If no actor is selected, stop */
     if (!gz.selected_actor.ptr) {
       gz.cull_view_state = CULLVIEW_STOP;
+      gz_log("out of memory, stopping cullview");
       return;
     }
     /* Now look through actor type list */
@@ -1623,6 +1670,12 @@ void gz_holl_view(void)
     holl_gfx_buf[0] = malloc(sizeof(*holl_gfx_buf[0]) * holl_gfx_cap);
     holl_gfx_buf[1] = malloc(sizeof(*holl_gfx_buf[1]) * holl_gfx_cap);
 
+    if (!holl_gfx_buf[0] || !holl_gfx_buf[1]) {
+      gz.holl_view_state = HOLLVIEW_STOPPING;
+      gz_log("out of memory, stopping hollview");
+      return;
+    }
+
     gz.holl_view_state = HOLLVIEW_ACTIVE;
   }
   if (enable && gz.holl_view_state == HOLLVIEW_ACTIVE) {
@@ -1869,9 +1922,9 @@ void gz_holl_view(void)
     gSPDisplayList((*p_gfx_p)++, holl_gfx);
   }
 
-  if (gz.holl_view_state == HOLLVIEW_BEGIN_STOP)
-    gz.holl_view_state = HOLLVIEW_STOP;
-  else if (gz.holl_view_state == HOLLVIEW_STOP) {
+  if (gz.holl_view_state == HOLLVIEW_STOP)
+    gz.holl_view_state = HOLLVIEW_STOPPING;
+  else if (gz.holl_view_state == HOLLVIEW_STOPPING) {
     release_mem(&holl_gfx_buf[0]);
     release_mem(&holl_gfx_buf[1]);
 
diff --git a/src/gz/gz_command.c b/src/gz/gz_command.c
index 6f0b9a95..7599f2cb 100644
--- a/src/gz/gz_command.c
+++ b/src/gz/gz_command.c
@@ -214,9 +214,15 @@ void command_savestate(void)
   if (!zu_in_game())
     gz_log("can not save here");
   else {
-    if (gz.state_buf[gz.state_slot])
+    if (gz.state_buf[gz.state_slot]) {
       free(gz.state_buf[gz.state_slot]);
+      gz.state_buf[gz.state_slot] = NULL;
+    }
     struct state_meta *state = malloc(368 * 1024);
+    if (!state) {
+      gz_log("out of memory, can't save");
+      return;
+    }
     state->z64_version = Z64_VERSION;
     state->state_version = SETTINGS_STATE_VERSION;
     state->scene_idx = z64_game.scene_index;
@@ -361,7 +367,7 @@ void command_hollview(void)
   if (gz.holl_view_state == HOLLVIEW_INACTIVE)
     gz.holl_view_state = HOLLVIEW_START;
   else
-    gz.holl_view_state = HOLLVIEW_BEGIN_STOP;
+    gz.holl_view_state = HOLLVIEW_STOP;
 }
 
 void command_resetlag(void)
diff --git a/src/gz/gz_scene.c b/src/gz/gz_scene.c
index 25ce5289..cce7ac29 100644
--- a/src/gz/gz_scene.c
+++ b/src/gz/gz_scene.c
@@ -96,7 +96,7 @@ static int holl_view_proc(struct menu_item *item,
   }
   else if (reason == MENU_CALLBACK_SWITCH_OFF) {
     if (gz.holl_view_state != HOLLVIEW_INACTIVE)
-      gz.holl_view_state = HOLLVIEW_BEGIN_STOP;
+      gz.holl_view_state = HOLLVIEW_STOP;
   }
   else if (reason == MENU_CALLBACK_THINK) {
     _Bool state = gz.holl_view_state == HOLLVIEW_START ||
diff --git a/src/gz/sys.c b/src/gz/sys.c
index cfd74b44..9b3495d0 100644
--- a/src/gz/sys.c
+++ b/src/gz/sys.c
@@ -2,6 +2,7 @@
 #include <stdio.h>
 #include <stdarg.h>
 #include <string.h>
+#include <stdint.h>
 #include "io.h"
 #include "sys.h"
 #include "fat.h"
@@ -718,3 +719,17 @@ void sys_reset(void)
     wd = 0;
   }
 }
+
+void *sbrk(intptr_t incr)
+{
+  extern char end[];
+  static void *brk = end;
+  if ((uintptr_t)brk + incr > 0x80800000) {
+    return (void *)-1;
+  }
+  else {
+    char *ret = brk;
+    brk += incr;
+    return ret;
+  }
+}