Fix restoring a diff image

4 files changed, 69 insertions, 65 deletions

diff --git a/src/backup.c b/src/backup.c
index 15628b8..186cff0 100644
--- a/src/backup.c
+++ b/src/backup.c
@@ -38,8 +38,6 @@ static int check_files(const options_t *const opts,
                        const resources_t *const res);
 static int allocate_buffers(const options_t *const opts, size_t out_buffer_size,
                             resources_t *const res);
-static int read_sector(FILE *const file, char *const buffer,
-                       uint32_t buffer_size, uint32_t sector_size);
 static int write_out_buffer(const char *const buffer, size_t size,
                             FILE *const file);
 
@@ -94,26 +92,6 @@ allocate_buffers(const options_t *const opts, size_t out_buffer_size,
 }
 
 static int
-read_sector(FILE *const file, char *const buffer, uint32_t buffer_size,
-            uint32_t sector_size)
-{
-    const size_t bytes_read = fread(buffer, 1U, buffer_size, file);
-
-    if (ferror(file)) {
-        print_error("cannot read from input file: %s", strerror(errno));
-        return 1;
-    } else if ((bytes_read % sector_size) != 0) {
-        print_error("data read from input file is not multiple of sector size");
-        return 1;
-    } else if (bytes_read != buffer_size) {
-        print_error("cannot read enough data from input file");
-        return 1;
-    }
-
-    return 0;
-}
-
-static int
 write_out_buffer(const char *const buffer, size_t size, FILE *const file)
 {
     const size_t bytes_written = fwrite(buffer, 1U, size, file);
@@ -145,29 +123,33 @@ backup(const options_t *const opts, resources_t *const res)
     size_t out_buffer_index = 0;
     uint64_t input_file_offset = 0;
 
-    while (!feof(res->in_file)) {
-        /* Read the sectors from the input file to the buffer */
-        if (read_sector(res->in_file, res->in_buffer, opts->buffer_size,
-                        opts->sector_size) != 0) {
+    for (;;) {
+        /* Read the sectors from the input and reference files into the buffers
+         */
+        const size_t in_sectors_read = file_read_sectors(
+            res->in_file, res->in_buffer, opts->buffer_size, opts->sector_size);
+        const size_t ref_sectors_read =
+            file_read_sectors(res->ref_file, res->ref_buffer, opts->buffer_size,
+                              opts->sector_size);
+
+        if ((in_sectors_read == 0) || (ref_sectors_read == 0)) {
+            break;
+        } else if (in_sectors_read != ref_sectors_read) {
+            print_error(
+                "cannot read equal amount of sectors from the input files");
             return 1;
         }
 
-        /* Read sectors from the reference file to the buffer */
-        if (read_sector(res->ref_file, res->ref_buffer, opts->buffer_size,
-                        opts->sector_size) != 0) {
-            return 1;
-        }
+        /* Process the sectors in the buffers */
+        for (size_t sector = 0; sector < in_sectors_read; ++sector) {
+            const size_t buffer_offset = sector * opts->sector_size;
 
-        /* Process sectors in the buffer */
-        for (size_t buffer_offset = 0; buffer_offset < opts->buffer_size;
-             buffer_offset += opts->sector_size) {
             if (memcmp(res->in_buffer + buffer_offset,
                        res->ref_buffer + buffer_offset,
                        opts->sector_size) != 0) {
                 /* Backup the changed sector */
                 if (out_buffer_index >= out_buffer_size) {
-                    /* The output buffer is full. Write it to the output file.
-                     */
+                    /* The output buffer is full. Write it to the output file */
                     if (write_out_buffer(res->out_buffer, out_buffer_index,
                                          res->out_file) != 1) {
                         return 1;
diff --git a/src/file.c b/src/file.c
index 634bb3b..57a847a 100644
--- a/src/file.c
+++ b/src/file.c
@@ -48,3 +48,20 @@ file_size(FILE *const file)
 
     return size;
 }
+
+size_t
+file_read_sectors(FILE *const file, char *const buffer, uint32_t buffer_size,
+                  uint32_t sector_size)
+{
+    const size_t bytes_read = fread(buffer, 1U, buffer_size, file);
+
+    if (ferror(file)) {
+        print_error("cannot read from file: %s", strerror(errno));
+        return 0;
+    } else if ((bytes_read % sector_size) != 0) {
+        print_error("data read from input file is not multiple of sector size");
+        return 0;
+    } else {
+        return (bytes_read / sector_size);
+    }
+}
diff --git a/src/file.h b/src/file.h
index 7c77b45..886098b 100644
--- a/src/file.h
+++ b/src/file.h
@@ -27,8 +27,11 @@
 #ifndef FILE_H
 #define FILE_H
 
+#include <stdint.h>
 #include <stdio.h>
 
 long file_size(FILE *const file);
+size_t file_read_sectors(FILE *const file, char *const buffer,
+                         uint32_t buffer_size, uint32_t sector_size);
 
 #endif /* FILE_H */
diff --git a/src/restore.c b/src/restore.c
index 1c2beae..6318766 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -63,7 +63,8 @@ is_reference_file_valid(resources_t *const res, uint32_t sector_size)
         return 1;
     }
 
-    uint64_t prev_out_offset;
+    uint64_t prev_out_offset = 0;
+    bool is_first_reading = true;
 
     /* Scan the reference file and check  */
     for (;;) {
@@ -78,18 +79,19 @@ is_reference_file_valid(resources_t *const res, uint32_t sector_size)
         } else if ((ref_read != 1U) || ferror(res->ref_file)) {
             print_error("cannot read from reference file: %s", strerror(errno));
             return false;
-        } else if (((out_offset <= prev_out_offset)) ||
-                   ((out_offset + sector_size) > out_size)) {
-            /* The offset must be higher than the previous one and it
-             * must point into the file
-             */
-            print_error("reference file is not valid");
+        } else if (!is_first_reading && (out_offset <= prev_out_offset)) {
+            print_error("a sector offset points behind the previous offset");
+            return false;
+        } else if ((out_offset + sector_size) > out_size) {
+            print_error(
+                "a sector offset points past the end of the output file");
             return false;
         } else if (fseek(res->ref_file, sector_size, SEEK_CUR) != 0) {
             print_error("cannot seek in reference file: %s", strerror(errno));
             return false;
         }
 
+        is_first_reading = false;
         prev_out_offset = out_offset;
     }
 
@@ -123,9 +125,9 @@ restore(const options_t *const opts, resources_t *const res)
 
     /* Allocate the buffer for data from the reference file */
     /* The reference buffer contains also the offsets */
-    const size_t ref_buffer_size =
-        ((opts->buffer_size / opts->sector_size) * sizeof(uint64_t)) +
-        opts->buffer_size;
+    const size_t ref_sector_size = sizeof(uint64_t) + opts->sector_size;
+    const size_t ref_buffer_sector_count = opts->buffer_size / ref_sector_size;
+    const size_t ref_buffer_size = ref_buffer_sector_count * ref_sector_size;
 
     if ((res->ref_buffer = (char *)malloc(ref_buffer_size)) == NULL) {
         print_error("cannot allocate buffer for reference file data");
@@ -134,34 +136,34 @@ restore(const options_t *const opts, resources_t *const res)
 
     /* Restore data from the differential image */
     for (;;) {
-        uint64_t out_offset;
 
         /* Read data of the offset and the next sector */
-        const size_t ref_read =
-            fread(res->ref_buffer, ref_buffer_size, 1U, res->ref_file);
+        const size_t ref_sectors_read = file_read_sectors(
+            res->ref_file, res->ref_buffer, ref_buffer_size, ref_sector_size);
 
-        if (feof(res->ref_file)) {
+        if (ref_sectors_read == 0) {
             break;
-        } else if ((ref_read != 1U) || ferror(res->ref_file)) {
-            print_error("cannot read from reference file: %s", strerror(errno));
-            return 1;
         }
 
-        /* Get offset */
-        out_offset = le64toh(*((uint64_t *)res->ref_buffer));
+        char *ref_buffer = res->ref_buffer;
 
-        if (fseek(res->out_file, out_offset, SEEK_SET) != 0) {
-            print_error("cannot seek in output file: %s", strerror(errno));
-            return 1;
-        }
+        for (size_t s = 0; s < ref_sectors_read; ++s) {
+            const uint64_t out_offset = le64toh(*((uint64_t *)ref_buffer));
+            ref_buffer += sizeof(uint64_t);
+
+            if (fseek(res->out_file, out_offset, SEEK_SET) != 0) {
+                print_error("cannot seek in output file: %s", strerror(errno));
+                return 1;
+            }
 
-        /* Write the sector data to the output file */
-        const size_t out_written = fwrite(res->ref_buffer + sizeof(out_offset),
-                                          opts->sector_size, 1U, res->out_file);
+            const size_t out_written =
+                fwrite(ref_buffer, opts->sector_size, 1U, res->out_file);
+            ref_buffer += opts->sector_size;
 
-        if (out_written != 1U) {
-            print_error("cannot write to output file: %s", strerror(errno));
-            return 1;
+            if (out_written != 1U) {
+                print_error("cannot write to output file: %s", strerror(errno));
+                return 1;
+            }
         }
     }