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