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

   1 /*
   2  * write.c
   3  *
   4  * Support for writing WIM files; write a WIM file, overwrite a WIM file, write
   5  * compressed file resources, etc.
   6  *
   7  * Copyright (C) 2010 Carl Thijssen
   8  * Copyright (C) 2012 Eric Biggers
   9  *
  10  * wimlib - Library for working with WIM files
  11  *
  12  * This library is free software; you can redistribute it and/or modify it under
  13  * the terms of the GNU Lesser General Public License as published by the Free
  14  * Software Foundation; either version 2.1 of the License, or (at your option) any
  15  * later version.
  16  *
  17  * This library is distributed in the hope that it will be useful, but WITHOUT ANY
  18  * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
  19  * PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.
  20  *
  21  * You should have received a copy of the GNU Lesser General Public License along
  22  * with this library; if not, write to the Free Software Foundation, Inc., 59
  23  * Temple Place, Suite 330, Boston, MA 02111-1307 USA
  24  */
  25 #include "wimlib_internal.h"
  26 #include "io.h"
  27 #include "lookup_table.h"
  28 #include "dentry.h"
  29 #include "sha1.h"
  30 #include "lzx.h"
  31 #include "xml.h"
  32 #include "xpress.h"
  33 #include <unistd.h>
  34
  35
  36
  37 /* Used for buffering FILE IO */
  38 #define BUFFER_SIZE 4096
  39
  40 /*
  41  * Copies bytes between two file streams.
  42  *
  43  * Copies @len bytes from @in to @out, at the current position in @out, and at
  44  * an offset of @in_offset in @in.
  45  */
  46 static int copy_between_files(FILE *in, off_t in_offset, FILE *out, size_t len)
  47 {
  48         u8 buf[BUFFER_SIZE];
  49         size_t n;
  50
  51         if (fseeko(in, in_offset, SEEK_SET) != 0) {
  52                 ERROR("Failed to seek to byte %"PRIu64" of input file: %m\n",
  53                                 in_offset);
  54                 return WIMLIB_ERR_READ;
  55         }
  56         /* To reduce memory usage and improve speed, read and write BUFFER_SIZE
  57          * bytes at a time. */
  58         while (len != 0) {
  59                 n = min(len, BUFFER_SIZE);
  60                 if (fread(buf, 1, n, in) != n) {
  61                         if (feof(in)) {
  62                                 ERROR("Unexpected EOF when copying data "
  63                                                 "between files\n");
  64                         } else {
  65                                 ERROR("Error copying data between files: %m\n");
  66                         }
  67                         return WIMLIB_ERR_READ;
  68                 }
  69
  70                 if (fwrite(buf, 1, n, out) != n) {
  71                         ERROR("Error copying data between files: %m\n");
  72                         return WIMLIB_ERR_WRITE;
  73                 }
  74                 len -= n;
  75         }
  76         return 0;
  77 }
  78
  79
  80 /*
  81  * Uncompresses a WIM file resource and writes it uncompressed to a file stream.
  82  *
  83  * @in:             The file stream that contains the file resource.
  84  * @size:           The size of the resource in the input file.
  85  * @original_size:  The original (uncompressed) size of the resource.
  86  * @offset:         The offset of the start of the resource in @in.
  87  * @input_ctype:    The compression type of the resource in @in.
  88  * @out:            The file stream to write the file resource to.
  89  */
  90 static int uncompress_resource(FILE *in, u64 size, u64 original_size,
  91                                off_t offset, int input_ctype, FILE *out)
  92 {
  93         int ret;
  94         u8 buf[WIM_CHUNK_SIZE];
  95         /* Determine how many compressed chunks the file is divided into. */
  96         u64 num_chunks;
  97         u64 i;
  98         u64 uncompressed_offset;
  99         u64 uncompressed_chunk_size;
 100
 101         num_chunks = (original_size + WIM_CHUNK_SIZE - 1) / WIM_CHUNK_SIZE;
 102
 103         for (i = 0; i < num_chunks; i++) {
 104
 105                 uncompressed_offset = i * WIM_CHUNK_SIZE;
 106                 uncompressed_chunk_size = min(WIM_CHUNK_SIZE,
 107                                         original_size - uncompressed_offset);
 108
 109                 ret = read_resource(in, size, original_size, offset, input_ctype,
 110                                         uncompressed_chunk_size,
 111                                         uncompressed_offset, buf);
 112                 if (ret != 0)
 113                         return ret;
 114
 115                 if (fwrite(buf, 1, uncompressed_chunk_size, out) !=
 116                                                 uncompressed_chunk_size) {
 117                         ERROR("Failed to write file resource: %m\n");
 118                         return WIMLIB_ERR_WRITE;
 119                 }
 120         }
 121         return 0;
 122 }
 123
 124 /*
 125  * Transfers a file resource between two files, writing it compressed.  The file
 126  * resource in the input file may be either compressed or uncompressed.
 127  * Alternatively, the input resource may be in-memory, but it must be
 128  * uncompressed.
 129  *
 130  * @in:             The file stream that contains the file resource.  Ignored
 131  *                      if uncompressed_resource != NULL.
 132  * @uncompressed_resource:      If this pointer is not NULL, it points to an
 133  *                                      array of @original_size bytes that are
 134  *                                      the uncompressed input resource.
 135  * @size:           The size of the resource in the input file.
 136  * @original_size:  The original (uncompressed) size of the resource.
 137  * @offset:         The offset of the start of the resource in @in.  Ignored
 138  *                      if uncompressed_resource != NULL.
 139  * @input_ctype:    The compression type of the resource in @in.  Ignored if
 140  *                      uncompressed_resource != NULL.
 141  * @out:            The file stream to write the file resource to.
 142  * @output_type:    The compression type to use when writing the resource to
 143  *                      @out.
 144  * @new_size_ret:   A location into which the new compressed size of the file
 145  *                      resource in returned.
 146  */
 147 static int recompress_resource(FILE *in, const u8 uncompressed_resource[],
 148                                         u64 size, u64 original_size,
 149                                         off_t offset, int input_ctype, FILE *out,
 150                                         int output_ctype, u64 *new_size_ret)
 151 {
 152         int ret;
 153         int (*compress)(const void *, uint, void *, uint *);
 154         if (output_ctype == WIM_COMPRESSION_TYPE_LZX)
 155                 compress = lzx_compress;
 156         else
 157                 compress = xpress_compress;
 158
 159         u8 uncompressed_buf[WIM_CHUNK_SIZE];
 160         u8 compressed_buf[WIM_CHUNK_SIZE - 1];
 161
 162         /* Determine how many compressed chunks the file needs to be divided
 163          * into. */
 164         u64 num_chunks = (original_size + WIM_CHUNK_SIZE - 1) / WIM_CHUNK_SIZE;
 165
 166         u64 num_chunk_entries = num_chunks - 1;
 167
 168         /* Size of the chunk entries--- 8 bytes for files over 4GB, otherwise 4
 169          * bytes */
 170         uint chunk_entry_size = (original_size >= (u64)1 << 32) ?  8 : 4;
 171
 172         /* Array in which to construct the chunk offset table. */
 173         u64 chunk_offsets[num_chunk_entries];
 174
 175         /* Offset of the start of the chunk table in the output file. */
 176         off_t chunk_tab_offset = ftello(out);
 177
 178         /* Total size of the chunk table (as written to the file) */
 179         u64 chunk_tab_size = chunk_entry_size * num_chunk_entries;
 180
 181         /* Reserve space for the chunk table. */
 182         if (fwrite(chunk_offsets, 1, chunk_tab_size, out) != chunk_tab_size) {
 183                 ERROR("Failed to write chunk offset table: %m\n");
 184                 return WIMLIB_ERR_WRITE;
 185         }
 186
 187         /* Read each chunk of the file, compress it, write it to the output
 188          * file, and update th chunk offset table. */
 189         u64 cur_chunk_offset = 0;
 190         for (u64 i = 0; i < num_chunks; i++) {
 191
 192                 u64 uncompressed_offset = i * WIM_CHUNK_SIZE;
 193                 u64 uncompressed_chunk_size = min(WIM_CHUNK_SIZE,
 194                                         original_size - uncompressed_offset);
 195
 196                 const u8 *uncompressed_p;
 197                 if (uncompressed_resource != NULL) {
 198                         uncompressed_p = uncompressed_resource +
 199                                                         uncompressed_offset;
 200
 201                 } else {
 202                         /* Read chunk i of the file into uncompressed_buf. */
 203                         ret = read_resource(in, size, original_size, offset, input_ctype,
 204                                                 uncompressed_chunk_size,
 205                                                 uncompressed_offset,
 206                                                 uncompressed_buf);
 207                         if (ret != 0)
 208                                 return ret;
 209                         uncompressed_p = uncompressed_buf;
 210                 }
 211
 212                 if (i != 0)
 213                         chunk_offsets[i - 1] = cur_chunk_offset;
 214
 215                 uint compressed_len;
 216
 217                 ret = compress(uncompressed_p, uncompressed_chunk_size,
 218                                compressed_buf, &compressed_len);
 219
 220                 /* if compress() returned nonzero, the compressed chunk would
 221                  * have been at least as large as the uncompressed chunk.  In
 222                  * this situation, the WIM format requires that the uncompressed
 223                  * chunk be written instead. */
 224                 const u8 *buf_to_write;
 225                 uint len_to_write;
 226                 if (ret == 0) {
 227                         buf_to_write = compressed_buf;
 228                         len_to_write = compressed_len;
 229                 } else {
 230                         buf_to_write = uncompressed_p;
 231                         len_to_write = uncompressed_chunk_size;
 232                 }
 233
 234                 if (fwrite(buf_to_write, 1, len_to_write, out) != len_to_write) {
 235                         ERROR("Failed to write compressed file resource: %m\n");
 236                         return WIMLIB_ERR_WRITE;
 237                 }
 238                 cur_chunk_offset += len_to_write;
 239         }
 240
 241         /* The chunk offset after the last chunk, plus the size of the chunk
 242          * table, gives the total compressed size of the resource. */
 243         *new_size_ret = cur_chunk_offset + chunk_tab_size;
 244
 245         /* Now that all entries of the chunk table are determined, rewind the
 246          * stream to where the chunk table was, and write it back out. */
 247
 248         if (fseeko(out, chunk_tab_offset, SEEK_SET) != 0) {
 249                 ERROR("Failed to seek to beginning of chunk table: %m\n");
 250                 return WIMLIB_ERR_READ;
 251         }
 252
 253         if (chunk_entry_size == 8) {
 254                 array_to_le64(chunk_offsets, num_chunk_entries);
 255
 256                 if (fwrite(chunk_offsets, 1, chunk_tab_size, out) !=
 257                                 chunk_tab_size) {
 258                         ERROR("Failed to write chunk table: %m\n");
 259                         return WIMLIB_ERR_WRITE;
 260                 }
 261         } else {
 262                 u32 chunk_entries_small[num_chunk_entries];
 263                 for (u64 i = 0; i < num_chunk_entries; i++)
 264                         chunk_entries_small[i] = to_le32(chunk_offsets[i]);
 265                 if (fwrite(chunk_entries_small, 1, chunk_tab_size, out) !=
 266                                 chunk_tab_size) {
 267                         ERROR("Failed to write chunk table: %m\n");
 268                         return WIMLIB_ERR_WRITE;
 269                 }
 270         }
 271
 272         if (fseeko(out, 0, SEEK_END) != 0) {
 273                 ERROR("Failed to seek to end of output file: %m\n");
 274                 return WIMLIB_ERR_WRITE;
 275         }
 276
 277         return 0;
 278 }
 279
 280 int write_resource_from_memory(const u8 resource[], int out_ctype,
 281                                u64 resource_original_size, FILE *out,
 282                                u64 *resource_size_ret)
 283 {
 284         if (out_ctype == WIM_COMPRESSION_TYPE_NONE) {
 285                 if (fwrite(resource, 1, resource_original_size, out) !=
 286                                         resource_original_size) {
 287                         ERROR("Failed to write resource of length "
 288                                         "%"PRIu64": %m\n",
 289                                         resource_original_size);
 290                         return WIMLIB_ERR_WRITE;
 291                 }
 292                 *resource_size_ret = resource_original_size;
 293                 return 0;
 294         } else {
 295                 return recompress_resource(NULL, resource, resource_original_size,
 296                                 resource_original_size, 0, 0, out, out_ctype,
 297                                                         resource_size_ret);
 298         }
 299 }
 300
 301
 302 /*
 303  * Transfers a file resource from a FILE* opened for reading to a FILE* opened
 304  * for writing, possibly changing the compression type.
 305  *
 306  * @in:                 The FILE* that contains the file resource.
 307  * @size:               The (compressed) size of the file resource.
 308  * @original_size:      The uncompressed size of the file resource.
 309  * @offset:             The offset of the file resource in the input file.
 310  * @input_ctype:        The compression type of the file resource in the input
 311  *                              file.
 312  * @out:                The FILE* for the output file.  The file resource is
 313  *                              written at the current position of @out.
 314  * @output_ctype:       The compression type to which the file resource will be
 315  *                              converted.
 316  * @output_res_entry:   A pointer to a resource entry that, upon successful
 317  *                              return of this function,  will have the size,
 318  *                              original size, offset, and flags fields filled
 319  *                              in for the file resource written to the output
 320  *                              file.
 321  */
 322 static int transfer_file_resource(FILE *in, u64 size, u64 original_size,
 323                                   off_t offset, int input_ctype, FILE *out,
 324                                   int output_ctype,
 325                                   struct resource_entry *output_res_entry)
 326 {
 327         int ret;
 328
 329         /* Handle zero-length files */
 330         if (original_size == 0) {
 331                 memset(output_res_entry, 0, sizeof(*output_res_entry));
 332                 return 0;
 333         }
 334
 335         /* Get current offset in the output file. */
 336         output_res_entry->offset = ftello(out);
 337         if (output_res_entry->offset == -1) {
 338                 ERROR("Failed to get output position: %m\n");
 339                 return WIMLIB_ERR_WRITE;
 340         }
 341
 342         if (output_ctype == input_ctype) {
 343                 /* The same compression types; simply copy the resource. */
 344
 345                 ret = copy_between_files(in, offset, out, size);
 346                 if (ret != 0)
 347                         return ret;
 348                 output_res_entry->size = size;
 349         } else {
 350                 /* Different compression types. */
 351
 352                 if (output_ctype == WIM_COMPRESSION_TYPE_NONE) {
 353                         /* Uncompress a compressed file resource */
 354                         ret = uncompress_resource(in, size,
 355                                                 original_size, offset,
 356                                                 input_ctype, out);
 357                         if (ret != 0)
 358                                 return ret;
 359                         output_res_entry->size = original_size;
 360                 } else {
 361                         u64 new_size;
 362                         /* Compress an uncompressed file resource, or compress a
 363                          * compressed file resource using a different
 364                          * compression type (the latter is currently unsupported
 365                          * since only LZX compression is supported. */
 366                         ret = recompress_resource(in, NULL, size, original_size,
 367                                                 offset, input_ctype, out,
 368                                                 output_ctype, &new_size);
 369                         if (ret != 0)
 370                                 return ret;
 371                         output_res_entry->size = new_size;
 372                 }
 373
 374         }
 375
 376         output_res_entry->original_size = original_size;
 377         if (output_ctype == WIM_COMPRESSION_TYPE_NONE)
 378                 output_res_entry->flags = 0;
 379         else
 380                 output_res_entry->flags = WIM_RESHDR_FLAG_COMPRESSED;
 381         return 0;
 382 }
 383
 384 /*
 385  * Writes a file resource to the output file.
 386  *
 387  * @dentry:  The dentry for the file resource.
 388  * @wim_p:  A pointer to the WIMStruct.  The fields of interest to this
 389  *      function are the input and output file streams and the lookup table.
 390  * @return zero on success, nonzero on failure.
 391  */
 392 static int write_file_resource(struct dentry *dentry, void *wim_p)
 393 {
 394         WIMStruct *w;
 395         FILE *out;
 396         FILE *in;
 397         struct lookup_table_entry *lte;
 398         int in_wim_ctype;
 399         int out_wim_ctype;
 400         int input_res_ctype;
 401         struct resource_entry *input_res_entry;
 402         struct resource_entry *output_res_entry;
 403         u64 len;
 404         int ret;
 405
 406         w = wim_p;
 407         out = w->out_fp;
 408
 409         /* Directories don't need file resources. */
 410         if (dentry_is_directory(dentry))
 411                 return 0;
 412
 413         /* Get the lookup entry for the file resource. */
 414         lte = wim_lookup_resource(w, dentry);
 415         if (!lte)
 416                 return 0;
 417
 418         /* No need to write file resources twice.  (This indicates file
 419          * resources that are part of a hard link set.) */
 420         if (++lte->out_refcnt != 1)
 421                 return 0;
 422
 423         out_wim_ctype = wimlib_get_compression_type(w);
 424         output_res_entry = &lte->output_resource_entry;
 425
 426         /* Figure out if we can read the resource from the WIM file, or
 427          * if we have to read it from the filesystem outside. */
 428         if (lte->file_on_disk) {
 429
 430                 /* Read from disk (uncompressed) */
 431
 432                 len = lte->resource_entry.original_size;
 433
 434                 in = fopen(lte->file_on_disk, "rb");
 435                 if (!in) {
 436                         ERROR("Failed to open the file `%s': %m\n",
 437                                         lte->file_on_disk);
 438                         return WIMLIB_ERR_OPEN;
 439                 }
 440
 441                 if (w->verbose)
 442                         puts(lte->file_on_disk);
 443
 444                 ret = transfer_file_resource(in, len, len, 0,
 445                                              WIM_COMPRESSION_TYPE_NONE, out,
 446                                              out_wim_ctype, output_res_entry);
 447                 fclose(in);
 448         } else {
 449
 450                 /* Read from input WIM (possibly compressed) */
 451
 452                 /* It may be a different WIM file, in the case of
 453                  * exporting images from one WIM file to another */
 454                 if (lte->other_wim_fp) {
 455                         /* Different WIM file. */
 456                         in = lte->other_wim_fp;
 457                         in_wim_ctype = lte->other_wim_ctype;
 458                 } else {
 459                         /* Same WIM file. */
 460                         in = w->fp;
 461                         in_wim_ctype = out_wim_ctype;
 462                 }
 463                 input_res_entry = &lte->resource_entry;
 464                 input_res_ctype = resource_compression_type(
 465                                         in_wim_ctype,
 466                                         input_res_entry->flags);
 467
 468                 ret = transfer_file_resource(in,
 469                                         input_res_entry->size,
 470                                         input_res_entry->original_size,
 471                                         input_res_entry->offset,
 472                                         input_res_ctype,
 473                                         out,
 474                                         out_wim_ctype,
 475                                         output_res_entry);
 476         }
 477         return ret;
 478 }
 479
 480 /* Reopens the FILE* for a WIM read-write. */
 481 static int reopen_rw(WIMStruct *w)
 482 {
 483         FILE *fp;
 484
 485         if (fclose(w->fp) != 0)
 486                 ERROR("Failed to close the file `%s': %m\n", w->filename);
 487         fp = fopen(w->filename, "r+b");
 488         if (!fp) {
 489                 ERROR("Failed to open `%s' for reading and writing: "
 490                                 "%m\n", w->filename);
 491                 return WIMLIB_ERR_OPEN;
 492         }
 493         w->fp = fp;
 494         return 0;
 495 }
 496
 497
 498
 499 /*
 500  * Writes a WIM file to the original file that it was read from, overwriting it.
 501  */
 502 WIMLIBAPI int wimlib_overwrite(WIMStruct *w, int flags)
 503 {
 504         const char *wimfile_name;
 505         size_t wim_name_len;
 506         int ret;
 507
 508         wimfile_name = w->filename;
 509
 510         DEBUG("Replacing WIM file `%s'\n", wimfile_name);
 511
 512         if (!wimfile_name)
 513                 return WIMLIB_ERR_NO_FILENAME;
 514
 515         /* Write the WIM to a temporary file. */
 516         /* XXX should the temporary file be somewhere else? */
 517         wim_name_len = strlen(wimfile_name);
 518         char tmpfile[wim_name_len + 10];
 519         memcpy(tmpfile, wimfile_name, wim_name_len);
 520         randomize_char_array_with_alnum(tmpfile + wim_name_len, 9);
 521         tmpfile[wim_name_len + 9] = '\0';
 522
 523         ret = wimlib_write(w, tmpfile, WIM_ALL_IMAGES, flags);
 524         if (ret != 0) {
 525                 ERROR("Failed to write the WIM file `%s'!\n", tmpfile);
 526                 return ret;
 527         }
 528
 529         DEBUG("Closing original WIM file.\n");
 530         /* Close the original WIM file that was opened for reading. */
 531         if (w->fp) {
 532                 if (fclose(w->fp) != 0) {
 533                         DEBUG("WARNING: Failed to close the file `%s'\n",
 534                                         wimfile_name);
 535                 }
 536                 w->fp = NULL;
 537         }
 538
 539         DEBUG("Renaming `%s' to `%s'\n", tmpfile, wimfile_name);
 540
 541         /* Rename the new file to the old file .*/
 542         if (rename(tmpfile, wimfile_name) != 0) {
 543                 ERROR("Failed to rename `%s' to `%s': %m\n", tmpfile,
 544                                                                 wimfile_name);
 545                 /* Remove temporary file. */
 546                 if (unlink(tmpfile) != 0)
 547                         ERROR("Failed to remove `%s': %m\n", tmpfile);
 548                 return WIMLIB_ERR_RENAME;
 549         }
 550
 551         return 0;
 552 }
 553
 554
 555 WIMLIBAPI int wimlib_overwrite_xml_and_header(WIMStruct *w, int flags)
 556 {
 557         int ret;
 558         FILE *fp;
 559         u8 *integrity_table = NULL;
 560         off_t xml_end;
 561         off_t xml_size;
 562         size_t bytes_written;
 563
 564         DEBUG("Overwriting XML and header of `%s', flags = %d\n",
 565                                 w->filename, flags);
 566         if (!w->filename)
 567                 return WIMLIB_ERR_NO_FILENAME;
 568
 569         ret = reopen_rw(w);
 570         if (ret != 0)
 571                 return ret;
 572
 573         fp = w->fp;
 574
 575         /* The old integrity table is still OK, as the SHA1 message digests in
 576          * the integrity table include neither the header nor the XML data.
 577          * Save it for later if it exists and an integrity table was required.
 578          * */
 579         if (flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY &&
 580                         w->hdr.integrity.offset != 0) {
 581                 DEBUG("Reading existing integrity table.\n");
 582                 integrity_table = MALLOC(w->hdr.integrity.size);
 583                 if (!integrity_table)
 584                         return WIMLIB_ERR_NOMEM;
 585
 586                 ret = read_uncompressed_resource(fp, w->hdr.integrity.offset,
 587                                                  w->hdr.integrity.original_size,
 588                                                  integrity_table);
 589                 if (ret != 0)
 590                         goto err;
 591                 DEBUG("Done reading existing integrity table.\n");
 592         }
 593
 594         DEBUG("Overwriting XML data.\n");
 595         /* Overwrite the XML data. */
 596         if (fseeko(fp, w->hdr.xml_res_entry.offset, SEEK_SET) != 0) {
 597                 ERROR("Failed to seek to byte %"PRIu64" for XML data: "
 598                                 "%m\n", w->hdr.xml_res_entry.offset);
 599                 ret = WIMLIB_ERR_WRITE;
 600                 goto err;
 601         }
 602         ret = write_xml_data(w->wim_info, WIM_ALL_IMAGES, fp);
 603         if (ret != 0)
 604                 goto err;
 605
 606         DEBUG("Updating XML resource entry.\n");
 607         /* Update the XML resource entry in the WIM header. */
 608         xml_end = ftello(fp);
 609         if (xml_end == -1) {
 610                 ret = WIMLIB_ERR_WRITE;
 611                 goto err;
 612         }
 613         xml_size = xml_end - w->hdr.xml_res_entry.offset;
 614         w->hdr.xml_res_entry.size = xml_size;
 615         w->hdr.xml_res_entry.original_size = xml_size;
 616
 617         if (flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) {
 618                 DEBUG("Writing integrity table.\n");
 619                 w->hdr.integrity.offset        = xml_end;
 620                 if (integrity_table) {
 621                         /* The existing integrity table was saved. */
 622                         bytes_written = fwrite(integrity_table, 1,
 623                                                w->hdr.integrity.size, fp);
 624                         if (bytes_written != w->hdr.integrity.size) {
 625                                 ERROR("Failed to write integrity table: %m\n");
 626                                 ret = WIMLIB_ERR_WRITE;
 627                                 goto err;
 628                         }
 629                         FREE(integrity_table);
 630                 } else {
 631                         /* There was no existing integrity table, so a new one
 632                          * must be calculated. */
 633                         ret = write_integrity_table(fp, WIM_HEADER_DISK_SIZE,
 634                                         w->hdr.lookup_table_res_entry.offset +
 635                                         w->hdr.lookup_table_res_entry.size,
 636                                         flags & WIMLIB_WRITE_FLAG_SHOW_PROGRESS);
 637                         if (ret != 0)
 638                                 goto err;
 639
 640                         off_t integrity_size           = ftello(fp) - xml_end;
 641                         w->hdr.integrity.size          = integrity_size;
 642                         w->hdr.integrity.original_size = integrity_size;
 643                         w->hdr.integrity.flags         = 0;
 644                 }
 645         } else {
 646                 DEBUG("Truncating file to end of XML data.\n");
 647                 /* No integrity table to write.  The file should be truncated
 648                  * because it's possible that the old file was longer (due to it
 649                  * including an integrity table, or due to its XML data being
 650                  * longer) */
 651                 if (fflush(fp) != 0) {
 652                         ERROR("Failed to flush stream for file `%s': %m\n",
 653                                         w->filename);
 654                         return WIMLIB_ERR_WRITE;
 655                 }
 656                 if (ftruncate(fileno(fp), xml_end) != 0) {
 657                         ERROR("Failed to truncate `%s' to %"PRIu64" "
 658                                         "bytes: %m\n",
 659                                         w->filename, xml_end);
 660                         return WIMLIB_ERR_WRITE;
 661                 }
 662                 memset(&w->hdr.integrity, 0, sizeof(struct resource_entry));
 663         }
 664
 665         DEBUG("Overwriting header.\n");
 666         /* Overwrite the header. */
 667         if (fseeko(fp, 0, SEEK_SET) != 0) {
 668                 ERROR("Failed to seek to beginning of `%s': %m\n",
 669                                 w->filename);
 670                 return WIMLIB_ERR_WRITE;
 671         }
 672
 673         ret = write_header(&w->hdr, fp);
 674         if (ret != 0)
 675                 return ret;;
 676
 677         DEBUG("Closing file.\n");
 678         if (fclose(fp) != 0) {
 679                 ERROR("Failed to close `%s': %m\n", w->filename);
 680                 return WIMLIB_ERR_WRITE;
 681         }
 682         w->fp = NULL;
 683         DEBUG("Done.\n");
 684         return 0;
 685 err:
 686         FREE(integrity_table);
 687         return ret;
 688 }
 689
 690 /* Write the metadata resource for the current image. */
 691 static int write_metadata_resource(WIMStruct *w)
 692 {
 693         FILE *out;
 694         u8 *buf;
 695         u8 *p;
 696         int ret;
 697         off_t subdir_offset;
 698         struct dentry *root;
 699         struct lookup_table_entry *lte;
 700         struct resource_entry *res_entry;
 701         off_t metadata_offset;
 702         u64 metadata_original_size;
 703         u64 metadata_compressed_size;
 704         int metadata_ctype;
 705         u8  hash[WIM_HASH_SIZE];
 706
 707         DEBUG("Writing metadata resource for image %u\n", w->current_image);
 708
 709         out = w->out_fp;
 710         root = wim_root_dentry(w);
 711         metadata_ctype = wimlib_get_compression_type(w);
 712         metadata_offset = ftello(out);
 713         if (metadata_offset == -1)
 714                 return WIMLIB_ERR_WRITE;
 715
 716         subdir_offset = 8 + root->length + 8;
 717         calculate_subdir_offsets(root, &subdir_offset);
 718         metadata_original_size = subdir_offset;
 719         buf = MALLOC(metadata_original_size);
 720         if (!buf) {
 721                 ERROR("Failed to allocate %"PRIu64" bytes for "
 722                                 "metadata resource\n", metadata_original_size);
 723                 return WIMLIB_ERR_NOMEM;
 724         }
 725         p = buf;
 726         #if 0
 727         /* Write the security data. */
 728         p = write_security_data(wim_security_data(w), p);
 729         #else
 730         p = put_u32(p, 8); /* Total length of security data. */
 731         p = put_u32(p, 0); /* Number of security data entries. */
 732         #endif
 733
 734         DEBUG("Writing dentry tree.\n");
 735         p = write_dentry_tree(root, p);
 736
 737         /* Like file resources, the lookup table entry for a metadata resource
 738          * uses for the hash code a SHA1 message digest of its uncompressed
 739          * contents. */
 740         sha1_buffer(buf, metadata_original_size, hash);
 741
 742         ret = write_resource_from_memory(buf,
 743                                          metadata_ctype,
 744                                          metadata_original_size,
 745                                          out,
 746                                          &metadata_compressed_size);
 747         FREE(buf);
 748         if (ret != 0)
 749                 return ret;
 750
 751         /* Update the lookup table entry, including the hash and output resource
 752          * entry fields, for this image's metadata resource.  */
 753         lte = wim_metadata_lookup_table_entry(w);
 754         res_entry = &lte->output_resource_entry;
 755         lte->out_refcnt++;
 756         if (memcmp(hash, lte->hash, WIM_HASH_SIZE) != 0) {
 757                 lookup_table_unlink(w->lookup_table, lte);
 758                 memcpy(lte->hash, hash, WIM_HASH_SIZE);
 759                 lookup_table_insert(w->lookup_table, lte);
 760         }
 761         res_entry->original_size = metadata_original_size;
 762         res_entry->offset        = metadata_offset;
 763         res_entry->size          = metadata_compressed_size;
 764         res_entry->flags         = WIM_RESHDR_FLAG_METADATA;
 765         if (metadata_ctype != WIM_COMPRESSION_TYPE_NONE)
 766                 res_entry->flags |= WIM_RESHDR_FLAG_COMPRESSED;
 767         return 0;
 768 }
 769
 770 /* Write the file resources for the current image. */
 771 static int write_file_resources(WIMStruct *w)
 772 {
 773
 774         DEBUG("Writing file resources for image %u\n", w->current_image);
 775         return for_dentry_in_tree(wim_root_dentry(w), write_file_resource, w);
 776 }
 777
 778 /* Write lookup table, xml data, lookup table, and rewrite header */
 779 static int finish_write(WIMStruct *w, int image, FILE *out, int flags)
 780 {
 781         off_t lookup_table_offset;
 782         off_t xml_data_offset;
 783         off_t lookup_table_size;
 784         off_t integrity_offset;
 785         off_t xml_data_size;
 786         off_t end_offset;
 787         off_t integrity_size;
 788         int ret;
 789         int i;
 790         struct wim_header hdr;
 791
 792         lookup_table_offset = ftello(out);
 793         if (lookup_table_offset == -1)
 794                 return WIMLIB_ERR_WRITE;
 795
 796         DEBUG("Writing lookup table.\n");
 797         /* Write the lookup table. */
 798         ret = write_lookup_table(w->lookup_table, out);
 799         if (ret != 0)
 800                 return ret;
 801
 802         DEBUG("Writing XML data.\n");
 803
 804         xml_data_offset = ftello(out);
 805         if (xml_data_offset == -1)
 806                 return WIMLIB_ERR_WRITE;
 807
 808         /* @hdr will be the header for the new WIM.  First copy all the data
 809          * from the header in the WIMStruct; then set all the fields that may
 810          * have changed, including the resource entries, boot index, and image
 811          * count.  */
 812         memcpy(&hdr, &w->hdr, sizeof(struct wim_header));
 813         lookup_table_size = xml_data_offset - lookup_table_offset;
 814         hdr.lookup_table_res_entry.offset        = lookup_table_offset;
 815         hdr.lookup_table_res_entry.size          = lookup_table_size;
 816         hdr.lookup_table_res_entry.original_size = lookup_table_size;
 817         hdr.lookup_table_res_entry.flags         = WIM_RESHDR_FLAG_METADATA;
 818
 819         ret = write_xml_data(w->wim_info, image, out);
 820         if (ret != 0)
 821                 return ret;
 822
 823         integrity_offset = ftello(out);
 824         if (integrity_offset == -1)
 825                 return WIMLIB_ERR_WRITE;
 826         xml_data_size = integrity_offset - xml_data_offset;
 827
 828         hdr.xml_res_entry.offset                 = xml_data_offset;
 829         hdr.xml_res_entry.size                   = xml_data_size;
 830         hdr.xml_res_entry.original_size          = xml_data_size;
 831         hdr.xml_res_entry.flags                  = 0;
 832
 833         if (flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) {
 834                 ret = write_integrity_table(out, WIM_HEADER_DISK_SIZE,
 835                                             xml_data_offset,
 836                                             flags & WIMLIB_WRITE_FLAG_SHOW_PROGRESS);
 837                 if (ret != 0)
 838                         return ret;
 839                 end_offset = ftello(out);
 840                 if (end_offset == -1)
 841                         return WIMLIB_ERR_WRITE;
 842                 integrity_size = end_offset - integrity_offset;
 843                 hdr.integrity.offset = integrity_offset;
 844                 hdr.integrity.size   = integrity_size;
 845                 hdr.integrity.original_size = integrity_size;
 846         } else {
 847                 hdr.integrity.offset        = 0;
 848                 hdr.integrity.size          = 0;
 849                 hdr.integrity.original_size = 0;
 850         }
 851         hdr.integrity.flags = 0;
 852
 853         DEBUG("Updating WIM header.\n");
 854
 855
 856         /*
 857          * In the WIM header, there is room for the resource entry for a
 858          * metadata resource labeled as the "boot metadata".  This entry should
 859          * be zeroed out if there is no bootable image (boot_idx 0).  Otherwise,
 860          * it should be a copy of the resource entry for the image that is
 861          * marked as bootable.  This is not well documented...
 862          */
 863         if (hdr.boot_idx == 0 || !w->image_metadata
 864                         || (image != WIM_ALL_IMAGES && image != hdr.boot_idx)) {
 865                 memset(&hdr.boot_metadata_res_entry, 0,
 866                        sizeof(struct resource_entry));
 867         } else {
 868                 memcpy(&hdr.boot_metadata_res_entry,
 869                        &w->image_metadata[hdr.boot_idx - 1].lookup_table_entry->
 870                                         output_resource_entry,
 871                                         sizeof(struct resource_entry));
 872         }
 873
 874         /* Set image count and boot index correctly for single image writes */
 875         if (image != WIM_ALL_IMAGES) {
 876                 hdr.image_count = 1;
 877                 if (hdr.boot_idx == image)
 878                         hdr.boot_idx = 1;
 879                 else
 880                         hdr.boot_idx = 0;
 881         }
 882
 883
 884         if (fseeko(out, 0, SEEK_SET) != 0)
 885                 return WIMLIB_ERR_WRITE;
 886
 887         return write_header(&hdr, out);
 888 }
 889
 890 /* Writes the WIM to a file.  */
 891 WIMLIBAPI int wimlib_write(WIMStruct *w, const char *path, int image, int flags)
 892 {
 893         int ret;
 894         const char *mode;
 895         FILE *out;
 896
 897         if (image != WIM_ALL_IMAGES &&
 898                         (image < 1 || image > w->hdr.image_count))
 899                 return WIMLIB_ERR_INVALID_IMAGE;
 900
 901         if (image == WIM_ALL_IMAGES)
 902                 DEBUG("Writing all images to `%s'\n", path);
 903         else
 904                 DEBUG("Writing image %d to `%s'\n", image, path);
 905
 906         /* checking the integrity requires going back over the file to read it.
 907          * XXX
 908          * (It also would be possible to keep a running sha1sum as the file
 909          * as written-- this would be faster, but a bit more complicated) */
 910         if (flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY)
 911                 mode = "w+b";
 912         else
 913                 mode = "wb";
 914
 915         out = fopen(path, mode);
 916         if (!out) {
 917                 ERROR("Failed to open the file `%s' for writing!\n",
 918                                 path);
 919                 return WIMLIB_ERR_OPEN;
 920         }
 921
 922         w->out_fp = out;
 923
 924         /* Write dummy header. It will be overwritten later. */
 925         ret = write_header(&w->hdr, out);
 926         if (ret != 0)
 927                 goto done;
 928
 929         for_lookup_table_entry(w->lookup_table, zero_out_refcnts, NULL);
 930
 931         ret = for_image(w, image, write_file_resources);
 932         if (ret != 0) {
 933                 ERROR("Failed to write file resources!\n");
 934                 goto done;
 935         }
 936
 937         ret = for_image(w, image, write_metadata_resource);
 938
 939         if (ret != 0) {
 940                 ERROR("Failed to write image metadata!\n");
 941                 goto done;
 942         }
 943
 944         ret = finish_write(w, image, out, flags);
 945
 946 done:
 947         DEBUG("Closing output file.\n");
 948         w->out_fp = NULL;
 949         if (fclose(out) != 0) {
 950                 ERROR("Failed to close the file `%s': %m\n", path);
 951                 ret = WIMLIB_ERR_WRITE;
 952         }
 953         return ret;
 954 }