+
+ *size_ret = size;
+ return 0;
+}
+
+/*
+ * Writes @size bytes of @contents to the named file @path.
+ *
+ * Returns 0 on success; WIMLIB_ERR_OPEN or WIMLIB_ERR_WRITE on failure.
+ */
+static int
+write_wimoverlay_dat(const wchar_t *path, const void *contents, u32 size)
+{
+ HANDLE h;
+ DWORD bytes_written;
+
+ h = CreateFile(path, GENERIC_WRITE, 0, NULL,
+ CREATE_ALWAYS, FILE_FLAG_BACKUP_SEMANTICS, NULL);
+ if (h == INVALID_HANDLE_VALUE) {
+ win32_error(GetLastError(),
+ L"\"%ls\": Can't open file for writing", path);
+ return WIMLIB_ERR_OPEN;
+ }
+
+ SetLastError(0);
+ if (!WriteFile(h, contents, size, &bytes_written, NULL) ||
+ bytes_written != size)
+ {
+ win32_error(GetLastError(),
+ L"\"%ls\": Can't write file", path);
+ CloseHandle(h);
+ return WIMLIB_ERR_WRITE;
+ }
+
+ if (!CloseHandle(h)) {
+ win32_error(GetLastError(),
+ L"\"%ls\": Can't close handle", path);
+ return WIMLIB_ERR_WRITE;
+ }
+
+ return 0;
+}
+
+/*
+ * Generates the contents of WimOverlay.dat in memory, with one entry added.
+ *
+ * @buf
+ * Buffer large enough to hold the new contents.
+ * @old_hdr
+ * Old contents of WimOverlay.dat, or NULL if it did not exist.
+ * @wim_path
+ * Absolute path to the WIM file. It must begin with a drive letter; for
+ * example, D:\install.wim.
+ * @wim_guid
+ * GUID of the WIM, from the WIM header.
+ * @image
+ * Number of the image in the WIM to specify in the new entry.
+ * @new_data_source_id
+ * Data source ID to use for the new entry.
+ * @part_info
+ * Partition information for the WIM file.
+ * @disk_info
+ * Disk information for the WIM file.
+ * @new_entry_2_size
+ * Size, in bytes, of the new location information structure ('struct
+ * WimOverlay_dat_entry_2').
+ *
+ * Returns a pointer one past the last byte of @buf filled in.
+ */
+static u8 *
+fill_in_wimoverlay_dat(u8 *buf,
+ const struct WimOverlay_dat_header *old_hdr,
+ const wchar_t *wim_path,
+ const u8 wim_guid[WIM_GUID_LEN],
+ int image,
+ u64 new_data_source_id,
+ const PARTITION_INFORMATION_EX *part_info,
+ const DRIVE_LAYOUT_INFORMATION_EX *disk_info,
+ u32 new_entry_2_size)
+{
+ struct WimOverlay_dat_header *new_hdr;
+ struct WimOverlay_dat_entry_1 *new_entry_1;
+ struct WimOverlay_dat_entry_2 *new_entry_2;
+ u32 entry_2_offset;
+ u8 *p = buf;
+
+ new_hdr = (struct WimOverlay_dat_header *)p;
+
+ /* Fill in new header */
+ new_hdr->magic = WIMOVERLAY_DAT_MAGIC;
+ new_hdr->wim_provider_version = WIM_PROVIDER_CURRENT_VERSION;
+ new_hdr->unknown_0x08 = 0x00000028;
+ new_hdr->num_entries_1 = (old_hdr ? old_hdr->num_entries_1 : 0) + 1;
+ new_hdr->num_entries_2 = (old_hdr ? old_hdr->num_entries_2 : 0) + 1;
+ new_hdr->unknown_0x14 = 0x00000000;
+
+ p += sizeof(struct WimOverlay_dat_header);
+
+ /* Copy WIM-specific information for old entries */
+ entry_2_offset = sizeof(struct WimOverlay_dat_header) +
+ (new_hdr->num_entries_1 * sizeof(struct WimOverlay_dat_entry_1));
+ if (old_hdr) {
+ for (u32 i = 0; i < old_hdr->num_entries_1; i++) {
+ new_entry_1 = (struct WimOverlay_dat_entry_1 *)p;
+
+ p = mempcpy(p, &old_hdr->entry_1s[i],
+ sizeof(struct WimOverlay_dat_entry_1));
+
+ new_entry_1->entry_2_offset = entry_2_offset;
+ entry_2_offset += new_entry_1->entry_2_length;
+ }
+ }
+
+ /* Generate WIM-specific information for new entry */
+ new_entry_1 = (struct WimOverlay_dat_entry_1 *)p;
+
+ new_entry_1->data_source_id = new_data_source_id;
+ new_entry_1->entry_2_offset = entry_2_offset;
+ new_entry_1->entry_2_length = new_entry_2_size;
+ new_entry_1->wim_type = WIM_BOOT_NOT_OS_WIM;
+ new_entry_1->wim_index = image;
+ BUILD_BUG_ON(sizeof(new_entry_1->guid) != WIM_GUID_LEN);
+ memcpy(new_entry_1->guid, wim_guid, WIM_GUID_LEN);
+
+ p += sizeof(struct WimOverlay_dat_entry_1);
+
+ /* Copy WIM location information for old entries */
+ if (old_hdr) {
+ for (u32 i = 0; i < old_hdr->num_entries_1; i++) {
+ wimlib_assert(new_hdr->entry_1s[i].entry_2_offset == p - buf);
+ wimlib_assert(old_hdr->entry_1s[i].entry_2_length ==
+ new_hdr->entry_1s[i].entry_2_length);
+ p = mempcpy(p,
+ ((const u8 *)old_hdr + old_hdr->entry_1s[i].entry_2_offset),
+ old_hdr->entry_1s[i].entry_2_length);
+ }
+ }
+
+ /* Generate WIM location information for new entry */
+ new_entry_2 = (struct WimOverlay_dat_entry_2 *)p;
+
+ new_entry_2->unknown_0x00 = 0x00000000;
+ new_entry_2->unknown_0x04 = 0x00000000;
+ new_entry_2->entry_2_length = new_entry_2_size;
+ new_entry_2->unknown_0x0C = 0x00000000;
+ new_entry_2->unknown_0x10 = 0x00000005;
+ new_entry_2->unknown_0x14 = 0x00000001;
+ new_entry_2->inner_struct_size = new_entry_2_size - 0x14;
+ new_entry_2->unknown_0x1C = 0x00000005;
+ new_entry_2->unknown_0x20 = 0x00000006;
+ new_entry_2->unknown_0x24 = 0x00000000;
+ new_entry_2->unknown_0x28 = 0x00000048;
+ new_entry_2->unknown_0x2C = 0x00000000;
+ new_entry_2->unknown_0x40 = 0x00000000;
+
+ if (part_info->PartitionStyle == PARTITION_STYLE_MBR) {
+ new_entry_2->partition.mbr.part_start_offset = part_info->StartingOffset.QuadPart;
+ new_entry_2->partition.mbr.padding = 0;
+ new_entry_2->partition_table_type = WIMOVERLAY_PARTITION_TYPE_MBR;
+ new_entry_2->disk.mbr.disk_id = disk_info->Mbr.Signature;
+ new_entry_2->disk.mbr.padding[0] = 0x00000000;
+ new_entry_2->disk.mbr.padding[1] = 0x00000000;
+ new_entry_2->disk.mbr.padding[2] = 0x00000000;
+ } else {
+ BUILD_BUG_ON(sizeof(new_entry_2->partition.gpt.part_unique_guid) !=
+ sizeof(part_info->Gpt.PartitionId));
+ memcpy(new_entry_2->partition.gpt.part_unique_guid,
+ &part_info->Gpt.PartitionId,
+ sizeof(part_info->Gpt.PartitionId));
+ new_entry_2->partition_table_type = WIMOVERLAY_PARTITION_TYPE_GPT;
+
+ BUILD_BUG_ON(sizeof(new_entry_2->disk.gpt.disk_guid) !=
+ sizeof(disk_info->Gpt.DiskId));
+ memcpy(new_entry_2->disk.gpt.disk_guid,
+ &disk_info->Gpt.DiskId,
+ sizeof(disk_info->Gpt.DiskId));
+
+ BUILD_BUG_ON(sizeof(new_entry_2->disk.gpt.disk_guid) !=
+ sizeof(new_entry_2->partition.gpt.part_unique_guid));
+ }
+ new_entry_2->unknown_0x58[0] = 0x00000000;
+ new_entry_2->unknown_0x58[1] = 0x00000000;
+ new_entry_2->unknown_0x58[2] = 0x00000000;
+ new_entry_2->unknown_0x58[3] = 0x00000000;
+
+ wimlib_assert(wim_path[2] == L'\\');
+ return mempcpy(new_entry_2->wim_file_name,
+ wim_path + 2,
+ new_entry_2_size - sizeof(struct WimOverlay_dat_entry_2));
+}
+
+/*
+ * Prepares the new contents of WimOverlay.dat in memory, with one entry added.
+ *
+ * @old_hdr
+ * Old contents of WimOverlay.dat, or NULL if it did not exist.
+ * @wim_path
+ * Absolute path to the WIM file. It must begin with a drive letter; for
+ * example, D:\install.wim.
+ * @wim_guid
+ * GUID of the WIM, from the WIM header.
+ * @image
+ * Number of the image in the WIM to specify in the new entry.
+ * @new_contents_ret
+ * Location into which to return the new contents as a malloc()ed buffer on
+ * success.
+ * @new_contents_size_ret
+ * Location into which to return the size, in bytes, of the new contents on
+ * success.
+ * @new_data_source_id_ret
+ * Location into which to return the data source ID of the new entry on
+ * success.
+ *
+ * Returns 0 on success, or a positive error code on failure.
+ */
+static int
+prepare_wimoverlay_dat(const struct WimOverlay_dat_header *old_hdr,
+ const wchar_t *wim_path,
+ const u8 wim_guid[WIM_GUID_LEN],
+ int image,
+ void **new_contents_ret,
+ u32 *new_contents_size_ret,
+ u64 *new_data_source_id_ret)
+{
+ int ret;
+ PARTITION_INFORMATION_EX part_info;
+ DRIVE_LAYOUT_INFORMATION_EX disk_info;
+ u64 new_data_source_id;
+ u32 new_entry_2_size;
+ u32 new_contents_size;
+ u8 *buf;
+ u8 *end;
+
+ ret = query_partition_and_disk_info(wim_path, &part_info, &disk_info);
+ if (ret)
+ return ret;
+
+ new_data_source_id = alloc_new_data_source_id(old_hdr);
+
+ new_entry_2_size = sizeof(struct WimOverlay_dat_entry_2) +
+ ((wcslen(wim_path) - 2 + 1) * sizeof(wchar_t));
+ ret = calculate_wimoverlay_dat_size(old_hdr, new_entry_2_size,
+ &new_contents_size);
+ if (ret)
+ return ret;
+
+ buf = MALLOC(new_contents_size);
+ if (!buf)
+ return WIMLIB_ERR_NOMEM;
+
+ end = fill_in_wimoverlay_dat(buf, old_hdr, wim_path, wim_guid, image,
+ new_data_source_id,
+ &part_info, &disk_info, new_entry_2_size);
+
+ wimlib_assert(end - buf == new_contents_size);
+
+ *new_contents_ret = buf;
+ *new_contents_size_ret = new_contents_size;
+ *new_data_source_id_ret = new_data_source_id;
+ return 0;
+}
+
+/*
+ * Reads and validates a WimOverlay.dat file.
+ *
+ * @path
+ * Path to the WimOverlay.dat file, such as
+ * C:\System Volume Information\WimOverlay.dat
+ * @contents_ret
+ * Location into which to return the contents as a malloc()ed buffer on
+ * success. This can be cast to 'struct WimOverlay_dat_header', and its
+ * contents are guaranteed to be valid. Alternatively, if the file does
+ * not exist, NULL will be returned here.
+ *
+ * Returns 0 on success, or a positive error code on failure.
+ */
+static int
+read_wimoverlay_dat(const wchar_t *path, void **contents_ret)
+{
+ HANDLE h;
+ BY_HANDLE_FILE_INFORMATION info;
+ int ret;
+ void *contents;
+ const struct WimOverlay_dat_header *hdr;
+ DWORD bytes_read;
+ bool already_retried = false;
+
+retry:
+ h = open_file(path, GENERIC_READ);
+ if (h == INVALID_HANDLE_VALUE) {
+ DWORD err = GetLastError();
+ if (err == ERROR_FILE_NOT_FOUND) {
+ *contents_ret = NULL;
+ return 0;
+ }
+ if (err == ERROR_PATH_NOT_FOUND &&
+ func_RtlCreateSystemVolumeInformationFolder)
+ {
+ wchar_t volume_root_path[] = L"\\??\\X:\\";
+
+ *(wcschr(volume_root_path, L'X')) = path[0];
+
+ UNICODE_STRING str = {
+ .Length = sizeof(volume_root_path) - sizeof(wchar_t),
+ .MaximumLength = sizeof(volume_root_path),
+ .Buffer = volume_root_path,
+ };
+ NTSTATUS status;
+ DWORD err2;
+
+ status = (*func_RtlCreateSystemVolumeInformationFolder)(&str);
+
+ err2 = (*func_RtlNtStatusToDosError)(status);
+ if (err2 == ERROR_SUCCESS) {
+ if (!already_retried) {
+ already_retried = true;
+ goto retry;
+ }
+ } else {
+ err = err2;
+ }
+ }
+ win32_error(err, L"\"%ls\": Can't open for reading", path);
+ return WIMLIB_ERR_OPEN;
+ }
+ if (!GetFileInformationByHandle(h, &info)) {
+ win32_error(GetLastError(), L"\"%ls\": Can't query metadata", path);
+ CloseHandle(h);
+ return WIMLIB_ERR_STAT;