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