chain, partiter, options: bugfix, new options

- small bugfix in partiter.c
- new options: hideall, mbrchs
- new partiter feature - rawindex - counts all (inc. empty) partitions in all
  cases

hideall is expansion of hide, that hides all partitions in a chain, not
only primary ones

mbrchs walks over all partitions and fixes chs values in mbr entries

rawindex keeps count of all partitions, including holes in EBRs

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

5 files changed, 158 insertions, 38 deletions

diff --git a/com32/chain/chain.c b/com32/chain/chain.c
index 7eed8808..19351f4a 100644
--- a/com32/chain/chain.c
+++ b/com32/chain/chain.c
@@ -288,7 +288,7 @@ enomem:
     error("Out of memory\n");
     return;
 }
-
+#if 0
 static void hide_unhide(const struct part_iter *_iter)
 {
     int i;
@@ -333,6 +333,111 @@ static void hide_unhide(const struct part_iter *_iter)
 bail:
     free(mbr);
 }
+#endif
+static int pem_sethide(struct disk_dos_part_entry *dp, int midx, int idx)
+{
+    static const uint16_t mask =
+	(1 << 0x01) | (1 << 0x04) | (1 << 0x06) |
+	(1 << 0x07) | (1 << 0x0b) | (1 << 0x0c) | (1 << 0x0e);
+    uint8_t t;
+
+    t = dp->ostype;
+    if ((t <= 0x1f) && ((mask >> (t & ~0x10u)) & 1)) {
+	/* It's a hideable partition type */
+	if (midx == idx)
+	    t &= (uint8_t)(~0x10u);	/* unhide */
+	else
+	    t |= 0x10u;	/* hide */
+    }
+    if (t != dp->ostype) {
+	dp->ostype = t;
+	return -1;
+    }
+    return 0;
+}
+
+static int pem_setchs(const struct disk_info *di,
+		     struct disk_dos_part_entry *dp)
+{
+    uint32_t ochs1, ochs2;
+
+    ochs1 = *(uint32_t *)dp->start;
+    ochs2 = *(uint32_t *)dp->end;
+
+    *(uint32_t *)dp->start =
+	lba2chs(di, dp->start_lba) |
+	(*(uint32_t *)dp->start & 0xFF000000);
+
+    *(uint32_t *)dp->end =
+	lba2chs(di, dp->start_lba + dp->length - 1) |
+	(*(uint32_t *)dp->end & 0xFF000000);
+
+    return
+	*(uint32_t *)dp->start != ochs1 ||
+	*(uint32_t *)dp->end != ochs2;
+}
+
+static int pe_mangle(const struct part_iter *_iter)
+{
+    int wb = 0, werr = 0;
+    uint32_t mbr_lba = 0;
+    struct part_iter *iter = NULL;
+    struct disk_dos_part_entry *dp;
+    struct disk_dos_mbr mbr;
+    int ridx;
+
+    if (_iter->type != typedos) {
+	error("Partition entry mangling ('hide[all]', 'mbrchs')\n"
+	      "is meaningful only for legacy partition scheme.");
+	goto bail;
+    }
+    if ((_iter->index < 1 || _iter->index > 4) && opt.hide == 1)
+	error("WARNING: option 'hide' specified with a non-primary partition.\n");
+
+    if (!(iter = pi_begin(&_iter->di, 1)))  /* turn on stepall */
+	goto bail;
+
+    memcpy(&mbr, iter->data, sizeof(struct disk_dos_mbr));
+
+    while (pi_next(&iter) && !werr) {
+	ridx = iter->rawindex;
+	if (ridx > 4) {
+	    if (opt.hide < 2 && !opt.mbrchs)
+		break;
+	    if (wb && !werr) {
+		werr |= disk_write_verify_sector(&_iter->di, mbr_lba, &mbr);
+		wb = false;
+	    }
+	    memcpy(&mbr, iter->data, sizeof(struct disk_dos_mbr));
+	    mbr_lba = iter->sub.dos.mbr_lba;
+	    dp = mbr.table;
+	} else
+	    dp = mbr.table + ridx - 1;
+	if (opt.hide == 2 ||
+	    (opt.hide == 1 && ridx <= 4)) {
+	    wb |= pem_sethide(dp, _iter->index, iter->index);
+	    if (_iter->index == iter->index) {
+		((struct disk_dos_part_entry *)_iter->record)->ostype =
+		    dp->ostype;
+	    }
+	}
+	if (opt.mbrchs) {
+	    wb |= pem_setchs(&_iter->di, dp);
+	    if (ridx > 4)
+		wb |= pem_setchs(&_iter->di, mbr.table + 1);
+	}
+    }
+    /* last write */
+    if (wb && !werr)
+	werr |= disk_write_verify_sector(&_iter->di, mbr_lba, &mbr);
+
+bail:
+    pi_del(&iter);
+    if (werr)
+	error("WARNING: failed to write E/MBR for partition\n"
+	      "mangling options ('hide[all]', 'mbrchs')\n");
+    return 0;
+}
 
 int find_dp(struct part_iter **_iter)
 {
@@ -605,9 +710,10 @@ int main(int argc, char *argv[])
     /* DOS kernels want the drive number in BL instead of DL. Indulge them. */
     opt.regs.ebx.b[0] = opt.regs.edx.b[0] = (uint8_t)iter->di.disk;
 
-    /* Do hide / unhide if appropriate */
-    if (opt.hide)
-	hide_unhide(iter);
+    /* Perform initial partition entry mangling */
+    if (opt.hide || opt.mbrchs)
+	pe_mangle(iter);
+/*	hide_unhide(iter);*/
 
     /* Load the boot file */
     if (opt.file) {
diff --git a/com32/chain/chain.h b/com32/chain/chain.h
index 7a63022a..7e303b8a 100644
--- a/com32/chain/chain.h
+++ b/com32/chain/chain.h
@@ -25,13 +25,14 @@ struct options {
     bool hand;
     bool hptr;
     bool swap;
-    bool hide;
+    int hide;
     bool sethid;
     bool setgeo;
     bool setdrv;
     bool sect;
     bool save;
     bool filebpb;
+    bool mbrchs;
     bool warn;
     uint16_t keeppxe;
     struct syslinux_rm_regs regs;
diff --git a/com32/chain/options.c b/com32/chain/options.c
index 35141248..f12a9755 100644
--- a/com32/chain/options.c
+++ b/com32/chain/options.c
@@ -79,6 +79,8 @@ Usage:\n\
     nohptr               Force ds:si and ds:bp to point to handover area\n\
     noswap               Swap drive numbers, if bootdisk is not fd0/hd0\n\
     nohide               Hide primary partitions, unhide selected partition\n\
+    nohideall            Hide *all* partitions, unhide selected partition\n\
+    nombrchs             Walk *all* partitions and fix E/MBRs' chs values\n\
     nokeeppxe            Keep the PXE and UNDI stacks in memory (PXELINUX)\n\
     nowarn               Wait for a keypress to continue chainloading\n\
                          - useful to see emited warnings\n\
@@ -238,10 +240,13 @@ int parse_args(int argc, char *argv[])
 	    opt.swap = true;
 	} else if (!strcmp(argv[i], "noswap")) {
 	    opt.swap = false;
+	} else if (!strcmp(argv[i], "nohide") ||
+		    !strcmp(argv[i], "nohideall")) {
+	    opt.hide = 0;
 	} else if (!strcmp(argv[i], "hide")) {
-	    opt.hide = true;
-	} else if (!strcmp(argv[i], "nohide")) {
-	    opt.hide = false;
+	    opt.hide = 1;
+	} else if (!strcmp(argv[i], "hideall")) {
+	    opt.hide = 2;
 	} else if (!strcmp(argv[i], "sethid") ||
 		   !strcmp(argv[i], "sethidden")) {
 	    opt.sethid = true;
@@ -301,6 +306,10 @@ int parse_args(int argc, char *argv[])
 	    opt.filebpb = true;
 	} else if (!strcmp(argv[i], "nofilebpb")) {
 	    opt.filebpb = false;
+	} else if (!strcmp(argv[i], "mbrchs")) {
+	    opt.mbrchs = true;
+	} else if (!strcmp(argv[i], "nombrchs")) {
+	    opt.mbrchs = false;
 	} else if (!strcmp(argv[i], "warn")) {
 	    opt.warn = true;
 	} else if (!strcmp(argv[i], "nowarn")) {
diff --git a/com32/chain/partiter.c b/com32/chain/partiter.c
index 0d1ea5e7..0fc542b6 100644
--- a/com32/chain/partiter.c
+++ b/com32/chain/partiter.c
@@ -111,6 +111,7 @@ static int iter_ctor(struct part_iter *iter, va_list *args)
 
     memcpy(&iter->di, di, sizeof(struct disk_info));
     iter->stepall = stepall;
+    iter->index0 = -1;
 
     return 0;
 }
@@ -155,7 +156,6 @@ static int iter_dos_ctor(struct part_iter *iter, va_list *args)
 
     memcpy(iter->data, mbr, sizeof(struct disk_dos_mbr));
 
-    iter->sub.dos.index0 = -1;
     iter->sub.dos.bebr_index0 = -1;
     iter->sub.dos.disk_sig = mbr->disk_sig;
 
@@ -206,7 +206,6 @@ static int iter_gpt_ctor(struct part_iter *iter, va_list *args)
 
     memcpy(iter->data, gptl, (size_t)siz);
 
-    iter->sub.gpt.index0 = -1;
     iter->sub.gpt.pe_count = (int)gpth->part_count;
     iter->sub.gpt.pe_size = (int)gpth->part_size;
     memcpy(&iter->sub.gpt.disk_guid, &gpth->disk_guid, sizeof(struct guid));
@@ -329,8 +328,8 @@ static int pi_dos_next_mbr(struct part_iter *iter, uint32_t *lba,
 {
     struct disk_dos_part_entry *dp;
 
-    while (++iter->sub.dos.index0 < 4) {
-	dp = ((struct disk_dos_mbr *)iter->data)->table + iter->sub.dos.index0;
+    while (++iter->index0 < 4) {
+	dp = ((struct disk_dos_mbr *)iter->data)->table + iter->index0;
 
 	if (notsane_primary(dp))
 	    goto bail;
@@ -341,7 +340,7 @@ static int pi_dos_next_mbr(struct part_iter *iter, uint32_t *lba,
 		goto bail;
 	    }
 	    /* record base EBR index */
-	    iter->sub.dos.bebr_index0 = iter->sub.dos.index0;
+	    iter->sub.dos.bebr_index0 = iter->index0;
 	}
 	if (!ost_is_nondata(dp->ostype) || iter->stepall) {
 	    *lba = dp->start_lba;
@@ -370,7 +369,7 @@ static int prep_base_ebr(struct part_iter *iter)
 	iter->sub.dos.ebr_start = 0;
 	iter->sub.dos.ebr_size = iter->sub.dos.bebr_size;
 
-	iter->sub.dos.index0--;
+	iter->index0--;
     }
     return 0;
 }
@@ -382,13 +381,14 @@ static int pi_dos_next_ebr(struct part_iter *iter, uint32_t *lba,
     uint32_t abs_ebr;
 
     if (prep_base_ebr(iter))
-	goto bail;
+	return -1;
 
-    if(++iter->sub.dos.index0 >= 1024)
+#if 0
+    if(++iter->index0 >= 1024)
 	/* that's one paranoid upper bound */
 	goto bail;
-
-    while (iter->sub.dos.ebr_size) {
+#endif
+    while (++iter->index0 < 1024 && iter->sub.dos.ebr_size) {
 
 	abs_ebr = iter->sub.dos.bebr_start + iter->sub.dos.ebr_start;
 
@@ -396,11 +396,11 @@ static int pi_dos_next_ebr(struct part_iter *iter, uint32_t *lba,
 	free(iter->data);
 	if (!(iter->data = disk_read_sectors(&iter->di, abs_ebr, 1))) {
 	    error("Couldn't load EBR.\n");
-	    goto bail;
+	    return -1;
 	}
 
 	if (notsane_logical(iter) || notsane_extended(iter))
-	    goto bail;
+	    return -1;
 
 	iter->sub.dos.mbr_lba = abs_ebr;
 	dp = ((struct disk_dos_mbr *)iter->data)->table;
@@ -414,10 +414,13 @@ static int pi_dos_next_ebr(struct part_iter *iter, uint32_t *lba,
 	    iter->sub.dos.ebr_size = 0;
 	}
 
+	if (!dp[0].ostype)
+	    iter->sub.dos.skipcnt++;
+
 	if (dp[0].ostype || iter->stepall) {
 	    *lba = abs_ebr;
 	    *_dp = dp;
-	    break;
+	    return 0;
 	}
 	/*
 	 * This way it's possible to continue, if some crazy soft left a "hole"
@@ -427,8 +430,10 @@ static int pi_dos_next_ebr(struct part_iter *iter, uint32_t *lba,
 	 * this place.
 	 */
     }
+#if 0
     return 0;
 bail:
+#endif
     return -1;
 }
 
@@ -438,28 +443,26 @@ static struct part_iter *pi_dos_next(struct part_iter *iter)
     struct disk_dos_part_entry *dos_part = NULL;
 
     /* look for primary partitions */
-    if (iter->sub.dos.index0 < 4 &&
+    if (iter->index0 < 4 &&
 	    pi_dos_next_mbr(iter, &start_lba, &dos_part))
 	goto bail;
 
     /* look for logical partitions */
-    if (iter->sub.dos.index0 >= 4 &&
+    if (iter->index0 >= 4 &&
 	    pi_dos_next_ebr(iter, &start_lba, &dos_part))
 	goto bail;
 
-    /* dos_part and start_lba are guaranteed to be valid here */
-
     /*
      * note special index handling, if we have stepall set -
      * this is made to keep index consistent with non-stepall
      * iterators
      */
 
-    if (iter->sub.dos.index0 >= 4 && !dos_part->ostype) {
+    if (iter->index0 >= 4 && !dos_part->ostype)
 	iter->index = -1;
-	iter->sub.dos.index0--;
-    } else
-	iter->index = iter->sub.dos.index0 + 1;
+    else
+	iter->index = iter->index0 - iter->sub.dos.skipcnt + 1;
+    iter->rawindex = iter->index0 + 1;
     iter->start_lba = start_lba;
     iter->record = (char *)dos_part;
 
@@ -478,7 +481,7 @@ static void gpt_conv_label(struct part_iter *iter)
     const int16_t *orig_lab;
 
     gp = (const struct disk_gpt_part_entry *)
-	(iter->data + iter->sub.gpt.index0 * iter->sub.gpt.pe_size);
+	(iter->data + iter->index0 * iter->sub.gpt.pe_size);
     orig_lab = (const int16_t *)gp->name;
 
     /* caveat: this is very crude conversion */
@@ -492,9 +495,9 @@ static struct part_iter *pi_gpt_next(struct part_iter *iter)
 {
     const struct disk_gpt_part_entry *gpt_part = NULL;
 
-    while (++iter->sub.gpt.index0 < iter->sub.gpt.pe_count) {
+    while (++iter->index0 < iter->sub.gpt.pe_count) {
 	gpt_part = (const struct disk_gpt_part_entry *)
-	    (iter->data + iter->sub.gpt.index0 * iter->sub.gpt.pe_size);
+	    (iter->data + iter->index0 * iter->sub.gpt.pe_size);
 
 	if (notsane_gpt(gpt_part))
 	    goto bail;
@@ -503,11 +506,12 @@ static struct part_iter *pi_gpt_next(struct part_iter *iter)
 	    break;
     }
     /* no more partitions ? */
-    if (iter->sub.gpt.index0 == iter->sub.gpt.pe_count) {
+    if (iter->index0 == iter->sub.gpt.pe_count) {
 	goto bail;
     }
     /* gpt_part is guaranteed to be valid here */
-    iter->index = iter->sub.gpt.index0 + 1;
+    iter->index = iter->index0 + 1;
+    iter->rawindex = iter->index0 + 1;
     iter->start_lba = gpt_part->lba_first;
     iter->record = (char *)gpt_part;
     memcpy(&iter->sub.gpt.part_guid, &gpt_part->uid, sizeof(struct guid));
diff --git a/com32/chain/partiter.h b/com32/chain/partiter.h
index 718f6033..694969ab 100644
--- a/com32/chain/partiter.h
+++ b/com32/chain/partiter.h
@@ -56,9 +56,11 @@ struct part_iter {
     char *record;
     uint64_t start_lba;
     int index;
+    int rawindex;
     struct disk_info di;
-    /* internal */
     int stepall;
+    /* internal */
+    int index0;
     union _sub {
 	struct _dos {
 	    uint32_t disk_sig;
@@ -68,17 +70,15 @@ struct part_iter {
 	    uint32_t ebr_size;
 	    uint32_t bebr_start;
 	    uint32_t bebr_size;
-	    int index0;
 	    int bebr_index0;
+	    int skipcnt;
 	} dos;
 	struct _gpt {
 	    struct guid disk_guid;
 	    struct guid part_guid;
 	    char part_label[PI_GPTLABSIZE/2+1];
-	    /* internal */
 	    int pe_count;
 	    int pe_size;
-	    int index0;
 	} gpt;
     } sub;
 };