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