From 78aab1bd135de9364ca5d936f08553d529868136 Mon Sep 17 00:00:00 2001
From: Peter Evans <pevans@activecampaign.com>
Date: Mon, 1 Jan 2018 17:11:03 -0600
Subject: [PATCH] Fix potential memory leaks in create()

---
 src/apple2.c     | 46 +++++++++++++++++++++++++++++++++++++++++++---
 src/vm_segment.c |  3 +++
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/src/apple2.c b/src/apple2.c
index 3ec7b80..e90c1dc 100644
--- a/src/apple2.c
+++ b/src/apple2.c
@@ -37,7 +37,20 @@ apple2_create(int width, int height)
         return NULL;
     }
 
+    // Forward set these to NULL in case we fail to build the machine
+    // properly; that way, we won't try to free garbage data
+    mach->sysfont = NULL;
+    mach->screen = NULL;
+    mach->drive1 = NULL;
+    mach->drive2 = NULL;
+
     mach->cpu = mos6502_create();
+    if (mach->cpu == NULL) {
+        log_critical("Could not create CPU!");
+        apple2_free(mach);
+        return NULL;
+    }
+
     mach->memory = mach->cpu->memory;
 
     // Our two drives -- we create both of them, even if we intend to
@@ -45,10 +58,17 @@ apple2_create(int width, int height)
     mach->drive1 = apple2dd_create();
     mach->drive2 = apple2dd_create();
 
+    if (mach->drive1 == NULL || mach->drive2 == NULL) {
+        log_critical("Could not create disk drives!");
+        apple2_free(mach);
+        return NULL;
+    }
+
     // Let's build our screen abstraction!
     mach->screen = vm_screen_create();
     if (mach->screen == NULL) {
-        log_critical("Screen creation failed!\n");
+        log_critical("Screen creation failed!");
+        apple2_free(mach);
         return NULL;
     }
 
@@ -56,7 +76,8 @@ apple2_create(int width, int height)
     // graphics.
     err = vm_screen_add_window(mach->screen, width, height);
     if (err != OK) {
-        log_critical("Window creation failed!\n");
+        log_critical("Window creation failed!");
+        apple2_free(mach);
         return NULL;
     }
 
@@ -72,6 +93,7 @@ apple2_create(int width, int height)
                                       7, 8,         // 7 pixels wide, 8 pixels tall
                                       0x7f);        // 7-bit values only
     if (mach->sysfont == NULL) {
+        apple2_free(mach);
         log_critical("Could not initialize apple2: bad font");
         return NULL;
     }
@@ -153,7 +175,25 @@ apple2_clear_strobe(apple2 *mach)
 void
 apple2_free(apple2 *mach)
 {
-    mos6502_free(mach->cpu);
+    if (mach->cpu) {
+        mos6502_free(mach->cpu);
+    }
+
+    if (mach->sysfont) {
+        vm_bitfont_free(mach->sysfont);
+    }
+
+    if (mach->drive1) {
+        apple2dd_free(mach->drive1);
+    }
+
+    if (mach->drive2) {
+        apple2dd_free(mach->drive2);
+    }
+
+    if (mach->screen) {
+        vm_screen_free(mach->screen);
+    }
 
     // NOTE: we do _NOT_ want to clear the memory field of mach, as it's
     // co-owned with the cpu struct that we just freed above.
diff --git a/src/vm_segment.c b/src/vm_segment.c
index 9f0c473..d895237 100644
--- a/src/vm_segment.c
+++ b/src/vm_segment.c
@@ -34,6 +34,7 @@ vm_segment_create(size_t size)
 
     segment->memory = malloc(sizeof(vm_8bit) * size);
     if (segment->memory == NULL) {
+        free(segment);
         log_critical("Couldn't allocate enough space for vm_segment");
         return NULL;
     }
@@ -45,12 +46,14 @@ vm_segment_create(size_t size)
     segment->read_table = malloc(sizeof(vm_segment_read_fn) * size);
     if (segment->read_table == NULL) {
         log_critical("Couldn't allocate enough space for segment read_table");
+        vm_segment_free(segment);
         return NULL;
     }
 
     segment->write_table = malloc(sizeof(vm_segment_write_fn) * size);
     if (segment->write_table == NULL) {
         log_critical("Couldn't allocate enough space for segment write_table");
+        vm_segment_free(segment);
         return NULL;
     }