From 59addd3efc166fbd79227319400c65a1f75f7d56 Mon Sep 17 00:00:00 2001
From: Natalia Portillo <claunia@claunia.com>
Date: Fri, 3 Oct 2025 00:57:14 +0100
Subject: [PATCH] Implement on-the-fly MD5 checksum calculation

---
 include/aaruformat/context.h | 3 +++
 src/close.c                  | 7 +++++++
 src/create.c                 | 6 ++++++
 src/write.c                  | 8 +++++++-
 4 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/include/aaruformat/context.h b/include/aaruformat/context.h
index 1856df3..cbf87de 100644
--- a/include/aaruformat/context.h
+++ b/include/aaruformat/context.h
@@ -22,6 +22,7 @@
 #include "crc64.h"
 #include "hash_map.h"
 #include "lru.h"
+#include "md5.h"
 #include "structs.h"
 #include "utarray.h"
 
@@ -212,6 +213,8 @@ typedef struct aaruformatContext
 
     bool     rewinded;            ///< True if stream has been rewound after open (write path).
     uint64_t last_written_block;  ///< Last written block number (write path).
+    bool     calculating_md5;     ///< True if whole-image MD5 being calculated on-the-fly.
+    md5_ctx  md5_context;         ///< Opaque MD5 context for streaming updates
 } aaruformatContext;
 
 /** \struct DumpHardwareEntriesWithData
diff --git a/src/close.c b/src/close.c
index 1f0be6a..7ba479a 100644
--- a/src/close.c
+++ b/src/close.c
@@ -500,6 +500,13 @@ int aaruf_close(void *context)
         uint64_t alignment_mask;
         uint64_t aligned_position;
 
+        // Finalize pending checksums
+        if(ctx->calculating_md5)
+        {
+            ctx->checksums.hasMd5 = true;
+            aaruf_md5_final(&ctx->md5_context, ctx->checksums.md5);
+        }
+
         // Write the checksums block
         bool has_checksums =
             ctx->checksums.hasMd5 || ctx->checksums.hasSha1 || ctx->checksums.hasSha256 || ctx->checksums.hasSpamSum;
diff --git a/src/create.c b/src/create.c
index 63d18eb..0e072d7 100644
--- a/src/create.c
+++ b/src/create.c
@@ -310,6 +310,12 @@ void *aaruf_create(const char *filepath, const uint32_t media_type, const uint32
     ctx->rewinded           = false;
     ctx->last_written_block = 0;
 
+    if(parsed_options.md5)
+    {
+        ctx->calculating_md5 = true;
+        aaruf_md5_init(&ctx->md5_context);
+    }
+
     // Is writing
     ctx->isWriting = true;
 
diff --git a/src/write.c b/src/write.c
index b87b37e..d449872 100644
--- a/src/write.c
+++ b/src/write.c
@@ -145,18 +145,24 @@ int32_t aaruf_write_sector(void *context, uint64_t sector_address, bool negative
         return AARUF_ERROR_SECTOR_OUT_OF_BOUNDS;
     }
 
-    // TODO: Check rewinded for disabling checksums
     if(!ctx->rewinded)
     {
         if(sector_address <= ctx->last_written_block)
         {
             TRACE("Rewinded");
             ctx->rewinded = true;
+
+            // Disable MD5 calculation
+            if(ctx->calculating_md5) ctx->calculating_md5 = false;
         }
         else
             ctx->last_written_block = sector_address;
     }
 
+    // Calculate MD5 on-the-fly if requested and sector is within user sectors (not negative or overflow)
+    if(ctx->calculating_md5 && !negative && sector_address <= ctx->imageInfo.Sectors)
+        aaruf_md5_update(&ctx->md5_context, data, length);
+
     // TODO: If optical disc check track
 
     // Close current block first