wimlib_split(): fix bug where no parts would be written if WIM contained no images
[wimlib] / src / split.c
1 /*
2  * split.c
3  *
4  * Split a WIM file into parts.
5  */
6
7 /*
8  * Copyright (C) 2012, 2013 Eric Biggers
9  *
10  * This file is free software; you can redistribute it and/or modify it under
11  * the terms of the GNU Lesser General Public License as published by the Free
12  * Software Foundation; either version 3 of the License, or (at your option) any
13  * later version.
14  *
15  * This file is distributed in the hope that it will be useful, but WITHOUT
16  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
17  * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
18  * details.
19  *
20  * You should have received a copy of the GNU Lesser General Public License
21  * along with this file; if not, see http://www.gnu.org/licenses/.
22  */
23
24 #ifdef HAVE_CONFIG_H
25 #  include "config.h"
26 #endif
27
28 #include "wimlib.h"
29 #include "wimlib/alloca.h"
30 #include "wimlib/blob_table.h"
31 #include "wimlib/error.h"
32 #include "wimlib/list.h"
33 #include "wimlib/metadata.h"
34 #include "wimlib/progress.h"
35 #include "wimlib/resource.h"
36 #include "wimlib/wim.h"
37 #include "wimlib/write.h"
38
39 struct swm_part_info {
40         struct list_head blob_list;
41         u64 size;
42 };
43
44 static void
45 copy_part_info(struct swm_part_info *dst, struct swm_part_info *src)
46 {
47         list_replace(&src->blob_list, &dst->blob_list);
48         dst->size = src->size;
49 }
50
51 struct swm_info {
52         struct swm_part_info *parts;
53         unsigned num_parts;
54         unsigned num_alloc_parts;
55         u64 total_bytes;
56         u64 max_part_size;
57 };
58
59 static int
60 write_split_wim(WIMStruct *orig_wim, const tchar *swm_name,
61                 struct swm_info *swm_info, int write_flags)
62 {
63         size_t swm_name_len;
64         tchar *swm_name_buf;
65         const tchar *dot;
66         tchar *swm_suffix;
67         size_t swm_base_name_len;
68
69         union wimlib_progress_info progress;
70         unsigned part_number;
71         int ret;
72         u8 guid[GUID_SIZE];
73
74         swm_name_len = tstrlen(swm_name);
75         swm_name_buf = alloca((swm_name_len + 20) * sizeof(tchar));
76         tstrcpy(swm_name_buf, swm_name);
77         dot = tstrchr(swm_name_buf, T('.'));
78         if (dot) {
79                 swm_base_name_len = dot - swm_name_buf;
80                 swm_suffix = alloca((tstrlen(dot) + 1) * sizeof(tchar));
81                 tstrcpy(swm_suffix, dot);
82         } else {
83                 swm_base_name_len = swm_name_len;
84                 swm_suffix = alloca(1 * sizeof(tchar));
85                 swm_suffix[0] = T('\0');
86         }
87
88         progress.split.completed_bytes = 0;
89         progress.split.total_bytes = 0;
90         for (part_number = 1; part_number <= swm_info->num_parts; part_number++)
91                 progress.split.total_bytes += swm_info->parts[part_number - 1].size;
92         progress.split.total_parts = swm_info->num_parts;
93
94         generate_guid(guid);
95
96         for (part_number = 1; part_number <= swm_info->num_parts; part_number++) {
97                 int part_write_flags;
98                 wimlib_progress_func_t progfunc;
99
100                 if (part_number != 1) {
101                         tsprintf(swm_name_buf + swm_base_name_len,
102                                  T("%u%"TS), part_number, swm_suffix);
103                 }
104
105                 progress.split.cur_part_number = part_number;
106                 progress.split.part_name = swm_name_buf;
107
108                 ret = call_progress(orig_wim->progfunc,
109                                     WIMLIB_PROGRESS_MSG_SPLIT_BEGIN_PART,
110                                     &progress,
111                                     orig_wim->progctx);
112                 if (ret)
113                         return ret;
114
115                 part_write_flags = write_flags;
116                 part_write_flags |= WIMLIB_WRITE_FLAG_USE_EXISTING_TOTALBYTES;
117                 if (part_number != 1)
118                         part_write_flags |= WIMLIB_WRITE_FLAG_NO_METADATA;
119
120                 progfunc = orig_wim->progfunc;
121                 orig_wim->progfunc = NULL;
122                 ret = write_wim_part(orig_wim,
123                                      progress.split.part_name,
124                                      WIMLIB_ALL_IMAGES,
125                                      part_write_flags,
126                                      1,
127                                      part_number,
128                                      swm_info->num_parts,
129                                      &swm_info->parts[part_number - 1].blob_list,
130                                      guid);
131                 orig_wim->progfunc = progfunc;
132                 if (ret)
133                         return ret;
134
135                 progress.split.completed_bytes += swm_info->parts[part_number - 1].size;
136
137                 ret = call_progress(orig_wim->progfunc,
138                                     WIMLIB_PROGRESS_MSG_SPLIT_END_PART,
139                                     &progress,
140                                     orig_wim->progctx);
141                 if (ret)
142                         return ret;
143         }
144         return 0;
145 }
146
147 static int
148 start_new_swm_part(struct swm_info *swm_info)
149 {
150         if (swm_info->num_parts == swm_info->num_alloc_parts) {
151                 struct swm_part_info *parts;
152                 size_t num_alloc_parts = swm_info->num_alloc_parts;
153
154                 num_alloc_parts += 8;
155                 parts = MALLOC(num_alloc_parts * sizeof(parts[0]));
156                 if (!parts)
157                         return WIMLIB_ERR_NOMEM;
158
159                 for (unsigned i = 0; i < swm_info->num_parts; i++)
160                         copy_part_info(&parts[i], &swm_info->parts[i]);
161
162                 FREE(swm_info->parts);
163                 swm_info->parts = parts;
164                 swm_info->num_alloc_parts = num_alloc_parts;
165         }
166         swm_info->num_parts++;
167         INIT_LIST_HEAD(&swm_info->parts[swm_info->num_parts - 1].blob_list);
168         swm_info->parts[swm_info->num_parts - 1].size = 0;
169         return 0;
170 }
171
172 static int
173 add_blob_to_swm(struct blob_descriptor *blob, void *_swm_info)
174 {
175         struct swm_info *swm_info = _swm_info;
176         u64 blob_stored_size;
177         int ret;
178
179         if (blob->blob_location == BLOB_IN_WIM)
180                 blob_stored_size = blob->rdesc->size_in_wim;
181         else
182                 blob_stored_size = blob->size;
183
184         /* Start the next part if adding this blob exceeds the maximum part
185          * size, UNLESS the blob is metadata or if no blobs at all have been
186          * added to the current part.  */
187         if ((swm_info->parts[swm_info->num_parts - 1].size +
188              blob_stored_size >= swm_info->max_part_size)
189             && !(blob->is_metadata ||
190                  swm_info->parts[swm_info->num_parts - 1].size == 0))
191         {
192                 ret = start_new_swm_part(swm_info);
193                 if (ret)
194                         return ret;
195         }
196         swm_info->parts[swm_info->num_parts - 1].size += blob_stored_size;
197         if (!blob->is_metadata) {
198                 list_add_tail(&blob->write_blobs_list,
199                               &swm_info->parts[swm_info->num_parts - 1].blob_list);
200         }
201         swm_info->total_bytes += blob_stored_size;
202         return 0;
203 }
204
205 /* API function documented in wimlib.h  */
206 WIMLIBAPI int
207 wimlib_split(WIMStruct *wim, const tchar *swm_name,
208              u64 part_size, int write_flags)
209 {
210         struct swm_info swm_info;
211         unsigned i;
212         int ret;
213
214         if (swm_name == NULL || swm_name[0] == T('\0') || part_size == 0)
215                 return WIMLIB_ERR_INVALID_PARAM;
216
217         if (write_flags & ~WIMLIB_WRITE_MASK_PUBLIC)
218                 return WIMLIB_ERR_INVALID_PARAM;
219
220         if (!wim_has_metadata(wim))
221                 return WIMLIB_ERR_METADATA_NOT_FOUND;
222
223         if (wim_has_solid_resources(wim)) {
224                 ERROR("Splitting of WIM containing solid resources is not supported.\n"
225                       "        Export it in non-solid format first.");
226                 return WIMLIB_ERR_UNSUPPORTED;
227         }
228
229         for (i = 0; i < wim->hdr.image_count; i++) {
230                 if (!is_image_unchanged_from_wim(wim->image_metadata[i], wim)) {
231                         ERROR("Only an unmodified, on-disk WIM file can be split.");
232                         return WIMLIB_ERR_UNSUPPORTED;
233                 }
234         }
235
236         memset(&swm_info, 0, sizeof(swm_info));
237         swm_info.max_part_size = part_size;
238
239         ret = start_new_swm_part(&swm_info);
240         if (ret)
241                 goto out_free_swm_info;
242
243         for (i = 0; i < wim->hdr.image_count; i++) {
244                 ret = add_blob_to_swm(wim->image_metadata[i]->metadata_blob,
245                                       &swm_info);
246                 if (ret)
247                         goto out_free_swm_info;
248         }
249
250         ret = for_blob_in_table_sorted_by_sequential_order(wim->blob_table,
251                                                            add_blob_to_swm,
252                                                            &swm_info);
253         if (ret)
254                 goto out_free_swm_info;
255
256         ret = write_split_wim(wim, swm_name, &swm_info, write_flags);
257 out_free_swm_info:
258         FREE(swm_info.parts);
259         return ret;
260 }