Fix setting header flag during overwrite

[wimlib] / src / write.c

diff --git a/src/write.c b/src/write.c
index 90dcaed4201d0d92ce6cf74dfc9691dd4bf11025..4c68921cc9164efb926e58325f0e0a99109ec6bc 100644 (file)
--- a/src/write.c
+++ b/src/write.c
@@ -377,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;
@@ -1207,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
@@ -1547,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,
@@ -1713,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);
@@ -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;
@@ -1825,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;
@@ -1907,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) {
@@ -2109,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")",
@@ -2121,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);
@@ -2148,6 +2161,7 @@ out_truncate:
                 * an error path. */
                (void)ttruncate(w->filename, old_wim_end);
        }
+out_unlock_wim:
        w->wim_locked = 0;
        return ret;
 }
@@ -2212,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))