Fix setting header flag during overwrite

diff --git a/include/wimlib/wim.h b/include/wimlib/wim.h
index 57adebde7522d7f780276865045215ccf3db700c..03e0184ce40bf2dbd0830f4fc196aacbe7caeec8 100644 (file)
--- a/include/wimlib/wim.h
+++ b/include/wimlib/wim.h
@@ -72,6 +72,11 @@ write_header(const struct wim_header *hdr, int out_fd);
 extern int
 init_header(struct wim_header *hdr, int ctype);
 
+extern int
+write_header_flags(u32 hdr_flags, int out_fd);
+
+extern int
+write_header_part_data(u16 part_number, u16 total_parts, int out_fd);
 
 extern int
 rename_wim_path(WIMStruct *wim, const tchar *from, const tchar *to);

diff --git a/src/header.c b/src/header.c
index 332d350fe07a04109f3ca51decf4e6dcfd67300f..5d0b683ba7d6ff361f28de36e84f466ff02fa2a9 100644 (file)
--- a/src/header.c
+++ b/src/header.c
@@ -228,6 +228,38 @@ write_header(const struct wim_header *hdr, int out_fd)
        return 0;
 }
 
+/* Update just the wim_flags field. */
+int
+write_header_flags(u32 hdr_flags, int out_fd)
+{
+       le32 flags = cpu_to_le32(hdr_flags);
+       if (full_pwrite(out_fd, &flags, sizeof(flags),
+                       offsetof(struct wim_header_disk, wim_flags)) != sizeof(flags))
+       {
+               return WIMLIB_ERR_WRITE;
+       } else {
+               return 0;
+       }
+
+}
+
+/* Update just the part_number and total_parts fields. */
+int
+write_header_part_data(u16 part_number, u16 total_parts, int out_fd)
+{
+       le16 part_data[2] = {
+               cpu_to_le16(part_number),
+               cpu_to_le16(total_parts),
+       };
+       if (full_pwrite(out_fd, &part_data, sizeof(part_data),
+                       offsetof(struct wim_header_disk, part_number)) != sizeof(part_data))
+       {
+               return WIMLIB_ERR_WRITE;
+       } else {
+               return 0;
+       }
+}
+
 /*
  * Initializes the header for a WIM file.
  */

diff --git a/src/split.c b/src/split.c
index 2636164b31cf83ee1ab746a7b236ee1e9ef2a065..25ca3cdb0ddc3747fd60627ab321342e0c57f838 100644 (file)
--- a/src/split.c
+++ b/src/split.c
@@ -222,6 +222,7 @@ wimlib_split(WIMStruct *w, const tchar *swm_name,
                int part_fd;
                le16 part_data[2];
                size_t bytes_written;
+               int ret2;
 
                if (i == 1) {
                        part_name = swm_name;
@@ -237,16 +238,14 @@ wimlib_split(WIMStruct *w, const tchar *swm_name,
                        ret = WIMLIB_ERR_OPEN;
                        goto out;
                }
-               part_data[0] = cpu_to_le16(i);
-               part_data[1] = cpu_to_le16(total_parts);
 
-               bytes_written = full_pwrite(part_fd, part_data,
-                                           sizeof(part_data), 40);
-               ret = close(part_fd);
-               if (bytes_written != sizeof(part_data) || ret != 0) {
+               ret = write_header_part_data(i, total_parts, part_fd);
+               ret2 = close(part_fd);
+               if (ret == 0 && ret2 != 0)
+                       ret = WIMLIB_ERR_WRITE;
+               if (ret) {
                        ERROR_WITH_ERRNO("Error updating header of `%"TS"'",
                                         part_name);
-                       ret = WIMLIB_ERR_WRITE;
                        goto out;
                }
        }

diff --git a/src/write.c b/src/write.c
index 54257e046f2effae9916681cfe4f60ff3be427a5..4c68921cc9164efb926e58325f0e0a99109ec6bc 100644 (file)
--- a/src/write.c
+++ b/src/write.c
@@ -1798,6 +1798,7 @@ finish_write(WIMStruct *w, int image, int write_flags,
                        struct wim_header checkpoint_hdr;
                        memcpy(&checkpoint_hdr, &hdr, sizeof(struct wim_header));
                        zero_resource_entry(&checkpoint_hdr.integrity);
+                       checkpoint_hdr.flags |= WIM_HDR_FLAG_WRITE_IN_PROGRESS;
                        ret = write_header(&checkpoint_hdr, w->out_fd);
                        if (ret)
                                goto out_close_wim;
@@ -2112,11 +2113,11 @@ overwrite_wim_inplace(WIMStruct *w, int write_flags,
                return ret;
        }
 
-       /* Write header with write in progress flag set. */
-       w->hdr.flags |= WIM_HDR_FLAG_WRITE_IN_PROGRESS;
-       ret = write_header(&w->hdr, w->out_fd);
-       w->hdr.flags &= ~WIM_HDR_FLAG_WRITE_IN_PROGRESS;
+       /* Set WIM_HDR_FLAG_WRITE_IN_PROGRESS flag in header. */
+       ret = write_header_flags(w->hdr.flags | WIM_HDR_FLAG_WRITE_IN_PROGRESS,
+                                w->out_fd);
        if (ret) {
+               ERROR_WITH_ERRNO("Error updating WIM header flags");
                close_wim_writable(w);
                goto out_unlock_wim;
        }