d1317cf9d44f0ce82b810e65b1f53f28852d43c7
[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 part of wimlib, a library for working with WIM files.
11  *
12  * wimlib is free software; you can redistribute it and/or modify it under the
13  * terms of the GNU General Public License as published by the Free
14  * Software Foundation; either version 3 of the License, or (at your option)
15  * any later version.
16  *
17  * wimlib 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
19  * A PARTICULAR PURPOSE. See the GNU General Public License for more
20  * details.
21  *
22  * You should have received a copy of the GNU General Public License
23  * along with wimlib; if not, see http://www.gnu.org/licenses/.
24  */
25
26 #ifdef HAVE_CONFIG_H
27 #  include "config.h"
28 #endif
29
30 #include "wimlib.h"
31 #include "wimlib/error.h"
32 #include "wimlib/list.h"
33 #include "wimlib/lookup_table.h"
34 #include "wimlib/metadata.h"
35 #include "wimlib/resource.h"
36 #include "wimlib/wim.h"
37 #include "wimlib/write.h"
38
39 #ifdef HAVE_ALLOCA_H
40 #  include <alloca.h>
41 #else
42 #  include <stdlib.h>
43 #endif
44
45 struct swm_part_info {
46         struct list_head stream_list;
47         u64 size;
48 };
49
50 static void
51 copy_part_info(struct swm_part_info *dst, struct swm_part_info *src)
52 {
53         list_transfer(&src->stream_list, &dst->stream_list);
54         dst->size = src->size;
55 }
56
57 struct swm_info {
58         struct swm_part_info *parts;
59         unsigned num_parts;
60         unsigned num_alloc_parts;
61         u64 total_bytes;
62         u64 max_part_size;
63 };
64
65 static int
66 write_split_wim(WIMStruct *orig_wim, const tchar *swm_name,
67                 struct swm_info *swm_info, int write_flags,
68                 wimlib_progress_func_t progress_func)
69 {
70         size_t swm_name_len;
71         tchar *swm_name_buf;
72         const tchar *dot;
73         tchar *swm_suffix;
74         size_t swm_base_name_len;
75
76         union wimlib_progress_info progress;
77         unsigned part_number;
78         int ret;
79         u8 guid[WIMLIB_GUID_LEN];
80
81         swm_name_len = tstrlen(swm_name);
82         swm_name_buf = alloca((swm_name_len + 20) * sizeof(tchar));
83         tstrcpy(swm_name_buf, swm_name);
84         dot = tstrchr(swm_name_buf, T('.'));
85         if (dot) {
86                 swm_base_name_len = dot - swm_name_buf;
87                 swm_suffix = alloca((tstrlen(dot) + 1) * sizeof(tchar));
88                 tstrcpy(swm_suffix, dot);
89         } else {
90                 swm_base_name_len = swm_name_len;
91                 swm_suffix = alloca(1 * sizeof(tchar));
92                 swm_suffix[0] = T('\0');
93         }
94
95         progress.split.completed_bytes = 0;
96         progress.split.total_bytes = 0;
97         for (part_number = 1; part_number <= swm_info->num_parts; part_number++)
98                 progress.split.total_bytes += swm_info->parts[part_number - 1].size;
99         progress.split.total_parts = swm_info->num_parts;
100         progress.split.part_name = swm_name_buf;
101
102         randomize_byte_array(guid, WIMLIB_GUID_LEN);
103
104         for (part_number = 1; part_number <= swm_info->num_parts; part_number++) {
105                 int part_write_flags;
106
107                 if (part_number != 1) {
108                         tsprintf(swm_name_buf + swm_base_name_len,
109                                  T("%u%"TS), part_number, swm_suffix);
110                 }
111
112                 progress.split.cur_part_number = part_number;
113                 if (progress_func) {
114                         progress_func(WIMLIB_PROGRESS_MSG_SPLIT_BEGIN_PART,
115                                       &progress);
116                 }
117
118                 part_write_flags = write_flags;
119                 part_write_flags |= WIMLIB_WRITE_FLAG_USE_EXISTING_TOTALBYTES;
120                 if (part_number != 1)
121                         part_write_flags |= WIMLIB_WRITE_FLAG_NO_METADATA;
122
123                 ret = write_wim_part(orig_wim,
124                                      swm_name_buf,
125                                      WIMLIB_ALL_IMAGES,
126                                      part_write_flags,
127                                      1,
128                                      NULL,
129                                      part_number,
130                                      swm_info->num_parts,
131                                      &swm_info->parts[part_number - 1].stream_list,
132                                      guid);
133                 if (ret)
134                         return ret;
135
136                 progress.split.completed_bytes += swm_info->parts[part_number - 1].size;
137                 if (progress_func) {
138                         progress_func(WIMLIB_PROGRESS_MSG_SPLIT_END_PART,
139                                       &progress);
140                 }
141         }
142         return 0;
143 }
144
145 static int
146 add_stream_to_swm(struct wim_lookup_table_entry *lte, void *_swm_info)
147 {
148         struct swm_info *swm_info = _swm_info;
149         u64 stream_size;
150
151         if (lte_is_partial(lte)) {
152                 ERROR("Splitting of WIM containing packed streams is not supported.\n"
153                       "        Export it in the default format first.");
154                 return WIMLIB_ERR_UNSUPPORTED;
155         }
156         if (lte->resource_location == RESOURCE_IN_WIM)
157                 stream_size = lte->rspec->size_in_wim;
158         else
159                 stream_size = lte->size;
160
161         /* - Start first part if no parts have been started so far;
162          * - Start next part if adding this stream exceeds maximum part size,
163          *   UNLESS the stream is metadata or if no streams at all have been
164          *   added to the current part.
165          */
166         if (swm_info->num_parts == 0 ||
167             ((swm_info->parts[swm_info->num_parts - 1].size +
168                         stream_size >= swm_info->max_part_size)
169              && !((lte->flags & WIM_RESHDR_FLAG_METADATA) ||
170                    swm_info->parts[swm_info->num_parts - 1].size == 0)))
171         {
172                 if (swm_info->num_parts == swm_info->num_alloc_parts) {
173                         struct swm_part_info *parts;
174                         size_t num_alloc_parts = swm_info->num_alloc_parts;
175
176                         num_alloc_parts += 8;
177                         parts = MALLOC(num_alloc_parts * sizeof(parts[0]));
178                         if (!parts)
179                                 return WIMLIB_ERR_NOMEM;
180
181                         for (unsigned i = 0; i < swm_info->num_parts; i++)
182                                 copy_part_info(&parts[i], &swm_info->parts[i]);
183
184                         FREE(swm_info->parts);
185                         swm_info->parts = parts;
186                         swm_info->num_alloc_parts = num_alloc_parts;
187                 }
188                 swm_info->num_parts++;
189                 INIT_LIST_HEAD(&swm_info->parts[swm_info->num_parts - 1].stream_list);
190                 swm_info->parts[swm_info->num_parts - 1].size = 0;
191         }
192         swm_info->parts[swm_info->num_parts - 1].size += stream_size;
193         if (!(lte->flags & WIM_RESHDR_FLAG_METADATA)) {
194                 list_add_tail(&lte->write_streams_list,
195                               &swm_info->parts[swm_info->num_parts - 1].stream_list);
196         }
197         swm_info->total_bytes += stream_size;
198         return 0;
199 }
200
201 /* API function documented in wimlib.h  */
202 WIMLIBAPI int
203 wimlib_split(WIMStruct *wim, const tchar *swm_name,
204              u64 part_size, int write_flags,
205              wimlib_progress_func_t progress_func)
206 {
207         struct swm_info swm_info;
208         unsigned i;
209         int ret;
210
211         if (swm_name == NULL || swm_name[0] == T('\0') || part_size == 0)
212                 return WIMLIB_ERR_INVALID_PARAM;
213
214         if (write_flags & ~WIMLIB_WRITE_MASK_PUBLIC)
215                 return WIMLIB_ERR_INVALID_PARAM;
216
217         if (!wim_has_metadata(wim))
218                 return WIMLIB_ERR_METADATA_NOT_FOUND;
219
220         memset(&swm_info, 0, sizeof(swm_info));
221         swm_info.max_part_size = part_size;
222
223         for (i = 0; i < wim->hdr.image_count; i++) {
224                 ret = add_stream_to_swm(wim->image_metadata[i]->metadata_lte,
225                                         &swm_info);
226                 if (ret)
227                         goto out_free_swm_info;
228         }
229
230         ret = for_lookup_table_entry_pos_sorted(wim->lookup_table,
231                                                 add_stream_to_swm,
232                                                 &swm_info);
233         if (ret)
234                 goto out_free_swm_info;
235
236         ret = write_split_wim(wim, swm_name, &swm_info, write_flags,
237                               progress_func);
238 out_free_swm_info:
239         FREE(swm_info.parts);
240         return ret;
241 }