]> wimlib.net Git - wimlib/blobdiff - src/write.c
Fix setting header flag during overwrite
[wimlib] / src / write.c
index 5218fee0251e8d45ff355d2c4b499487edd1cce1..4c68921cc9164efb926e58325f0e0a99109ec6bc 100644 (file)
  * along with wimlib; if not, see http://www.gnu.org/licenses/.
  */
 
-#include "config.h"
+#ifdef HAVE_CONFIG_H
+#  include "config.h"
+#endif
 
 #if defined(HAVE_SYS_FILE_H) && defined(HAVE_FLOCK)
-/* On BSD, this should be included before "list.h" so that "list.h" can
+/* On BSD, this should be included before "wimlib/list.h" so that "wimlib/list.h" can
  * overwrite the LIST_HEAD macro. */
 #  include <sys/file.h>
 #endif
 
+#include "wimlib/endianness.h"
+#include "wimlib/error.h"
+#include "wimlib/file_io.h"
+#include "wimlib/header.h"
+#include "wimlib/integrity.h"
+#include "wimlib/lookup_table.h"
+#include "wimlib/metadata.h"
+#include "wimlib/resource.h"
+#include "wimlib/write.h"
+#include "wimlib/xml.h"
+
 #ifdef __WIN32__
-#  include "win32.h"
+#  include "wimlib/win32.h" /* win32_get_number_of_processors() */
 #endif
 
-#include "list.h"
-#include "wimlib_internal.h"
-#include "buffer_io.h"
-#include "dentry.h"
-#include "lookup_table.h"
-#include "xml.h"
-
 #ifdef ENABLE_MULTITHREADED_COMPRESSION
 #  include <pthread.h>
 #endif
  * array of offsets.) */
 struct chunk_table {
        off_t file_offset;
-       u64 num_chunks;
        u64 original_resource_size;
-       u64 bytes_per_chunk_entry;
+       u64 num_chunks;
        u64 table_disk_size;
-       u64 cur_offset;
-       u64 *cur_offset_p;
-       u64 offsets[0];
+       unsigned bytes_per_chunk_entry;
+       void *cur_offset_p;
+       union {
+               u32 cur_offset_u32;
+               u64 cur_offset_u64;
+       };
+       /* Beginning of chunk offsets, in either 32-bit or 64-bit little endian
+        * integers, including the first offset of 0, which will not be written.
+        * */
+       u8 offsets[] _aligned_attribute(8);
 };
 
 /*
@@ -95,11 +107,12 @@ begin_wim_resource_chunk_tab(const struct wim_lookup_table_entry *lte,
                             struct chunk_table **chunk_tab_ret)
 {
        u64 size = wim_resource_size(lte);
-       u64 num_chunks = (size + WIM_CHUNK_SIZE - 1) / WIM_CHUNK_SIZE;
+       u64 num_chunks = wim_resource_chunks(lte);
+       unsigned bytes_per_chunk_entry = (size > (1ULL << 32)) ? 8 : 4;
        size_t alloc_size = sizeof(struct chunk_table) + num_chunks * sizeof(u64);
        struct chunk_table *chunk_tab = CALLOC(1, alloc_size);
 
-       DEBUG("Begin chunk table for stream with size %"PRIu64, size);
+       DEBUG("Beginning chunk table for stream with size %"PRIu64, size);
 
        if (!chunk_tab) {
                ERROR("Failed to allocate chunk table for %"PRIu64" byte "
@@ -109,13 +122,15 @@ begin_wim_resource_chunk_tab(const struct wim_lookup_table_entry *lte,
        chunk_tab->file_offset = file_offset;
        chunk_tab->num_chunks = num_chunks;
        chunk_tab->original_resource_size = size;
-       chunk_tab->bytes_per_chunk_entry = (size >= (1ULL << 32)) ? 8 : 4;
+       chunk_tab->bytes_per_chunk_entry = bytes_per_chunk_entry;
        chunk_tab->table_disk_size = chunk_tab->bytes_per_chunk_entry *
                                     (num_chunks - 1);
-       chunk_tab->cur_offset = 0;
        chunk_tab->cur_offset_p = chunk_tab->offsets;
 
-       if (full_write(out_fd, chunk_tab,
+       /* We don't know the correct offsets yet; this just writes zeroes to
+        * reserve space for the table, so we can go back to it later after
+        * we've written the compressed chunks following it. */
+       if (full_write(out_fd, chunk_tab->offsets,
                       chunk_tab->table_disk_size) != chunk_tab->table_disk_size)
        {
                ERROR_WITH_ERRNO("Failed to write chunk table in compressed "
@@ -127,6 +142,22 @@ begin_wim_resource_chunk_tab(const struct wim_lookup_table_entry *lte,
        return 0;
 }
 
+/* Add the offset for the next chunk to the chunk table being constructed for a
+ * compressed stream. */
+static void
+chunk_tab_record_chunk(struct chunk_table *chunk_tab, unsigned out_chunk_size)
+{
+       if (chunk_tab->bytes_per_chunk_entry == 4) {
+               *(le32*)chunk_tab->cur_offset_p = cpu_to_le32(chunk_tab->cur_offset_u32);
+               chunk_tab->cur_offset_p = (le32*)chunk_tab->cur_offset_p + 1;
+               chunk_tab->cur_offset_u32 += out_chunk_size;
+       } else {
+               *(le64*)chunk_tab->cur_offset_p = cpu_to_le64(chunk_tab->cur_offset_u64);
+               chunk_tab->cur_offset_p = (le64*)chunk_tab->cur_offset_p + 1;
+               chunk_tab->cur_offset_u64 += out_chunk_size;
+       }
+}
+
 /*
  * compress_func_t- Pointer to a function to compresses a chunk
  *                  of a WIM resource.  This may be either
@@ -196,8 +227,7 @@ write_wim_resource_chunk(const void * restrict chunk,
                        out_chunk = chunk;
                        out_chunk_size = chunk_size;
                }
-               *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
-               chunk_tab->cur_offset += out_chunk_size;
+               chunk_tab_record_chunk(chunk_tab, out_chunk_size);
        } else {
                /* Write uncompressed */
                out_chunk = chunk;
@@ -223,15 +253,8 @@ finish_wim_resource_chunk_tab(struct chunk_table *chunk_tab,
 {
        size_t bytes_written;
 
-       if (chunk_tab->bytes_per_chunk_entry == 8) {
-               array_cpu_to_le64(chunk_tab->offsets, chunk_tab->num_chunks);
-       } else {
-               for (u64 i = 0; i < chunk_tab->num_chunks; i++)
-                       ((u32*)chunk_tab->offsets)[i] =
-                               cpu_to_le32(chunk_tab->offsets[i]);
-       }
        bytes_written = full_pwrite(out_fd,
-                                   (u8*)chunk_tab->offsets + chunk_tab->bytes_per_chunk_entry,
+                                   chunk_tab->offsets + chunk_tab->bytes_per_chunk_entry,
                                    chunk_tab->table_disk_size,
                                    chunk_tab->file_offset);
        if (bytes_written != chunk_tab->table_disk_size) {
@@ -239,7 +262,10 @@ finish_wim_resource_chunk_tab(struct chunk_table *chunk_tab,
                                 "file resource");
                return WIMLIB_ERR_WRITE;
        }
-       *compressed_size_p = chunk_tab->cur_offset + chunk_tab->table_disk_size;
+       if (chunk_tab->bytes_per_chunk_entry == 4)
+               *compressed_size_p = chunk_tab->cur_offset_u32 + chunk_tab->table_disk_size;
+       else
+               *compressed_size_p = chunk_tab->cur_offset_u64 + chunk_tab->table_disk_size;
        return 0;
 }
 
@@ -351,7 +377,7 @@ write_wim_resource(struct wim_lookup_table_entry *lte,
        if (!(flags & WIMLIB_RESOURCE_FLAG_RECOMPRESS) &&
            lte->resource_location == RESOURCE_IN_WIM &&
            out_ctype != WIMLIB_COMPRESSION_TYPE_NONE &&
-           wimlib_get_compression_type(lte->wim) == out_ctype)
+           lte->wim->compression_type == out_ctype)
        {
                flags |= WIMLIB_RESOURCE_FLAG_RAW;
                write_ctx.doing_sha = false;
@@ -785,10 +811,8 @@ static int
 write_wim_chunks(struct message *msg, int out_fd,
                 struct chunk_table *chunk_tab)
 {
-       for (unsigned i = 0; i < msg->num_chunks; i++) {
-               *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
-               chunk_tab->cur_offset += msg->out_chunks[i].iov_len;
-       }
+       for (unsigned i = 0; i < msg->num_chunks; i++)
+               chunk_tab_record_chunk(chunk_tab, msg->out_chunks[i].iov_len);
        if (full_writev(out_fd, msg->out_chunks,
                        msg->num_chunks) != msg->total_out_bytes)
        {
@@ -1183,7 +1207,7 @@ main_thread_process_next_stream(struct wim_lookup_table_entry *lte, void *_ctx)
            ctx->out_ctype == WIMLIB_COMPRESSION_TYPE_NONE ||
            (lte->resource_location == RESOURCE_IN_WIM &&
             !(ctx->write_resource_flags & WIMLIB_RESOURCE_FLAG_RECOMPRESS) &&
-            wimlib_get_compression_type(lte->wim) == ctx->out_ctype))
+            lte->wim->compression_type == ctx->out_ctype))
        {
                /* Stream is too small or isn't being compressed.  Process it by
                 * the main thread when we have a chance.  We can't necessarily
@@ -1200,7 +1224,7 @@ main_thread_process_next_stream(struct wim_lookup_table_entry *lte, void *_ctx)
 }
 
 static long
-get_default_num_threads()
+get_default_num_threads(void)
 {
 #ifdef __WIN32__
        return win32_get_number_of_processors();
@@ -1415,7 +1439,7 @@ write_stream_list(struct list_head *stream_list,
        progress.write_streams._private          = 0;
 
 #ifdef ENABLE_MULTITHREADED_COMPRESSION
-       if (total_compression_bytes >= 1000000 && num_threads != 1)
+       if (total_compression_bytes >= 2000000 && num_threads != 1)
                ret = write_stream_list_parallel(stream_list,
                                                 lookup_table,
                                                 out_fd,
@@ -1523,10 +1547,10 @@ lte_overwrite_prepare_2(struct wim_lookup_table_entry *lte, void *_args)
                if (lte->resource_entry.offset +
                    lte->resource_entry.size > args->end_offset)
                {
-               #ifdef ENABLE_ERROR_MESSAGES
-                       ERROR("The following resource is after the XML data:");
-                       print_lookup_table_entry(lte, stderr);
-               #endif
+                       if (wimlib_print_errors) {
+                               ERROR("The following resource is after the XML data:");
+                               print_lookup_table_entry(lte, stderr);
+                       }
                        return WIMLIB_ERR_RESOURCE_ORDER;
                }
                copy_resource_entry(&lte->output_resource_entry,
@@ -1689,7 +1713,7 @@ write_wim_streams(WIMStruct *wim, int image, int write_flags,
        return write_stream_list(&stream_list,
                                 wim->lookup_table,
                                 wim->out_fd,
-                                wimlib_get_compression_type(wim),
+                                wim->compression_type,
                                 write_flags,
                                 num_threads,
                                 progress_func);
@@ -1704,9 +1728,6 @@ write_wim_streams(WIMStruct *wim, int image, int write_flags,
  *     (public)  WIMLIB_WRITE_FLAG_CHECK_INTEGRITY:
  *             Include an integrity table.
  *
- *     (public)  WIMLIB_WRITE_FLAG_SHOW_PROGRESS:
- *             Show progress information when (if) writing the integrity table.
- *
  *     (private) WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE:
  *             Don't write the lookup table.
  *
@@ -1777,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;
@@ -1804,6 +1826,7 @@ finish_write(WIMStruct *w, int image, int write_flags,
                zero_resource_entry(&hdr.integrity);
        }
 
+       hdr.flags &= ~WIM_HDR_FLAG_WRITE_IN_PROGRESS;
        ret = write_header(&hdr, w->out_fd);
        if (ret)
                goto out_close_wim;
@@ -1886,7 +1909,9 @@ begin_write(WIMStruct *w, const tchar *path, int write_flags)
        if (ret)
                return ret;
        /* Write dummy header. It will be overwritten later. */
+       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;
        if (ret)
                return ret;
        if (lseek(w->out_fd, WIM_HEADER_DISK_SIZE, SEEK_SET) == -1) {
@@ -2088,11 +2113,20 @@ overwrite_wim_inplace(WIMStruct *w, int write_flags,
                return ret;
        }
 
+       /* 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;
+       }
+
        if (lseek(w->out_fd, old_wim_end, SEEK_SET) == -1) {
                ERROR_WITH_ERRNO("Can't seek to end of WIM");
                close_wim_writable(w);
-               w->wim_locked = 0;
-               return WIMLIB_ERR_WRITE;
+               ret = WIMLIB_ERR_WRITE;
+               goto out_unlock_wim;
        }
 
        DEBUG("Writing newly added streams (offset = %"PRIu64")",
@@ -2100,7 +2134,7 @@ overwrite_wim_inplace(WIMStruct *w, int write_flags,
        ret = write_stream_list(&stream_list,
                                w->lookup_table,
                                w->out_fd,
-                               wimlib_get_compression_type(w),
+                               w->compression_type,
                                write_flags,
                                num_threads,
                                progress_func);
@@ -2127,6 +2161,7 @@ out_truncate:
                 * an error path. */
                (void)ttruncate(w->filename, old_wim_end);
        }
+out_unlock_wim:
        w->wim_locked = 0;
        return ret;
 }
@@ -2191,15 +2226,16 @@ wimlib_overwrite(WIMStruct *w, int write_flags,
                 unsigned num_threads,
                 wimlib_progress_func_t progress_func)
 {
+       int ret;
+
        write_flags &= WIMLIB_WRITE_MASK_PUBLIC;
 
        if (!w->filename)
                return WIMLIB_ERR_NO_FILENAME;
 
-       if (w->hdr.total_parts != 1) {
-               ERROR("Cannot modify a split WIM");
-               return WIMLIB_ERR_SPLIT_UNSUPPORTED;
-       }
+       ret = can_modify_wim(w);
+       if (ret)
+               return ret;
 
        if ((!w->deletion_occurred || (write_flags & WIMLIB_WRITE_FLAG_SOFT_DELETE))
            && !(write_flags & WIMLIB_WRITE_FLAG_REBUILD))