Close out_fp in finish_write()
[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 <unistd.h>
34
35
36 /* Reopens the FILE* for a WIM read-write. */
37 static int reopen_rw(WIMStruct *w)
38 {
39         FILE *fp;
40
41         if (fclose(w->fp) != 0)
42                 ERROR_WITH_ERRNO("Failed to close the file `%s'", w->filename);
43         w->fp = NULL;
44         fp = fopen(w->filename, "r+b");
45         if (!fp) {
46                 ERROR_WITH_ERRNO("Failed to open `%s' for reading and writing",
47                                  w->filename);
48                 return WIMLIB_ERR_OPEN;
49         }
50         w->fp = fp;
51         return 0;
52 }
53
54
55
56 /*
57  * Writes a WIM file to the original file that it was read from, overwriting it.
58  */
59 WIMLIBAPI int wimlib_overwrite(WIMStruct *w, int write_flags)
60 {
61         const char *wimfile_name;
62         size_t wim_name_len;
63         int ret;
64
65         if (!w)
66                 return WIMLIB_ERR_INVALID_PARAM;
67
68         write_flags &= ~WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE;
69
70         wimfile_name = w->filename;
71
72         DEBUG("Replacing WIM file `%s'.", wimfile_name);
73
74         if (!wimfile_name)
75                 return WIMLIB_ERR_NO_FILENAME;
76
77         /* Write the WIM to a temporary file. */
78         /* XXX should the temporary file be somewhere else? */
79         wim_name_len = strlen(wimfile_name);
80         char tmpfile[wim_name_len + 10];
81         memcpy(tmpfile, wimfile_name, wim_name_len);
82         randomize_char_array_with_alnum(tmpfile + wim_name_len, 9);
83         tmpfile[wim_name_len + 9] = '\0';
84
85         ret = wimlib_write(w, tmpfile, WIM_ALL_IMAGES, write_flags);
86         if (ret != 0) {
87                 ERROR("Failed to write the WIM file `%s'", tmpfile);
88                 if (unlink(tmpfile) != 0)
89                         WARNING("Failed to remove `%s'", tmpfile);
90                 return ret;
91         }
92
93         DEBUG("Closing original WIM file.");
94         /* Close the original WIM file that was opened for reading. */
95         if (w->fp) {
96                 if (fclose(w->fp) != 0) {
97                         WARNING("Failed to close the file `%s'", wimfile_name);
98                 }
99                 w->fp = NULL;
100         }
101
102         DEBUG("Renaming `%s' to `%s'", tmpfile, wimfile_name);
103
104         /* Rename the new file to the old file .*/
105         if (rename(tmpfile, wimfile_name) != 0) {
106                 ERROR_WITH_ERRNO("Failed to rename `%s' to `%s'",
107                                  tmpfile, wimfile_name);
108                 /* Remove temporary file. */
109                 if (unlink(tmpfile) != 0)
110                         ERROR_WITH_ERRNO("Failed to remove `%s'", tmpfile);
111                 return WIMLIB_ERR_RENAME;
112         }
113
114         return 0;
115 }
116
117 static int check_resource_offset(struct lookup_table_entry *lte, void *arg)
118 {
119         u64 xml_data_offset = *(u64*)arg;
120         if (lte->resource_entry.offset > xml_data_offset) {
121                 ERROR("The following resource is *after* the XML data:");
122                 print_lookup_table_entry(lte);
123                 return WIMLIB_ERR_RESOURCE_ORDER;
124         }
125         return 0;
126 }
127
128 WIMLIBAPI int wimlib_overwrite_xml_and_header(WIMStruct *w, int write_flags)
129 {
130         int ret;
131         FILE *fp;
132         u8 *integrity_table = NULL;
133         off_t xml_end;
134         off_t xml_size;
135         size_t bytes_written;
136
137         DEBUG("Overwriting XML and header of `%s', write_flags = %#x",
138               w->filename, write_flags);
139
140         if (!w->filename)
141                 return WIMLIB_ERR_NO_FILENAME;
142
143         write_flags &= ~WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE;
144
145         /* Make sure that the integrity table (if present) is after the XML
146          * data, and that there are no stream resources, metadata resources, or
147          * lookup tables after the XML data.  Otherwise, these data would be
148          * destroyed by this function. */
149         if (w->hdr.integrity.offset != 0 &&
150             w->hdr.integrity.offset < w->hdr.xml_res_entry.offset) {
151                 ERROR("Didn't expect the integrity table to be before the XML data");
152                 return WIMLIB_ERR_RESOURCE_ORDER;
153         }
154
155         if (w->hdr.lookup_table_res_entry.offset >
156             w->hdr.xml_res_entry.offset) {
157                 ERROR("Didn't expect the lookup table to be after the XML data");
158                 return WIMLIB_ERR_RESOURCE_ORDER;
159         }
160
161         ret = for_lookup_table_entry(w->lookup_table, check_resource_offset,
162                                      &w->hdr.xml_res_entry.offset);
163         if (ret != 0)
164                 return ret;
165
166         ret = reopen_rw(w);
167         if (ret != 0)
168                 return ret;
169
170         fp = w->fp;
171
172         /* The old integrity table is still OK, as the SHA1 message digests in
173          * the integrity table include neither the header nor the XML data.
174          * Save it for later if it exists and an integrity table was required.
175          * */
176         if ((write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY)
177              && w->hdr.integrity.offset != 0)
178         {
179                 DEBUG("Reading existing integrity table.");
180                 integrity_table = MALLOC(w->hdr.integrity.size);
181                 if (!integrity_table)
182                         return WIMLIB_ERR_NOMEM;
183
184                 ret = read_uncompressed_resource(fp, w->hdr.integrity.offset,
185                                                  w->hdr.integrity.original_size,
186                                                  integrity_table);
187                 if (ret != 0)
188                         goto err;
189                 DEBUG("Done reading existing integrity table.");
190         }
191
192         DEBUG("Overwriting XML data.");
193         /* Overwrite the XML data. */
194         if (fseeko(fp, w->hdr.xml_res_entry.offset, SEEK_SET) != 0) {
195                 ERROR_WITH_ERRNO("Failed to seek to byte %"PRIu64" "
196                                  "for XML data", w->hdr.xml_res_entry.offset);
197                 ret = WIMLIB_ERR_WRITE;
198                 goto err;
199         }
200         ret = write_xml_data(w->wim_info, WIM_ALL_IMAGES, fp, 0);
201         if (ret != 0)
202                 goto err;
203
204         DEBUG("Updating XML resource entry.");
205         /* Update the XML resource entry in the WIM header. */
206         xml_end = ftello(fp);
207         if (xml_end == -1) {
208                 ret = WIMLIB_ERR_WRITE;
209                 goto err;
210         }
211         xml_size = xml_end - w->hdr.xml_res_entry.offset;
212         w->hdr.xml_res_entry.size = xml_size;
213         w->hdr.xml_res_entry.original_size = xml_size;
214         /* XML data offset is unchanged. */
215
216         if (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) {
217                 DEBUG("Writing integrity table.");
218                 w->hdr.integrity.offset = xml_end;
219                 if (integrity_table) {
220                         /* The existing integrity table was saved. */
221                         bytes_written = fwrite(integrity_table, 1,
222                                                w->hdr.integrity.size, fp);
223                         if (bytes_written != w->hdr.integrity.size) {
224                                 ERROR_WITH_ERRNO("Failed to write integrity "
225                                                  "table");
226                                 ret = WIMLIB_ERR_WRITE;
227                                 goto err;
228                         }
229                         FREE(integrity_table);
230                 } else {
231                         /* There was no existing integrity table, so a new one
232                          * must be calculated. */
233                         ret = write_integrity_table(fp, WIM_HEADER_DISK_SIZE,
234                                         w->hdr.lookup_table_res_entry.offset +
235                                         w->hdr.lookup_table_res_entry.size,
236                                         write_flags & WIMLIB_WRITE_FLAG_SHOW_PROGRESS);
237                         if (ret != 0)
238                                 return ret;
239
240                         off_t end_integrity = ftello(fp);
241                         if (end_integrity == -1)
242                                 return WIMLIB_ERR_WRITE;
243
244                         off_t integrity_size           = end_integrity - xml_end;
245                         w->hdr.integrity.size          = integrity_size;
246                         w->hdr.integrity.original_size = integrity_size;
247                         w->hdr.integrity.flags         = 0;
248                 }
249         } else {
250                 DEBUG("Truncating file to end of XML data.");
251                 /* No integrity table to write.  The file should be truncated
252                  * because it's possible that the old file was longer (due to it
253                  * including an integrity table, or due to its XML data being
254                  * longer) */
255                 if (fflush(fp) != 0) {
256                         ERROR_WITH_ERRNO("Failed to flush stream for file `%s'",
257                                          w->filename);
258                         return WIMLIB_ERR_WRITE;
259                 }
260                 if (ftruncate(fileno(fp), xml_end) != 0) {
261                         ERROR_WITH_ERRNO("Failed to truncate `%s' to %"PRIu64" "
262                                          "bytes", w->filename, xml_end);
263                         return WIMLIB_ERR_WRITE;
264                 }
265                 memset(&w->hdr.integrity, 0, sizeof(struct resource_entry));
266         }
267
268         DEBUG("Overwriting header.");
269         /* Overwrite the header. */
270         if (fseeko(fp, 0, SEEK_SET) != 0) {
271                 ERROR_WITH_ERRNO("Failed to seek to beginning of `%s'",
272                                  w->filename);
273                 return WIMLIB_ERR_WRITE;
274         }
275
276         ret = write_header(&w->hdr, fp);
277         if (ret != 0)
278                 return ret;
279
280         DEBUG("Closing `%s'.", w->filename);
281         if (fclose(fp) != 0) {
282                 ERROR_WITH_ERRNO("Failed to close `%s'", w->filename);
283                 return WIMLIB_ERR_WRITE;
284         }
285         w->fp = NULL;
286         DEBUG("Done.");
287         return 0;
288 err:
289         FREE(integrity_table);
290         return ret;
291 }
292
293
294 /* Write the file resources for the current image. */
295 static int write_file_resources(WIMStruct *w)
296 {
297         DEBUG("Writing file resources for image %u.", w->current_image);
298         return for_dentry_in_tree(wim_root_dentry(w), write_dentry_resources, w);
299 }
300
301 /*
302  * Write the lookup table, xml data, and integrity table, then overwrite the WIM
303  * header.
304  */
305 int finish_write(WIMStruct *w, int image, int write_flags)
306 {
307         off_t lookup_table_offset;
308         off_t xml_data_offset;
309         off_t lookup_table_size;
310         off_t integrity_offset;
311         off_t xml_data_size;
312         off_t end_offset;
313         off_t integrity_size;
314         int ret;
315         struct wim_header hdr;
316         FILE *out = w->out_fp;
317
318         if (!(write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE)) {
319                 /* Write the lookup table. */
320                 lookup_table_offset = ftello(out);
321                 if (lookup_table_offset == -1)
322                         return WIMLIB_ERR_WRITE;
323
324                 DEBUG("Writing lookup table (offset %"PRIu64")",
325                       lookup_table_offset);
326                 ret = write_lookup_table(w->lookup_table, out);
327                 if (ret != 0)
328                         return ret;
329         }
330
331         xml_data_offset = ftello(out);
332         if (xml_data_offset == -1)
333                 return WIMLIB_ERR_WRITE;
334
335         /* @hdr will be the header for the new WIM.  First copy all the data
336          * from the header in the WIMStruct; then set all the fields that may
337          * have changed, including the resource entries, boot index, and image
338          * count.  */
339         memcpy(&hdr, &w->hdr, sizeof(struct wim_header));
340         if (!(write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE)) {
341                 lookup_table_size = xml_data_offset - lookup_table_offset;
342                 hdr.lookup_table_res_entry.offset = lookup_table_offset;
343                 hdr.lookup_table_res_entry.size = lookup_table_size;
344         }
345         hdr.lookup_table_res_entry.original_size = hdr.lookup_table_res_entry.size;
346         hdr.lookup_table_res_entry.flags = WIM_RESHDR_FLAG_METADATA;
347
348         DEBUG("Writing XML data (offset %"PRIu64")", xml_data_offset);
349         ret = write_xml_data(w->wim_info, image, out,
350                              (write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE) ?
351                                 wim_info_get_total_bytes(w->wim_info) : 0);
352         if (ret != 0)
353                 return ret;
354
355         integrity_offset = ftello(out);
356         if (integrity_offset == -1)
357                 return WIMLIB_ERR_WRITE;
358         xml_data_size = integrity_offset - xml_data_offset;
359
360         hdr.xml_res_entry.offset                 = xml_data_offset;
361         hdr.xml_res_entry.size                   = xml_data_size;
362         hdr.xml_res_entry.original_size          = xml_data_size;
363         hdr.xml_res_entry.flags                  = 0;
364
365         if (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) {
366                 ret = write_integrity_table(out, WIM_HEADER_DISK_SIZE,
367                                             xml_data_offset,
368                                             write_flags & WIMLIB_WRITE_FLAG_SHOW_PROGRESS);
369                 if (ret != 0)
370                         return ret;
371                 end_offset = ftello(out);
372                 if (end_offset == -1)
373                         return WIMLIB_ERR_WRITE;
374                 integrity_size              = end_offset - integrity_offset;
375                 hdr.integrity.offset        = integrity_offset;
376                 hdr.integrity.size          = integrity_size;
377                 hdr.integrity.original_size = integrity_size;
378         } else {
379                 hdr.integrity.offset        = 0;
380                 hdr.integrity.size          = 0;
381                 hdr.integrity.original_size = 0;
382         }
383         hdr.integrity.flags = 0;
384
385         DEBUG("Updating WIM header.");
386
387         /*
388          * In the WIM header, there is room for the resource entry for a
389          * metadata resource labeled as the "boot metadata".  This entry should
390          * be zeroed out if there is no bootable image (boot_idx 0).  Otherwise,
391          * it should be a copy of the resource entry for the image that is
392          * marked as bootable.  This is not well documented...
393          */
394         if (hdr.boot_idx == 0 || !w->image_metadata
395                         || (image != WIM_ALL_IMAGES && image != hdr.boot_idx)) {
396                 memset(&hdr.boot_metadata_res_entry, 0,
397                        sizeof(struct resource_entry));
398         } else {
399                 memcpy(&hdr.boot_metadata_res_entry,
400                        &w->image_metadata[
401                           hdr.boot_idx - 1].metadata_lte->output_resource_entry,
402                        sizeof(struct resource_entry));
403         }
404
405         /* Set image count and boot index correctly for single image writes */
406         if (image != WIM_ALL_IMAGES) {
407                 hdr.image_count = 1;
408                 if (hdr.boot_idx == image)
409                         hdr.boot_idx = 1;
410                 else
411                         hdr.boot_idx = 0;
412         }
413
414
415         if (fseeko(out, 0, SEEK_SET) != 0)
416                 return WIMLIB_ERR_WRITE;
417
418         ret = write_header(&hdr, out);
419         if (ret != 0)
420                 return ret;
421
422         DEBUG("Closing output file.");
423         wimlib_assert(w->out_fp != NULL);
424         if (fclose(w->out_fp) != 0) {
425                 ERROR_WITH_ERRNO("Failed to close the WIM file");
426                 ret = WIMLIB_ERR_WRITE;
427         }
428         w->out_fp = NULL;
429         return ret;
430 }
431
432 /* Open file stream and write dummy header for WIM. */
433 int begin_write(WIMStruct *w, const char *path, int write_flags)
434 {
435         const char *mode;
436         DEBUG("Opening `%s' for new WIM", path);
437
438         /* checking the integrity requires going back over the file to read it.
439          * XXX
440          * (It also would be possible to keep a running sha1sum as the file is
441          * written-- this would be faster, but a bit more complicated) */
442         if (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY)
443                 mode = "w+b";
444         else
445                 mode = "wb";
446
447         if (w->out_fp)
448                 fclose(w->out_fp);
449
450         w->out_fp = fopen(path, mode);
451         if (!w->out_fp) {
452                 ERROR_WITH_ERRNO("Failed to open the file `%s' for writing",
453                                  path);
454                 return WIMLIB_ERR_OPEN;
455         }
456
457         /* Write dummy header. It will be overwritten later. */
458         return write_header(&w->hdr, w->out_fp);
459 }
460
461 /* Writes a stand-alone WIM to a file.  */
462 WIMLIBAPI int wimlib_write(WIMStruct *w, const char *path,
463                            int image, int write_flags)
464 {
465         int ret;
466
467         if (!w || !path)
468                 return WIMLIB_ERR_INVALID_PARAM;
469
470         write_flags &= ~WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE;
471
472         if (image != WIM_ALL_IMAGES &&
473              (image < 1 || image > w->hdr.image_count))
474                 return WIMLIB_ERR_INVALID_IMAGE;
475
476
477         if (w->hdr.total_parts != 1) {
478                 ERROR("Cannot call wimlib_write() on part of a split WIM");
479                 return WIMLIB_ERR_SPLIT_UNSUPPORTED;
480         }
481
482         if (image == WIM_ALL_IMAGES)
483                 DEBUG("Writing all images to `%s'.", path);
484         else
485                 DEBUG("Writing image %d to `%s'.", image, path);
486
487         ret = begin_write(w, path, write_flags);
488         if (ret != 0)
489                 return ret;
490
491         for_lookup_table_entry(w->lookup_table, lte_zero_out_refcnt, NULL);
492
493         w->write_flags = write_flags;
494
495         ret = for_image(w, image, write_file_resources);
496         if (ret != 0) {
497                 ERROR("Failed to write WIM file resources to `%s'", path);
498                 return ret;
499         }
500
501         ret = for_image(w, image, write_metadata_resource);
502
503         if (ret != 0) {
504                 /*ERROR("Failed to write WIM image metadata to `%s'", path);*/
505                 return ret;
506         }
507
508         return finish_write(w, image, write_flags);
509 }