8451b391db9d3e1824e59eec5e29d0732223d10b
[wimlib] / 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) 2012, 2013 Eric Biggers
10  *
11  * This file is part of wimlib, a library for working with WIM files.
12  *
13  * wimlib is free software; you can redistribute it and/or modify it under the
14  * terms of the GNU General Public License as published by the Free
15  * Software Foundation; either version 3 of the License, or (at your option)
16  * any later version.
17  *
18  * wimlib is distributed in the hope that it will be useful, but WITHOUT ANY
19  * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
20  * A PARTICULAR PURPOSE. See the GNU General Public License for more
21  * details.
22  *
23  * You should have received a copy of the GNU General Public License
24  * along with wimlib; if not, see http://www.gnu.org/licenses/.
25  */
26
27 #ifdef HAVE_CONFIG_H
28 #  include "config.h"
29 #endif
30
31 #if defined(HAVE_SYS_FILE_H) && defined(HAVE_FLOCK)
32 /* On BSD, this should be included before "wimlib/list.h" so that "wimlib/list.h" can
33  * overwrite the LIST_HEAD macro. */
34 #  include <sys/file.h>
35 #endif
36
37 #include "wimlib/endianness.h"
38 #include "wimlib/error.h"
39 #include "wimlib/file_io.h"
40 #include "wimlib/header.h"
41 #include "wimlib/integrity.h"
42 #include "wimlib/lookup_table.h"
43 #include "wimlib/metadata.h"
44 #include "wimlib/resource.h"
45 #include "wimlib/write.h"
46 #include "wimlib/xml.h"
47
48 #ifdef __WIN32__
49 #  include "wimlib/win32.h" /* win32_get_number_of_processors() */
50 #endif
51
52 #ifdef ENABLE_MULTITHREADED_COMPRESSION
53 #  include <pthread.h>
54 #endif
55
56 #include <unistd.h>
57 #include <fcntl.h>
58 #include <errno.h>
59
60 #ifdef WITH_NTFS_3G
61 #  include <time.h>
62 #  include <ntfs-3g/attrib.h>
63 #  include <ntfs-3g/inode.h>
64 #  include <ntfs-3g/dir.h>
65 #endif
66
67 #ifdef HAVE_ALLOCA_H
68 #  include <alloca.h>
69 #else
70 #  include <stdlib.h>
71 #endif
72
73 #include <limits.h>
74
75 #ifndef __WIN32__
76 #  include <sys/uio.h> /* for `struct iovec' */
77 #endif
78
79 /* Chunk table that's located at the beginning of each compressed resource in
80  * the WIM.  (This is not the on-disk format; the on-disk format just has an
81  * array of offsets.) */
82 struct chunk_table {
83         off_t file_offset;
84         u64 num_chunks;
85         u64 original_resource_size;
86         u64 bytes_per_chunk_entry;
87         u64 table_disk_size;
88         u64 cur_offset;
89         u64 *cur_offset_p;
90         union {
91                 u64 offsets[0];
92                 u32 u32_offsets[0];
93         };
94 };
95
96 /*
97  * Allocates and initializes a chunk table, and reserves space for it in the
98  * output file.
99  */
100 static int
101 begin_wim_resource_chunk_tab(const struct wim_lookup_table_entry *lte,
102                              int out_fd,
103                              off_t file_offset,
104                              struct chunk_table **chunk_tab_ret)
105 {
106         u64 size = wim_resource_size(lte);
107         u64 num_chunks = (size + WIM_CHUNK_SIZE - 1) / WIM_CHUNK_SIZE;
108         size_t alloc_size = sizeof(struct chunk_table) + num_chunks * sizeof(u64);
109         struct chunk_table *chunk_tab = CALLOC(1, alloc_size);
110
111         DEBUG("Begin chunk table for stream with size %"PRIu64, size);
112
113         if (!chunk_tab) {
114                 ERROR("Failed to allocate chunk table for %"PRIu64" byte "
115                       "resource", size);
116                 return WIMLIB_ERR_NOMEM;
117         }
118         chunk_tab->file_offset = file_offset;
119         chunk_tab->num_chunks = num_chunks;
120         chunk_tab->original_resource_size = size;
121         chunk_tab->bytes_per_chunk_entry = (size >= (1ULL << 32)) ? 8 : 4;
122         chunk_tab->table_disk_size = chunk_tab->bytes_per_chunk_entry *
123                                      (num_chunks - 1);
124         chunk_tab->cur_offset = 0;
125         chunk_tab->cur_offset_p = chunk_tab->offsets;
126
127         if (full_write(out_fd, chunk_tab,
128                        chunk_tab->table_disk_size) != chunk_tab->table_disk_size)
129         {
130                 ERROR_WITH_ERRNO("Failed to write chunk table in compressed "
131                                  "file resource");
132                 FREE(chunk_tab);
133                 return WIMLIB_ERR_WRITE;
134         }
135         *chunk_tab_ret = chunk_tab;
136         return 0;
137 }
138
139 /*
140  * compress_func_t- Pointer to a function to compresses a chunk
141  *                  of a WIM resource.  This may be either
142  *                  wimlib_xpress_compress() (xpress-compress.c) or
143  *                  wimlib_lzx_compress() (lzx-compress.c).
144  *
145  * @chunk:        Uncompressed data of the chunk.
146  * @chunk_size:   Size of the uncompressed chunk, in bytes.
147  * @out:          Pointer to output buffer of size at least (@chunk_size - 1) bytes.
148  *
149  * Returns the size of the compressed data written to @out in bytes, or 0 if the
150  * data could not be compressed to (@chunk_size - 1) bytes or fewer.
151  *
152  * As a special requirement, the compression code is optimized for the WIM
153  * format and therefore requires (@chunk_size <= 32768).
154  *
155  * As another special requirement, the compression code will read up to 8 bytes
156  * off the end of the @chunk array for performance reasons.  The values of these
157  * bytes will not affect the output of the compression, but the calling code
158  * must make sure that the buffer holding the uncompressed chunk is actually at
159  * least (@chunk_size + 8) bytes, or at least that these extra bytes are in
160  * mapped memory that will not cause a memory access violation if accessed.
161  */
162 typedef unsigned (*compress_func_t)(const void *chunk, unsigned chunk_size,
163                                     void *out);
164
165 static compress_func_t
166 get_compress_func(int out_ctype)
167 {
168         if (out_ctype == WIMLIB_COMPRESSION_TYPE_LZX)
169                 return wimlib_lzx_compress;
170         else
171                 return wimlib_xpress_compress;
172 }
173
174 /*
175  * Writes a chunk of a WIM resource to an output file.
176  *
177  * @chunk:        Uncompressed data of the chunk.
178  * @chunk_size:   Size of the chunk (<= WIM_CHUNK_SIZE)
179  * @out_fd:       File descriptor to write the chunk to.
180  * @compress:     Compression function to use (NULL if writing uncompressed
181  *                      data).
182  * @chunk_tab:    Pointer to chunk table being created.  It is updated with the
183  *                      offset of the chunk we write.
184  *
185  * Returns 0 on success; nonzero on failure.
186  */
187 static int
188 write_wim_resource_chunk(const void * restrict chunk,
189                          unsigned chunk_size,
190                          int out_fd,
191                          compress_func_t compress,
192                          struct chunk_table * restrict chunk_tab)
193 {
194         const void *out_chunk;
195         unsigned out_chunk_size;
196         if (compress) {
197                 void *compressed_chunk = alloca(chunk_size);
198
199                 out_chunk_size = (*compress)(chunk, chunk_size, compressed_chunk);
200                 if (out_chunk_size) {
201                         /* Write compressed */
202                         out_chunk = compressed_chunk;
203                 } else {
204                         /* Write uncompressed */
205                         out_chunk = chunk;
206                         out_chunk_size = chunk_size;
207                 }
208                 *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
209                 chunk_tab->cur_offset += out_chunk_size;
210         } else {
211                 /* Write uncompressed */
212                 out_chunk = chunk;
213                 out_chunk_size = chunk_size;
214         }
215         if (full_write(out_fd, out_chunk, out_chunk_size) != out_chunk_size) {
216                 ERROR_WITH_ERRNO("Failed to write WIM resource chunk");
217                 return WIMLIB_ERR_WRITE;
218         }
219         return 0;
220 }
221
222 /*
223  * Finishes a WIM chunk table and writes it to the output file at the correct
224  * offset.
225  *
226  * The final size of the full compressed resource is returned in the
227  * @compressed_size_p.
228  */
229 static int
230 finish_wim_resource_chunk_tab(struct chunk_table *chunk_tab,
231                               int out_fd, u64 *compressed_size_p)
232 {
233         size_t bytes_written;
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                         chunk_tab->u32_offsets[i] = cpu_to_le32(chunk_tab->offsets[i]);
240         }
241         bytes_written = full_pwrite(out_fd,
242                                     (u8*)chunk_tab->offsets + chunk_tab->bytes_per_chunk_entry,
243                                     chunk_tab->table_disk_size,
244                                     chunk_tab->file_offset);
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         *compressed_size_p = chunk_tab->cur_offset + chunk_tab->table_disk_size;
251         return 0;
252 }
253
254 static int
255 seek_and_truncate(int out_fd, off_t offset)
256 {
257         if (lseek(out_fd, offset, SEEK_SET) == -1 ||
258             ftruncate(out_fd, offset))
259         {
260                 ERROR_WITH_ERRNO("Failed to truncate output WIM file");
261                 return WIMLIB_ERR_WRITE;
262         } else {
263                 return 0;
264         }
265 }
266
267 static int
268 finalize_and_check_sha1(SHA_CTX * restrict sha_ctx,
269                         struct wim_lookup_table_entry * restrict lte)
270 {
271         u8 md[SHA1_HASH_SIZE];
272         sha1_final(md, sha_ctx);
273         if (lte->unhashed) {
274                 copy_hash(lte->hash, md);
275         } else if (!hashes_equal(md, lte->hash)) {
276                 ERROR("WIM resource has incorrect hash!");
277                 if (lte_filename_valid(lte)) {
278                         ERROR("We were reading it from \"%"TS"\"; maybe "
279                               "it changed while we were reading it.",
280                               lte->file_on_disk);
281                 }
282                 return WIMLIB_ERR_INVALID_RESOURCE_HASH;
283         }
284         return 0;
285 }
286
287
288 struct write_resource_ctx {
289         compress_func_t compress;
290         struct chunk_table *chunk_tab;
291         int out_fd;
292         SHA_CTX sha_ctx;
293         bool doing_sha;
294 };
295
296 static int
297 write_resource_cb(const void *restrict chunk, size_t chunk_size,
298                   void *restrict _ctx)
299 {
300         struct write_resource_ctx *ctx = _ctx;
301
302         if (ctx->doing_sha)
303                 sha1_update(&ctx->sha_ctx, chunk, chunk_size);
304         return write_wim_resource_chunk(chunk, chunk_size,
305                                         ctx->out_fd, ctx->compress,
306                                         ctx->chunk_tab);
307 }
308
309 /*
310  * Write a resource to an output WIM.
311  *
312  * @lte:  Lookup table entry for the resource, which could be in another WIM,
313  *        in an external file, or in another location.
314  *
315  * @out_fd:  File descriptor opened to the output WIM.
316  *
317  * @out_ctype:  One of the WIMLIB_COMPRESSION_TYPE_* constants to indicate
318  *              which compression algorithm to use.
319  *
320  * @out_res_entry:  On success, this is filled in with the offset, flags,
321  *                  compressed size, and uncompressed size of the resource
322  *                  in the output WIM.
323  *
324  * @flags:  WIMLIB_RESOURCE_FLAG_RECOMPRESS to force data to be recompressed
325  *          even if it could otherwise be copied directly from the input.
326  *
327  * Additional notes:  The SHA1 message digest of the uncompressed data is
328  * calculated (except when doing a raw copy --- see below).  If the @unhashed
329  * flag is set on the lookup table entry, this message digest is simply copied
330  * to it; otherwise, the message digest is compared with the existing one, and
331  * the function will fail if they do not match.
332  */
333 int
334 write_wim_resource(struct wim_lookup_table_entry *lte,
335                    int out_fd, int out_ctype,
336                    struct resource_entry *out_res_entry,
337                    int flags)
338 {
339         struct write_resource_ctx write_ctx;
340         u64 read_size;
341         u64 new_size;
342         off_t offset;
343         int ret;
344
345         flags &= ~WIMLIB_RESOURCE_FLAG_RECOMPRESS;
346
347         /* Get current position in output WIM */
348         offset = filedes_offset(out_fd);
349         if (offset == -1) {
350                 ERROR_WITH_ERRNO("Can't get position in output WIM");
351                 return WIMLIB_ERR_WRITE;
352         }
353
354         /* If we are not forcing the data to be recompressed, and the input
355          * resource is located in a WIM with the same compression type as that
356          * desired other than no compression, we can simply copy the compressed
357          * data without recompressing it.  This also means we must skip
358          * calculating the SHA1, as we never will see the uncompressed data. */
359         if (!(flags & WIMLIB_RESOURCE_FLAG_RECOMPRESS) &&
360             lte->resource_location == RESOURCE_IN_WIM &&
361             out_ctype != WIMLIB_COMPRESSION_TYPE_NONE &&
362             wimlib_get_compression_type(lte->wim) == out_ctype)
363         {
364                 flags |= WIMLIB_RESOURCE_FLAG_RAW;
365                 write_ctx.doing_sha = false;
366                 read_size = lte->resource_entry.size;
367         } else {
368                 write_ctx.doing_sha = true;
369                 sha1_init(&write_ctx.sha_ctx);
370                 read_size = lte->resource_entry.original_size;
371         }
372
373         /* Initialize the chunk table and set the compression function if
374          * compressing the resource. */
375         if (out_ctype == WIMLIB_COMPRESSION_TYPE_NONE ||
376             (flags & WIMLIB_RESOURCE_FLAG_RAW)) {
377                 write_ctx.compress = NULL;
378                 write_ctx.chunk_tab = NULL;
379         } else {
380                 write_ctx.compress = get_compress_func(out_ctype);
381                 ret = begin_wim_resource_chunk_tab(lte, out_fd,
382                                                    offset,
383                                                    &write_ctx.chunk_tab);
384                 if (ret)
385                         return ret;
386         }
387
388         /* Write the entire resource by reading the entire resource and feeding
389          * the data through the write_resource_cb function. */
390         write_ctx.out_fd = out_fd;
391 try_write_again:
392         ret = read_resource_prefix(lte, read_size,
393                                    write_resource_cb, &write_ctx, flags);
394         if (ret)
395                 goto out_free_chunk_tab;
396
397         /* Verify SHA1 message digest of the resource, or set the hash for the
398          * first time. */
399         if (write_ctx.doing_sha) {
400                 ret = finalize_and_check_sha1(&write_ctx.sha_ctx, lte);
401                 if (ret)
402                         goto out_free_chunk_tab;
403         }
404
405         out_res_entry->flags = lte->resource_entry.flags;
406         out_res_entry->original_size = wim_resource_size(lte);
407         out_res_entry->offset = offset;
408         if (flags & WIMLIB_RESOURCE_FLAG_RAW) {
409                 /* Doing a raw write:  The new compressed size is the same as
410                  * the compressed size in the other WIM. */
411                 new_size = lte->resource_entry.size;
412         } else if (out_ctype == WIMLIB_COMPRESSION_TYPE_NONE) {
413                 /* Using WIMLIB_COMPRESSION_TYPE_NONE:  The new compressed size
414                  * is the original size. */
415                 new_size = lte->resource_entry.original_size;
416                 out_res_entry->flags &= ~WIM_RESHDR_FLAG_COMPRESSED;
417         } else {
418                 /* Using a different compression type:  Call
419                  * finish_wim_resource_chunk_tab() and it will provide the new
420                  * compressed size. */
421                 ret = finish_wim_resource_chunk_tab(write_ctx.chunk_tab, out_fd,
422                                                     &new_size);
423                 if (ret)
424                         goto out_free_chunk_tab;
425                 if (new_size >= wim_resource_size(lte)) {
426                         /* Oops!  We compressed the resource to larger than the original
427                          * size.  Write the resource uncompressed instead. */
428                         DEBUG("Compressed %"PRIu64" => %"PRIu64" bytes; "
429                               "writing uncompressed instead",
430                               wim_resource_size(lte), new_size);
431                         ret = seek_and_truncate(out_fd, offset);
432                         if (ret)
433                                 goto out_free_chunk_tab;
434                         write_ctx.compress = NULL;
435                         write_ctx.doing_sha = false;
436                         out_ctype = WIMLIB_COMPRESSION_TYPE_NONE;
437                         goto try_write_again;
438                 }
439                 out_res_entry->flags |= WIM_RESHDR_FLAG_COMPRESSED;
440         }
441         out_res_entry->size = new_size;
442         ret = 0;
443 out_free_chunk_tab:
444         FREE(write_ctx.chunk_tab);
445         return ret;
446 }
447
448 #ifdef ENABLE_MULTITHREADED_COMPRESSION
449
450 /* Blocking shared queue (solves the producer-consumer problem) */
451 struct shared_queue {
452         unsigned size;
453         unsigned front;
454         unsigned back;
455         unsigned filled_slots;
456         void **array;
457         pthread_mutex_t lock;
458         pthread_cond_t msg_avail_cond;
459         pthread_cond_t space_avail_cond;
460 };
461
462 static int
463 shared_queue_init(struct shared_queue *q, unsigned size)
464 {
465         wimlib_assert(size != 0);
466         q->array = CALLOC(sizeof(q->array[0]), size);
467         if (!q->array)
468                 goto err;
469         q->filled_slots = 0;
470         q->front = 0;
471         q->back = size - 1;
472         q->size = size;
473         if (pthread_mutex_init(&q->lock, NULL)) {
474                 ERROR_WITH_ERRNO("Failed to initialize mutex");
475                 goto err;
476         }
477         if (pthread_cond_init(&q->msg_avail_cond, NULL)) {
478                 ERROR_WITH_ERRNO("Failed to initialize condition variable");
479                 goto err_destroy_lock;
480         }
481         if (pthread_cond_init(&q->space_avail_cond, NULL)) {
482                 ERROR_WITH_ERRNO("Failed to initialize condition variable");
483                 goto err_destroy_msg_avail_cond;
484         }
485         return 0;
486 err_destroy_msg_avail_cond:
487         pthread_cond_destroy(&q->msg_avail_cond);
488 err_destroy_lock:
489         pthread_mutex_destroy(&q->lock);
490 err:
491         return WIMLIB_ERR_NOMEM;
492 }
493
494 static void
495 shared_queue_destroy(struct shared_queue *q)
496 {
497         FREE(q->array);
498         pthread_mutex_destroy(&q->lock);
499         pthread_cond_destroy(&q->msg_avail_cond);
500         pthread_cond_destroy(&q->space_avail_cond);
501 }
502
503 static void
504 shared_queue_put(struct shared_queue *q, void *obj)
505 {
506         pthread_mutex_lock(&q->lock);
507         while (q->filled_slots == q->size)
508                 pthread_cond_wait(&q->space_avail_cond, &q->lock);
509
510         q->back = (q->back + 1) % q->size;
511         q->array[q->back] = obj;
512         q->filled_slots++;
513
514         pthread_cond_broadcast(&q->msg_avail_cond);
515         pthread_mutex_unlock(&q->lock);
516 }
517
518 static void *
519 shared_queue_get(struct shared_queue *q)
520 {
521         void *obj;
522
523         pthread_mutex_lock(&q->lock);
524         while (q->filled_slots == 0)
525                 pthread_cond_wait(&q->msg_avail_cond, &q->lock);
526
527         obj = q->array[q->front];
528         q->array[q->front] = NULL;
529         q->front = (q->front + 1) % q->size;
530         q->filled_slots--;
531
532         pthread_cond_broadcast(&q->space_avail_cond);
533         pthread_mutex_unlock(&q->lock);
534         return obj;
535 }
536
537 struct compressor_thread_params {
538         struct shared_queue *res_to_compress_queue;
539         struct shared_queue *compressed_res_queue;
540         compress_func_t compress;
541 };
542
543 #define MAX_CHUNKS_PER_MSG 2
544
545 struct message {
546         struct wim_lookup_table_entry *lte;
547         u8 *uncompressed_chunks[MAX_CHUNKS_PER_MSG];
548         u8 *compressed_chunks[MAX_CHUNKS_PER_MSG];
549         unsigned uncompressed_chunk_sizes[MAX_CHUNKS_PER_MSG];
550         struct iovec out_chunks[MAX_CHUNKS_PER_MSG];
551         size_t total_out_bytes;
552         unsigned num_chunks;
553         struct list_head list;
554         bool complete;
555         u64 begin_chunk;
556 };
557
558 static void
559 compress_chunks(struct message *msg, compress_func_t compress)
560 {
561         msg->total_out_bytes = 0;
562         for (unsigned i = 0; i < msg->num_chunks; i++) {
563                 unsigned len = compress(msg->uncompressed_chunks[i],
564                                         msg->uncompressed_chunk_sizes[i],
565                                         msg->compressed_chunks[i]);
566                 void *out_chunk;
567                 unsigned out_len;
568                 if (len) {
569                         /* To be written compressed */
570                         out_chunk = msg->compressed_chunks[i];
571                         out_len = len;
572                 } else {
573                         /* To be written uncompressed */
574                         out_chunk = msg->uncompressed_chunks[i];
575                         out_len = msg->uncompressed_chunk_sizes[i];
576                 }
577                 msg->out_chunks[i].iov_base = out_chunk;
578                 msg->out_chunks[i].iov_len = out_len;
579                 msg->total_out_bytes += out_len;
580         }
581 }
582
583 /* Compressor thread routine.  This is a lot simpler than the main thread
584  * routine: just repeatedly get a group of chunks from the
585  * res_to_compress_queue, compress them, and put them in the
586  * compressed_res_queue.  A NULL pointer indicates that the thread should stop.
587  * */
588 static void *
589 compressor_thread_proc(void *arg)
590 {
591         struct compressor_thread_params *params = arg;
592         struct shared_queue *res_to_compress_queue = params->res_to_compress_queue;
593         struct shared_queue *compressed_res_queue = params->compressed_res_queue;
594         compress_func_t compress = params->compress;
595         struct message *msg;
596
597         DEBUG("Compressor thread ready");
598         while ((msg = shared_queue_get(res_to_compress_queue)) != NULL) {
599                 compress_chunks(msg, compress);
600                 shared_queue_put(compressed_res_queue, msg);
601         }
602         DEBUG("Compressor thread terminating");
603         return NULL;
604 }
605 #endif /* ENABLE_MULTITHREADED_COMPRESSION */
606
607 static void
608 do_write_streams_progress(union wimlib_progress_info *progress,
609                           wimlib_progress_func_t progress_func,
610                           uint64_t size_added,
611                           bool stream_discarded)
612 {
613         if (stream_discarded) {
614                 progress->write_streams.total_bytes -= size_added;
615                 if (progress->write_streams._private != ~(uint64_t)0 &&
616                     progress->write_streams._private > progress->write_streams.total_bytes)
617                 {
618                         progress->write_streams._private = progress->write_streams.total_bytes;
619                 }
620         } else {
621                 progress->write_streams.completed_bytes += size_added;
622         }
623         progress->write_streams.completed_streams++;
624         if (progress_func &&
625             progress->write_streams.completed_bytes >= progress->write_streams._private)
626         {
627                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_STREAMS,
628                               progress);
629                 if (progress->write_streams._private == progress->write_streams.total_bytes) {
630                         progress->write_streams._private = ~(uint64_t)0;
631                 } else {
632                         progress->write_streams._private =
633                                 min(progress->write_streams.total_bytes,
634                                     progress->write_streams.completed_bytes +
635                                         progress->write_streams.total_bytes / 100);
636                 }
637         }
638 }
639
640 struct serial_write_stream_ctx {
641         int out_fd;
642         int out_ctype;
643         int write_resource_flags;
644 };
645
646 static int
647 serial_write_stream(struct wim_lookup_table_entry *lte, void *_ctx)
648 {
649         struct serial_write_stream_ctx *ctx = _ctx;
650         return write_wim_resource(lte, ctx->out_fd,
651                                   ctx->out_ctype, &lte->output_resource_entry,
652                                   ctx->write_resource_flags);
653 }
654
655 /* Write a list of streams, taking into account that some streams may be
656  * duplicates that are checksummed and discarded on the fly, and also delegating
657  * the actual writing of a stream to a function @write_stream_cb, which is
658  * passed the context @write_stream_ctx. */
659 static int
660 do_write_stream_list(struct list_head *stream_list,
661                      struct wim_lookup_table *lookup_table,
662                      int (*write_stream_cb)(struct wim_lookup_table_entry *, void *),
663                      void *write_stream_ctx,
664                      wimlib_progress_func_t progress_func,
665                      union wimlib_progress_info *progress)
666 {
667         int ret = 0;
668         struct wim_lookup_table_entry *lte;
669         bool stream_discarded;
670
671         /* For each stream in @stream_list ... */
672         while (!list_empty(stream_list)) {
673                 stream_discarded = false;
674                 lte = container_of(stream_list->next,
675                                    struct wim_lookup_table_entry,
676                                    write_streams_list);
677                 list_del(&lte->write_streams_list);
678                 if (lte->unhashed && !lte->unique_size) {
679                         /* Unhashed stream that shares a size with some other
680                          * stream in the WIM we are writing.  The stream must be
681                          * checksummed to know if we need to write it or not. */
682                         struct wim_lookup_table_entry *tmp;
683                         u32 orig_refcnt = lte->out_refcnt;
684
685                         ret = hash_unhashed_stream(lte, lookup_table, &tmp);
686                         if (ret)
687                                 break;
688                         if (tmp != lte) {
689                                 lte = tmp;
690                                 /* We found a duplicate stream. */
691                                 if (orig_refcnt != tmp->out_refcnt) {
692                                         /* We have already written, or are going
693                                          * to write, the duplicate stream.  So
694                                          * just skip to the next stream. */
695                                         DEBUG("Discarding duplicate stream of length %"PRIu64,
696                                               wim_resource_size(lte));
697                                         lte->no_progress = 0;
698                                         stream_discarded = true;
699                                         goto skip_to_progress;
700                                 }
701                         }
702                 }
703
704                 /* Here, @lte is either a hashed stream or an unhashed stream
705                  * with a unique size.  In either case we know that the stream
706                  * has to be written.  In either case the SHA1 message digest
707                  * will be calculated over the stream while writing it; however,
708                  * in the former case this is done merely to check the data,
709                  * while in the latter case this is done because we do not have
710                  * the SHA1 message digest yet.  */
711                 wimlib_assert(lte->out_refcnt != 0);
712                 lte->deferred = 0;
713                 lte->no_progress = 0;
714                 ret = (*write_stream_cb)(lte, write_stream_ctx);
715                 if (ret)
716                         break;
717                 /* In parallel mode, some streams are deferred for later,
718                  * serialized processing; ignore them here. */
719                 if (lte->deferred)
720                         continue;
721                 if (lte->unhashed) {
722                         list_del(&lte->unhashed_list);
723                         lookup_table_insert(lookup_table, lte);
724                         lte->unhashed = 0;
725                 }
726         skip_to_progress:
727                 if (!lte->no_progress) {
728                         do_write_streams_progress(progress,
729                                                   progress_func,
730                                                   wim_resource_size(lte),
731                                                   stream_discarded);
732                 }
733         }
734         return ret;
735 }
736
737 static int
738 do_write_stream_list_serial(struct list_head *stream_list,
739                             struct wim_lookup_table *lookup_table,
740                             int out_fd,
741                             int out_ctype,
742                             int write_resource_flags,
743                             wimlib_progress_func_t progress_func,
744                             union wimlib_progress_info *progress)
745 {
746         struct serial_write_stream_ctx ctx = {
747                 .out_fd = out_fd,
748                 .out_ctype = out_ctype,
749                 .write_resource_flags = write_resource_flags,
750         };
751         return do_write_stream_list(stream_list,
752                                     lookup_table,
753                                     serial_write_stream,
754                                     &ctx,
755                                     progress_func,
756                                     progress);
757 }
758
759 static inline int
760 write_flags_to_resource_flags(int write_flags)
761 {
762         int resource_flags = 0;
763
764         if (write_flags & WIMLIB_WRITE_FLAG_RECOMPRESS)
765                 resource_flags |= WIMLIB_RESOURCE_FLAG_RECOMPRESS;
766         return resource_flags;
767 }
768
769 static int
770 write_stream_list_serial(struct list_head *stream_list,
771                          struct wim_lookup_table *lookup_table,
772                          int out_fd,
773                          int out_ctype,
774                          int write_resource_flags,
775                          wimlib_progress_func_t progress_func,
776                          union wimlib_progress_info *progress)
777 {
778         DEBUG("Writing stream list (serial version)");
779         progress->write_streams.num_threads = 1;
780         if (progress_func)
781                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_STREAMS, progress);
782         return do_write_stream_list_serial(stream_list,
783                                            lookup_table,
784                                            out_fd,
785                                            out_ctype,
786                                            write_resource_flags,
787                                            progress_func,
788                                            progress);
789 }
790
791 #ifdef ENABLE_MULTITHREADED_COMPRESSION
792 static int
793 write_wim_chunks(struct message *msg, int out_fd,
794                  struct chunk_table *chunk_tab)
795 {
796         for (unsigned i = 0; i < msg->num_chunks; i++) {
797                 *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
798                 chunk_tab->cur_offset += msg->out_chunks[i].iov_len;
799         }
800         if (full_writev(out_fd, msg->out_chunks,
801                         msg->num_chunks) != msg->total_out_bytes)
802         {
803                 ERROR_WITH_ERRNO("Failed to write WIM chunks");
804                 return WIMLIB_ERR_WRITE;
805         }
806         return 0;
807 }
808
809 struct main_writer_thread_ctx {
810         struct list_head *stream_list;
811         struct wim_lookup_table *lookup_table;
812         int out_fd;
813         int out_ctype;
814         int write_resource_flags;
815         struct shared_queue *res_to_compress_queue;
816         struct shared_queue *compressed_res_queue;
817         size_t num_messages;
818         wimlib_progress_func_t progress_func;
819         union wimlib_progress_info *progress;
820
821         struct list_head available_msgs;
822         struct list_head outstanding_streams;
823         struct list_head serial_streams;
824         size_t num_outstanding_messages;
825
826         SHA_CTX next_sha_ctx;
827         u64 next_chunk;
828         u64 next_num_chunks;
829         struct wim_lookup_table_entry *next_lte;
830
831         struct message *msgs;
832         struct message *next_msg;
833         struct chunk_table *cur_chunk_tab;
834 };
835
836 static int
837 init_message(struct message *msg)
838 {
839         for (size_t i = 0; i < MAX_CHUNKS_PER_MSG; i++) {
840                 msg->compressed_chunks[i] = MALLOC(WIM_CHUNK_SIZE);
841                 msg->uncompressed_chunks[i] = MALLOC(WIM_CHUNK_SIZE);
842                 if (msg->compressed_chunks[i] == NULL ||
843                     msg->uncompressed_chunks[i] == NULL)
844                         return WIMLIB_ERR_NOMEM;
845         }
846         return 0;
847 }
848
849 static void
850 destroy_message(struct message *msg)
851 {
852         for (size_t i = 0; i < MAX_CHUNKS_PER_MSG; i++) {
853                 FREE(msg->compressed_chunks[i]);
854                 FREE(msg->uncompressed_chunks[i]);
855         }
856 }
857
858 static void
859 free_messages(struct message *msgs, size_t num_messages)
860 {
861         if (msgs) {
862                 for (size_t i = 0; i < num_messages; i++)
863                         destroy_message(&msgs[i]);
864                 FREE(msgs);
865         }
866 }
867
868 static struct message *
869 allocate_messages(size_t num_messages)
870 {
871         struct message *msgs;
872
873         msgs = CALLOC(num_messages, sizeof(struct message));
874         if (!msgs)
875                 return NULL;
876         for (size_t i = 0; i < num_messages; i++) {
877                 if (init_message(&msgs[i])) {
878                         free_messages(msgs, num_messages);
879                         return NULL;
880                 }
881         }
882         return msgs;
883 }
884
885 static void
886 main_writer_thread_destroy_ctx(struct main_writer_thread_ctx *ctx)
887 {
888         while (ctx->num_outstanding_messages--)
889                 shared_queue_get(ctx->compressed_res_queue);
890         free_messages(ctx->msgs, ctx->num_messages);
891         FREE(ctx->cur_chunk_tab);
892 }
893
894 static int
895 main_writer_thread_init_ctx(struct main_writer_thread_ctx *ctx)
896 {
897         /* Pre-allocate all the buffers that will be needed to do the chunk
898          * compression. */
899         ctx->msgs = allocate_messages(ctx->num_messages);
900         if (!ctx->msgs)
901                 return WIMLIB_ERR_NOMEM;
902
903         /* Initially, all the messages are available to use. */
904         INIT_LIST_HEAD(&ctx->available_msgs);
905         for (size_t i = 0; i < ctx->num_messages; i++)
906                 list_add_tail(&ctx->msgs[i].list, &ctx->available_msgs);
907
908         /* outstanding_streams is the list of streams that currently have had
909          * chunks sent off for compression.
910          *
911          * The first stream in outstanding_streams is the stream that is
912          * currently being written.
913          *
914          * The last stream in outstanding_streams is the stream that is
915          * currently being read and having chunks fed to the compressor threads.
916          * */
917         INIT_LIST_HEAD(&ctx->outstanding_streams);
918         ctx->num_outstanding_messages = 0;
919
920         ctx->next_msg = NULL;
921
922         /* Resources that don't need any chunks compressed are added to this
923          * list and written directly by the main thread. */
924         INIT_LIST_HEAD(&ctx->serial_streams);
925
926         ctx->cur_chunk_tab = NULL;
927
928         return 0;
929 }
930
931 static int
932 receive_compressed_chunks(struct main_writer_thread_ctx *ctx)
933 {
934         struct message *msg;
935         struct wim_lookup_table_entry *cur_lte;
936         int ret;
937
938         wimlib_assert(!list_empty(&ctx->outstanding_streams));
939         wimlib_assert(ctx->num_outstanding_messages != 0);
940
941         cur_lte = container_of(ctx->outstanding_streams.next,
942                                struct wim_lookup_table_entry,
943                                being_compressed_list);
944
945         /* Get the next message from the queue and process it.
946          * The message will contain 1 or more data chunks that have been
947          * compressed. */
948         msg = shared_queue_get(ctx->compressed_res_queue);
949         msg->complete = true;
950         --ctx->num_outstanding_messages;
951
952         /* Is this the next chunk in the current resource?  If it's not
953          * (i.e., an earlier chunk in a same or different resource
954          * hasn't been compressed yet), do nothing, and keep this
955          * message around until all earlier chunks are received.
956          *
957          * Otherwise, write all the chunks we can. */
958         while (cur_lte != NULL &&
959                !list_empty(&cur_lte->msg_list)
960                && (msg = container_of(cur_lte->msg_list.next,
961                                       struct message,
962                                       list))->complete)
963         {
964                 list_move(&msg->list, &ctx->available_msgs);
965                 if (msg->begin_chunk == 0) {
966                         /* This is the first set of chunks.  Leave space
967                          * for the chunk table in the output file. */
968                         off_t cur_offset = filedes_offset(ctx->out_fd);
969                         if (cur_offset == -1)
970                                 return WIMLIB_ERR_WRITE;
971                         ret = begin_wim_resource_chunk_tab(cur_lte,
972                                                            ctx->out_fd,
973                                                            cur_offset,
974                                                            &ctx->cur_chunk_tab);
975                         if (ret)
976                                 return ret;
977                 }
978
979                 /* Write the compressed chunks from the message. */
980                 ret = write_wim_chunks(msg, ctx->out_fd, ctx->cur_chunk_tab);
981                 if (ret)
982                         return ret;
983
984                 /* Was this the last chunk of the stream?  If so, finish
985                  * it. */
986                 if (list_empty(&cur_lte->msg_list) &&
987                     msg->begin_chunk + msg->num_chunks == ctx->cur_chunk_tab->num_chunks)
988                 {
989                         u64 res_csize;
990                         off_t offset;
991
992                         ret = finish_wim_resource_chunk_tab(ctx->cur_chunk_tab,
993                                                             ctx->out_fd,
994                                                             &res_csize);
995                         if (ret)
996                                 return ret;
997
998                         list_del(&cur_lte->being_compressed_list);
999
1000                         /* Grab the offset of this stream in the output file
1001                          * from the chunk table before we free it. */
1002                         offset = ctx->cur_chunk_tab->file_offset;
1003
1004                         FREE(ctx->cur_chunk_tab);
1005                         ctx->cur_chunk_tab = NULL;
1006
1007                         if (res_csize >= wim_resource_size(cur_lte)) {
1008                                 /* Oops!  We compressed the resource to
1009                                  * larger than the original size.  Write
1010                                  * the resource uncompressed instead. */
1011                                 DEBUG("Compressed %"PRIu64" => %"PRIu64" bytes; "
1012                                       "writing uncompressed instead",
1013                                       wim_resource_size(cur_lte), res_csize);
1014                                 ret = seek_and_truncate(ctx->out_fd, offset);
1015                                 if (ret)
1016                                         return ret;
1017                                 ret = write_wim_resource(cur_lte,
1018                                                          ctx->out_fd,
1019                                                          WIMLIB_COMPRESSION_TYPE_NONE,
1020                                                          &cur_lte->output_resource_entry,
1021                                                          ctx->write_resource_flags);
1022                                 if (ret)
1023                                         return ret;
1024                         } else {
1025                                 cur_lte->output_resource_entry.size =
1026                                         res_csize;
1027
1028                                 cur_lte->output_resource_entry.original_size =
1029                                         cur_lte->resource_entry.original_size;
1030
1031                                 cur_lte->output_resource_entry.offset =
1032                                         offset;
1033
1034                                 cur_lte->output_resource_entry.flags =
1035                                         cur_lte->resource_entry.flags |
1036                                                 WIM_RESHDR_FLAG_COMPRESSED;
1037                         }
1038
1039                         do_write_streams_progress(ctx->progress,
1040                                                   ctx->progress_func,
1041                                                   wim_resource_size(cur_lte),
1042                                                   false);
1043
1044                         /* Since we just finished writing a stream, write any
1045                          * streams that have been added to the serial_streams
1046                          * list for direct writing by the main thread (e.g.
1047                          * resources that don't need to be compressed because
1048                          * the desired compression type is the same as the
1049                          * previous compression type). */
1050                         if (!list_empty(&ctx->serial_streams)) {
1051                                 ret = do_write_stream_list_serial(&ctx->serial_streams,
1052                                                                   ctx->lookup_table,
1053                                                                   ctx->out_fd,
1054                                                                   ctx->out_ctype,
1055                                                                   ctx->write_resource_flags,
1056                                                                   ctx->progress_func,
1057                                                                   ctx->progress);
1058                                 if (ret)
1059                                         return ret;
1060                         }
1061
1062                         /* Advance to the next stream to write. */
1063                         if (list_empty(&ctx->outstanding_streams)) {
1064                                 cur_lte = NULL;
1065                         } else {
1066                                 cur_lte = container_of(ctx->outstanding_streams.next,
1067                                                        struct wim_lookup_table_entry,
1068                                                        being_compressed_list);
1069                         }
1070                 }
1071         }
1072         return 0;
1073 }
1074
1075 /* Called when the main thread has read a new chunk of data. */
1076 static int
1077 main_writer_thread_cb(const void *chunk, size_t chunk_size, void *_ctx)
1078 {
1079         struct main_writer_thread_ctx *ctx = _ctx;
1080         int ret;
1081         struct message *next_msg;
1082         u64 next_chunk_in_msg;
1083
1084         /* Update SHA1 message digest for the stream currently being read by the
1085          * main thread. */
1086         sha1_update(&ctx->next_sha_ctx, chunk, chunk_size);
1087
1088         /* We send chunks of data to the compressor chunks in batches which we
1089          * refer to as "messages".  @next_msg is the message that is currently
1090          * being prepared to send off.  If it is NULL, that indicates that we
1091          * need to start a new message. */
1092         next_msg = ctx->next_msg;
1093         if (!next_msg) {
1094                 /* We need to start a new message.  First check to see if there
1095                  * is a message available in the list of available messages.  If
1096                  * so, we can just take one.  If not, all the messages (there is
1097                  * a fixed number of them, proportional to the number of
1098                  * threads) have been sent off to the compressor threads, so we
1099                  * receive messages from the compressor threads containing
1100                  * compressed chunks of data.
1101                  *
1102                  * We may need to receive multiple messages before one is
1103                  * actually available to use because messages received that are
1104                  * *not* for the very next set of chunks to compress must be
1105                  * buffered until it's time to write those chunks. */
1106                 while (list_empty(&ctx->available_msgs)) {
1107                         ret = receive_compressed_chunks(ctx);
1108                         if (ret)
1109                                 return ret;
1110                 }
1111
1112                 next_msg = container_of(ctx->available_msgs.next,
1113                                         struct message, list);
1114                 list_del(&next_msg->list);
1115                 next_msg->complete = false;
1116                 next_msg->begin_chunk = ctx->next_chunk;
1117                 next_msg->num_chunks = min(MAX_CHUNKS_PER_MSG,
1118                                            ctx->next_num_chunks - ctx->next_chunk);
1119                 ctx->next_msg = next_msg;
1120         }
1121
1122         /* Fill in the next chunk to compress */
1123         next_chunk_in_msg = ctx->next_chunk - next_msg->begin_chunk;
1124
1125         next_msg->uncompressed_chunk_sizes[next_chunk_in_msg] = chunk_size;
1126         memcpy(next_msg->uncompressed_chunks[next_chunk_in_msg],
1127                chunk, chunk_size);
1128         ctx->next_chunk++;
1129         if (++next_chunk_in_msg == next_msg->num_chunks) {
1130                 /* Send off an array of chunks to compress */
1131                 list_add_tail(&next_msg->list, &ctx->next_lte->msg_list);
1132                 shared_queue_put(ctx->res_to_compress_queue, next_msg);
1133                 ++ctx->num_outstanding_messages;
1134                 ctx->next_msg = NULL;
1135         }
1136         return 0;
1137 }
1138
1139 static int
1140 main_writer_thread_finish(void *_ctx)
1141 {
1142         struct main_writer_thread_ctx *ctx = _ctx;
1143         int ret;
1144         while (ctx->num_outstanding_messages != 0) {
1145                 ret = receive_compressed_chunks(ctx);
1146                 if (ret)
1147                         return ret;
1148         }
1149         wimlib_assert(list_empty(&ctx->outstanding_streams));
1150         return do_write_stream_list_serial(&ctx->serial_streams,
1151                                            ctx->lookup_table,
1152                                            ctx->out_fd,
1153                                            ctx->out_ctype,
1154                                            ctx->write_resource_flags,
1155                                            ctx->progress_func,
1156                                            ctx->progress);
1157 }
1158
1159 static int
1160 submit_stream_for_compression(struct wim_lookup_table_entry *lte,
1161                               struct main_writer_thread_ctx *ctx)
1162 {
1163         int ret;
1164
1165         /* Read the entire stream @lte, feeding its data chunks to the
1166          * compressor threads.  Also SHA1-sum the stream; this is required in
1167          * the case that @lte is unhashed, and a nice additional verification
1168          * when @lte is already hashed. */
1169         sha1_init(&ctx->next_sha_ctx);
1170         ctx->next_chunk = 0;
1171         ctx->next_num_chunks = wim_resource_chunks(lte);
1172         ctx->next_lte = lte;
1173         INIT_LIST_HEAD(&lte->msg_list);
1174         list_add_tail(&lte->being_compressed_list, &ctx->outstanding_streams);
1175         ret = read_resource_prefix(lte, wim_resource_size(lte),
1176                                    main_writer_thread_cb, ctx, 0);
1177         if (ret == 0) {
1178                 wimlib_assert(ctx->next_chunk == ctx->next_num_chunks);
1179                 ret = finalize_and_check_sha1(&ctx->next_sha_ctx, lte);
1180         }
1181         return ret;
1182 }
1183
1184 static int
1185 main_thread_process_next_stream(struct wim_lookup_table_entry *lte, void *_ctx)
1186 {
1187         struct main_writer_thread_ctx *ctx = _ctx;
1188         int ret;
1189
1190         if (wim_resource_size(lte) < 1000 ||
1191             ctx->out_ctype == WIMLIB_COMPRESSION_TYPE_NONE ||
1192             (lte->resource_location == RESOURCE_IN_WIM &&
1193              !(ctx->write_resource_flags & WIMLIB_RESOURCE_FLAG_RECOMPRESS) &&
1194              wimlib_get_compression_type(lte->wim) == ctx->out_ctype))
1195         {
1196                 /* Stream is too small or isn't being compressed.  Process it by
1197                  * the main thread when we have a chance.  We can't necessarily
1198                  * process it right here, as the main thread could be in the
1199                  * middle of writing a different stream. */
1200                 list_add_tail(&lte->write_streams_list, &ctx->serial_streams);
1201                 lte->deferred = 1;
1202                 ret = 0;
1203         } else {
1204                 ret = submit_stream_for_compression(lte, ctx);
1205         }
1206         lte->no_progress = 1;
1207         return ret;
1208 }
1209
1210 static long
1211 get_default_num_threads(void)
1212 {
1213 #ifdef __WIN32__
1214         return win32_get_number_of_processors();
1215 #else
1216         return sysconf(_SC_NPROCESSORS_ONLN);
1217 #endif
1218 }
1219
1220 /* Equivalent to write_stream_list_serial(), except this takes a @num_threads
1221  * parameter and will perform compression using that many threads.  Falls
1222  * back to write_stream_list_serial() on certain errors, such as a failure to
1223  * create the number of threads requested.
1224  *
1225  * High level description of the algorithm for writing compressed streams in
1226  * parallel:  We perform compression on chunks of size WIM_CHUNK_SIZE bytes
1227  * rather than on full files.  The currently executing thread becomes the main
1228  * thread and is entirely in charge of reading the data to compress (which may
1229  * be in any location understood by the resource code--- such as in an external
1230  * file being captured, or in another WIM file from which an image is being
1231  * exported) and actually writing the compressed data to the output file.
1232  * Additional threads are "compressor threads" and all execute the
1233  * compressor_thread_proc, where they repeatedly retrieve buffers of data from
1234  * the main thread, compress them, and hand them back to the main thread.
1235  *
1236  * Certain streams, such as streams that do not need to be compressed (e.g.
1237  * input compression type same as output compression type) or streams of very
1238  * small size are placed in a list (main_writer_thread_ctx.serial_list) and
1239  * handled entirely by the main thread at an appropriate time.
1240  *
1241  * At any given point in time, multiple streams may be having chunks compressed
1242  * concurrently.  The stream that the main thread is currently *reading* may be
1243  * later in the list that the stream that the main thread is currently
1244  * *writing*.
1245  */
1246 static int
1247 write_stream_list_parallel(struct list_head *stream_list,
1248                            struct wim_lookup_table *lookup_table,
1249                            int out_fd,
1250                            int out_ctype,
1251                            int write_resource_flags,
1252                            wimlib_progress_func_t progress_func,
1253                            union wimlib_progress_info *progress,
1254                            unsigned num_threads)
1255 {
1256         int ret;
1257         struct shared_queue res_to_compress_queue;
1258         struct shared_queue compressed_res_queue;
1259         pthread_t *compressor_threads = NULL;
1260
1261         if (num_threads == 0) {
1262                 long nthreads = get_default_num_threads();
1263                 if (nthreads < 1 || nthreads > UINT_MAX) {
1264                         WARNING("Could not determine number of processors! Assuming 1");
1265                         goto out_serial;
1266                 } else if (nthreads == 1) {
1267                         goto out_serial_quiet;
1268                 } else {
1269                         num_threads = nthreads;
1270                 }
1271         }
1272
1273         DEBUG("Writing stream list (parallel version, num_threads=%u)",
1274               num_threads);
1275
1276         progress->write_streams.num_threads = num_threads;
1277
1278         static const size_t MESSAGES_PER_THREAD = 2;
1279         size_t queue_size = (size_t)(num_threads * MESSAGES_PER_THREAD);
1280
1281         DEBUG("Initializing shared queues (queue_size=%zu)", queue_size);
1282
1283         ret = shared_queue_init(&res_to_compress_queue, queue_size);
1284         if (ret)
1285                 goto out_serial;
1286
1287         ret = shared_queue_init(&compressed_res_queue, queue_size);
1288         if (ret)
1289                 goto out_destroy_res_to_compress_queue;
1290
1291         struct compressor_thread_params params;
1292         params.res_to_compress_queue = &res_to_compress_queue;
1293         params.compressed_res_queue = &compressed_res_queue;
1294         params.compress = get_compress_func(out_ctype);
1295
1296         compressor_threads = MALLOC(num_threads * sizeof(pthread_t));
1297         if (!compressor_threads) {
1298                 ret = WIMLIB_ERR_NOMEM;
1299                 goto out_destroy_compressed_res_queue;
1300         }
1301
1302         for (unsigned i = 0; i < num_threads; i++) {
1303                 DEBUG("pthread_create thread %u of %u", i + 1, num_threads);
1304                 ret = pthread_create(&compressor_threads[i], NULL,
1305                                      compressor_thread_proc, &params);
1306                 if (ret != 0) {
1307                         ret = -1;
1308                         ERROR_WITH_ERRNO("Failed to create compressor "
1309                                          "thread %u of %u",
1310                                          i + 1, num_threads);
1311                         num_threads = i;
1312                         goto out_join;
1313                 }
1314         }
1315
1316         if (progress_func)
1317                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_STREAMS, progress);
1318
1319         struct main_writer_thread_ctx ctx;
1320         ctx.stream_list           = stream_list;
1321         ctx.lookup_table          = lookup_table;
1322         ctx.out_fd                = out_fd;
1323         ctx.out_ctype             = out_ctype;
1324         ctx.res_to_compress_queue = &res_to_compress_queue;
1325         ctx.compressed_res_queue  = &compressed_res_queue;
1326         ctx.num_messages          = queue_size;
1327         ctx.write_resource_flags  = write_resource_flags;
1328         ctx.progress_func         = progress_func;
1329         ctx.progress              = progress;
1330         ret = main_writer_thread_init_ctx(&ctx);
1331         if (ret)
1332                 goto out_join;
1333         ret = do_write_stream_list(stream_list, lookup_table,
1334                                    main_thread_process_next_stream,
1335                                    &ctx, progress_func, progress);
1336         if (ret)
1337                 goto out_destroy_ctx;
1338
1339         /* The main thread has finished reading all streams that are going to be
1340          * compressed in parallel, and it now needs to wait for all remaining
1341          * chunks to be compressed so that the remaining streams can actually be
1342          * written to the output file.  Furthermore, any remaining streams that
1343          * had processing deferred to the main thread need to be handled.  These
1344          * tasks are done by the main_writer_thread_finish() function. */
1345         ret = main_writer_thread_finish(&ctx);
1346 out_destroy_ctx:
1347         main_writer_thread_destroy_ctx(&ctx);
1348 out_join:
1349         for (unsigned i = 0; i < num_threads; i++)
1350                 shared_queue_put(&res_to_compress_queue, NULL);
1351
1352         for (unsigned i = 0; i < num_threads; i++) {
1353                 if (pthread_join(compressor_threads[i], NULL)) {
1354                         WARNING_WITH_ERRNO("Failed to join compressor "
1355                                            "thread %u of %u",
1356                                            i + 1, num_threads);
1357                 }
1358         }
1359         FREE(compressor_threads);
1360 out_destroy_compressed_res_queue:
1361         shared_queue_destroy(&compressed_res_queue);
1362 out_destroy_res_to_compress_queue:
1363         shared_queue_destroy(&res_to_compress_queue);
1364         if (ret >= 0 && ret != WIMLIB_ERR_NOMEM)
1365                 return ret;
1366 out_serial:
1367         WARNING("Falling back to single-threaded compression");
1368 out_serial_quiet:
1369         return write_stream_list_serial(stream_list,
1370                                         lookup_table,
1371                                         out_fd,
1372                                         out_ctype,
1373                                         write_resource_flags,
1374                                         progress_func,
1375                                         progress);
1376
1377 }
1378 #endif
1379
1380 /*
1381  * Write a list of streams to a WIM (@out_fd) using the compression type
1382  * @out_ctype and up to @num_threads compressor threads.
1383  */
1384 static int
1385 write_stream_list(struct list_head *stream_list,
1386                   struct wim_lookup_table *lookup_table,
1387                   int out_fd, int out_ctype, int write_flags,
1388                   unsigned num_threads, wimlib_progress_func_t progress_func)
1389 {
1390         struct wim_lookup_table_entry *lte;
1391         size_t num_streams = 0;
1392         u64 total_bytes = 0;
1393         u64 total_compression_bytes = 0;
1394         union wimlib_progress_info progress;
1395         int ret;
1396         int write_resource_flags;
1397
1398         if (list_empty(stream_list))
1399                 return 0;
1400
1401         write_resource_flags = write_flags_to_resource_flags(write_flags);
1402
1403         /* Calculate the total size of the streams to be written.  Note: this
1404          * will be the uncompressed size, as we may not know the compressed size
1405          * yet, and also this will assume that every unhashed stream will be
1406          * written (which will not necessarily be the case). */
1407         list_for_each_entry(lte, stream_list, write_streams_list) {
1408                 num_streams++;
1409                 total_bytes += wim_resource_size(lte);
1410                 if (out_ctype != WIMLIB_COMPRESSION_TYPE_NONE
1411                        && (wim_resource_compression_type(lte) != out_ctype ||
1412                            (write_resource_flags & WIMLIB_RESOURCE_FLAG_RECOMPRESS)))
1413                 {
1414                         total_compression_bytes += wim_resource_size(lte);
1415                 }
1416         }
1417         progress.write_streams.total_bytes       = total_bytes;
1418         progress.write_streams.total_streams     = num_streams;
1419         progress.write_streams.completed_bytes   = 0;
1420         progress.write_streams.completed_streams = 0;
1421         progress.write_streams.num_threads       = num_threads;
1422         progress.write_streams.compression_type  = out_ctype;
1423         progress.write_streams._private          = 0;
1424
1425 #ifdef ENABLE_MULTITHREADED_COMPRESSION
1426         if (total_compression_bytes >= 2000000 && num_threads != 1)
1427                 ret = write_stream_list_parallel(stream_list,
1428                                                  lookup_table,
1429                                                  out_fd,
1430                                                  out_ctype,
1431                                                  write_resource_flags,
1432                                                  progress_func,
1433                                                  &progress,
1434                                                  num_threads);
1435         else
1436 #endif
1437                 ret = write_stream_list_serial(stream_list,
1438                                                lookup_table,
1439                                                out_fd,
1440                                                out_ctype,
1441                                                write_resource_flags,
1442                                                progress_func,
1443                                                &progress);
1444         return ret;
1445 }
1446
1447 struct stream_size_table {
1448         struct hlist_head *array;
1449         size_t num_entries;
1450         size_t capacity;
1451 };
1452
1453 static int
1454 init_stream_size_table(struct stream_size_table *tab, size_t capacity)
1455 {
1456         tab->array = CALLOC(capacity, sizeof(tab->array[0]));
1457         if (!tab->array)
1458                 return WIMLIB_ERR_NOMEM;
1459         tab->num_entries = 0;
1460         tab->capacity = capacity;
1461         return 0;
1462 }
1463
1464 static void
1465 destroy_stream_size_table(struct stream_size_table *tab)
1466 {
1467         FREE(tab->array);
1468 }
1469
1470 static int
1471 stream_size_table_insert(struct wim_lookup_table_entry *lte, void *_tab)
1472 {
1473         struct stream_size_table *tab = _tab;
1474         size_t pos;
1475         struct wim_lookup_table_entry *same_size_lte;
1476         struct hlist_node *tmp;
1477
1478         pos = hash_u64(wim_resource_size(lte)) % tab->capacity;
1479         lte->unique_size = 1;
1480         hlist_for_each_entry(same_size_lte, tmp, &tab->array[pos], hash_list_2) {
1481                 if (wim_resource_size(same_size_lte) == wim_resource_size(lte)) {
1482                         lte->unique_size = 0;
1483                         same_size_lte->unique_size = 0;
1484                         break;
1485                 }
1486         }
1487
1488         hlist_add_head(&lte->hash_list_2, &tab->array[pos]);
1489         tab->num_entries++;
1490         return 0;
1491 }
1492
1493
1494 struct lte_overwrite_prepare_args {
1495         WIMStruct *wim;
1496         off_t end_offset;
1497         struct list_head stream_list;
1498         struct stream_size_table stream_size_tab;
1499 };
1500
1501 /* First phase of preparing streams for an in-place overwrite.  This is called
1502  * on all streams, both hashed and unhashed, except the metadata resources. */
1503 static int
1504 lte_overwrite_prepare(struct wim_lookup_table_entry *lte, void *_args)
1505 {
1506         struct lte_overwrite_prepare_args *args = _args;
1507
1508         wimlib_assert(!(lte->resource_entry.flags & WIM_RESHDR_FLAG_METADATA));
1509         if (lte->resource_location != RESOURCE_IN_WIM || lte->wim != args->wim)
1510                 list_add_tail(&lte->write_streams_list, &args->stream_list);
1511         lte->out_refcnt = lte->refcnt;
1512         stream_size_table_insert(lte, &args->stream_size_tab);
1513         return 0;
1514 }
1515
1516 /* Second phase of preparing streams for an in-place overwrite.  This is called
1517  * on existing metadata resources and hashed streams, but not unhashed streams.
1518  *
1519  * NOTE: lte->output_resource_entry is in union with lte->hash_list_2, so
1520  * lte_overwrite_prepare_2() must be called after lte_overwrite_prepare(), as
1521  * the latter uses lte->hash_list_2, while the former expects to set
1522  * lte->output_resource_entry. */
1523 static int
1524 lte_overwrite_prepare_2(struct wim_lookup_table_entry *lte, void *_args)
1525 {
1526         struct lte_overwrite_prepare_args *args = _args;
1527
1528         if (lte->resource_location == RESOURCE_IN_WIM && lte->wim == args->wim) {
1529                 /* We can't do an in place overwrite on the WIM if there are
1530                  * streams after the XML data. */
1531                 if (lte->resource_entry.offset +
1532                     lte->resource_entry.size > args->end_offset)
1533                 {
1534                 #ifdef ENABLE_ERROR_MESSAGES
1535                         ERROR("The following resource is after the XML data:");
1536                         print_lookup_table_entry(lte, stderr);
1537                 #endif
1538                         return WIMLIB_ERR_RESOURCE_ORDER;
1539                 }
1540                 copy_resource_entry(&lte->output_resource_entry,
1541                                     &lte->resource_entry);
1542         }
1543         return 0;
1544 }
1545
1546 /* Given a WIM that we are going to overwrite in place with zero or more
1547  * additional streams added, construct a list the list of new unique streams
1548  * ('struct wim_lookup_table_entry's) that must be written, plus any unhashed
1549  * streams that need to be added but may be identical to other hashed or
1550  * unhashed streams.  These unhashed streams are checksummed while the streams
1551  * are being written.  To aid this process, the member @unique_size is set to 1
1552  * on streams that have a unique size and therefore must be written.
1553  *
1554  * The out_refcnt member of each 'struct wim_lookup_table_entry' is set to
1555  * indicate the number of times the stream is referenced in only the streams
1556  * that are being written; this may still be adjusted later when unhashed
1557  * streams are being resolved.
1558  */
1559 static int
1560 prepare_streams_for_overwrite(WIMStruct *wim, off_t end_offset,
1561                               struct list_head *stream_list)
1562 {
1563         int ret;
1564         struct lte_overwrite_prepare_args args;
1565         unsigned i;
1566
1567         args.wim = wim;
1568         args.end_offset = end_offset;
1569         ret = init_stream_size_table(&args.stream_size_tab,
1570                                      wim->lookup_table->capacity);
1571         if (ret)
1572                 return ret;
1573
1574         INIT_LIST_HEAD(&args.stream_list);
1575         for (i = 0; i < wim->hdr.image_count; i++) {
1576                 struct wim_image_metadata *imd;
1577                 struct wim_lookup_table_entry *lte;
1578
1579                 imd = wim->image_metadata[i];
1580                 image_for_each_unhashed_stream(lte, imd)
1581                         lte_overwrite_prepare(lte, &args);
1582         }
1583         for_lookup_table_entry(wim->lookup_table, lte_overwrite_prepare, &args);
1584         list_transfer(&args.stream_list, stream_list);
1585
1586         for (i = 0; i < wim->hdr.image_count; i++) {
1587                 ret = lte_overwrite_prepare_2(wim->image_metadata[i]->metadata_lte,
1588                                               &args);
1589                 if (ret)
1590                         goto out_destroy_stream_size_table;
1591         }
1592         ret = for_lookup_table_entry(wim->lookup_table,
1593                                      lte_overwrite_prepare_2, &args);
1594 out_destroy_stream_size_table:
1595         destroy_stream_size_table(&args.stream_size_tab);
1596         return ret;
1597 }
1598
1599
1600 struct find_streams_ctx {
1601         struct list_head stream_list;
1602         struct stream_size_table stream_size_tab;
1603 };
1604
1605 static void
1606 inode_find_streams_to_write(struct wim_inode *inode,
1607                             struct wim_lookup_table *table,
1608                             struct list_head *stream_list,
1609                             struct stream_size_table *tab)
1610 {
1611         struct wim_lookup_table_entry *lte;
1612         for (unsigned i = 0; i <= inode->i_num_ads; i++) {
1613                 lte = inode_stream_lte(inode, i, table);
1614                 if (lte) {
1615                         if (lte->out_refcnt == 0) {
1616                                 if (lte->unhashed)
1617                                         stream_size_table_insert(lte, tab);
1618                                 list_add_tail(&lte->write_streams_list, stream_list);
1619                         }
1620                         lte->out_refcnt += inode->i_nlink;
1621                 }
1622         }
1623 }
1624
1625 static int
1626 image_find_streams_to_write(WIMStruct *w)
1627 {
1628         struct find_streams_ctx *ctx;
1629         struct wim_image_metadata *imd;
1630         struct wim_inode *inode;
1631         struct wim_lookup_table_entry *lte;
1632
1633         ctx = w->private;
1634         imd = wim_get_current_image_metadata(w);
1635
1636         image_for_each_unhashed_stream(lte, imd)
1637                 lte->out_refcnt = 0;
1638
1639         /* Go through this image's inodes to find any streams that have not been
1640          * found yet. */
1641         image_for_each_inode(inode, imd) {
1642                 inode_find_streams_to_write(inode, w->lookup_table,
1643                                             &ctx->stream_list,
1644                                             &ctx->stream_size_tab);
1645         }
1646         return 0;
1647 }
1648
1649 /* Given a WIM that from which one or all of the images is being written, build
1650  * the list of unique streams ('struct wim_lookup_table_entry's) that must be
1651  * written, plus any unhashed streams that need to be written but may be
1652  * identical to other hashed or unhashed streams being written.  These unhashed
1653  * streams are checksummed while the streams are being written.  To aid this
1654  * process, the member @unique_size is set to 1 on streams that have a unique
1655  * size and therefore must be written.
1656  *
1657  * The out_refcnt member of each 'struct wim_lookup_table_entry' is set to
1658  * indicate the number of times the stream is referenced in only the streams
1659  * that are being written; this may still be adjusted later when unhashed
1660  * streams are being resolved.
1661  */
1662 static int
1663 prepare_stream_list(WIMStruct *wim, int image, struct list_head *stream_list)
1664 {
1665         int ret;
1666         struct find_streams_ctx ctx;
1667
1668         for_lookup_table_entry(wim->lookup_table, lte_zero_out_refcnt, NULL);
1669         ret = init_stream_size_table(&ctx.stream_size_tab,
1670                                      wim->lookup_table->capacity);
1671         if (ret)
1672                 return ret;
1673         for_lookup_table_entry(wim->lookup_table, stream_size_table_insert,
1674                                &ctx.stream_size_tab);
1675         INIT_LIST_HEAD(&ctx.stream_list);
1676         wim->private = &ctx;
1677         ret = for_image(wim, image, image_find_streams_to_write);
1678         destroy_stream_size_table(&ctx.stream_size_tab);
1679         if (ret == 0)
1680                 list_transfer(&ctx.stream_list, stream_list);
1681         return ret;
1682 }
1683
1684 /* Writes the streams for the specified @image in @wim to @wim->out_fd.
1685  */
1686 static int
1687 write_wim_streams(WIMStruct *wim, int image, int write_flags,
1688                   unsigned num_threads,
1689                   wimlib_progress_func_t progress_func)
1690 {
1691         int ret;
1692         struct list_head stream_list;
1693
1694         ret = prepare_stream_list(wim, image, &stream_list);
1695         if (ret)
1696                 return ret;
1697         return write_stream_list(&stream_list,
1698                                  wim->lookup_table,
1699                                  wim->out_fd,
1700                                  wimlib_get_compression_type(wim),
1701                                  write_flags,
1702                                  num_threads,
1703                                  progress_func);
1704 }
1705
1706 /*
1707  * Finish writing a WIM file: write the lookup table, xml data, and integrity
1708  * table (optional), then overwrite the WIM header.
1709  *
1710  * write_flags is a bitwise OR of the following:
1711  *
1712  *      (public)  WIMLIB_WRITE_FLAG_CHECK_INTEGRITY:
1713  *              Include an integrity table.
1714  *
1715  *      (private) WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE:
1716  *              Don't write the lookup table.
1717  *
1718  *      (private) WIMLIB_WRITE_FLAG_REUSE_INTEGRITY_TABLE:
1719  *              When (if) writing the integrity table, re-use entries from the
1720  *              existing integrity table, if possible.
1721  *
1722  *      (private) WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML:
1723  *              After writing the XML data but before writing the integrity
1724  *              table, write a temporary WIM header and flush the stream so that
1725  *              the WIM is less likely to become corrupted upon abrupt program
1726  *              termination.
1727  *
1728  *      (private) WIMLIB_WRITE_FLAG_FSYNC:
1729  *              fsync() the output file before closing it.
1730  *
1731  */
1732 int
1733 finish_write(WIMStruct *w, int image, int write_flags,
1734              wimlib_progress_func_t progress_func)
1735 {
1736         int ret;
1737         struct wim_header hdr;
1738
1739         /* @hdr will be the header for the new WIM.  First copy all the data
1740          * from the header in the WIMStruct; then set all the fields that may
1741          * have changed, including the resource entries, boot index, and image
1742          * count.  */
1743         memcpy(&hdr, &w->hdr, sizeof(struct wim_header));
1744
1745         /* Set image count and boot index correctly for single image writes */
1746         if (image != WIMLIB_ALL_IMAGES) {
1747                 hdr.image_count = 1;
1748                 if (hdr.boot_idx == image)
1749                         hdr.boot_idx = 1;
1750                 else
1751                         hdr.boot_idx = 0;
1752         }
1753
1754         /* In the WIM header, there is room for the resource entry for a
1755          * metadata resource labeled as the "boot metadata".  This entry should
1756          * be zeroed out if there is no bootable image (boot_idx 0).  Otherwise,
1757          * it should be a copy of the resource entry for the image that is
1758          * marked as bootable.  This is not well documented...  */
1759         if (hdr.boot_idx == 0) {
1760                 zero_resource_entry(&hdr.boot_metadata_res_entry);
1761         } else {
1762                 copy_resource_entry(&hdr.boot_metadata_res_entry,
1763                             &w->image_metadata[ hdr.boot_idx- 1
1764                                         ]->metadata_lte->output_resource_entry);
1765         }
1766
1767         if (!(write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE)) {
1768                 ret = write_lookup_table(w, image, &hdr.lookup_table_res_entry);
1769                 if (ret)
1770                         goto out_close_wim;
1771         }
1772
1773         ret = write_xml_data(w->wim_info, image, w->out_fd,
1774                              (write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE) ?
1775                               wim_info_get_total_bytes(w->wim_info) : 0,
1776                              &hdr.xml_res_entry);
1777         if (ret)
1778                 goto out_close_wim;
1779
1780         if (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) {
1781                 if (write_flags & WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML) {
1782                         struct wim_header checkpoint_hdr;
1783                         memcpy(&checkpoint_hdr, &hdr, sizeof(struct wim_header));
1784                         zero_resource_entry(&checkpoint_hdr.integrity);
1785                         ret = write_header(&checkpoint_hdr, w->out_fd);
1786                         if (ret)
1787                                 goto out_close_wim;
1788                 }
1789
1790                 off_t old_lookup_table_end;
1791                 off_t new_lookup_table_end;
1792                 if (write_flags & WIMLIB_WRITE_FLAG_REUSE_INTEGRITY_TABLE) {
1793                         old_lookup_table_end = w->hdr.lookup_table_res_entry.offset +
1794                                                w->hdr.lookup_table_res_entry.size;
1795                 } else {
1796                         old_lookup_table_end = 0;
1797                 }
1798                 new_lookup_table_end = hdr.lookup_table_res_entry.offset +
1799                                        hdr.lookup_table_res_entry.size;
1800
1801                 ret = write_integrity_table(w->out_fd,
1802                                             &hdr.integrity,
1803                                             new_lookup_table_end,
1804                                             old_lookup_table_end,
1805                                             progress_func);
1806                 if (ret)
1807                         goto out_close_wim;
1808         } else {
1809                 zero_resource_entry(&hdr.integrity);
1810         }
1811
1812         ret = write_header(&hdr, w->out_fd);
1813         if (ret)
1814                 goto out_close_wim;
1815
1816         if (write_flags & WIMLIB_WRITE_FLAG_FSYNC) {
1817                 if (fsync(w->out_fd)) {
1818                         ERROR_WITH_ERRNO("Error syncing data to WIM file");
1819                         ret = WIMLIB_ERR_WRITE;
1820                 }
1821         }
1822 out_close_wim:
1823         if (close(w->out_fd)) {
1824                 ERROR_WITH_ERRNO("Failed to close the output WIM file");
1825                 if (ret == 0)
1826                         ret = WIMLIB_ERR_WRITE;
1827         }
1828         w->out_fd = -1;
1829         return ret;
1830 }
1831
1832 #if defined(HAVE_SYS_FILE_H) && defined(HAVE_FLOCK)
1833 int
1834 lock_wim(WIMStruct *w, int fd)
1835 {
1836         int ret = 0;
1837         if (fd != -1 && !w->wim_locked) {
1838                 ret = flock(fd, LOCK_EX | LOCK_NB);
1839                 if (ret != 0) {
1840                         if (errno == EWOULDBLOCK) {
1841                                 ERROR("`%"TS"' is already being modified or has been "
1842                                       "mounted read-write\n"
1843                                       "        by another process!", w->filename);
1844                                 ret = WIMLIB_ERR_ALREADY_LOCKED;
1845                         } else {
1846                                 WARNING_WITH_ERRNO("Failed to lock `%"TS"'",
1847                                                    w->filename);
1848                                 ret = 0;
1849                         }
1850                 } else {
1851                         w->wim_locked = 1;
1852                 }
1853         }
1854         return ret;
1855 }
1856 #endif
1857
1858 static int
1859 open_wim_writable(WIMStruct *w, const tchar *path, int open_flags)
1860 {
1861         w->out_fd = topen(path, open_flags | O_BINARY, 0644);
1862         if (w->out_fd == -1) {
1863                 ERROR_WITH_ERRNO("Failed to open `%"TS"' for writing", path);
1864                 return WIMLIB_ERR_OPEN;
1865         }
1866         return 0;
1867 }
1868
1869
1870 void
1871 close_wim_writable(WIMStruct *w)
1872 {
1873         if (w->out_fd != -1) {
1874                 if (close(w->out_fd))
1875                         WARNING_WITH_ERRNO("Failed to close output WIM");
1876                 w->out_fd = -1;
1877         }
1878 }
1879
1880 /* Open file stream and write dummy header for WIM. */
1881 int
1882 begin_write(WIMStruct *w, const tchar *path, int write_flags)
1883 {
1884         int ret;
1885         int open_flags = O_TRUNC | O_CREAT;
1886         if (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY)
1887                 open_flags |= O_RDWR;
1888         else
1889                 open_flags |= O_WRONLY;
1890         ret = open_wim_writable(w, path, open_flags);
1891         if (ret)
1892                 return ret;
1893         /* Write dummy header. It will be overwritten later. */
1894         ret = write_header(&w->hdr, w->out_fd);
1895         if (ret)
1896                 return ret;
1897         if (lseek(w->out_fd, WIM_HEADER_DISK_SIZE, SEEK_SET) == -1) {
1898                 ERROR_WITH_ERRNO("Failed to seek to end of WIM header");
1899                 return WIMLIB_ERR_WRITE;
1900         }
1901         return 0;
1902 }
1903
1904 /* Writes a stand-alone WIM to a file.  */
1905 WIMLIBAPI int
1906 wimlib_write(WIMStruct *w, const tchar *path,
1907              int image, int write_flags, unsigned num_threads,
1908              wimlib_progress_func_t progress_func)
1909 {
1910         int ret;
1911
1912         if (!path)
1913                 return WIMLIB_ERR_INVALID_PARAM;
1914
1915         write_flags &= WIMLIB_WRITE_MASK_PUBLIC;
1916
1917         if (image != WIMLIB_ALL_IMAGES &&
1918              (image < 1 || image > w->hdr.image_count))
1919                 return WIMLIB_ERR_INVALID_IMAGE;
1920
1921         if (w->hdr.total_parts != 1) {
1922                 ERROR("Cannot call wimlib_write() on part of a split WIM");
1923                 return WIMLIB_ERR_SPLIT_UNSUPPORTED;
1924         }
1925
1926         ret = begin_write(w, path, write_flags);
1927         if (ret)
1928                 goto out_close_wim;
1929
1930         ret = write_wim_streams(w, image, write_flags, num_threads,
1931                                 progress_func);
1932         if (ret)
1933                 goto out_close_wim;
1934
1935         if (progress_func)
1936                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_METADATA_BEGIN, NULL);
1937
1938         ret = for_image(w, image, write_metadata_resource);
1939         if (ret)
1940                 goto out_close_wim;
1941
1942         if (progress_func)
1943                 progress_func(WIMLIB_PROGRESS_MSG_WRITE_METADATA_END, NULL);
1944
1945         ret = finish_write(w, image, write_flags, progress_func);
1946         /* finish_write() closed the WIM for us */
1947         goto out;
1948 out_close_wim:
1949         close_wim_writable(w);
1950 out:
1951         DEBUG("wimlib_write(path=%"TS") = %d", path, ret);
1952         return ret;
1953 }
1954
1955 static bool
1956 any_images_modified(WIMStruct *w)
1957 {
1958         for (int i = 0; i < w->hdr.image_count; i++)
1959                 if (w->image_metadata[i]->modified)
1960                         return true;
1961         return false;
1962 }
1963
1964 /*
1965  * Overwrite a WIM, possibly appending streams to it.
1966  *
1967  * A WIM looks like (or is supposed to look like) the following:
1968  *
1969  *                   Header (212 bytes)
1970  *                   Streams and metadata resources (variable size)
1971  *                   Lookup table (variable size)
1972  *                   XML data (variable size)
1973  *                   Integrity table (optional) (variable size)
1974  *
1975  * If we are not adding any streams or metadata resources, the lookup table is
1976  * unchanged--- so we only need to overwrite the XML data, integrity table, and
1977  * header.  This operation is potentially unsafe if the program is abruptly
1978  * terminated while the XML data or integrity table are being overwritten, but
1979  * before the new header has been written.  To partially alleviate this problem,
1980  * a special flag (WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML) is passed to
1981  * finish_write() to cause a temporary WIM header to be written after the XML
1982  * data has been written.  This may prevent the WIM from becoming corrupted if
1983  * the program is terminated while the integrity table is being calculated (but
1984  * no guarantees, due to write re-ordering...).
1985  *
1986  * If we are adding new streams or images (metadata resources), the lookup table
1987  * needs to be changed, and those streams need to be written.  In this case, we
1988  * try to perform a safe update of the WIM file by writing the streams *after*
1989  * the end of the previous WIM, then writing the new lookup table, XML data, and
1990  * (optionally) integrity table following the new streams.  This will produce a
1991  * layout like the following:
1992  *
1993  *                   Header (212 bytes)
1994  *                   (OLD) Streams and metadata resources (variable size)
1995  *                   (OLD) Lookup table (variable size)
1996  *                   (OLD) XML data (variable size)
1997  *                   (OLD) Integrity table (optional) (variable size)
1998  *                   (NEW) Streams and metadata resources (variable size)
1999  *                   (NEW) Lookup table (variable size)
2000  *                   (NEW) XML data (variable size)
2001  *                   (NEW) Integrity table (optional) (variable size)
2002  *
2003  * At all points, the WIM is valid as nothing points to the new data yet.  Then,
2004  * the header is overwritten to point to the new lookup table, XML data, and
2005  * integrity table, to produce the following layout:
2006  *
2007  *                   Header (212 bytes)
2008  *                   Streams and metadata resources (variable size)
2009  *                   Nothing (variable size)
2010  *                   More Streams and metadata resources (variable size)
2011  *                   Lookup table (variable size)
2012  *                   XML data (variable size)
2013  *                   Integrity table (optional) (variable size)
2014  *
2015  * This method allows an image to be appended to a large WIM very quickly, and
2016  * is is crash-safe except in the case of write re-ordering, but the
2017  * disadvantage is that a small hole is left in the WIM where the old lookup
2018  * table, xml data, and integrity table were.  (These usually only take up a
2019  * small amount of space compared to the streams, however.)
2020  */
2021 static int
2022 overwrite_wim_inplace(WIMStruct *w, int write_flags,
2023                       unsigned num_threads,
2024                       wimlib_progress_func_t progress_func)
2025 {
2026         int ret;
2027         struct list_head stream_list;
2028         off_t old_wim_end;
2029         u64 old_lookup_table_end, old_xml_begin, old_xml_end;
2030         int open_flags;
2031
2032         DEBUG("Overwriting `%"TS"' in-place", w->filename);
2033
2034         /* Make sure that the integrity table (if present) is after the XML
2035          * data, and that there are no stream resources, metadata resources, or
2036          * lookup tables after the XML data.  Otherwise, these data would be
2037          * overwritten. */
2038         old_xml_begin = w->hdr.xml_res_entry.offset;
2039         old_xml_end = old_xml_begin + w->hdr.xml_res_entry.size;
2040         old_lookup_table_end = w->hdr.lookup_table_res_entry.offset +
2041                                w->hdr.lookup_table_res_entry.size;
2042         if (w->hdr.integrity.offset != 0 && w->hdr.integrity.offset < old_xml_end) {
2043                 ERROR("Didn't expect the integrity table to be before the XML data");
2044                 return WIMLIB_ERR_RESOURCE_ORDER;
2045         }
2046
2047         if (old_lookup_table_end > old_xml_begin) {
2048                 ERROR("Didn't expect the lookup table to be after the XML data");
2049                 return WIMLIB_ERR_RESOURCE_ORDER;
2050         }
2051
2052         /* Set @old_wim_end, which indicates the point beyond which we don't
2053          * allow any file and metadata resources to appear without returning
2054          * WIMLIB_ERR_RESOURCE_ORDER (due to the fact that we would otherwise
2055          * overwrite these resources). */
2056         if (!w->deletion_occurred && !any_images_modified(w)) {
2057                 /* If no images have been modified and no images have been
2058                  * deleted, a new lookup table does not need to be written.  We
2059                  * shall write the new XML data and optional integrity table
2060                  * immediately after the lookup table.  Note that this may
2061                  * overwrite an existing integrity table. */
2062                 DEBUG("Skipping writing lookup table "
2063                       "(no images modified or deleted)");
2064                 old_wim_end = old_lookup_table_end;
2065                 write_flags |= WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE |
2066                                WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML;
2067         } else if (w->hdr.integrity.offset) {
2068                 /* Old WIM has an integrity table; begin writing new streams
2069                  * after it. */
2070                 old_wim_end = w->hdr.integrity.offset + w->hdr.integrity.size;
2071         } else {
2072                 /* No existing integrity table; begin writing new streams after
2073                  * the old XML data. */
2074                 old_wim_end = old_xml_end;
2075         }
2076
2077         ret = prepare_streams_for_overwrite(w, old_wim_end, &stream_list);
2078         if (ret)
2079                 return ret;
2080
2081         open_flags = 0;
2082         if (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY)
2083                 open_flags |= O_RDWR;
2084         else
2085                 open_flags |= O_WRONLY;
2086         ret = open_wim_writable(w, w->filename, open_flags);
2087         if (ret)
2088                 return ret;
2089
2090         ret = lock_wim(w, w->out_fd);
2091         if (ret) {
2092                 close_wim_writable(w);
2093                 return ret;
2094         }
2095
2096         if (lseek(w->out_fd, old_wim_end, SEEK_SET) == -1) {
2097                 ERROR_WITH_ERRNO("Can't seek to end of WIM");
2098                 close_wim_writable(w);
2099                 w->wim_locked = 0;
2100                 return WIMLIB_ERR_WRITE;
2101         }
2102
2103         DEBUG("Writing newly added streams (offset = %"PRIu64")",
2104               old_wim_end);
2105         ret = write_stream_list(&stream_list,
2106                                 w->lookup_table,
2107                                 w->out_fd,
2108                                 wimlib_get_compression_type(w),
2109                                 write_flags,
2110                                 num_threads,
2111                                 progress_func);
2112         if (ret)
2113                 goto out_truncate;
2114
2115         for (int i = 0; i < w->hdr.image_count; i++) {
2116                 if (w->image_metadata[i]->modified) {
2117                         select_wim_image(w, i + 1);
2118                         ret = write_metadata_resource(w);
2119                         if (ret)
2120                                 goto out_truncate;
2121                 }
2122         }
2123         write_flags |= WIMLIB_WRITE_FLAG_REUSE_INTEGRITY_TABLE;
2124         ret = finish_write(w, WIMLIB_ALL_IMAGES, write_flags,
2125                            progress_func);
2126 out_truncate:
2127         close_wim_writable(w);
2128         if (ret != 0 && !(write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE)) {
2129                 WARNING("Truncating `%"TS"' to its original size (%"PRIu64" bytes)",
2130                         w->filename, old_wim_end);
2131                 /* Return value of truncate() is ignored because this is already
2132                  * an error path. */
2133                 (void)ttruncate(w->filename, old_wim_end);
2134         }
2135         w->wim_locked = 0;
2136         return ret;
2137 }
2138
2139 static int
2140 overwrite_wim_via_tmpfile(WIMStruct *w, int write_flags,
2141                           unsigned num_threads,
2142                           wimlib_progress_func_t progress_func)
2143 {
2144         size_t wim_name_len;
2145         int ret;
2146
2147         DEBUG("Overwriting `%"TS"' via a temporary file", w->filename);
2148
2149         /* Write the WIM to a temporary file in the same directory as the
2150          * original WIM. */
2151         wim_name_len = tstrlen(w->filename);
2152         tchar tmpfile[wim_name_len + 10];
2153         tmemcpy(tmpfile, w->filename, wim_name_len);
2154         randomize_char_array_with_alnum(tmpfile + wim_name_len, 9);
2155         tmpfile[wim_name_len + 9] = T('\0');
2156
2157         ret = wimlib_write(w, tmpfile, WIMLIB_ALL_IMAGES,
2158                            write_flags | WIMLIB_WRITE_FLAG_FSYNC,
2159                            num_threads, progress_func);
2160         if (ret) {
2161                 ERROR("Failed to write the WIM file `%"TS"'", tmpfile);
2162                 goto out_unlink;
2163         }
2164
2165         close_wim(w);
2166
2167         DEBUG("Renaming `%"TS"' to `%"TS"'", tmpfile, w->filename);
2168         /* Rename the new file to the old file .*/
2169         if (trename(tmpfile, w->filename) != 0) {
2170                 ERROR_WITH_ERRNO("Failed to rename `%"TS"' to `%"TS"'",
2171                                  tmpfile, w->filename);
2172                 ret = WIMLIB_ERR_RENAME;
2173                 goto out_unlink;
2174         }
2175
2176         if (progress_func) {
2177                 union wimlib_progress_info progress;
2178                 progress.rename.from = tmpfile;
2179                 progress.rename.to = w->filename;
2180                 progress_func(WIMLIB_PROGRESS_MSG_RENAME, &progress);
2181         }
2182         goto out;
2183 out_unlink:
2184         /* Remove temporary file. */
2185         if (tunlink(tmpfile) != 0)
2186                 WARNING_WITH_ERRNO("Failed to remove `%"TS"'", tmpfile);
2187 out:
2188         return ret;
2189 }
2190
2191 /*
2192  * Writes a WIM file to the original file that it was read from, overwriting it.
2193  */
2194 WIMLIBAPI int
2195 wimlib_overwrite(WIMStruct *w, int write_flags,
2196                  unsigned num_threads,
2197                  wimlib_progress_func_t progress_func)
2198 {
2199         write_flags &= WIMLIB_WRITE_MASK_PUBLIC;
2200
2201         if (!w->filename)
2202                 return WIMLIB_ERR_NO_FILENAME;
2203
2204         if (w->hdr.total_parts != 1) {
2205                 ERROR("Cannot modify a split WIM");
2206                 return WIMLIB_ERR_SPLIT_UNSUPPORTED;
2207         }
2208
2209         if ((!w->deletion_occurred || (write_flags & WIMLIB_WRITE_FLAG_SOFT_DELETE))
2210             && !(write_flags & WIMLIB_WRITE_FLAG_REBUILD))
2211         {
2212                 int ret;
2213                 ret = overwrite_wim_inplace(w, write_flags, num_threads,
2214                                             progress_func);
2215                 if (ret == WIMLIB_ERR_RESOURCE_ORDER)
2216                         WARNING("Falling back to re-building entire WIM");
2217                 else
2218                         return ret;
2219         }
2220         return overwrite_wim_via_tmpfile(w, write_flags, num_threads,
2221                                          progress_func);
2222 }