wimlib.net Git - wimlib/blob - src/write.c

   1 /*
   2  * write.c
   3  *
   4  * Support for writing WIM files; write a WIM file, overwrite a WIM file, write
   5  * compressed file resources, etc.
   6  */
   7
   8 /*
   9  * Copyright (C) 2010 Carl Thijssen
  10  * Copyright (C) 2012 Eric Biggers
  11  *
  12  * This file is part of wimlib, a library for working with WIM files.
  13  *
  14  * wimlib is free software; you can redistribute it and/or modify it under the
  15  * terms of the GNU General Public License as published by the Free
  16  * Software Foundation; either version 3 of the License, or (at your option)
  17  * any later version.
  18  *
  19  * wimlib is distributed in the hope that it will be useful, but WITHOUT ANY
  20  * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  21  * A PARTICULAR PURPOSE. See the GNU General Public License for more
  22  * details.
  23  *
  24  * You should have received a copy of the GNU General Public License
  25  * along with wimlib; if not, see http://www.gnu.org/licenses/.
  26  */
  27
  28 #include "config.h"
  29
  30 #if defined(HAVE_SYS_FILE_H) && defined(HAVE_FLOCK)
  31 /* On BSD, this should be included before "list.h" so that "list.h" can
  32  * overwrite the LIST_HEAD macro. */
  33 #include <sys/file.h>
  34 #endif
  35
  36 #include "list.h"
  37 #include "wimlib_internal.h"
  38 #include "io.h"
  39 #include "dentry.h"
  40 #include "lookup_table.h"
  41 #include "xml.h"
  42 #include "lzx.h"
  43 #include "xpress.h"
  44
  45 #ifdef ENABLE_MULTITHREADED_COMPRESSION
  46 #include <pthread.h>
  47 #endif
  48
  49 #include <unistd.h>
  50 #include <errno.h>
  51
  52 #ifdef WITH_NTFS_3G
  53 #include <time.h>
  54 #include <ntfs-3g/attrib.h>
  55 #include <ntfs-3g/inode.h>
  56 #include <ntfs-3g/dir.h>
  57 #endif
  58
  59 #ifdef HAVE_ALLOCA_H
  60 #include <alloca.h>
  61 #else
  62 #include <stdlib.h>
  63 #endif
  64
  65
  66 static int do_fflush(FILE *fp)
  67 {
  68         int ret = fflush(fp);
  69         if (ret != 0) {
  70                 ERROR_WITH_ERRNO("Failed to flush data to output WIM file");
  71                 return WIMLIB_ERR_WRITE;
  72         }
  73         return 0;
  74 }
  75
  76 static int fflush_and_ftruncate(FILE *fp, off_t size)
  77 {
  78         int ret;
  79
  80         ret = do_fflush(fp);
  81         if (ret != 0)
  82                 return ret;
  83         ret = ftruncate(fileno(fp), size);
  84         if (ret != 0) {
  85                 ERROR_WITH_ERRNO("Failed to truncate output WIM file to "
  86                                  "%"PRIu64" bytes", size);
  87                 return WIMLIB_ERR_WRITE;
  88         }
  89         return 0;
  90 }
  91
  92 /* Chunk table that's located at the beginning of each compressed resource in
  93  * the WIM.  (This is not the on-disk format; the on-disk format just has an
  94  * array of offsets.) */
  95 struct chunk_table {
  96         off_t file_offset;
  97         u64 num_chunks;
  98         u64 original_resource_size;
  99         u64 bytes_per_chunk_entry;
 100         u64 table_disk_size;
 101         u64 cur_offset;
 102         u64 *cur_offset_p;
 103         u64 offsets[0];
 104 };
 105
 106 /*
 107  * Allocates and initializes a chunk table, and reserves space for it in the
 108  * output file.
 109  */
 110 static int
 111 begin_wim_resource_chunk_tab(const struct lookup_table_entry *lte,
 112                              FILE *out_fp,
 113                              off_t file_offset,
 114                              struct chunk_table **chunk_tab_ret)
 115 {
 116         u64 size = wim_resource_size(lte);
 117         u64 num_chunks = (size + WIM_CHUNK_SIZE - 1) / WIM_CHUNK_SIZE;
 118         size_t alloc_size = sizeof(struct chunk_table) + num_chunks * sizeof(u64);
 119         struct chunk_table *chunk_tab = CALLOC(1, alloc_size);
 120         int ret;
 121
 122         if (!chunk_tab) {
 123                 ERROR("Failed to allocate chunk table for %"PRIu64" byte "
 124                       "resource", size);
 125                 ret = WIMLIB_ERR_NOMEM;
 126                 goto out;
 127         }
 128         chunk_tab->file_offset = file_offset;
 129         chunk_tab->num_chunks = num_chunks;
 130         chunk_tab->original_resource_size = size;
 131         chunk_tab->bytes_per_chunk_entry = (size >= (1ULL << 32)) ? 8 : 4;
 132         chunk_tab->table_disk_size = chunk_tab->bytes_per_chunk_entry *
 133                                      (num_chunks - 1);
 134         chunk_tab->cur_offset = 0;
 135         chunk_tab->cur_offset_p = chunk_tab->offsets;
 136
 137         if (fwrite(chunk_tab, 1, chunk_tab->table_disk_size, out_fp) !=
 138                    chunk_tab->table_disk_size) {
 139                 ERROR_WITH_ERRNO("Failed to write chunk table in compressed "
 140                                  "file resource");
 141                 ret = WIMLIB_ERR_WRITE;
 142                 goto out;
 143         }
 144
 145         ret = 0;
 146 out:
 147         *chunk_tab_ret = chunk_tab;
 148         return ret;
 149 }
 150
 151 /*
 152  * Pointer to function to compresses a chunk of a WIM resource.
 153  *
 154  * @chunk:              Uncompressed data of the chunk.
 155  * @chunk_size:         Size of the uncompressed chunk in bytes.
 156  * @compressed_chunk:   Pointer to output buffer of size at least
 157  *                              (@chunk_size - 1) bytes.
 158  * @compressed_chunk_len_ret:   Pointer to an unsigned int into which the size
 159  *                                      of the compressed chunk will be
 160  *                                      returned.
 161  *
 162  * Returns zero if compressed succeeded, and nonzero if the chunk could not be
 163  * compressed to any smaller than @chunk_size.  This function cannot fail for
 164  * any other reasons.
 165  */
 166 typedef int (*compress_func_t)(const void *, unsigned, void *, unsigned *);
 167
 168 compress_func_t get_compress_func(int out_ctype)
 169 {
 170         if (out_ctype == WIMLIB_COMPRESSION_TYPE_LZX)
 171                 return lzx_compress;
 172         else
 173                 return xpress_compress;
 174 }
 175
 176 /*
 177  * Writes a chunk of a WIM resource to an output file.
 178  *
 179  * @chunk:        Uncompressed data of the chunk.
 180  * @chunk_size:   Size of the chunk (<= WIM_CHUNK_SIZE)
 181  * @out_fp:       FILE * to write tho chunk to.
 182  * @out_ctype:    Compression type to use when writing the chunk (ignored if no
 183  *                      chunk table provided)
 184  * @chunk_tab:    Pointer to chunk table being created.  It is updated with the
 185  *                      offset of the chunk we write.
 186  *
 187  * Returns 0 on success; nonzero on failure.
 188  */
 189 static int write_wim_resource_chunk(const u8 chunk[], unsigned chunk_size,
 190                                     FILE *out_fp, compress_func_t compress,
 191                                     struct chunk_table *chunk_tab)
 192 {
 193         const u8 *out_chunk;
 194         unsigned out_chunk_size;
 195         if (chunk_tab) {
 196                 u8 *compressed_chunk = alloca(chunk_size);
 197                 int ret;
 198
 199                 ret = compress(chunk, chunk_size, compressed_chunk,
 200                                &out_chunk_size);
 201                 if (ret == 0) {
 202                         out_chunk = compressed_chunk;
 203                 } else {
 204                         out_chunk = chunk;
 205                         out_chunk_size = chunk_size;
 206                 }
 207                 *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
 208                 chunk_tab->cur_offset += out_chunk_size;
 209         } else {
 210                 out_chunk = chunk;
 211                 out_chunk_size = chunk_size;
 212         }
 213         if (fwrite(out_chunk, 1, out_chunk_size, out_fp) != out_chunk_size) {
 214                 ERROR_WITH_ERRNO("Failed to write WIM resource chunk");
 215                 return WIMLIB_ERR_WRITE;
 216         }
 217         return 0;
 218 }
 219
 220 /*
 221  * Finishes a WIM chunk tale and writes it to the output file at the correct
 222  * offset.
 223  *
 224  * The final size of the full compressed resource is returned in the
 225  * @compressed_size_p.
 226  */
 227 static int
 228 finish_wim_resource_chunk_tab(struct chunk_table *chunk_tab,
 229                               FILE *out_fp, u64 *compressed_size_p)
 230 {
 231         size_t bytes_written;
 232         if (fseeko(out_fp, chunk_tab->file_offset, SEEK_SET) != 0) {
 233                 ERROR_WITH_ERRNO("Failed to seek to byte %"PRIu64" of output "
 234                                  "WIM file", chunk_tab->file_offset);
 235                 return WIMLIB_ERR_WRITE;
 236         }
 237
 238         if (chunk_tab->bytes_per_chunk_entry == 8) {
 239                 array_cpu_to_le64(chunk_tab->offsets, chunk_tab->num_chunks);
 240         } else {
 241                 for (u64 i = 0; i < chunk_tab->num_chunks; i++)
 242                         ((u32*)chunk_tab->offsets)[i] =
 243                                 cpu_to_le32(chunk_tab->offsets[i]);
 244         }
 245         bytes_written = fwrite((u8*)chunk_tab->offsets +
 246                                         chunk_tab->bytes_per_chunk_entry,
 247                                1, chunk_tab->table_disk_size, out_fp);
 248         if (bytes_written != chunk_tab->table_disk_size) {
 249                 ERROR_WITH_ERRNO("Failed to write chunk table in compressed "
 250                                  "file resource");
 251                 return WIMLIB_ERR_WRITE;
 252         }
 253         if (fseeko(out_fp, 0, SEEK_END) != 0) {
 254                 ERROR_WITH_ERRNO("Failed to seek to end of output WIM file");
 255                 return WIMLIB_ERR_WRITE;
 256         }
 257         *compressed_size_p = chunk_tab->cur_offset + chunk_tab->table_disk_size;
 258         return 0;
 259 }
 260
 261 /* Prepare for multiple reads to a resource by caching a FILE * or NTFS
 262  * attribute pointer in the lookup table entry. */
 263 static int prepare_resource_for_read(struct lookup_table_entry *lte
 264
 265                                         #ifdef WITH_NTFS_3G
 266                                         , ntfs_inode **ni_ret
 267                                         #endif
 268                 )
 269 {
 270         if (lte->resource_location == RESOURCE_IN_FILE_ON_DISK
 271              && !lte->file_on_disk_fp)
 272         {
 273                 wimlib_assert(lte->file_on_disk);
 274                 lte->file_on_disk_fp = fopen(lte->file_on_disk, "rb");
 275                 if (!lte->file_on_disk_fp) {
 276                         ERROR_WITH_ERRNO("Failed to open the file `%s' for "
 277                                          "reading", lte->file_on_disk);
 278                         return WIMLIB_ERR_OPEN;
 279                 }
 280         }
 281 #ifdef WITH_NTFS_3G
 282         else if (lte->resource_location == RESOURCE_IN_NTFS_VOLUME
 283                   && !lte->attr)
 284         {
 285                 struct ntfs_location *loc = lte->ntfs_loc;
 286                 ntfs_inode *ni;
 287                 wimlib_assert(loc);
 288                 ni = ntfs_pathname_to_inode(*loc->ntfs_vol_p, NULL, loc->path_utf8);
 289                 if (!ni) {
 290                         ERROR_WITH_ERRNO("Failed to open inode `%s' in NTFS "
 291                                          "volume", loc->path_utf8);
 292                         return WIMLIB_ERR_NTFS_3G;
 293                 }
 294                 lte->attr = ntfs_attr_open(ni,
 295                                            loc->is_reparse_point ? AT_REPARSE_POINT : AT_DATA,
 296                                            (ntfschar*)loc->stream_name_utf16,
 297                                            loc->stream_name_utf16_num_chars);
 298                 if (!lte->attr) {
 299                         ERROR_WITH_ERRNO("Failed to open attribute of `%s' in "
 300                                          "NTFS volume", loc->path_utf8);
 301                         ntfs_inode_close(ni);
 302                         return WIMLIB_ERR_NTFS_3G;
 303                 }
 304                 *ni_ret = ni;
 305         }
 306 #endif
 307         return 0;
 308 }
 309
 310 /* Undo prepare_resource_for_read() by closing the cached FILE * or NTFS
 311  * attribute. */
 312 static void end_wim_resource_read(struct lookup_table_entry *lte
 313                                 #ifdef WITH_NTFS_3G
 314                                         , ntfs_inode *ni
 315                                 #endif
 316                                         )
 317 {
 318         if (lte->resource_location == RESOURCE_IN_FILE_ON_DISK
 319             && lte->file_on_disk_fp) {
 320                 fclose(lte->file_on_disk_fp);
 321                 lte->file_on_disk_fp = NULL;
 322         }
 323 #ifdef WITH_NTFS_3G
 324         else if (lte->resource_location == RESOURCE_IN_NTFS_VOLUME) {
 325                 if (lte->attr) {
 326                         ntfs_attr_close(lte->attr);
 327                         lte->attr = NULL;
 328                 }
 329                 if (ni)
 330                         ntfs_inode_close(ni);
 331         }
 332 #endif
 333 }
 334
 335 /*
 336  * Writes a WIM resource to a FILE * opened for writing.  The resource may be
 337  * written uncompressed or compressed depending on the @out_ctype parameter.
 338  *
 339  * If by chance the resource compresses to more than the original size (this may
 340  * happen with random data or files than are pre-compressed), the resource is
 341  * instead written uncompressed (and this is reflected in the @out_res_entry by
 342  * removing the WIM_RESHDR_FLAG_COMPRESSED flag).
 343  *
 344  * @lte:        The lookup table entry for the WIM resource.
 345  * @out_fp:     The FILE * to write the resource to.
 346  * @out_ctype:  The compression type of the resource to write.  Note: if this is
 347  *                      the same as the compression type of the WIM resource we
 348  *                      need to read, we simply copy the data (i.e. we do not
 349  *                      uncompress it, then compress it again).
 350  * @out_res_entry:  If non-NULL, a resource entry that is filled in with the
 351  *                  offset, original size, compressed size, and compression flag
 352  *                  of the output resource.
 353  *
 354  * Returns 0 on success; nonzero on failure.
 355  */
 356 int write_wim_resource(struct lookup_table_entry *lte,
 357                        FILE *out_fp, int out_ctype,
 358                        struct resource_entry *out_res_entry,
 359                        int flags)
 360 {
 361         u64 bytes_remaining;
 362         u64 original_size;
 363         u64 old_compressed_size;
 364         u64 new_compressed_size;
 365         u64 offset;
 366         int ret;
 367         struct chunk_table *chunk_tab = NULL;
 368         bool raw;
 369         off_t file_offset;
 370         compress_func_t compress = NULL;
 371 #ifdef WITH_NTFS_3G
 372         ntfs_inode *ni = NULL;
 373 #endif
 374
 375         wimlib_assert(lte);
 376
 377         /* Original size of the resource */
 378         original_size = wim_resource_size(lte);
 379
 380         /* Compressed size of the resource (as it exists now) */
 381         old_compressed_size = wim_resource_compressed_size(lte);
 382
 383         /* Current offset in output file */
 384         file_offset = ftello(out_fp);
 385         if (file_offset == -1) {
 386                 ERROR_WITH_ERRNO("Failed to get offset in output "
 387                                  "stream");
 388                 return WIMLIB_ERR_WRITE;
 389         }
 390
 391         /* Are the compression types the same?  If so, do a raw copy (copy
 392          * without decompressing and recompressing the data). */
 393         raw = (wim_resource_compression_type(lte) == out_ctype
 394                && out_ctype != WIMLIB_COMPRESSION_TYPE_NONE
 395                && !(flags & WIMLIB_RESOURCE_FLAG_RECOMPRESS));
 396
 397         if (raw) {
 398                 flags |= WIMLIB_RESOURCE_FLAG_RAW;
 399                 bytes_remaining = old_compressed_size;
 400         } else {
 401                 flags &= ~WIMLIB_RESOURCE_FLAG_RAW;
 402                 bytes_remaining = original_size;
 403         }
 404
 405         /* Empty resource; nothing needs to be done, so just return success. */
 406         if (bytes_remaining == 0)
 407                 return 0;
 408
 409         /* Buffer for reading chunks for the resource */
 410         u8 buf[min(WIM_CHUNK_SIZE, bytes_remaining)];
 411
 412         /* If we are writing a compressed resource and not doing a raw copy, we
 413          * need to initialize the chunk table */
 414         if (out_ctype != WIMLIB_COMPRESSION_TYPE_NONE && !raw) {
 415                 ret = begin_wim_resource_chunk_tab(lte, out_fp, file_offset,
 416                                                    &chunk_tab);
 417                 if (ret != 0)
 418                         goto out;
 419         }
 420
 421         /* If the WIM resource is in an external file, open a FILE * to it so we
 422          * don't have to open a temporary one in read_wim_resource() for each
 423          * chunk. */
 424 #ifdef WITH_NTFS_3G
 425         ret = prepare_resource_for_read(lte, &ni);
 426 #else
 427         ret = prepare_resource_for_read(lte);
 428 #endif
 429         if (ret != 0)
 430                 goto out;
 431
 432         /* If we aren't doing a raw copy, we will compute the SHA1 message
 433          * digest of the resource as we read it, and verify it's the same as the
 434          * hash given in the lookup table entry once we've finished reading the
 435          * resource. */
 436         SHA_CTX ctx;
 437         if (!raw) {
 438                 sha1_init(&ctx);
 439                 compress = get_compress_func(out_ctype);
 440         }
 441         offset = 0;
 442
 443         /* While there are still bytes remaining in the WIM resource, read a
 444          * chunk of the resource, update SHA1, then write that chunk using the
 445          * desired compression type. */
 446         do {
 447                 u64 to_read = min(bytes_remaining, WIM_CHUNK_SIZE);
 448                 ret = read_wim_resource(lte, buf, to_read, offset, flags);
 449                 if (ret != 0)
 450                         goto out_fclose;
 451                 if (!raw)
 452                         sha1_update(&ctx, buf, to_read);
 453                 ret = write_wim_resource_chunk(buf, to_read, out_fp,
 454                                                compress, chunk_tab);
 455                 if (ret != 0)
 456                         goto out_fclose;
 457                 bytes_remaining -= to_read;
 458                 offset += to_read;
 459         } while (bytes_remaining);
 460
 461         /* Raw copy:  The new compressed size is the same as the old compressed
 462          * size
 463          *
 464          * Using WIMLIB_COMPRESSION_TYPE_NONE:  The new compressed size is the
 465          * original size
 466          *
 467          * Using a different compression type:  Call
 468          * finish_wim_resource_chunk_tab() and it will provide the new
 469          * compressed size.
 470          */
 471         if (raw) {
 472                 new_compressed_size = old_compressed_size;
 473         } else {
 474                 if (out_ctype == WIMLIB_COMPRESSION_TYPE_NONE)
 475                         new_compressed_size = original_size;
 476                 else {
 477                         ret = finish_wim_resource_chunk_tab(chunk_tab, out_fp,
 478                                                             &new_compressed_size);
 479                         if (ret != 0)
 480                                 goto out_fclose;
 481                 }
 482         }
 483
 484         /* Verify SHA1 message digest of the resource, unless we are doing a raw
 485          * write (in which case we never even saw the uncompressed data).  Or,
 486          * if the hash we had before is all 0's, just re-set it to be the new
 487          * hash. */
 488         if (!raw) {
 489                 u8 md[SHA1_HASH_SIZE];
 490                 sha1_final(md, &ctx);
 491                 if (is_zero_hash(lte->hash)) {
 492                         copy_hash(lte->hash, md);
 493                 } else if (!hashes_equal(md, lte->hash)) {
 494                         ERROR("WIM resource has incorrect hash!");
 495                         if (lte->resource_location == RESOURCE_IN_FILE_ON_DISK) {
 496                                 ERROR("We were reading it from `%s'; maybe it changed "
 497                                       "while we were reading it.",
 498                                       lte->file_on_disk);
 499                         }
 500                         ret = WIMLIB_ERR_INVALID_RESOURCE_HASH;
 501                         goto out_fclose;
 502                 }
 503         }
 504
 505         if (!raw && new_compressed_size >= original_size &&
 506             out_ctype != WIMLIB_COMPRESSION_TYPE_NONE)
 507         {
 508                 /* Oops!  We compressed the resource to larger than the original
 509                  * size.  Write the resource uncompressed instead. */
 510                 if (fseeko(out_fp, file_offset, SEEK_SET) != 0) {
 511                         ERROR_WITH_ERRNO("Failed to seek to byte %"PRIu64" "
 512                                          "of output WIM file", file_offset);
 513                         ret = WIMLIB_ERR_WRITE;
 514                         goto out_fclose;
 515                 }
 516                 ret = write_wim_resource(lte, out_fp, WIMLIB_COMPRESSION_TYPE_NONE,
 517                                          out_res_entry, flags);
 518                 if (ret != 0)
 519                         goto out_fclose;
 520
 521                 ret = fflush_and_ftruncate(out_fp, file_offset + out_res_entry->size);
 522                 if (ret != 0)
 523                         goto out_fclose;
 524         } else {
 525                 if (out_res_entry) {
 526                         out_res_entry->size          = new_compressed_size;
 527                         out_res_entry->original_size = original_size;
 528                         out_res_entry->offset        = file_offset;
 529                         out_res_entry->flags         = lte->resource_entry.flags
 530                                                         & ~WIM_RESHDR_FLAG_COMPRESSED;
 531                         if (out_ctype != WIMLIB_COMPRESSION_TYPE_NONE)
 532                                 out_res_entry->flags |= WIM_RESHDR_FLAG_COMPRESSED;
 533                 }
 534         }
 535         ret = 0;
 536 out_fclose:
 537 #ifdef WITH_NTFS_3G
 538         end_wim_resource_read(lte, ni);
 539 #else
 540         end_wim_resource_read(lte);
 541 #endif
 542 out:
 543         FREE(chunk_tab);
 544         return ret;
 545 }
 546
 547 #ifdef ENABLE_MULTITHREADED_COMPRESSION
 548 struct shared_queue {
 549         unsigned size;
 550         unsigned front;
 551         unsigned back;
 552         unsigned filled_slots;
 553         void **array;
 554         pthread_mutex_t lock;
 555         pthread_cond_t msg_avail_cond;
 556         pthread_cond_t space_avail_cond;
 557 };
 558
 559 static int shared_queue_init(struct shared_queue *q, unsigned size)
 560 {
 561         q->array = CALLOC(sizeof(q->array[0]), size);
 562         if (!q->array)
 563                 return WIMLIB_ERR_NOMEM;
 564         q->filled_slots = 0;
 565         q->front = 0;
 566         q->back = size - 1;
 567         q->size = size;
 568         pthread_mutex_init(&q->lock, NULL);
 569         pthread_cond_init(&q->msg_avail_cond, NULL);
 570         pthread_cond_init(&q->space_avail_cond, NULL);
 571         return 0;
 572 }
 573
 574 static void shared_queue_destroy(struct shared_queue *q)
 575 {
 576         FREE(q->array);
 577         pthread_mutex_destroy(&q->lock);
 578         pthread_cond_destroy(&q->msg_avail_cond);
 579         pthread_cond_destroy(&q->space_avail_cond);
 580 }
 581
 582 static void shared_queue_put(struct shared_queue *q, void *obj)
 583 {
 584         pthread_mutex_lock(&q->lock);
 585         while (q->filled_slots == q->size)
 586                 pthread_cond_wait(&q->space_avail_cond, &q->lock);
 587
 588         q->back = (q->back + 1) % q->size;
 589         q->array[q->back] = obj;
 590         q->filled_slots++;
 591
 592         pthread_cond_broadcast(&q->msg_avail_cond);
 593         pthread_mutex_unlock(&q->lock);
 594 }
 595
 596 static void *shared_queue_get(struct shared_queue *q)
 597 {
 598         void *obj;
 599
 600         pthread_mutex_lock(&q->lock);
 601         while (q->filled_slots == 0)
 602                 pthread_cond_wait(&q->msg_avail_cond, &q->lock);
 603
 604         obj = q->array[q->front];
 605         q->array[q->front] = NULL;
 606         q->front = (q->front + 1) % q->size;
 607         q->filled_slots--;
 608
 609         pthread_cond_broadcast(&q->space_avail_cond);
 610         pthread_mutex_unlock(&q->lock);
 611         return obj;
 612 }
 613
 614 struct compressor_thread_params {
 615         struct shared_queue *res_to_compress_queue;
 616         struct shared_queue *compressed_res_queue;
 617         compress_func_t compress;
 618 };
 619
 620 #define MAX_CHUNKS_PER_MSG 2
 621
 622 struct message {
 623         struct lookup_table_entry *lte;
 624         u8 *uncompressed_chunks[MAX_CHUNKS_PER_MSG];
 625         u8 *out_compressed_chunks[MAX_CHUNKS_PER_MSG];
 626         u8 *compressed_chunks[MAX_CHUNKS_PER_MSG];
 627         unsigned uncompressed_chunk_sizes[MAX_CHUNKS_PER_MSG];
 628         unsigned compressed_chunk_sizes[MAX_CHUNKS_PER_MSG];
 629         unsigned num_chunks;
 630         struct list_head list;
 631         bool complete;
 632         u64 begin_chunk;
 633 };
 634
 635 static void compress_chunks(struct message *msg, compress_func_t compress)
 636 {
 637         for (unsigned i = 0; i < msg->num_chunks; i++) {
 638                 DEBUG2("compress chunk %u of %u", i, msg->num_chunks);
 639                 int ret = compress(msg->uncompressed_chunks[i],
 640                                    msg->uncompressed_chunk_sizes[i],
 641                                    msg->compressed_chunks[i],
 642                                    &msg->compressed_chunk_sizes[i]);
 643                 if (ret == 0) {
 644                         msg->out_compressed_chunks[i] = msg->compressed_chunks[i];
 645                 } else {
 646                         msg->out_compressed_chunks[i] = msg->uncompressed_chunks[i];
 647                         msg->compressed_chunk_sizes[i] = msg->uncompressed_chunk_sizes[i];
 648                 }
 649         }
 650 }
 651
 652 static void *compressor_thread_proc(void *arg)
 653 {
 654         struct compressor_thread_params *params = arg;
 655         struct shared_queue *res_to_compress_queue = params->res_to_compress_queue;
 656         struct shared_queue *compressed_res_queue = params->compressed_res_queue;
 657         compress_func_t compress = params->compress;
 658         struct message *msg;
 659
 660         DEBUG("Compressor thread ready");
 661         while ((msg = shared_queue_get(res_to_compress_queue)) != NULL) {
 662                 compress_chunks(msg, compress);
 663                 shared_queue_put(compressed_res_queue, msg);
 664         }
 665         DEBUG("Compressor thread terminating");
 666         return NULL;
 667 }
 668 #endif
 669
 670 static int do_write_stream_list(struct list_head *my_resources,
 671                                 FILE *out_fp,
 672                                 int out_ctype,
 673                                 wimlib_progress_func_t progress_func,
 674                                 union wimlib_progress_info *progress,
 675                                 int write_resource_flags)
 676 {
 677         int ret;
 678         struct lookup_table_entry *lte, *tmp;
 679
 680         list_for_each_entry_safe(lte, tmp, my_resources, staging_list) {
 681                 ret = write_wim_resource(lte,
 682                                          out_fp,
 683                                          out_ctype,
 684                                          &lte->output_resource_entry,
 685                                          write_resource_flags);
 686                 if (ret != 0)
 687                         return ret;
 688                 list_del(&lte->staging_list);
 689                 progress->write_streams.completed_bytes +=
 690                         wim_resource_size(lte);
 691                 progress->write_streams.completed_streams++;
 692                 if (progress_func) {
 693                         progress_func(WIMLIB_PROGRESS_MSG_WRITE_STREAMS,
 694                                       progress);
 695                 }
 696         }
 697         return 0;
 698 }
 699
 700 static int write_stream_list_serial(struct list_head *stream_list,
 701                                     FILE *out_fp,
 702                                     int out_ctype,
 703                                     int write_flags,
 704                                     wimlib_progress_func_t progress_func,
 705                                     union wimlib_progress_info *progress)
 706 {
 707         int write_resource_flags;
 708
 709         if (write_flags & WIMLIB_WRITE_FLAG_RECOMPRESS)
 710                 write_resource_flags = WIMLIB_RESOURCE_FLAG_RECOMPRESS;
 711         else
 712                 write_resource_flags = 0;
 713         progress->write_streams.num_threads = 1;
 714         if (progress_func)
 715                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_STREAMS, progress);
 716         return do_write_stream_list(stream_list, out_fp,
 717                                     out_ctype, progress_func,
 718                                     progress, write_resource_flags);
 719 }
 720
 721 #ifdef ENABLE_MULTITHREADED_COMPRESSION
 722 static int write_wim_chunks(struct message *msg, FILE *out_fp,
 723                             struct chunk_table *chunk_tab)
 724 {
 725         for (unsigned i = 0; i < msg->num_chunks; i++) {
 726                 unsigned chunk_csize = msg->compressed_chunk_sizes[i];
 727
 728                 DEBUG2("Write wim chunk %u of %u (csize = %u)",
 729                       i, msg->num_chunks, chunk_csize);
 730
 731                 if (fwrite(msg->out_compressed_chunks[i], 1, chunk_csize, out_fp)
 732                     != chunk_csize)
 733                 {
 734                         ERROR_WITH_ERRNO("Failed to write WIM chunk");
 735                         return WIMLIB_ERR_WRITE;
 736                 }
 737
 738                 *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
 739                 chunk_tab->cur_offset += chunk_csize;
 740         }
 741         return 0;
 742 }
 743
 744 /*
 745  * This function is executed by the main thread when the resources are being
 746  * compressed in parallel.  The main thread is in change of all reading of the
 747  * uncompressed data and writing of the compressed data.  The compressor threads
 748  * *only* do compression from/to in-memory buffers.
 749  *
 750  * Each unit of work given to a compressor thread is up to MAX_CHUNKS_PER_MSG
 751  * chunks of compressed data to compress, represented in a `struct message'.
 752  * Each message is passed from the main thread to a worker thread through the
 753  * res_to_compress_queue, and it is passed back through the
 754  * compressed_res_queue.
 755  */
 756 static int main_writer_thread_proc(struct list_head *stream_list,
 757                                    FILE *out_fp,
 758                                    int out_ctype,
 759                                    struct shared_queue *res_to_compress_queue,
 760                                    struct shared_queue *compressed_res_queue,
 761                                    size_t queue_size,
 762                                    int write_flags,
 763                                    wimlib_progress_func_t progress_func,
 764                                    union wimlib_progress_info *progress)
 765 {
 766         int ret;
 767
 768         struct message msgs[queue_size];
 769         ZERO_ARRAY(msgs);
 770
 771         // Initially, all the messages are available to use.
 772         LIST_HEAD(available_msgs);
 773         for (size_t i = 0; i < ARRAY_LEN(msgs); i++)
 774                 list_add(&msgs[i].list, &available_msgs);
 775
 776         // outstanding_resources is the list of resources that currently have
 777         // had chunks sent off for compression.
 778         //
 779         // The first stream in outstanding_resources is the stream that is
 780         // currently being written (cur_lte).
 781         //
 782         // The last stream in outstanding_resources is the stream that is
 783         // currently being read and chunks fed to the compressor threads
 784         // (next_lte).
 785         //
 786         // Depending on the number of threads and the sizes of the resource,
 787         // the outstanding streams list may contain streams between cur_lte and
 788         // next_lte that have all their chunks compressed or being compressed,
 789         // but haven't been written yet.
 790         //
 791         LIST_HEAD(outstanding_resources);
 792         struct list_head *next_resource = stream_list->next;
 793         struct lookup_table_entry *next_lte = container_of(next_resource,
 794                                                            struct lookup_table_entry,
 795                                                            staging_list);
 796         next_resource = next_resource->next;
 797         u64 next_chunk = 0;
 798         u64 next_num_chunks = wim_resource_chunks(next_lte);
 799         INIT_LIST_HEAD(&next_lte->msg_list);
 800         list_add_tail(&next_lte->staging_list, &outstanding_resources);
 801
 802         // As in write_wim_resource(), each resource we read is checksummed.
 803         SHA_CTX next_sha_ctx;
 804         sha1_init(&next_sha_ctx);
 805         u8 next_hash[SHA1_HASH_SIZE];
 806
 807         // Resources that don't need any chunks compressed are added to this
 808         // list and written directly by the main thread.
 809         LIST_HEAD(my_resources);
 810
 811         struct lookup_table_entry *cur_lte = next_lte;
 812         struct chunk_table *cur_chunk_tab = NULL;
 813         struct message *msg;
 814
 815 #ifdef WITH_NTFS_3G
 816         ntfs_inode *ni = NULL;
 817 #endif
 818
 819 #ifdef WITH_NTFS_3G
 820         ret = prepare_resource_for_read(next_lte, &ni);
 821 #else
 822         ret = prepare_resource_for_read(next_lte);
 823 #endif
 824         if (ret != 0)
 825                 goto out;
 826
 827         DEBUG("Initializing buffers for uncompressed "
 828               "and compressed data (%zu bytes needed)",
 829               queue_size * MAX_CHUNKS_PER_MSG * WIM_CHUNK_SIZE * 2);
 830
 831         // Pre-allocate all the buffers that will be needed to do the chunk
 832         // compression.
 833         for (size_t i = 0; i < ARRAY_LEN(msgs); i++) {
 834                 for (size_t j = 0; j < MAX_CHUNKS_PER_MSG; j++) {
 835                         msgs[i].compressed_chunks[j] = MALLOC(WIM_CHUNK_SIZE);
 836                         msgs[i].uncompressed_chunks[j] = MALLOC(WIM_CHUNK_SIZE);
 837                         if (msgs[i].compressed_chunks[j] == NULL ||
 838                             msgs[i].uncompressed_chunks[j] == NULL)
 839                         {
 840                                 ERROR("Could not allocate enough memory for "
 841                                       "multi-threaded compression");
 842                                 ret = WIMLIB_ERR_NOMEM;
 843                                 goto out;
 844                         }
 845                 }
 846         }
 847
 848         // This loop is executed until all resources have been written, except
 849         // possibly a few that have been added to the @my_resources list for
 850         // writing later.
 851         while (1) {
 852                 // Send chunks to the compressor threads until either (a) there
 853                 // are no more messages available since they were all sent off,
 854                 // or (b) there are no more resources that need to be
 855                 // compressed.
 856                 while (!list_empty(&available_msgs) && next_lte != NULL) {
 857
 858                         // Get a message from the available messages
 859                         // list
 860                         msg = container_of(available_msgs.next,
 861                                            struct message,
 862                                            list);
 863
 864                         // ... and delete it from the available messages
 865                         // list
 866                         list_del(&msg->list);
 867
 868                         // Initialize the message with the chunks to
 869                         // compress.
 870                         msg->num_chunks = min(next_num_chunks - next_chunk,
 871                                               MAX_CHUNKS_PER_MSG);
 872                         msg->lte = next_lte;
 873                         msg->complete = false;
 874                         msg->begin_chunk = next_chunk;
 875
 876                         unsigned size = WIM_CHUNK_SIZE;
 877                         for (unsigned i = 0; i < msg->num_chunks; i++) {
 878
 879                                 // Read chunk @next_chunk of the stream into the
 880                                 // message so that a compressor thread can
 881                                 // compress it.
 882
 883                                 if (next_chunk == next_num_chunks - 1 &&
 884                                      wim_resource_size(next_lte) % WIM_CHUNK_SIZE != 0)
 885                                 {
 886                                         size = wim_resource_size(next_lte) % WIM_CHUNK_SIZE;
 887                                 }
 888
 889
 890                                 DEBUG2("Read resource (size=%u, offset=%zu)",
 891                                       size, next_chunk * WIM_CHUNK_SIZE);
 892
 893                                 msg->uncompressed_chunk_sizes[i] = size;
 894
 895                                 ret = read_wim_resource(next_lte,
 896                                                         msg->uncompressed_chunks[i],
 897                                                         size,
 898                                                         next_chunk * WIM_CHUNK_SIZE,
 899                                                         0);
 900                                 if (ret != 0)
 901                                         goto out;
 902                                 sha1_update(&next_sha_ctx,
 903                                             msg->uncompressed_chunks[i], size);
 904                                 next_chunk++;
 905                         }
 906
 907                         // Send the compression request
 908                         list_add_tail(&msg->list, &next_lte->msg_list);
 909                         shared_queue_put(res_to_compress_queue, msg);
 910                         DEBUG2("Compression request sent");
 911
 912                         if (next_chunk != next_num_chunks)
 913                                 // More chunks to send for this resource
 914                                 continue;
 915
 916                         // Done sending compression requests for a resource!
 917                         // Check the SHA1 message digest.
 918                         DEBUG2("Finalize SHA1 md (next_num_chunks=%zu)", next_num_chunks);
 919                         sha1_final(next_hash, &next_sha_ctx);
 920                         if (!hashes_equal(next_lte->hash, next_hash)) {
 921                                 ERROR("WIM resource has incorrect hash!");
 922                                 if (next_lte->resource_location == RESOURCE_IN_FILE_ON_DISK) {
 923                                         ERROR("We were reading it from `%s'; maybe it changed "
 924                                               "while we were reading it.",
 925                                               next_lte->file_on_disk);
 926                                 }
 927                                 ret = WIMLIB_ERR_INVALID_RESOURCE_HASH;
 928                                 goto out;
 929                         }
 930
 931                         // Advance to the next resource.
 932                         //
 933                         // If the next resource needs no compression, just write
 934                         // it with this thread (not now though--- we could be in
 935                         // the middle of writing another resource.)  Keep doing
 936                         // this until we either get to the end of the resources
 937                         // list, or we get to a resource that needs compression.
 938
 939                         while (1) {
 940                                 if (next_resource == stream_list) {
 941                                         next_lte = NULL;
 942                                         break;
 943                                 }
 944                         #ifdef WITH_NTFS_3G
 945                                 end_wim_resource_read(next_lte, ni);
 946                                 ni = NULL;
 947                         #else
 948                                 end_wim_resource_read(next_lte);
 949                         #endif
 950
 951                                 next_lte = container_of(next_resource,
 952                                                         struct lookup_table_entry,
 953                                                         staging_list);
 954                                 next_resource = next_resource->next;
 955                                 if ((!(write_flags & WIMLIB_WRITE_FLAG_RECOMPRESS)
 956                                       && next_lte->resource_location == RESOURCE_IN_WIM
 957                                       && wimlib_get_compression_type(next_lte->wim) == out_ctype)
 958                                     || wim_resource_size(next_lte) == 0)
 959                                 {
 960                                         list_add_tail(&next_lte->staging_list,
 961                                                       &my_resources);
 962                                 } else {
 963                                         list_add_tail(&next_lte->staging_list,
 964                                                       &outstanding_resources);
 965                                         next_chunk = 0;
 966                                         next_num_chunks = wim_resource_chunks(next_lte);
 967                                         sha1_init(&next_sha_ctx);
 968                                         INIT_LIST_HEAD(&next_lte->msg_list);
 969                                 #ifdef WITH_NTFS_3G
 970                                         ret = prepare_resource_for_read(next_lte, &ni);
 971                                 #else
 972                                         ret = prepare_resource_for_read(next_lte);
 973                                 #endif
 974                                         if (ret != 0)
 975                                                 goto out;
 976                                         DEBUG2("Updated next_lte");
 977                                         break;
 978                                 }
 979                         }
 980                 }
 981
 982                 // If there are no outstanding resources, there are no more
 983                 // resources that need to be written.
 984                 if (list_empty(&outstanding_resources)) {
 985                         DEBUG("No outstanding resources! Done");
 986                         ret = 0;
 987                         goto out;
 988                 }
 989
 990                 // Get the next message from the queue and process it.
 991                 // The message will contain 1 or more data chunks that have been
 992                 // compressed.
 993                 DEBUG2("Waiting for message");
 994                 msg = shared_queue_get(compressed_res_queue);
 995                 msg->complete = true;
 996
 997                 DEBUG2("Received msg (begin_chunk=%"PRIu64")", msg->begin_chunk);
 998
 999                 list_for_each_entry(msg, &cur_lte->msg_list, list) {
1000                         DEBUG2("complete=%d", msg->complete);
1001                 }
1002
1003                 // Is this the next chunk in the current resource?  If it's not
1004                 // (i.e., an earlier chunk in a same or different resource
1005                 // hasn't been compressed yet), do nothing, and keep this
1006                 // message around until all earlier chunks are received.
1007                 //
1008                 // Otherwise, write all the chunks we can.
1009                 while (!list_empty(&cur_lte->msg_list)
1010                         && (msg = container_of(cur_lte->msg_list.next,
1011                                                struct message,
1012                                                list))->complete)
1013                 {
1014                         DEBUG2("Complete msg (begin_chunk=%"PRIu64")", msg->begin_chunk);
1015                         if (msg->begin_chunk == 0) {
1016                                 DEBUG2("Begin chunk tab");
1017
1018                                 // This is the first set of chunks.  Leave space
1019                                 // for the chunk table in the output file.
1020                                 off_t cur_offset = ftello(out_fp);
1021                                 if (cur_offset == -1) {
1022                                         ret = WIMLIB_ERR_WRITE;
1023                                         goto out;
1024                                 }
1025                                 ret = begin_wim_resource_chunk_tab(cur_lte,
1026                                                                    out_fp,
1027                                                                    cur_offset,
1028                                                                    &cur_chunk_tab);
1029                                 if (ret != 0)
1030                                         goto out;
1031                         }
1032
1033                         // Write the compressed chunks from the message.
1034                         ret = write_wim_chunks(msg, out_fp, cur_chunk_tab);
1035                         if (ret != 0)
1036                                 goto out;
1037
1038                         list_del(&msg->list);
1039
1040                         // This message is available to use for different chunks
1041                         // now.
1042                         list_add(&msg->list, &available_msgs);
1043
1044                         // Was this the last chunk of the stream?  If so,
1045                         // finish it.
1046                         if (list_empty(&cur_lte->msg_list) &&
1047                             msg->begin_chunk + msg->num_chunks == cur_chunk_tab->num_chunks)
1048                         {
1049                                 DEBUG2("Finish wim chunk tab");
1050                                 u64 res_csize;
1051                                 ret = finish_wim_resource_chunk_tab(cur_chunk_tab,
1052                                                                     out_fp,
1053                                                                     &res_csize);
1054                                 if (ret != 0)
1055                                         goto out;
1056
1057                                 progress->write_streams.completed_bytes +=
1058                                                 wim_resource_size(cur_lte);
1059                                 progress->write_streams.completed_streams++;
1060
1061                                 if (progress_func) {
1062                                         progress_func(WIMLIB_PROGRESS_MSG_WRITE_STREAMS,
1063                                                       progress);
1064                                 }
1065
1066                                 cur_lte->output_resource_entry.size =
1067                                         res_csize;
1068
1069                                 cur_lte->output_resource_entry.original_size =
1070                                         cur_lte->resource_entry.original_size;
1071
1072                                 cur_lte->output_resource_entry.offset =
1073                                         cur_chunk_tab->file_offset;
1074
1075                                 cur_lte->output_resource_entry.flags =
1076                                         cur_lte->resource_entry.flags |
1077                                                 WIM_RESHDR_FLAG_COMPRESSED;
1078
1079                                 FREE(cur_chunk_tab);
1080                                 cur_chunk_tab = NULL;
1081
1082                                 struct list_head *next = cur_lte->staging_list.next;
1083                                 list_del(&cur_lte->staging_list);
1084
1085                                 if (next == &outstanding_resources) {
1086                                         DEBUG("No more outstanding resources");
1087                                         ret = 0;
1088                                         goto out;
1089                                 } else {
1090                                         cur_lte = container_of(cur_lte->staging_list.next,
1091                                                                struct lookup_table_entry,
1092                                                                staging_list);
1093                                 }
1094
1095                                 // Since we just finished writing a stream,
1096                                 // write any streams that have been added to the
1097                                 // my_resources list for direct writing by the
1098                                 // main thread (e.g. resources that don't need
1099                                 // to be compressed because the desired
1100                                 // compression type is the same as the previous
1101                                 // compression type).
1102                                 ret = do_write_stream_list(&my_resources,
1103                                                            out_fp,
1104                                                            out_ctype,
1105                                                            progress_func,
1106                                                            progress,
1107                                                            0);
1108                                 if (ret != 0)
1109                                         goto out;
1110                         }
1111                 }
1112         }
1113
1114 out:
1115 #ifdef WITH_NTFS_3G
1116         end_wim_resource_read(cur_lte, ni);
1117 #else
1118         end_wim_resource_read(cur_lte);
1119 #endif
1120         if (ret == 0) {
1121                 ret = do_write_stream_list(&my_resources, out_fp,
1122                                            out_ctype, progress_func,
1123                                            progress, 0);
1124         } else {
1125                 size_t num_available_msgs = 0;
1126                 struct list_head *cur;
1127
1128                 list_for_each(cur, &available_msgs) {
1129                         num_available_msgs++;
1130                 }
1131
1132                 while (num_available_msgs < ARRAY_LEN(msgs)) {
1133                         shared_queue_get(compressed_res_queue);
1134                         num_available_msgs++;
1135                 }
1136         }
1137
1138         for (size_t i = 0; i < ARRAY_LEN(msgs); i++) {
1139                 for (size_t j = 0; j < MAX_CHUNKS_PER_MSG; j++) {
1140                         FREE(msgs[i].compressed_chunks[j]);
1141                         FREE(msgs[i].uncompressed_chunks[j]);
1142                 }
1143         }
1144
1145         if (cur_chunk_tab != NULL)
1146                 FREE(cur_chunk_tab);
1147         return ret;
1148 }
1149
1150
1151 static int write_stream_list_parallel(struct list_head *stream_list,
1152                                       FILE *out_fp,
1153                                       int out_ctype,
1154                                       int write_flags,
1155                                       unsigned num_threads,
1156                                       wimlib_progress_func_t progress_func,
1157                                       union wimlib_progress_info *progress)
1158 {
1159         int ret;
1160         struct shared_queue res_to_compress_queue;
1161         struct shared_queue compressed_res_queue;
1162         pthread_t *compressor_threads = NULL;
1163
1164         if (num_threads == 0) {
1165                 long nthreads = sysconf(_SC_NPROCESSORS_ONLN);
1166                 if (nthreads < 1) {
1167                         WARNING("Could not determine number of processors! Assuming 1");
1168                         goto out_serial;
1169                 } else {
1170                         num_threads = nthreads;
1171                 }
1172         }
1173
1174         progress->write_streams.num_threads = num_threads;
1175         wimlib_assert(stream_list->next != stream_list);
1176
1177         static const double MESSAGES_PER_THREAD = 2.0;
1178         size_t queue_size = (size_t)(num_threads * MESSAGES_PER_THREAD);
1179
1180         DEBUG("Initializing shared queues (queue_size=%zu)", queue_size);
1181
1182         ret = shared_queue_init(&res_to_compress_queue, queue_size);
1183         if (ret != 0)
1184                 goto out_serial;
1185
1186         ret = shared_queue_init(&compressed_res_queue, queue_size);
1187         if (ret != 0)
1188                 goto out_destroy_res_to_compress_queue;
1189
1190         struct compressor_thread_params params;
1191         params.res_to_compress_queue = &res_to_compress_queue;
1192         params.compressed_res_queue = &compressed_res_queue;
1193         params.compress = get_compress_func(out_ctype);
1194
1195         compressor_threads = MALLOC(num_threads * sizeof(pthread_t));
1196
1197         for (unsigned i = 0; i < num_threads; i++) {
1198                 DEBUG("pthread_create thread %u", i);
1199                 ret = pthread_create(&compressor_threads[i], NULL,
1200                                      compressor_thread_proc, &params);
1201                 if (ret != 0) {
1202                         ret = -1;
1203                         ERROR_WITH_ERRNO("Failed to create compressor "
1204                                          "thread %u", i);
1205                         num_threads = i;
1206                         goto out_join;
1207                 }
1208         }
1209
1210         if (progress_func)
1211                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_STREAMS, progress);
1212
1213         ret = main_writer_thread_proc(stream_list,
1214                                       out_fp,
1215                                       out_ctype,
1216                                       &res_to_compress_queue,
1217                                       &compressed_res_queue,
1218                                       queue_size,
1219                                       write_flags,
1220                                       progress_func,
1221                                       progress);
1222 out_join:
1223         for (unsigned i = 0; i < num_threads; i++)
1224                 shared_queue_put(&res_to_compress_queue, NULL);
1225
1226         for (unsigned i = 0; i < num_threads; i++) {
1227                 if (pthread_join(compressor_threads[i], NULL)) {
1228                         WARNING("Failed to join compressor thread %u: %s",
1229                                 i, strerror(errno));
1230                 }
1231         }
1232         FREE(compressor_threads);
1233         shared_queue_destroy(&compressed_res_queue);
1234 out_destroy_res_to_compress_queue:
1235         shared_queue_destroy(&res_to_compress_queue);
1236         if (ret >= 0 && ret != WIMLIB_ERR_NOMEM)
1237                 return ret;
1238 out_serial:
1239         WARNING("Falling back to single-threaded compression");
1240         return write_stream_list_serial(stream_list,
1241                                         out_fp,
1242                                         out_ctype,
1243                                         write_flags,
1244                                         progress_func,
1245                                         progress);
1246
1247 }
1248 #endif
1249
1250 /*
1251  * Write a list of streams to a WIM (@out_fp) using the compression type
1252  * @out_ctype and up to @num_threads compressor threads.
1253  */
1254 static int write_stream_list(struct list_head *stream_list, FILE *out_fp,
1255                              int out_ctype, int write_flags,
1256                              unsigned num_threads,
1257                              wimlib_progress_func_t progress_func)
1258 {
1259         struct lookup_table_entry *lte;
1260         size_t num_streams = 0;
1261         u64 total_bytes = 0;
1262         bool compression_needed = false;
1263         union wimlib_progress_info progress;
1264         int ret;
1265
1266         list_for_each_entry(lte, stream_list, staging_list) {
1267                 num_streams++;
1268                 total_bytes += wim_resource_size(lte);
1269                 if (!compression_needed
1270                     &&
1271                     (out_ctype != WIMLIB_COMPRESSION_TYPE_NONE
1272                        && (lte->resource_location != RESOURCE_IN_WIM
1273                            || wimlib_get_compression_type(lte->wim) != out_ctype
1274                            || (write_flags & WIMLIB_WRITE_FLAG_REBUILD)))
1275                     && wim_resource_size(lte) != 0)
1276                         compression_needed = true;
1277         }
1278         progress.write_streams.total_bytes       = total_bytes;
1279         progress.write_streams.total_streams     = num_streams;
1280         progress.write_streams.completed_bytes   = 0;
1281         progress.write_streams.completed_streams = 0;
1282         progress.write_streams.num_threads       = num_threads;
1283         progress.write_streams.compression_type  = out_ctype;
1284
1285         if (num_streams == 0) {
1286                 ret = 0;
1287                 goto out;
1288         }
1289
1290 #ifdef ENABLE_MULTITHREADED_COMPRESSION
1291         if (compression_needed && total_bytes >= 1000000 && num_threads != 1) {
1292                 ret = write_stream_list_parallel(stream_list,
1293                                                  out_fp,
1294                                                  out_ctype,
1295                                                  write_flags,
1296                                                  num_threads,
1297                                                  progress_func,
1298                                                  &progress);
1299         }
1300         else
1301 #endif
1302         {
1303                 ret = write_stream_list_serial(stream_list,
1304                                                out_fp,
1305                                                out_ctype,
1306                                                write_flags,
1307                                                progress_func,
1308                                                &progress);
1309         }
1310 out:
1311         return ret;
1312 }
1313
1314
1315 static int dentry_find_streams_to_write(struct dentry *dentry,
1316                                         void *wim)
1317 {
1318         WIMStruct *w = wim;
1319         struct list_head *stream_list = w->private;
1320         struct lookup_table_entry *lte;
1321         for (unsigned i = 0; i <= dentry->d_inode->num_ads; i++) {
1322                 lte = inode_stream_lte(dentry->d_inode, i, w->lookup_table);
1323                 if (lte && ++lte->out_refcnt == 1)
1324                         list_add_tail(&lte->staging_list, stream_list);
1325         }
1326         return 0;
1327 }
1328
1329 static int find_streams_to_write(WIMStruct *w)
1330 {
1331         return for_dentry_in_tree(wim_root_dentry(w),
1332                                   dentry_find_streams_to_write, w);
1333 }
1334
1335 static int write_wim_streams(WIMStruct *w, int image, int write_flags,
1336                              unsigned num_threads,
1337                              wimlib_progress_func_t progress_func)
1338 {
1339
1340         for_lookup_table_entry(w->lookup_table, lte_zero_out_refcnt, NULL);
1341         LIST_HEAD(stream_list);
1342         w->private = &stream_list;
1343         for_image(w, image, find_streams_to_write);
1344         return write_stream_list(&stream_list, w->out_fp,
1345                                  wimlib_get_compression_type(w), write_flags,
1346                                  num_threads, progress_func);
1347 }
1348
1349 /*
1350  * Finish writing a WIM file: write the lookup table, xml data, and integrity
1351  * table (optional), then overwrite the WIM header.
1352  *
1353  * write_flags is a bitwise OR of the following:
1354  *
1355  *      (public)  WIMLIB_WRITE_FLAG_CHECK_INTEGRITY:
1356  *              Include an integrity table.
1357  *
1358  *      (public)  WIMLIB_WRITE_FLAG_SHOW_PROGRESS:
1359  *              Show progress information when (if) writing the integrity table.
1360  *
1361  *      (private) WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE:
1362  *              Don't write the lookup table.
1363  *
1364  *      (private) WIMLIB_WRITE_FLAG_REUSE_INTEGRITY_TABLE:
1365  *              When (if) writing the integrity table, re-use entries from the
1366  *              existing integrity table, if possible.
1367  *
1368  *      (private) WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML:
1369  *              After writing the XML data but before writing the integrity
1370  *              table, write a temporary WIM header and flush the stream so that
1371  *              the WIM is less likely to become corrupted upon abrupt program
1372  *              termination.
1373  *
1374  *      (private) WIMLIB_WRITE_FLAG_FSYNC:
1375  *              fsync() the output file before closing it.
1376  *
1377  */
1378 int finish_write(WIMStruct *w, int image, int write_flags,
1379                  wimlib_progress_func_t progress_func)
1380 {
1381         int ret;
1382         struct wim_header hdr;
1383         FILE *out = w->out_fp;
1384
1385         /* @hdr will be the header for the new WIM.  First copy all the data
1386          * from the header in the WIMStruct; then set all the fields that may
1387          * have changed, including the resource entries, boot index, and image
1388          * count.  */
1389         memcpy(&hdr, &w->hdr, sizeof(struct wim_header));
1390
1391         if (!(write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE)) {
1392                 ret = write_lookup_table(w->lookup_table, out, &hdr.lookup_table_res_entry);
1393                 if (ret != 0)
1394                         goto out;
1395         }
1396
1397         ret = write_xml_data(w->wim_info, image, out,
1398                              (write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE) ?
1399                               wim_info_get_total_bytes(w->wim_info) : 0,
1400                              &hdr.xml_res_entry);
1401         if (ret != 0)
1402                 goto out;
1403
1404         if (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) {
1405                 if (write_flags & WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML) {
1406                         struct wim_header checkpoint_hdr;
1407                         memcpy(&checkpoint_hdr, &hdr, sizeof(struct wim_header));
1408                         memset(&checkpoint_hdr.integrity, 0, sizeof(struct resource_entry));
1409                         if (fseeko(out, 0, SEEK_SET) != 0) {
1410                                 ret = WIMLIB_ERR_WRITE;
1411                                 goto out;
1412                         }
1413                         ret = write_header(&checkpoint_hdr, out);
1414                         if (ret != 0)
1415                                 goto out;
1416
1417                         if (fflush(out) != 0) {
1418                                 ERROR_WITH_ERRNO("Can't write data to WIM");
1419                                 ret = WIMLIB_ERR_WRITE;
1420                                 goto out;
1421                         }
1422
1423                         if (fseeko(out, 0, SEEK_END) != 0) {
1424                                 ret = WIMLIB_ERR_WRITE;
1425                                 goto out;
1426                         }
1427                 }
1428
1429                 off_t old_lookup_table_end;
1430                 off_t new_lookup_table_end;
1431                 if (write_flags & WIMLIB_WRITE_FLAG_REUSE_INTEGRITY_TABLE) {
1432                         old_lookup_table_end = w->hdr.lookup_table_res_entry.offset +
1433                                                w->hdr.lookup_table_res_entry.size;
1434                 } else {
1435                         old_lookup_table_end = 0;
1436                 }
1437                 new_lookup_table_end = hdr.lookup_table_res_entry.offset +
1438                                        hdr.lookup_table_res_entry.size;
1439
1440                 ret = write_integrity_table(out,
1441                                             &hdr.integrity,
1442                                             new_lookup_table_end,
1443                                             old_lookup_table_end,
1444                                             progress_func);
1445                 if (ret != 0)
1446                         goto out;
1447         } else {
1448                 memset(&hdr.integrity, 0, sizeof(struct resource_entry));
1449         }
1450
1451         /*
1452          * In the WIM header, there is room for the resource entry for a
1453          * metadata resource labeled as the "boot metadata".  This entry should
1454          * be zeroed out if there is no bootable image (boot_idx 0).  Otherwise,
1455          * it should be a copy of the resource entry for the image that is
1456          * marked as bootable.  This is not well documented...
1457          */
1458         if (hdr.boot_idx == 0 || !w->image_metadata
1459                         || (image != WIMLIB_ALL_IMAGES && image != hdr.boot_idx)) {
1460                 memset(&hdr.boot_metadata_res_entry, 0,
1461                        sizeof(struct resource_entry));
1462         } else {
1463                 memcpy(&hdr.boot_metadata_res_entry,
1464                        &w->image_metadata[
1465                           hdr.boot_idx - 1].metadata_lte->output_resource_entry,
1466                        sizeof(struct resource_entry));
1467         }
1468
1469         /* Set image count and boot index correctly for single image writes */
1470         if (image != WIMLIB_ALL_IMAGES) {
1471                 hdr.image_count = 1;
1472                 if (hdr.boot_idx == image)
1473                         hdr.boot_idx = 1;
1474                 else
1475                         hdr.boot_idx = 0;
1476         }
1477
1478         if (fseeko(out, 0, SEEK_SET) != 0) {
1479                 ret = WIMLIB_ERR_WRITE;
1480                 goto out;
1481         }
1482
1483         ret = write_header(&hdr, out);
1484         if (ret != 0)
1485                 goto out;
1486
1487         if (write_flags & WIMLIB_WRITE_FLAG_FSYNC) {
1488                 if (fflush(out) != 0
1489                     || fsync(fileno(out)) != 0)
1490                 {
1491                         ERROR_WITH_ERRNO("Error flushing data to WIM file");
1492                         ret = WIMLIB_ERR_WRITE;
1493                 }
1494         }
1495 out:
1496         if (fclose(out) != 0) {
1497                 ERROR_WITH_ERRNO("Failed to close the WIM file");
1498                 if (ret == 0)
1499                         ret = WIMLIB_ERR_WRITE;
1500         }
1501         w->out_fp = NULL;
1502         return ret;
1503 }
1504
1505 #if defined(HAVE_SYS_FILE_H) && defined(HAVE_FLOCK)
1506 int lock_wim(FILE *fp, const char *path)
1507 {
1508         int ret = 0;
1509         if (fp) {
1510                 ret = flock(fileno(fp), LOCK_EX | LOCK_NB);
1511                 if (ret != 0) {
1512                         if (errno == EWOULDBLOCK) {
1513                                 ERROR("`%s' is already being modified or has been "
1514                                       "mounted read-write\n"
1515                                       "        by another process!", path);
1516                                 ret = WIMLIB_ERR_ALREADY_LOCKED;
1517                         } else {
1518                                 WARNING("Failed to lock `%s': %s",
1519                                         path, strerror(errno));
1520                                 ret = 0;
1521                         }
1522                 }
1523         }
1524         return ret;
1525 }
1526 #endif
1527
1528 static int open_wim_writable(WIMStruct *w, const char *path,
1529                              bool trunc, bool readable)
1530 {
1531         const char *mode;
1532         if (trunc)
1533                 if (readable)
1534                         mode = "w+b";
1535                 else
1536                         mode = "wb";
1537         else
1538                 mode = "r+b";
1539
1540         wimlib_assert(w->out_fp == NULL);
1541         w->out_fp = fopen(path, mode);
1542         if (w->out_fp) {
1543                 return 0;
1544         } else {
1545                 ERROR_WITH_ERRNO("Failed to open `%s' for writing", path);
1546                 return WIMLIB_ERR_OPEN;
1547         }
1548 }
1549
1550
1551 void close_wim_writable(WIMStruct *w)
1552 {
1553         if (w->out_fp) {
1554                 if (fclose(w->out_fp) != 0) {
1555                         WARNING("Failed to close output WIM: %s",
1556                                 strerror(errno));
1557                 }
1558                 w->out_fp = NULL;
1559         }
1560 }
1561
1562 /* Open file stream and write dummy header for WIM. */
1563 int begin_write(WIMStruct *w, const char *path, int write_flags)
1564 {
1565         int ret;
1566         ret = open_wim_writable(w, path, true,
1567                                 (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) != 0);
1568         if (ret != 0)
1569                 return ret;
1570         /* Write dummy header. It will be overwritten later. */
1571         return write_header(&w->hdr, w->out_fp);
1572 }
1573
1574 /* Writes a stand-alone WIM to a file.  */
1575 WIMLIBAPI int wimlib_write(WIMStruct *w, const char *path,
1576                            int image, int write_flags, unsigned num_threads,
1577                            wimlib_progress_func_t progress_func)
1578 {
1579         int ret;
1580
1581         if (!w || !path)
1582                 return WIMLIB_ERR_INVALID_PARAM;
1583
1584         write_flags &= WIMLIB_WRITE_MASK_PUBLIC;
1585
1586         if (image != WIMLIB_ALL_IMAGES &&
1587              (image < 1 || image > w->hdr.image_count))
1588                 return WIMLIB_ERR_INVALID_IMAGE;
1589
1590         if (w->hdr.total_parts != 1) {
1591                 ERROR("Cannot call wimlib_write() on part of a split WIM");
1592                 return WIMLIB_ERR_SPLIT_UNSUPPORTED;
1593         }
1594
1595         ret = begin_write(w, path, write_flags);
1596         if (ret != 0)
1597                 goto out;
1598
1599         ret = write_wim_streams(w, image, write_flags, num_threads,
1600                                 progress_func);
1601         if (ret != 0)
1602                 goto out;
1603
1604         if (progress_func)
1605                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_METADATA_BEGIN, NULL);
1606
1607         ret = for_image(w, image, write_metadata_resource);
1608         if (ret != 0)
1609                 goto out;
1610
1611         if (progress_func)
1612                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_METADATA_END, NULL);
1613
1614         ret = finish_write(w, image, write_flags, progress_func);
1615 out:
1616         close_wim_writable(w);
1617         return ret;
1618 }
1619
1620 static int lte_overwrite_prepare(struct lookup_table_entry *lte,
1621                                  void *ignore)
1622 {
1623         memcpy(&lte->output_resource_entry, &lte->resource_entry,
1624                sizeof(struct resource_entry));
1625         lte->out_refcnt = 0;
1626         return 0;
1627 }
1628
1629 static int check_resource_offset(struct lookup_table_entry *lte, void *arg)
1630 {
1631         off_t end_offset = *(u64*)arg;
1632
1633         wimlib_assert(lte->out_refcnt <= lte->refcnt);
1634         if (lte->out_refcnt < lte->refcnt) {
1635                 if (lte->resource_entry.offset + lte->resource_entry.size > end_offset) {
1636                         ERROR("The following resource is after the XML data:");
1637                         print_lookup_table_entry(lte);
1638                         return WIMLIB_ERR_RESOURCE_ORDER;
1639                 }
1640         }
1641         return 0;
1642 }
1643
1644 static int find_new_streams(struct lookup_table_entry *lte, void *arg)
1645 {
1646         if (lte->out_refcnt == lte->refcnt)
1647                 list_add(&lte->staging_list, (struct list_head*)arg);
1648         else
1649                 lte->out_refcnt = lte->refcnt;
1650         return 0;
1651 }
1652
1653 /*
1654  * Overwrite a WIM, possibly appending streams to it.
1655  *
1656  * A WIM looks like (or is supposed to look like) the following:
1657  *
1658  *                   Header (212 bytes)
1659  *                   Streams and metadata resources (variable size)
1660  *                   Lookup table (variable size)
1661  *                   XML data (variable size)
1662  *                   Integrity table (optional) (variable size)
1663  *
1664  * If we are not adding any streams or metadata resources, the lookup table is
1665  * unchanged--- so we only need to overwrite the XML data, integrity table, and
1666  * header.  This operation is potentially unsafe if the program is abruptly
1667  * terminated while the XML data or integrity table are being overwritten, but
1668  * before the new header has been written.  To partially alleviate this problem,
1669  * a special flag (WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML) is passed to
1670  * finish_write() to cause a temporary WIM header to be written after the XML
1671  * data has been written.  This may prevent the WIM from becoming corrupted if
1672  * the program is terminated while the integrity table is being calculated (but
1673  * no guarantees, due to write re-ordering...).
1674  *
1675  * If we are adding new streams or images (metadata resources), the lookup table
1676  * needs to be changed, and those streams need to be written.  In this case, we
1677  * try to perform a safe update of the WIM file by writing the streams *after*
1678  * the end of the previous WIM, then writing the new lookup table, XML data, and
1679  * (optionally) integrity table following the new streams.  This will produce a
1680  * layout like the following:
1681  *
1682  *                   Header (212 bytes)
1683  *                   (OLD) Streams and metadata resources (variable size)
1684  *                   (OLD) Lookup table (variable size)
1685  *                   (OLD) XML data (variable size)
1686  *                   (OLD) Integrity table (optional) (variable size)
1687  *                   (NEW) Streams and metadata resources (variable size)
1688  *                   (NEW) Lookup table (variable size)
1689  *                   (NEW) XML data (variable size)
1690  *                   (NEW) Integrity table (optional) (variable size)
1691  *
1692  * At all points, the WIM is valid as nothing points to the new data yet.  Then,
1693  * the header is overwritten to point to the new lookup table, XML data, and
1694  * integrity table, to produce the following layout:
1695  *
1696  *                   Header (212 bytes)
1697  *                   Streams and metadata resources (variable size)
1698  *                   Nothing (variable size)
1699  *                   More Streams and metadata resources (variable size)
1700  *                   Lookup table (variable size)
1701  *                   XML data (variable size)
1702  *                   Integrity table (optional) (variable size)
1703  *
1704  * This method allows an image to be appended to a large WIM very quickly, and
1705  * is is crash-safe except in the case of write re-ordering, but the
1706  * disadvantage is that a small hole is left in the WIM where the old lookup
1707  * table, xml data, and integrity table were.  (These usually only take up a
1708  * small amount of space compared to the streams, however.
1709  */
1710 static int overwrite_wim_inplace(WIMStruct *w, int write_flags,
1711                                  unsigned num_threads,
1712                                  wimlib_progress_func_t progress_func,
1713                                  int modified_image_idx)
1714 {
1715         int ret;
1716         struct list_head stream_list;
1717         off_t old_wim_end;
1718
1719         DEBUG("Overwriting `%s' in-place", w->filename);
1720
1721         /* Make sure that the integrity table (if present) is after the XML
1722          * data, and that there are no stream resources, metadata resources, or
1723          * lookup tables after the XML data.  Otherwise, these data would be
1724          * overwritten. */
1725         if (w->hdr.integrity.offset != 0 &&
1726             w->hdr.integrity.offset < w->hdr.xml_res_entry.offset) {
1727                 ERROR("Didn't expect the integrity table to be before the XML data");
1728                 return WIMLIB_ERR_RESOURCE_ORDER;
1729         }
1730
1731         if (w->hdr.lookup_table_res_entry.offset > w->hdr.xml_res_entry.offset) {
1732                 ERROR("Didn't expect the lookup table to be after the XML data");
1733                 return WIMLIB_ERR_RESOURCE_ORDER;
1734         }
1735
1736         DEBUG("Identifying newly added streams");
1737         for_lookup_table_entry(w->lookup_table, lte_overwrite_prepare, NULL);
1738         INIT_LIST_HEAD(&stream_list);
1739         for (int i = modified_image_idx; i < w->hdr.image_count; i++) {
1740                 DEBUG("Identifiying streams in image %d", i + 1);
1741                 w->private = &stream_list;
1742                 for_dentry_in_tree(w->image_metadata[i].root_dentry,
1743                                    dentry_find_streams_to_write, w);
1744         }
1745
1746         if (w->hdr.integrity.offset)
1747                 old_wim_end = w->hdr.integrity.offset + w->hdr.integrity.size;
1748         else
1749                 old_wim_end = w->hdr.xml_res_entry.offset + w->hdr.xml_res_entry.size;
1750
1751         ret = for_lookup_table_entry(w->lookup_table, check_resource_offset,
1752                                      &old_wim_end);
1753         if (ret != 0)
1754                 return ret;
1755
1756         if (modified_image_idx == w->hdr.image_count && !w->deletion_occurred) {
1757                 /* If no images have been modified and no images have been
1758                  * deleted, a new lookup table does not need to be written. */
1759                 old_wim_end = w->hdr.lookup_table_res_entry.offset +
1760                               w->hdr.lookup_table_res_entry.size;
1761                 write_flags |= WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE |
1762                                WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML;
1763         }
1764
1765         INIT_LIST_HEAD(&stream_list);
1766         for_lookup_table_entry(w->lookup_table, find_new_streams,
1767                                &stream_list);
1768
1769         ret = open_wim_writable(w, w->filename, false,
1770                                 (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) != 0);
1771         if (ret != 0)
1772                 return ret;
1773
1774         ret = lock_wim(w->out_fp, w->filename);
1775         if (ret != 0) {
1776                 fclose(w->out_fp);
1777                 w->out_fp = NULL;
1778                 return ret;
1779         }
1780
1781         if (fseeko(w->out_fp, old_wim_end, SEEK_SET) != 0) {
1782                 ERROR_WITH_ERRNO("Can't seek to end of WIM");
1783                 ret = WIMLIB_ERR_WRITE;
1784                 goto out_ftruncate;
1785         }
1786
1787         if (!list_empty(&stream_list)) {
1788                 DEBUG("Writing newly added streams (offset = %"PRIu64")",
1789                       old_wim_end);
1790                 ret = write_stream_list(&stream_list, w->out_fp,
1791                                         wimlib_get_compression_type(w),
1792                                         write_flags, num_threads,
1793                                         progress_func);
1794                 if (ret != 0)
1795                         goto out_ftruncate;
1796         } else {
1797                 DEBUG("No new streams were added");
1798         }
1799
1800         for (int i = modified_image_idx; i < w->hdr.image_count; i++) {
1801                 select_wim_image(w, i + 1);
1802                 ret = write_metadata_resource(w);
1803                 if (ret != 0)
1804                         goto out_ftruncate;
1805         }
1806         write_flags |= WIMLIB_WRITE_FLAG_REUSE_INTEGRITY_TABLE;
1807         ret = finish_write(w, WIMLIB_ALL_IMAGES, write_flags,
1808                            progress_func);
1809 out_ftruncate:
1810         close_wim_writable(w);
1811         if (ret != 0) {
1812                 WARNING("Truncating `%s' to its original size (%"PRIu64" bytes)",
1813                         w->filename, old_wim_end);
1814                 truncate(w->filename, old_wim_end);
1815         }
1816         return ret;
1817 }
1818
1819 static int overwrite_wim_via_tmpfile(WIMStruct *w, int write_flags,
1820                                      unsigned num_threads,
1821                                      wimlib_progress_func_t progress_func)
1822 {
1823         size_t wim_name_len;
1824         int ret;
1825
1826         DEBUG("Overwriting `%s' via a temporary file", w->filename);
1827
1828         /* Write the WIM to a temporary file in the same directory as the
1829          * original WIM. */
1830         wim_name_len = strlen(w->filename);
1831         char tmpfile[wim_name_len + 10];
1832         memcpy(tmpfile, w->filename, wim_name_len);
1833         randomize_char_array_with_alnum(tmpfile + wim_name_len, 9);
1834         tmpfile[wim_name_len + 9] = '\0';
1835
1836         ret = wimlib_write(w, tmpfile, WIMLIB_ALL_IMAGES,
1837                            write_flags | WIMLIB_WRITE_FLAG_FSYNC,
1838                            num_threads, progress_func);
1839         if (ret != 0) {
1840                 ERROR("Failed to write the WIM file `%s'", tmpfile);
1841                 goto err;
1842         }
1843
1844         DEBUG("Renaming `%s' to `%s'", tmpfile, w->filename);
1845
1846         /* Rename the new file to the old file .*/
1847         if (rename(tmpfile, w->filename) != 0) {
1848                 ERROR_WITH_ERRNO("Failed to rename `%s' to `%s'",
1849                                  tmpfile, w->filename);
1850                 ret = WIMLIB_ERR_RENAME;
1851                 goto err;
1852         }
1853
1854         if (progress_func) {
1855                 union wimlib_progress_info progress;
1856                 progress.rename.from = tmpfile;
1857                 progress.rename.to = w->filename;
1858                 progress_func(WIMLIB_PROGRESS_MSG_RENAME, &progress);
1859         }
1860
1861         /* Close the original WIM file that was opened for reading. */
1862         if (w->fp != NULL) {
1863                 fclose(w->fp);
1864                 w->fp = NULL;
1865         }
1866
1867         /* Re-open the WIM read-only. */
1868         w->fp = fopen(w->filename, "rb");
1869         if (w->fp == NULL) {
1870                 ret = WIMLIB_ERR_REOPEN;
1871                 WARNING("Failed to re-open `%s' read-only: %s",
1872                         w->filename, strerror(errno));
1873                 FREE(w->filename);
1874                 w->filename = NULL;
1875         }
1876         return ret;
1877 err:
1878         /* Remove temporary file. */
1879         if (unlink(tmpfile) != 0)
1880                 WARNING("Failed to remove `%s': %s", tmpfile, strerror(errno));
1881         return ret;
1882 }
1883
1884 /*
1885  * Writes a WIM file to the original file that it was read from, overwriting it.
1886  */
1887 WIMLIBAPI int wimlib_overwrite(WIMStruct *w, int write_flags,
1888                                unsigned num_threads,
1889                                wimlib_progress_func_t progress_func)
1890 {
1891         if (!w)
1892                 return WIMLIB_ERR_INVALID_PARAM;
1893
1894         write_flags &= WIMLIB_WRITE_MASK_PUBLIC;
1895
1896         if (!w->filename)
1897                 return WIMLIB_ERR_NO_FILENAME;
1898
1899         if (w->hdr.total_parts != 1) {
1900                 ERROR("Cannot modify a split WIM");
1901                 return WIMLIB_ERR_SPLIT_UNSUPPORTED;
1902         }
1903
1904         if ((!w->deletion_occurred || (write_flags & WIMLIB_WRITE_FLAG_SOFT_DELETE))
1905             && !(write_flags & WIMLIB_WRITE_FLAG_REBUILD))
1906         {
1907                 int i, modified_image_idx;
1908                 for (i = 0; i < w->hdr.image_count && !w->image_metadata[i].modified; i++)
1909                         ;
1910                 modified_image_idx = i;
1911                 for (; i < w->hdr.image_count && w->image_metadata[i].modified &&
1912                         !w->image_metadata[i].has_been_mounted_rw; i++)
1913                         ;
1914                 if (i == w->hdr.image_count) {
1915                         return overwrite_wim_inplace(w, write_flags, num_threads,
1916                                                      progress_func,
1917                                                      modified_image_idx);
1918                 }
1919         }
1920         return overwrite_wim_via_tmpfile(w, write_flags, num_threads,
1921                                          progress_func);
1922 }