chain module: make option 'save' truly separate

Previously, option 'save' has been executed as a part of bpb mangling.
This was rather inflexible.

Apart from that, all mangle functions test their respective opt.* options
internally.

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

4 files changed, 54 insertions, 41 deletions

diff --git a/com32/chain/Makefile b/com32/chain/Makefile
index 054c7681..36710ffb 100644
--- a/com32/chain/Makefile
+++ b/com32/chain/Makefile
@@ -16,7 +16,8 @@ topdir = ../..
 include ../MCONFIG
 
 OBJS = chain.o partiter.o utility.o options.o mangle.o
-GCCEXTRA = -Wextra -Wconversion -pedantic -Wno-error -DDEBUG
+#GCCEXTRA = -DDEBUG
+#GCCEXTRA += -Wextra -Wconversion -pedantic -Wno-error
 
 all: chain.c32
 
diff --git a/com32/chain/chain.c b/com32/chain/chain.c
index e983dcba..3bd2a318 100644
--- a/com32/chain/chain.c
+++ b/com32/chain/chain.c
@@ -662,6 +662,7 @@ int main(int argc, char *argv[])
 
     void *file_area = NULL;
     void *sect_area = NULL;
+    void *sbck_area = NULL;
     struct disk_dos_part_entry *hand_area = NULL;
 
     struct data_area data[3], bdata[3];
@@ -749,6 +750,14 @@ int main(int argc, char *argv[])
 	    }
 	    sect_area = (void *)data[ndata].data;
 
+	    if (opt.save) {
+		if (!(sbck_area = malloc(SECTOR))) {
+		    error("Couldn't allocate cmp-buf for option 'save'.\n");
+		    goto bail;
+		}
+		memcpy(sbck_area, data->data, data->size);
+	    }
+
 	    sidx = ndata;
 	    ndata++;
 	}
@@ -793,32 +802,33 @@ int main(int argc, char *argv[])
     /* Do file related stuff */
 
     if (fidx >= 0) {
-	if (opt.isolinux && manglef_isolinux(data + fidx))
+	if (manglef_isolinux(data + fidx))
 	    goto bail;
 
-	if (opt.grldr && manglef_grldr(iter))
+	if (manglef_grldr(iter))
 	    goto bail;
 
-	if (opt.grub && manglef_grub(iter, data + fidx))
+	if (manglef_grub(iter, data + fidx))
 	    goto bail;
-
 #if 0
-	if (opt.drmk && manglef_drmk(data + fidx))
+	if (manglef_drmk(data + fidx))
 	    goto bail;
 #endif
-
-	if (opt.filebpb && manglef_bpb(iter, data + fidx))
+	if (manglef_bpb(iter, data + fidx))
 	    goto bail;
     }
 
     /* Do sector related stuff */
 
     if (sidx >= 0) {
-	if (try_mangles_bpb(iter, data + sidx))
+	if (mangles_bpb(iter, data + sidx))
+	    goto bail;
+
+	if (mangles_save(iter, data + sidx, sbck_area))
 	    goto bail;
 
-	/* This *must* be after BPB mangling */
-	if (opt.cmldr && mangles_cmldr(data + sidx))
+	/* This *must* be after last BPB saving */
+	if (mangles_cmldr(data + sidx))
 	    goto bail;
     }
 
@@ -851,6 +861,7 @@ bail:
     /* Free allocated areas */
     free(file_area);
     free(sect_area);
+    free(sbck_area);
     free(hand_area);
     return 255;
 }
diff --git a/com32/chain/mangle.c b/com32/chain/mangle.c
index 5aea2f68..3ddff716 100644
--- a/com32/chain/mangle.c
+++ b/com32/chain/mangle.c
@@ -24,6 +24,9 @@ int manglef_isolinux(struct data_area *data)
     uint32_t *checksum, *chkhead, *chktail;
     uint32_t file_lba = 0;
 
+    if (!opt.isolinux)
+	return 0;
+
     sdi = syslinux_derivative_info();
 
     if (sdi->c.filesystem != SYSLINUX_FS_ISOLINUX) {
@@ -92,6 +95,9 @@ bail:
  */
 int manglef_grldr(const struct part_iter *iter)
 {
+    if (!opt.grldr)
+	return 0;
+
     opt.regs.edx.b[1] = (uint8_t)(iter->index - 1);
     return 0;
 }
@@ -136,6 +142,9 @@ int manglef_grub(const struct part_iter *iter, struct data_area *data)
 	char codestart[1];
     } __attribute__ ((packed)) *stage2;
 
+    if (!opt.grub)
+	return 0;
+
     if (data->size < sizeof(struct grub_stage2_patch_area)) {
 	error("The file specified by grub=<loader> is too small to be stage2 of GRUB Legacy.\n");
 	goto bail;
@@ -202,7 +211,7 @@ bail:
 /*
  * Adjust BPB of a BPB-compatible file
  */
-int manglef_bpb(const struct part_iter *iter, struct data_area *data)
+int mangles_bpb(const struct part_iter *iter, struct data_area *data)
 {
     /* BPB: hidden sectors */
     if (opt.sethid) {
@@ -223,7 +232,6 @@ int manglef_bpb(const struct part_iter *iter, struct data_area *data)
 		*(uint32_t *)((char *)data->data + 0x18) = 0x00020012;
 	}
     }
-
     /* BPB: drive */
     if (opt.setdrv)
 	*(uint8_t *)((char *)data->data + opt.drvoff) = (uint8_t)
@@ -232,42 +240,28 @@ int manglef_bpb(const struct part_iter *iter, struct data_area *data)
     return 0;
 }
 
-/*
- * Adjust BPB of a sector
- */
-int try_mangles_bpb(const struct part_iter *iter, struct data_area *data)
+int manglef_bpb(const struct part_iter *iter, struct data_area *data)
 {
-    void *cmp_buf = NULL;
-
-    if (!(opt.setdrv || opt.setgeo || opt.sethid))
+    if (!opt.filebpb)
 	return 0;
 
-#if 0
-    /* Turn this off for now. It's hard to find a reason to
-     * BPB-mangle sector 0 of whatever there is, but it's
-     * "potentially" useful (fixing fdd's sector ?).
-     */
-    if (!iter->index)
-	return 0;
-#endif
-
-    if (!(cmp_buf = malloc(data->size))) {
-	error("Could not allocate sector-compare buffer.\n");
-	goto bail;
-    }
-
-    memcpy(cmp_buf, data->data, data->size);
+    return mangles_bpb(iter, data);
+}
 
-    manglef_bpb(iter, data);
+int mangles_save(const struct part_iter *iter, const struct data_area *data, void *org)
+{
+    if (!opt.save)
+	return 0;
 
-    if (opt.save && memcmp(cmp_buf, data->data, data->size)) {
-	if (disk_write_verify_sector(&iter->di, iter->start_lba, data->data)) {
-	    error("Cannot write updated boot sector.\n");
+    if (memcmp(org, data->data, data->size)) {
+	if (disk_write_sector(&iter->di, iter->start_lba, data->data)) {
+	    error("Cannot write the updated sector.\n");
 	    goto bail;
 	}
+	/* function is ready do be called again */
+	memcpy(org, data->data, data->size);
     }
 
-    free(cmp_buf);
     return 0;
 
 bail:
@@ -281,6 +275,9 @@ bail:
  */
 int mangles_cmldr(struct data_area *data)
 {
+    if (!opt.cmldr)
+	return 0;
+
     memcpy((char *)data->data + 3, cmldr_signature, sizeof(cmldr_signature));
     return 0;
 }
@@ -297,6 +294,9 @@ int manglef_drmk(struct data_area *data)
      * We only really need 4 new, usable bytes at the end.
      */
 
+    if (!opt.drmk)
+	return 0;
+
     uint32_t tsize = (data->size + 19) & 0xfffffff0;
     opt.regs.ss = opt.regs.fs = opt.regs.gs = 0;	/* Used before initialized */
     if (!realloc(data->data, tsize)) {
diff --git a/com32/chain/mangle.h b/com32/chain/mangle.h
index 71a347ae..006fb3b2 100644
--- a/com32/chain/mangle.h
+++ b/com32/chain/mangle.h
@@ -8,7 +8,8 @@ int manglef_isolinux(struct data_area *data);
 int manglef_grldr(const struct part_iter *iter);
 int manglef_grub(const struct part_iter *iter, struct data_area *data);
 int manglef_bpb(const struct part_iter *iter, struct data_area *data);
-int try_mangles_bpb(const struct part_iter *iter, struct data_area *data);
+int mangles_bpb(const struct part_iter *iter, struct data_area *data);
+int mangles_save(const struct part_iter *iter, const struct data_area *data, void *org);
 int mangles_cmldr(struct data_area *data);
 #if 0
 int manglef_drmk(struct data_area *data);