1/*
2 * Copyright (C) 2010 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#define _FILE_OFFSET_BITS 64
18#define _LARGEFILE64_SOURCE 1
19
20#include <fcntl.h>
21#include <inttypes.h>
22#include <limits.h>
23#include <stdbool.h>
24#include <stddef.h>
25#include <stdlib.h>
26#include <string.h>
27#include <sys/stat.h>
28#include <sys/types.h>
29#include <unistd.h>
30#include <zlib.h>
31
32#include "defs.h"
33#include "output_file.h"
34#include "sparse_crc32.h"
35#include "sparse_format.h"
36
37#ifndef USE_MINGW
38#include <sys/mman.h>
39#define O_BINARY 0
40#else
41#define ftruncate64 ftruncate
42#endif
43
44#if defined(__APPLE__) && defined(__MACH__)
45#define lseek64 lseek
46#define ftruncate64 ftruncate
47#define mmap64 mmap
48#define off64_t off_t
49#endif
50
51#define min(a, b) \
52	({ typeof(a) _a = (a); typeof(b) _b = (b); (_a < _b) ? _a : _b; })
53
54#define SPARSE_HEADER_MAJOR_VER 1
55#define SPARSE_HEADER_MINOR_VER 0
56#define SPARSE_HEADER_LEN       (sizeof(sparse_header_t))
57#define CHUNK_HEADER_LEN (sizeof(chunk_header_t))
58
59#define container_of(inner, outer_t, elem) \
60	((outer_t *)((char *)inner - offsetof(outer_t, elem)))
61
62struct output_file_ops {
63	int (*open)(struct output_file *, int fd);
64	int (*skip)(struct output_file *, int64_t);
65	int (*pad)(struct output_file *, int64_t);
66	int (*write)(struct output_file *, void *, int);
67	void (*close)(struct output_file *);
68};
69
70struct sparse_file_ops {
71	int (*write_data_chunk)(struct output_file *out, unsigned int len,
72			void *data);
73	int (*write_fill_chunk)(struct output_file *out, unsigned int len,
74			uint32_t fill_val);
75	int (*write_skip_chunk)(struct output_file *out, int64_t len);
76	int (*write_end_chunk)(struct output_file *out);
77};
78
79struct output_file {
80	int64_t cur_out_ptr;
81	unsigned int chunk_cnt;
82	uint32_t crc32;
83	struct output_file_ops *ops;
84	struct sparse_file_ops *sparse_ops;
85	int use_crc;
86	unsigned int block_size;
87	int64_t len;
88	char *zero_buf;
89	uint32_t *fill_buf;
90	char *buf;
91};
92
93struct output_file_gz {
94	struct output_file out;
95	gzFile gz_fd;
96};
97
98#define to_output_file_gz(_o) \
99	container_of((_o), struct output_file_gz, out)
100
101struct output_file_normal {
102	struct output_file out;
103	int fd;
104};
105
106#define to_output_file_normal(_o) \
107	container_of((_o), struct output_file_normal, out)
108
109struct output_file_callback {
110	struct output_file out;
111	void *priv;
112	int (*write)(void *priv, const void *buf, int len);
113};
114
115#define to_output_file_callback(_o) \
116	container_of((_o), struct output_file_callback, out)
117
118static int file_open(struct output_file *out, int fd)
119{
120	struct output_file_normal *outn = to_output_file_normal(out);
121
122	outn->fd = fd;
123	return 0;
124}
125
126static int file_skip(struct output_file *out, int64_t cnt)
127{
128	off64_t ret;
129	struct output_file_normal *outn = to_output_file_normal(out);
130
131	ret = lseek64(outn->fd, cnt, SEEK_CUR);
132	if (ret < 0) {
133		error_errno("lseek64");
134		return -1;
135	}
136	return 0;
137}
138
139static int file_pad(struct output_file *out, int64_t len)
140{
141	int ret;
142	struct output_file_normal *outn = to_output_file_normal(out);
143
144	ret = ftruncate64(outn->fd, len);
145	if (ret < 0) {
146		return -errno;
147	}
148
149	return 0;
150}
151
152static int file_write(struct output_file *out, void *data, int len)
153{
154	int ret;
155	struct output_file_normal *outn = to_output_file_normal(out);
156
157	ret = write(outn->fd, data, len);
158	if (ret < 0) {
159		error_errno("write");
160		return -1;
161	} else if (ret < len) {
162		error("incomplete write");
163		return -1;
164	}
165
166	return 0;
167}
168
169static void file_close(struct output_file *out)
170{
171	struct output_file_normal *outn = to_output_file_normal(out);
172
173	free(outn);
174}
175
176static struct output_file_ops file_ops = {
177	.open = file_open,
178	.skip = file_skip,
179	.pad = file_pad,
180	.write = file_write,
181	.close = file_close,
182};
183
184static int gz_file_open(struct output_file *out, int fd)
185{
186	struct output_file_gz *outgz = to_output_file_gz(out);
187
188	outgz->gz_fd = gzdopen(fd, "wb9");
189	if (!outgz->gz_fd) {
190		error_errno("gzopen");
191		return -errno;
192	}
193
194	return 0;
195}
196
197
198static int gz_file_skip(struct output_file *out, int64_t cnt)
199{
200	off64_t ret;
201	struct output_file_gz *outgz = to_output_file_gz(out);
202
203	ret = gzseek(outgz->gz_fd, cnt, SEEK_CUR);
204	if (ret < 0) {
205		error_errno("gzseek");
206		return -1;
207	}
208	return 0;
209}
210
211static int gz_file_pad(struct output_file *out, int64_t len)
212{
213	off64_t ret;
214	struct output_file_gz *outgz = to_output_file_gz(out);
215
216	ret = gztell(outgz->gz_fd);
217	if (ret < 0) {
218		return -1;
219	}
220
221	if (ret >= len) {
222		return 0;
223	}
224
225	ret = gzseek(outgz->gz_fd, len - 1, SEEK_SET);
226	if (ret < 0) {
227		return -1;
228	}
229
230	gzwrite(outgz->gz_fd, "", 1);
231
232	return 0;
233}
234
235static int gz_file_write(struct output_file *out, void *data, int len)
236{
237	int ret;
238	struct output_file_gz *outgz = to_output_file_gz(out);
239
240	ret = gzwrite(outgz->gz_fd, data, len);
241	if (ret < 0) {
242		error_errno("gzwrite");
243		return -1;
244	} else if (ret < len) {
245		error("incomplete gzwrite");
246		return -1;
247	}
248
249	return 0;
250}
251
252static void gz_file_close(struct output_file *out)
253{
254	struct output_file_gz *outgz = to_output_file_gz(out);
255
256	gzclose(outgz->gz_fd);
257	free(outgz);
258}
259
260static struct output_file_ops gz_file_ops = {
261	.open = gz_file_open,
262	.skip = gz_file_skip,
263	.pad = gz_file_pad,
264	.write = gz_file_write,
265	.close = gz_file_close,
266};
267
268static int callback_file_open(struct output_file *out __unused, int fd __unused)
269{
270	return 0;
271}
272
273static int callback_file_skip(struct output_file *out, int64_t off)
274{
275	struct output_file_callback *outc = to_output_file_callback(out);
276	int to_write;
277	int ret;
278
279	while (off > 0) {
280		to_write = min(off, (int64_t)INT_MAX);
281		ret = outc->write(outc->priv, NULL, to_write);
282		if (ret < 0) {
283			return ret;
284		}
285		off -= to_write;
286	}
287
288	return 0;
289}
290
291static int callback_file_pad(struct output_file *out __unused, int64_t len __unused)
292{
293	return -1;
294}
295
296static int callback_file_write(struct output_file *out, void *data, int len)
297{
298	struct output_file_callback *outc = to_output_file_callback(out);
299
300	return outc->write(outc->priv, data, len);
301}
302
303static void callback_file_close(struct output_file *out)
304{
305	struct output_file_callback *outc = to_output_file_callback(out);
306
307	free(outc);
308}
309
310static struct output_file_ops callback_file_ops = {
311	.open = callback_file_open,
312	.skip = callback_file_skip,
313	.pad = callback_file_pad,
314	.write = callback_file_write,
315	.close = callback_file_close,
316};
317
318int read_all(int fd, void *buf, size_t len)
319{
320	size_t total = 0;
321	int ret;
322	char *ptr = buf;
323
324	while (total < len) {
325		ret = read(fd, ptr, len - total);
326
327		if (ret < 0)
328			return -errno;
329
330		if (ret == 0)
331			return -EINVAL;
332
333		ptr += ret;
334		total += ret;
335	}
336
337	return 0;
338}
339
340static int write_sparse_skip_chunk(struct output_file *out, int64_t skip_len)
341{
342	chunk_header_t chunk_header;
343	int ret;
344
345	if (skip_len % out->block_size) {
346		error("don't care size %"PRIi64" is not a multiple of the block size %u",
347				skip_len, out->block_size);
348		return -1;
349	}
350
351	/* We are skipping data, so emit a don't care chunk. */
352	chunk_header.chunk_type = CHUNK_TYPE_DONT_CARE;
353	chunk_header.reserved1 = 0;
354	chunk_header.chunk_sz = skip_len / out->block_size;
355	chunk_header.total_sz = CHUNK_HEADER_LEN;
356	ret = out->ops->write(out, &chunk_header, sizeof(chunk_header));
357	if (ret < 0)
358		return -1;
359
360	out->cur_out_ptr += skip_len;
361	out->chunk_cnt++;
362
363	return 0;
364}
365
366static int write_sparse_fill_chunk(struct output_file *out, unsigned int len,
367		uint32_t fill_val)
368{
369	chunk_header_t chunk_header;
370	int rnd_up_len, count;
371	int ret;
372
373	/* Round up the fill length to a multiple of the block size */
374	rnd_up_len = ALIGN(len, out->block_size);
375
376	/* Finally we can safely emit a chunk of data */
377	chunk_header.chunk_type = CHUNK_TYPE_FILL;
378	chunk_header.reserved1 = 0;
379	chunk_header.chunk_sz = rnd_up_len / out->block_size;
380	chunk_header.total_sz = CHUNK_HEADER_LEN + sizeof(fill_val);
381	ret = out->ops->write(out, &chunk_header, sizeof(chunk_header));
382
383	if (ret < 0)
384		return -1;
385	ret = out->ops->write(out, &fill_val, sizeof(fill_val));
386	if (ret < 0)
387		return -1;
388
389	if (out->use_crc) {
390		count = out->block_size / sizeof(uint32_t);
391		while (count--)
392			out->crc32 = sparse_crc32(out->crc32, &fill_val, sizeof(uint32_t));
393	}
394
395	out->cur_out_ptr += rnd_up_len;
396	out->chunk_cnt++;
397
398	return 0;
399}
400
401static int write_sparse_data_chunk(struct output_file *out, unsigned int len,
402		void *data)
403{
404	chunk_header_t chunk_header;
405	int rnd_up_len, zero_len;
406	int ret;
407
408	/* Round up the data length to a multiple of the block size */
409	rnd_up_len = ALIGN(len, out->block_size);
410	zero_len = rnd_up_len - len;
411
412	/* Finally we can safely emit a chunk of data */
413	chunk_header.chunk_type = CHUNK_TYPE_RAW;
414	chunk_header.reserved1 = 0;
415	chunk_header.chunk_sz = rnd_up_len / out->block_size;
416	chunk_header.total_sz = CHUNK_HEADER_LEN + rnd_up_len;
417	ret = out->ops->write(out, &chunk_header, sizeof(chunk_header));
418
419	if (ret < 0)
420		return -1;
421	ret = out->ops->write(out, data, len);
422	if (ret < 0)
423		return -1;
424	if (zero_len) {
425		ret = out->ops->write(out, out->zero_buf, zero_len);
426		if (ret < 0)
427			return -1;
428	}
429
430	if (out->use_crc) {
431		out->crc32 = sparse_crc32(out->crc32, data, len);
432		if (zero_len)
433			out->crc32 = sparse_crc32(out->crc32, out->zero_buf, zero_len);
434	}
435
436	out->cur_out_ptr += rnd_up_len;
437	out->chunk_cnt++;
438
439	return 0;
440}
441
442int write_sparse_end_chunk(struct output_file *out)
443{
444	chunk_header_t chunk_header;
445	int ret;
446
447	if (out->use_crc) {
448		chunk_header.chunk_type = CHUNK_TYPE_CRC32;
449		chunk_header.reserved1 = 0;
450		chunk_header.chunk_sz = 0;
451		chunk_header.total_sz = CHUNK_HEADER_LEN + 4;
452
453		ret = out->ops->write(out, &chunk_header, sizeof(chunk_header));
454		if (ret < 0) {
455			return ret;
456		}
457		out->ops->write(out, &out->crc32, 4);
458		if (ret < 0) {
459			return ret;
460		}
461
462		out->chunk_cnt++;
463	}
464
465	return 0;
466}
467
468static struct sparse_file_ops sparse_file_ops = {
469		.write_data_chunk = write_sparse_data_chunk,
470		.write_fill_chunk = write_sparse_fill_chunk,
471		.write_skip_chunk = write_sparse_skip_chunk,
472		.write_end_chunk = write_sparse_end_chunk,
473};
474
475static int write_normal_data_chunk(struct output_file *out, unsigned int len,
476		void *data)
477{
478	int ret;
479	unsigned int rnd_up_len = ALIGN(len, out->block_size);
480
481	ret = out->ops->write(out, data, len);
482	if (ret < 0) {
483		return ret;
484	}
485
486	if (rnd_up_len > len) {
487		ret = out->ops->skip(out, rnd_up_len - len);
488	}
489
490	return ret;
491}
492
493static int write_normal_fill_chunk(struct output_file *out, unsigned int len,
494		uint32_t fill_val)
495{
496	int ret;
497	unsigned int i;
498	unsigned int write_len;
499
500	/* Initialize fill_buf with the fill_val */
501	for (i = 0; i < out->block_size / sizeof(uint32_t); i++) {
502		out->fill_buf[i] = fill_val;
503	}
504
505	while (len) {
506		write_len = min(len, out->block_size);
507		ret = out->ops->write(out, out->fill_buf, write_len);
508		if (ret < 0) {
509			return ret;
510		}
511
512		len -= write_len;
513	}
514
515	return 0;
516}
517
518static int write_normal_skip_chunk(struct output_file *out, int64_t len)
519{
520	return out->ops->skip(out, len);
521}
522
523int write_normal_end_chunk(struct output_file *out)
524{
525	return out->ops->pad(out, out->len);
526}
527
528static struct sparse_file_ops normal_file_ops = {
529		.write_data_chunk = write_normal_data_chunk,
530		.write_fill_chunk = write_normal_fill_chunk,
531		.write_skip_chunk = write_normal_skip_chunk,
532		.write_end_chunk = write_normal_end_chunk,
533};
534
535void output_file_close(struct output_file *out)
536{
537	out->sparse_ops->write_end_chunk(out);
538	out->ops->close(out);
539}
540
541static int output_file_init(struct output_file *out, int block_size,
542		int64_t len, bool sparse, int chunks, bool crc)
543{
544	int ret;
545
546	out->len = len;
547	out->block_size = block_size;
548	out->cur_out_ptr = 0ll;
549	out->chunk_cnt = 0;
550	out->crc32 = 0;
551	out->use_crc = crc;
552
553	out->zero_buf = calloc(block_size, 1);
554	if (!out->zero_buf) {
555		error_errno("malloc zero_buf");
556		return -ENOMEM;
557	}
558
559	out->fill_buf = calloc(block_size, 1);
560	if (!out->fill_buf) {
561		error_errno("malloc fill_buf");
562		ret = -ENOMEM;
563		goto err_fill_buf;
564	}
565
566	if (sparse) {
567		out->sparse_ops = &sparse_file_ops;
568	} else {
569		out->sparse_ops = &normal_file_ops;
570	}
571
572	if (sparse) {
573		sparse_header_t sparse_header = {
574				.magic = SPARSE_HEADER_MAGIC,
575				.major_version = SPARSE_HEADER_MAJOR_VER,
576				.minor_version = SPARSE_HEADER_MINOR_VER,
577				.file_hdr_sz = SPARSE_HEADER_LEN,
578				.chunk_hdr_sz = CHUNK_HEADER_LEN,
579				.blk_sz = out->block_size,
580				.total_blks = out->len / out->block_size,
581				.total_chunks = chunks,
582				.image_checksum = 0
583		};
584
585		if (out->use_crc) {
586			sparse_header.total_chunks++;
587		}
588
589		ret = out->ops->write(out, &sparse_header, sizeof(sparse_header));
590		if (ret < 0) {
591			goto err_write;
592		}
593	}
594
595	return 0;
596
597err_write:
598	free(out->fill_buf);
599err_fill_buf:
600	free(out->zero_buf);
601	return ret;
602}
603
604static struct output_file *output_file_new_gz(void)
605{
606	struct output_file_gz *outgz = calloc(1, sizeof(struct output_file_gz));
607	if (!outgz) {
608		error_errno("malloc struct outgz");
609		return NULL;
610	}
611
612	outgz->out.ops = &gz_file_ops;
613
614	return &outgz->out;
615}
616
617static struct output_file *output_file_new_normal(void)
618{
619	struct output_file_normal *outn = calloc(1, sizeof(struct output_file_normal));
620	if (!outn) {
621		error_errno("malloc struct outn");
622		return NULL;
623	}
624
625	outn->out.ops = &file_ops;
626
627	return &outn->out;
628}
629
630struct output_file *output_file_open_callback(int (*write)(void *, const void *, int),
631		void *priv, unsigned int block_size, int64_t len,
632		int gz __unused, int sparse, int chunks, int crc)
633{
634	int ret;
635	struct output_file_callback *outc;
636
637	outc = calloc(1, sizeof(struct output_file_callback));
638	if (!outc) {
639		error_errno("malloc struct outc");
640		return NULL;
641	}
642
643	outc->out.ops = &callback_file_ops;
644	outc->priv = priv;
645	outc->write = write;
646
647	ret = output_file_init(&outc->out, block_size, len, sparse, chunks, crc);
648	if (ret < 0) {
649		free(outc);
650		return NULL;
651	}
652
653	return &outc->out;
654}
655
656struct output_file *output_file_open_fd(int fd, unsigned int block_size, int64_t len,
657		int gz, int sparse, int chunks, int crc)
658{
659	int ret;
660	struct output_file *out;
661
662	if (gz) {
663		out = output_file_new_gz();
664	} else {
665		out = output_file_new_normal();
666	}
667	if (!out) {
668		return NULL;
669	}
670
671	out->ops->open(out, fd);
672
673	ret = output_file_init(out, block_size, len, sparse, chunks, crc);
674	if (ret < 0) {
675		free(out);
676		return NULL;
677	}
678
679	return out;
680}
681
682/* Write a contiguous region of data blocks from a memory buffer */
683int write_data_chunk(struct output_file *out, unsigned int len, void *data)
684{
685	return out->sparse_ops->write_data_chunk(out, len, data);
686}
687
688/* Write a contiguous region of data blocks with a fill value */
689int write_fill_chunk(struct output_file *out, unsigned int len,
690		uint32_t fill_val)
691{
692	return out->sparse_ops->write_fill_chunk(out, len, fill_val);
693}
694
695int write_fd_chunk(struct output_file *out, unsigned int len,
696		int fd, int64_t offset)
697{
698	int ret;
699	int64_t aligned_offset;
700	int aligned_diff;
701	int buffer_size;
702	char *ptr;
703
704	aligned_offset = offset & ~(4096 - 1);
705	aligned_diff = offset - aligned_offset;
706	buffer_size = len + aligned_diff;
707
708#ifndef USE_MINGW
709	char *data = mmap64(NULL, buffer_size, PROT_READ, MAP_SHARED, fd,
710			aligned_offset);
711	if (data == MAP_FAILED) {
712		return -errno;
713	}
714	ptr = data + aligned_diff;
715#else
716	off64_t pos;
717	char *data = malloc(len);
718	if (!data) {
719		return -errno;
720	}
721	pos = lseek64(fd, offset, SEEK_SET);
722	if (pos < 0) {
723                free(data);
724		return -errno;
725	}
726	ret = read_all(fd, data, len);
727	if (ret < 0) {
728                free(data);
729		return ret;
730	}
731	ptr = data;
732#endif
733
734	ret = out->sparse_ops->write_data_chunk(out, len, ptr);
735
736#ifndef USE_MINGW
737	munmap(data, buffer_size);
738#else
739	free(data);
740#endif
741
742	return ret;
743}
744
745/* Write a contiguous region of data blocks from a file */
746int write_file_chunk(struct output_file *out, unsigned int len,
747		const char *file, int64_t offset)
748{
749	int ret;
750
751	int file_fd = open(file, O_RDONLY | O_BINARY);
752	if (file_fd < 0) {
753		return -errno;
754	}
755
756	ret = write_fd_chunk(out, len, file_fd, offset);
757
758	close(file_fd);
759
760	return ret;
761}
762
763int write_skip_chunk(struct output_file *out, int64_t len)
764{
765	return out->sparse_ops->write_skip_chunk(out, len);
766}
767