From 2271d6bde1fe9f68d0ad2d20a8e82a1704e8b877 Mon Sep 17 00:00:00 2001
From: Aiken Harris <harraiken@noreply.codingworkshop.git>
Date: Wed, 1 Oct 2025 23:06:15 +0200
Subject: [PATCH] Preserve BPB when writing custom VBR

---
 tools/diskimg.c | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/tools/diskimg.c b/tools/diskimg.c
index 3d288db..e4465f5 100644
--- a/tools/diskimg.c
+++ b/tools/diskimg.c
@@ -4,6 +4,7 @@
  * FILE:        tools/diskimg.c
  * DESCRIPTION: Disk Image manipulation tool
  * DEVELOPERS:  Rafal Kupiec <belliash@codingworkshop.eu.org>
+ *              Aiken Harris <harraiken91@gmail.com>
  */
 
 #include "xtchain.h"
@@ -164,7 +165,8 @@ int main(int argc, char **argv)
     MBR_PARTITION Partition = {0};
     char Zero[SECTOR_SIZE] = {0};
     uint8_t Mbr[SECTOR_SIZE] = {0};
-    uint8_t Vbr[SECTOR_SIZE];
+    uint8_t Vbr[SECTOR_SIZE] = {0};
+    uint8_t ImageVbr[SECTOR_SIZE] = {0};
     const char *FileName = NULL;
     const char *MbrFile = NULL;
     const char *VbrFile = NULL;
@@ -321,11 +323,36 @@ int main(int argc, char **argv)
     /* Write VBR to the start of the partition, if provided */
     if(VbrFile)
     {
-        if (LoadSector(VbrFile, Vbr) != 0) {
+        /* Read the VBR file into memory */
+        if(LoadSector(VbrFile, Vbr) != 0)
+        {
             fclose(File);
             return 1;
         }
-        /* Seek to the start of the partition and write VBR */
+
+        /* Read the existing VBR from the formatted partition to get the correct BPB */
+        fseek(File, Partition.StartLBA * SECTOR_SIZE, SEEK_SET);
+        if(fread(ImageVbr, 1, SECTOR_SIZE, File) != SECTOR_SIZE)
+        {
+            /* Failed to read VBR from disk image */
+            perror("Failed to read BPB from disk image");
+            fclose(File);
+            return 1;
+        }
+
+        /* Copy the BPB from the image's VBR to VBR buffer */
+        if(FatFormat == 32)
+        {
+            /* For FAT32, BPB is larger (up to offset 89) */
+            memcpy(&Vbr[3], &ImageVbr[3], 87);
+        }
+        else
+        {
+            /* For FAT16, BPB is smaller (up to offset 61) */
+            memcpy(&Vbr[3], &ImageVbr[3], 59);
+        }
+
+        /* Write the final, merged VBR to the start of the partition */
         fseek(File, Partition.StartLBA * SECTOR_SIZE, SEEK_SET);
         if(fwrite(Vbr, 1, SECTOR_SIZE, File) != SECTOR_SIZE)
         {
@@ -336,6 +363,7 @@ int main(int argc, char **argv)
         }
     }
 
+    /* Close file */
     fclose(File);
 
     /* Copy files if requested */