From: Eric Biggers <ebiggers3@gmail.com>
Date: Thu, 23 May 2013 01:48:14 +0000 (-0500)
Subject: Fix setting header flag during overwrite
X-Git-Tag: v1.4.1~15
X-Git-Url: https://wimlib.net/git/?p=wimlib;a=commitdiff_plain;h=58daaca96dc8b23e345342e844c58b1972e5fedf;hp=1da527afb0ea8fb58838b6ade3a93aea198191da

Fix setting header flag during overwrite
---

diff --git a/include/wimlib/wim.h b/include/wimlib/wim.h
index 57adebde..03e0184c 100644
--- 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 332d350f..5d0b683b 100644
--- 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 2636164b..25ca3cdb 100644
--- 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 54257e04..4c68921c 100644
--- 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;
 	}