com32/chain/partiter: make iterators not autofree after fin/err

This patch changes iterator behaviour to not free themselves
automatically after finished iteration or error. This allows
us to be able to always:

- check their status through ->status field
- access last valid data

It will allow simplification of pentry_mangle() function in
further commits.

Signed-off-by: Michal Soltys <soltys@ziu.info>

3 files changed, 57 insertions, 28 deletions

diff --git a/com32/chain/chain.c b/com32/chain/chain.c
index bd8cd0c6..de9e0ec6 100644
--- a/com32/chain/chain.c
+++ b/com32/chain/chain.c
@@ -117,7 +117,7 @@ static int find_by_guid(const struct guid *gpt_guid,
 	    goto ok;
 	}
 	/* disk guid doesn't match, maybe partition guid will */
-	while (pi_next(&boot_part)) {
+	while (!pi_next(&boot_part)) {
 	    if (!memcmp(&boot_part->sub.gpt.part_guid, gpt_guid, sizeof(*gpt_guid)))
 		goto ok;
 	}
@@ -149,7 +149,7 @@ static int find_by_label(const char *label, struct part_iter **_boot_part)
 	    continue;
 	}
 	/* Check for a matching partition */
-	while (pi_next(&boot_part)) {
+	while (!pi_next(&boot_part)) {
 	    if (!strcmp(label, boot_part->sub.gpt.part_label))
 		goto ok;
 	}
@@ -400,7 +400,7 @@ static int pentry_mangle(struct part_iter *_iter)
 
     memcpy(&mbr, iter->data, sizeof(struct disk_dos_mbr));
 
-    while (pi_next(&iter) && !werr) {
+    while (!pi_next(&iter) && !werr) {
 	ridx = iter->rawindex;
 	if (ridx > 4) {
 	    if (opt.hide < 2 && !opt.mbrchs)
@@ -442,7 +442,7 @@ bail:
 
 int find_dp(struct part_iter **_iter)
 {
-    struct part_iter *iter;
+    struct part_iter *iter = NULL;
     struct disk_info diskinfo;
     struct guid gpt_guid;
     uint64_t fs_lba;
@@ -506,12 +506,12 @@ int find_dp(struct part_iter **_iter)
 
 	/* 'fs' => we should lookup the syslinux partition number and use it */
 	if (!strcmp(opt.drivename, "fs")) {
-	    while (pi_next(&iter)) {
+	    while (!pi_next(&iter)) {
 		if (iter->start_lba == fs_lba)
 		    break;
 	    }
 	    /* broken part structure or other problems */
-	    if (!iter) {
+	    if (iter->status) {
 		error("Can't find myself on the drive I booted from.\n");
 		goto bail;
 	    }
@@ -530,8 +530,8 @@ int find_dp(struct part_iter **_iter)
 	do {
 	    if (iter->index == partition)
 		break;
-	} while (pi_next(&iter));
-	if (!iter) {
+	} while (!pi_next(&iter));
+	if (iter->status) {
 	    error("Requested disk / partition combination not found.\n");
 	    goto bail;
 	}
@@ -546,6 +546,7 @@ int find_dp(struct part_iter **_iter)
     return 0;
 
 bail:
+    pi_del(&iter);
     return -1;
 }
 
diff --git a/com32/chain/partiter.c b/com32/chain/partiter.c
index ef74b686..80b26fbf 100644
--- a/com32/chain/partiter.c
+++ b/com32/chain/partiter.c
@@ -331,12 +331,15 @@ static int pi_dos_next_mbr(struct part_iter *iter, uint32_t *lba,
     while (++iter->index0 < 4) {
 	dp = ((struct disk_dos_mbr *)iter->data)->table + iter->index0;
 
-	if (notsane_primary(dp))
+	if (notsane_primary(dp)) {
+	    iter->status = PI_INSANE;
 	    goto bail;
+	}
 
 	if (ost_is_ext(dp->ostype)) {
 	    if (iter->sub.dos.bebr_index0 >= 0) {
 		error("You have more than 1 extended partition.\n");
+		iter->status = PI_INSANE;
 		goto bail;
 	    }
 	    /* record base EBR index */
@@ -382,19 +385,24 @@ static int pi_dos_next_ebr(struct part_iter *iter, uint32_t *lba,
 {
     struct disk_dos_part_entry *dp;
 
-    if (prep_base_ebr(iter))
+    if (prep_base_ebr(iter)) {
+	iter->status = PI_DONE;
 	return -1;
+    }
 
     while (++iter->index0 < 1024 && iter->sub.dos.nebr_lba) {
 	free(iter->data);
 	if (!(iter->data =
 		    disk_read_sectors(&iter->di, iter->sub.dos.nebr_lba, 1))) {
 	    error("Couldn't load EBR.\n");
+	    iter->status = PI_ERRLOAD;
 	    return -1;
 	}
 
-	if (notsane_logical(iter) || notsane_extended(iter))
+	if (notsane_logical(iter) || notsane_extended(iter)) {
+	    iter->status = PI_INSANE;
 	    return -1;
+	}
 
 	dp = ((struct disk_dos_mbr *)iter->data)->table;
 
@@ -427,6 +435,7 @@ static int pi_dos_next_ebr(struct part_iter *iter, uint32_t *lba,
 	 * this place.
 	 */
     }
+    iter->status = PI_DONE;
     return -1;
 }
 
@@ -435,6 +444,9 @@ static struct part_iter *pi_dos_next(struct part_iter *iter)
     uint32_t start_lba = 0;
     struct disk_dos_part_entry *dos_part = NULL;
 
+    if (iter->status)
+	goto bail;
+
     /* look for primary partitions */
     if (iter->index0 < 4 &&
 	    pi_dos_next_mbr(iter, &start_lba, &dos_part))
@@ -465,7 +477,7 @@ static struct part_iter *pi_dos_next(struct part_iter *iter)
 
     return iter;
 bail:
-    return pi_del(&iter);
+    return NULL;
 }
 
 static void gpt_conv_label(struct part_iter *iter)
@@ -488,18 +500,24 @@ static struct part_iter *pi_gpt_next(struct part_iter *iter)
 {
     const struct disk_gpt_part_entry *gpt_part = NULL;
 
+    if (iter->status)
+	goto bail;
+
     while (++iter->index0 < iter->sub.gpt.pe_count) {
 	gpt_part = (const struct disk_gpt_part_entry *)
 	    (iter->data + iter->index0 * iter->sub.gpt.pe_size);
 
-	if (notsane_gpt(gpt_part))
+	if (notsane_gpt(gpt_part)) {
+	    iter->status = PI_INSANE;
 	    goto bail;
+	}
 
 	if (!guid_is0(&gpt_part->type) || iter->stepall)
 	    break;
     }
     /* no more partitions ? */
     if (iter->index0 == iter->sub.gpt.pe_count) {
+	iter->status = PI_DONE;
 	goto bail;
     }
     /* gpt_part is guaranteed to be valid here */
@@ -516,12 +534,13 @@ static struct part_iter *pi_gpt_next(struct part_iter *iter)
 
     return iter;
 bail:
-    return pi_del(&iter);
+    return NULL;
 }
 
 static struct part_iter *pi_raw_next(struct part_iter *iter)
 {
-    return pi_del(&iter);
+    iter->status = PI_DONE;
+    return NULL;
 }
 
 static int check_crc(uint32_t crc_match, const uint8_t *buf, unsigned int siz)
@@ -571,19 +590,23 @@ static int gpt_check_hdr_crc(const struct disk_info * const diskinfo, struct dis
  */
 
 
-struct part_iter *pi_next(struct part_iter **_iter)
+int pi_next(struct part_iter **_iter)
 {
-    struct part_iter *iter = *_iter;
+    struct part_iter *iter;
+
+    if(!_iter || !*_iter)
+	return 0;
+    iter = *_iter;
 #ifdef DEBUG
-    if (!iter)
-	return NULL;
     if (inv_type(iter->type)) {
 	error("This is not a valid iterator.\n");
-	return NULL;
+	return 0;
     }
 #endif
-    *_iter = iter->type->next(iter);
-    return *_iter;
+    if ((iter = iter->type->next(iter))) {
+	*_iter = iter;
+    }
+    return (*_iter)->status;
 }
 
 /**
@@ -638,25 +661,24 @@ bail:
  *
  **/
 
-void *pi_del(struct part_iter **_iter)
+void pi_del(struct part_iter **_iter)
 {
     struct part_iter *iter;
 
     if(!_iter || !*_iter)
-	return NULL;
+	return;
     iter = *_iter;
 
 #ifdef DEBUG
     if (inv_type(iter->type)) {
 	error("This is not a valid iterator.\n");
-	return NULL;
+	return;
     }
 #endif
 
     iter->type->dtor(iter);
     free(iter);
     *_iter = NULL;
-    return NULL;
 }
 
 /**
diff --git a/com32/chain/partiter.h b/com32/chain/partiter.h
index 22397bd4..5e17a9a4 100644
--- a/com32/chain/partiter.h
+++ b/com32/chain/partiter.h
@@ -39,6 +39,11 @@
 #include <stdint.h>
 #include <syslinux/disk.h>
 
+#define PI_ERRLOAD 3
+#define PI_INSANE 2
+#define PI_DONE 1
+#define PI_OK 0
+
 struct itertype;
 struct part_iter;
 
@@ -59,6 +64,7 @@ struct part_iter {
     int rawindex;
     struct disk_info di;
     int stepall;
+    int status;
     /* internal */
     int index0;
     union _sub {
@@ -90,8 +96,8 @@ extern const struct itertype * const typeraw;
 
 struct part_iter *pi_begin(const struct disk_info *, int stepall);
 struct part_iter *pi_new(const struct itertype *, ...);
-void *pi_del(struct part_iter **);
-struct part_iter *pi_next(struct part_iter **);
+void pi_del(struct part_iter **);
+int pi_next(struct part_iter **);
 
 #endif