Fix memory leaks in cgpt

DriveOpen mallocs a couple of buffers, but DriveClose only freed them if a
particular flag was passed causing GptSave to be called. Move the free calls
out of GptSave so that the buffers are always freed, and add DriveClose calls
to a couple of cgpt functions that are missing them.

BUG=chromium:940543
TEST=precq passes, manual testing with cgpt/valgrind shows that memory leaks
     for cgpt find, show are fixed.
BRANCH=none

Change-Id: I58aeddfa6b8b4715ba4f8e064e95a660371a01c9
Reviewed-on: https://chromium-review.googlesource.com/1516413
Commit-Ready: Fletcher Woodruff <fletcherw@chromium.org>
Tested-by: Fletcher Woodruff <fletcherw@chromium.org>
Reviewed-by: Julius Werner <jwerner@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>

4 files changed, 19 insertions, 19 deletions

diff --git a/cgpt/cgpt_common.c b/cgpt/cgpt_common.c
index a5550b4c..c129d92d 100644
--- a/cgpt/cgpt_common.c
+++ b/cgpt/cgpt_common.c
@@ -286,18 +286,6 @@ static int GptSave(struct drive *drive) {
     }
   }
 
-  if (drive->gpt.primary_header)
-    free(drive->gpt.primary_header);
-  drive->gpt.primary_header = 0;
-  if (drive->gpt.primary_entries)
-    free(drive->gpt.primary_entries);
-  drive->gpt.primary_entries = 0;
-  if (drive->gpt.secondary_header)
-    free(drive->gpt.secondary_header);
-  drive->gpt.secondary_header = 0;
-  if (drive->gpt.secondary_entries)
-    free(drive->gpt.secondary_entries);
-  drive->gpt.secondary_entries = 0;
   return errors ? -1 : 0;
 }
 
@@ -388,6 +376,15 @@ int DriveClose(struct drive *drive, int update_as_needed) {
     }
   }
 
+  free(drive->gpt.primary_header);
+  drive->gpt.primary_header = NULL;
+  free(drive->gpt.primary_entries);
+  drive->gpt.primary_entries = NULL;
+  free(drive->gpt.secondary_header);
+  drive->gpt.secondary_header = NULL;
+  free(drive->gpt.secondary_entries);
+  drive->gpt.secondary_entries = NULL;
+
   // Sync early! Only sync file descriptor here, and leave the whole system sync
   // outside cgpt because whole system sync would trigger tons of disk accesses
   // and timeout tests.
diff --git a/cgpt/cgpt_legacy.c b/cgpt/cgpt_legacy.c
index d5df34d5..cf84821c 100644
--- a/cgpt/cgpt_legacy.c
+++ b/cgpt/cgpt_legacy.c
@@ -23,7 +23,7 @@ int CgptLegacy(CgptLegacyParams *params) {
   if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
     Error("GptSanityCheck() returned %d: %s\n",
           gpt_retval, GptError(gpt_retval));
-    return CGPT_FAILED;
+    goto bad;
   }
 
   h1 = (GptHeader *)drive.gpt.primary_header;
@@ -40,7 +40,7 @@ int CgptLegacy(CgptLegacyParams *params) {
         !(drive.gpt.valid_entries & MASK_SECONDARY) ||
         drive.gpt.ignored & MASK_SECONDARY) {
       Error("Refusing to mark primary GPT ignored unless secondary is valid.");
-      return CGPT_FAILED;
+      goto bad;
     }
     memset(h1, 0, sizeof(*h1));
     memcpy(h1->signature, GPT_HEADER_SIGNATURE_IGNORED,
@@ -58,4 +58,8 @@ int CgptLegacy(CgptLegacyParams *params) {
 
   // Write it all out
   return DriveClose(&drive, 1);
+
+bad:
+  (void) DriveClose(&drive, 0);
+  return CGPT_FAILED;
 }
diff --git a/cgpt/cgpt_prioritize.c b/cgpt/cgpt_prioritize.c
index c4e812c2..b59005f7 100644
--- a/cgpt/cgpt_prioritize.c
+++ b/cgpt/cgpt_prioritize.c
@@ -113,11 +113,12 @@ int CgptPrioritize(CgptPrioritizeParams *params) {
   if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
     Error("GptSanityCheck() returned %d: %s\n",
           gpt_retval, GptError(gpt_retval));
-    return CGPT_FAILED;
+    goto bad;
   }
 
   if (CGPT_OK != CheckValid(&drive)) {
     Error("please run 'cgpt repair' before reordering the priority.\n");
+    (void) DriveClose(&drive, 0);
     return CGPT_OK;
   }
 
diff --git a/cgpt/cgpt_show.c b/cgpt/cgpt_show.c
index 2b882778..7240a6a7 100644
--- a/cgpt/cgpt_show.c
+++ b/cgpt/cgpt_show.c
@@ -385,9 +385,7 @@ int CgptShow(CgptShowParams *params) {
                            params->drive_size))
     return CGPT_FAILED;
 
-  if (GptShow(&drive, params))
-    return CGPT_FAILED;
-
+  int ret = GptShow(&drive, params);
   DriveClose(&drive, 0);
-  return CGPT_OK;
+  return ret;
 }