xref: /freebsd/usr.bin/mkimg/image.c (revision 0b3105a37d7adcadcb720112fed4dc4e8040be99) 
1 /*-
2  * Copyright (c) 2014 Juniper Networks, Inc.
3  * All rights reserved.
4  *
5  * Redistribution and use in source and binary forms, with or without
6  * modification, are permitted provided that the following conditions
7  * are met:
8  * 1. Redistributions of source code must retain the above copyright
9  *    notice, this list of conditions and the following disclaimer.
10  * 2. Redistributions in binary form must reproduce the above copyright
11  *    notice, this list of conditions and the following disclaimer in the
12  *    documentation and/or other materials provided with the distribution.
13  *
14  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
15  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
16  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
17  * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
18  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
19  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
20  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
21  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
22  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
23  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
24  * SUCH DAMAGE.
25  */
26 
27 #include <sys/cdefs.h>
28 __FBSDID("$FreeBSD$");
29 
30 #include <sys/mman.h>
31 #include <sys/queue.h>
32 #include <sys/stat.h>
33 #include <sys/types.h>
34 #include <assert.h>
35 #include <errno.h>
36 #include <limits.h>
37 #include <paths.h>
38 #include <stdint.h>
39 #include <stdio.h>
40 #include <stdlib.h>
41 #include <string.h>
42 #include <unistd.h>
43 
44 #include "image.h"
45 #include "mkimg.h"
46 
47 struct chunk {
48 	STAILQ_ENTRY(chunk) ch_list;
49 	size_t	ch_size;		/* Size of chunk in bytes. */
50 	lba_t	ch_block;		/* Block address in image. */
51 	union {
52 		struct {
53 			off_t	ofs;	/* Offset in backing file. */
54 			int	fd;	/* FD of backing file. */
55 		} file;
56 		struct {
57 			void	*ptr;	/* Pointer to data in memory */
58 		} mem;
59 	} ch_u;
60 	u_int	ch_type;
61 #define	CH_TYPE_ZEROES		0	/* Chunk is a gap (no data). */
62 #define	CH_TYPE_FILE		1	/* File-backed chunk. */
63 #define	CH_TYPE_MEMORY		2	/* Memory-backed chunk */
64 };
65 
66 static STAILQ_HEAD(chunk_head, chunk) image_chunks;
67 static u_int image_nchunks;
68 
69 static char image_swap_file[PATH_MAX];
70 static int image_swap_fd = -1;
71 static u_int image_swap_pgsz;
72 static off_t image_swap_size;
73 
74 static lba_t image_size;
75 
76 static int
77 is_empty_sector(void *buf)
78 {
79 	uint64_t *p = buf;
80 	size_t n, max;
81 
82 	assert(((uintptr_t)p & 3) == 0);
83 
84 	max = secsz / sizeof(uint64_t);
85 	for (n = 0; n < max; n++) {
86 		if (p[n] != 0UL)
87 			return (0);
88 	}
89 	return (1);
90 }
91 
92 /*
93  * Swap file handlng.
94  */
95 
96 static off_t
97 image_swap_alloc(size_t size)
98 {
99 	off_t ofs;
100 	size_t unit;
101 
102 	unit = (secsz > image_swap_pgsz) ? secsz : image_swap_pgsz;
103 	assert((unit & (unit - 1)) == 0);
104 
105 	size = (size + unit - 1) & ~(unit - 1);
106 
107 	ofs = image_swap_size;
108 	image_swap_size += size;
109 	if (ftruncate(image_swap_fd, image_swap_size) == -1) {
110 		image_swap_size = ofs;
111 		ofs = -1LL;
112 	}
113 	return (ofs);
114 }
115 
116 /*
117  * Image chunk handling.
118  */
119 
120 static struct chunk *
121 image_chunk_find(lba_t blk)
122 {
123 	static struct chunk *last = NULL;
124 	struct chunk *ch;
125 
126 	ch = (last != NULL && last->ch_block <= blk)
127 	    ? last : STAILQ_FIRST(&image_chunks);
128 	while (ch != NULL) {
129 		if (ch->ch_block <= blk &&
130 		    (lba_t)(ch->ch_block + (ch->ch_size / secsz)) > blk) {
131 			last = ch;
132 			break;
133 		}
134 		ch = STAILQ_NEXT(ch, ch_list);
135 	}
136 	return (ch);
137 }
138 
139 static size_t
140 image_chunk_grow(struct chunk *ch, size_t sz)
141 {
142 	size_t dsz, newsz;
143 
144 	newsz = ch->ch_size + sz;
145 	if (newsz > ch->ch_size) {
146 		ch->ch_size = newsz;
147 		return (0);
148 	}
149 	/* We would overflow -- create new chunk for remainder. */
150 	dsz = SIZE_MAX - ch->ch_size;
151 	assert(dsz < sz);
152 	ch->ch_size = SIZE_MAX;
153 	return (sz - dsz);
154 }
155 
156 static struct chunk *
157 image_chunk_memory(struct chunk *ch, lba_t blk)
158 {
159 	struct chunk *new;
160 	void *ptr;
161 
162 	ptr = calloc(1, secsz);
163 	if (ptr == NULL)
164 		return (NULL);
165 
166 	if (ch->ch_block < blk) {
167 		new = malloc(sizeof(*new));
168 		if (new == NULL) {
169 			free(ptr);
170 			return (NULL);
171 		}
172 		memcpy(new, ch, sizeof(*new));
173 		ch->ch_size = (blk - ch->ch_block) * secsz;
174 		new->ch_block = blk;
175 		new->ch_size -= ch->ch_size;
176 		STAILQ_INSERT_AFTER(&image_chunks, ch, new, ch_list);
177 		image_nchunks++;
178 		ch = new;
179 	}
180 
181 	if (ch->ch_size > secsz) {
182 		new = malloc(sizeof(*new));
183 		if (new == NULL) {
184 			free(ptr);
185 			return (NULL);
186 		}
187 		memcpy(new, ch, sizeof(*new));
188 		ch->ch_size = secsz;
189 		new->ch_block++;
190 		new->ch_size -= secsz;
191 		STAILQ_INSERT_AFTER(&image_chunks, ch, new, ch_list);
192 		image_nchunks++;
193 	}
194 
195 	ch->ch_type = CH_TYPE_MEMORY;
196 	ch->ch_u.mem.ptr = ptr;
197 	return (ch);
198 }
199 
200 static int
201 image_chunk_skipto(lba_t to)
202 {
203 	struct chunk *ch;
204 	lba_t from;
205 	size_t sz;
206 
207 	ch = STAILQ_LAST(&image_chunks, chunk, ch_list);
208 	from = (ch != NULL) ? ch->ch_block + (ch->ch_size / secsz) : 0LL;
209 
210 	assert(from <= to);
211 
212 	/* Nothing to do? */
213 	if (from == to)
214 		return (0);
215 	/* Avoid bugs due to overflows. */
216 	if ((uintmax_t)(to - from) > (uintmax_t)(SIZE_MAX / secsz))
217 		return (EFBIG);
218 	sz = (to - from) * secsz;
219 	if (ch != NULL && ch->ch_type == CH_TYPE_ZEROES) {
220 		sz = image_chunk_grow(ch, sz);
221 		if (sz == 0)
222 			return (0);
223 		from = ch->ch_block + (ch->ch_size / secsz);
224 	}
225 	ch = malloc(sizeof(*ch));
226 	if (ch == NULL)
227 		return (ENOMEM);
228 	memset(ch, 0, sizeof(*ch));
229 	ch->ch_block = from;
230 	ch->ch_size = sz;
231 	ch->ch_type = CH_TYPE_ZEROES;
232 	STAILQ_INSERT_TAIL(&image_chunks, ch, ch_list);
233 	image_nchunks++;
234 	return (0);
235 }
236 
237 static int
238 image_chunk_append(lba_t blk, size_t sz, off_t ofs, int fd)
239 {
240 	struct chunk *ch;
241 
242 	ch = STAILQ_LAST(&image_chunks, chunk, ch_list);
243 	if (ch != NULL && ch->ch_type == CH_TYPE_FILE) {
244 		if (fd == ch->ch_u.file.fd &&
245 		    blk == (lba_t)(ch->ch_block + (ch->ch_size / secsz)) &&
246 		    ofs == (off_t)(ch->ch_u.file.ofs + ch->ch_size)) {
247 			sz = image_chunk_grow(ch, sz);
248 			if (sz == 0)
249 				return (0);
250 			blk = ch->ch_block + (ch->ch_size / secsz);
251 			ofs = ch->ch_u.file.ofs + ch->ch_size;
252 		}
253 	}
254 	ch = malloc(sizeof(*ch));
255 	if (ch == NULL)
256 		return (ENOMEM);
257 	memset(ch, 0, sizeof(*ch));
258 	ch->ch_block = blk;
259 	ch->ch_size = sz;
260 	ch->ch_type = CH_TYPE_FILE;
261 	ch->ch_u.file.ofs = ofs;
262 	ch->ch_u.file.fd = fd;
263 	STAILQ_INSERT_TAIL(&image_chunks, ch, ch_list);
264 	image_nchunks++;
265 	return (0);
266 }
267 
268 static int
269 image_chunk_copyin(lba_t blk, void *buf, size_t sz, off_t ofs, int fd)
270 {
271 	uint8_t *p = buf;
272 	int error;
273 
274 	error = 0;
275 	sz = (sz + secsz - 1) & ~(secsz - 1);
276 	while (!error && sz > 0) {
277 		if (is_empty_sector(p))
278 			error = image_chunk_skipto(blk + 1);
279 		else
280 			error = image_chunk_append(blk, secsz, ofs, fd);
281 		blk++;
282 		p += secsz;
283 		sz -= secsz;
284 		ofs += secsz;
285 	}
286 	return (error);
287 }
288 
289 /*
290  * File mapping support.
291  */
292 
293 static void *
294 image_file_map(int fd, off_t ofs, size_t sz)
295 {
296 	void *ptr;
297 	size_t unit;
298 	int flags, prot;
299 
300 	unit = (secsz > image_swap_pgsz) ? secsz : image_swap_pgsz;
301 	assert((unit & (unit - 1)) == 0);
302 
303 	flags = MAP_NOCORE | MAP_NOSYNC | MAP_SHARED;
304 	/* Allow writing to our swap file only. */
305 	prot = PROT_READ | ((fd == image_swap_fd) ? PROT_WRITE : 0);
306 	sz = (sz + unit - 1) & ~(unit - 1);
307 	ptr = mmap(NULL, sz, prot, flags, fd, ofs);
308 	return ((ptr == MAP_FAILED) ? NULL : ptr);
309 }
310 
311 static int
312 image_file_unmap(void *buffer, size_t sz)
313 {
314 	size_t unit;
315 
316 	unit = (secsz > image_swap_pgsz) ? secsz : image_swap_pgsz;
317 	sz = (sz + unit - 1) & ~(unit - 1);
318 	munmap(buffer, sz);
319 	return (0);
320 }
321 
322 /*
323  * Input/source file handling.
324  */
325 
326 static int
327 image_copyin_stream(lba_t blk, int fd, uint64_t *sizep)
328 {
329 	char *buffer;
330 	uint64_t bytesize;
331 	off_t swofs;
332 	size_t iosz;
333 	ssize_t rdsz;
334 	int error;
335 
336 	/*
337 	 * This makes sure we're doing I/O in multiples of the page
338 	 * size as well as of the sector size. 2MB is the minimum
339 	 * by virtue of secsz at least 512 bytes and the page size
340 	 * at least 4K bytes.
341 	 */
342 	iosz = secsz * image_swap_pgsz;
343 
344 	bytesize = 0;
345 	do {
346 		swofs = image_swap_alloc(iosz);
347 		if (swofs == -1LL)
348 			return (errno);
349 		buffer = image_file_map(image_swap_fd, swofs, iosz);
350 		if (buffer == NULL)
351 			return (errno);
352 		rdsz = read(fd, buffer, iosz);
353 		if (rdsz > 0)
354 			error = image_chunk_copyin(blk, buffer, rdsz, swofs,
355 			    image_swap_fd);
356 		else if (rdsz < 0)
357 			error = errno;
358 		else
359 			error = 0;
360 		image_file_unmap(buffer, iosz);
361 		/* XXX should we relinguish unused swap space? */
362 		if (error)
363 			return (error);
364 
365 		bytesize += rdsz;
366 		blk += (rdsz + secsz - 1) / secsz;
367 	} while (rdsz > 0);
368 
369 	if (sizep != NULL)
370 		*sizep = bytesize;
371 	return (0);
372 }
373 
374 static int
375 image_copyin_mapped(lba_t blk, int fd, uint64_t *sizep)
376 {
377 	off_t cur, data, end, hole, pos;
378 	void *buf;
379 	uint64_t bytesize;
380 	size_t iosz, sz;
381 	int error;
382 
383 	/*
384 	 * We'd like to know the size of the file and we must
385 	 * be able to seek in order to mmap(2). If this isn't
386 	 * possible, then treat the file as a stream/pipe.
387 	 */
388 	end = lseek(fd, 0L, SEEK_END);
389 	if (end == -1L)
390 		return (image_copyin_stream(blk, fd, sizep));
391 
392 	/*
393 	 * We need the file opened for the duration and our
394 	 * caller is going to close the file. Make a dup(2)
395 	 * so that control the faith of the descriptor.
396 	 */
397 	fd = dup(fd);
398 	if (fd == -1)
399 		return (errno);
400 
401 	iosz = secsz * image_swap_pgsz;
402 
403 	bytesize = 0;
404 	cur = pos = 0;
405 	error = 0;
406 	while (!error && cur < end) {
407 		hole = lseek(fd, cur, SEEK_HOLE);
408 		if (hole == -1)
409 			hole = end;
410 		data = lseek(fd, cur, SEEK_DATA);
411 		if (data == -1)
412 			data = end;
413 
414 		/*
415 		 * Treat the entire file as data if sparse files
416 		 * are not supported by the underlying file system.
417 		 */
418 		if (hole == end && data == end)
419 			data = cur;
420 
421 		if (cur == hole && data > hole) {
422 			hole = pos;
423 			pos = data & ~((uint64_t)secsz - 1);
424 
425 			blk += (pos - hole) / secsz;
426 			error = image_chunk_skipto(blk);
427 
428 			bytesize += pos - hole;
429 			cur = data;
430 		} else if (cur == data && hole > data) {
431 			data = pos;
432 			pos = (hole + secsz - 1) & ~((uint64_t)secsz - 1);
433 
434 			while (data < pos) {
435 				sz = (pos - data > (off_t)iosz)
436 				    ? iosz : (size_t)(pos - data);
437 
438 				buf = image_file_map(fd, data, sz);
439 				if (buf != NULL) {
440 					error = image_chunk_copyin(blk, buf,
441 					    sz, data, fd);
442 					image_file_unmap(buf, sz);
443 				} else
444 					error = errno;
445 
446 				blk += sz / secsz;
447 				bytesize += sz;
448 				data += sz;
449 			}
450 			cur = hole;
451 		} else {
452 			/*
453 			 * I don't know what this means or whether it
454 			 * can happen at all...
455 			 */
456 			error = EDOOFUS;
457 			break;
458 		}
459 	}
460 	if (error)
461 		close(fd);
462 	if (!error && sizep != NULL)
463 		*sizep = bytesize;
464 	return (error);
465 }
466 
467 int
468 image_copyin(lba_t blk, int fd, uint64_t *sizep)
469 {
470 	struct stat sb;
471 	int error;
472 
473 	error = image_chunk_skipto(blk);
474 	if (!error) {
475 		if (fstat(fd, &sb) == -1 || !S_ISREG(sb.st_mode))
476 			error = image_copyin_stream(blk, fd, sizep);
477 		else
478 			error = image_copyin_mapped(blk, fd, sizep);
479 	}
480 	return (error);
481 }
482 
483 /*
484  * Output/sink file handling.
485  */
486 
487 int
488 image_copyout(int fd)
489 {
490 	int error;
491 
492 	error = image_copyout_region(fd, 0, image_size);
493 	if (!error)
494 		error = image_copyout_done(fd);
495 	return (error);
496 }
497 
498 int
499 image_copyout_done(int fd)
500 {
501 	off_t ofs;
502 	int error;
503 
504 	ofs = lseek(fd, 0L, SEEK_CUR);
505 	if (ofs == -1)
506 		return (0);
507 	error = (ftruncate(fd, ofs) == -1) ? errno : 0;
508 	return (error);
509 }
510 
511 static int
512 image_copyout_memory(int fd, size_t size, void *ptr)
513 {
514 
515 	if (write(fd, ptr, size) == -1)
516 		return (errno);
517 	return (0);
518 }
519 
520 int
521 image_copyout_zeroes(int fd, size_t count)
522 {
523 	static uint8_t *zeroes = NULL;
524 	size_t sz;
525 	int error;
526 
527 	if (lseek(fd, (off_t)count, SEEK_CUR) != -1)
528 		return (0);
529 
530 	/*
531 	 * If we can't seek, we must write.
532 	 */
533 
534 	if (zeroes == NULL) {
535 		zeroes = calloc(1, secsz);
536 		if (zeroes == NULL)
537 			return (ENOMEM);
538 	}
539 
540 	while (count > 0) {
541 		sz = (count > secsz) ? secsz : count;
542 		error = image_copyout_memory(fd, sz, zeroes);
543 		if (error)
544 			return (error);
545 		count -= sz;
546 	}
547 	return (0);
548 }
549 
550 static int
551 image_copyout_file(int fd, size_t size, int ifd, off_t iofs)
552 {
553 	void *buf;
554 	size_t iosz, sz;
555 	int error;
556 
557 	iosz = secsz * image_swap_pgsz;
558 
559 	while (size > 0) {
560 		sz = (size > iosz) ? iosz : size;
561 		buf = image_file_map(ifd, iofs, sz);
562 		if (buf == NULL)
563 			return (errno);
564 		error = image_copyout_memory(fd, sz, buf);
565 		image_file_unmap(buf, sz);
566 		if (error)
567 			return (error);
568 		size -= sz;
569 		iofs += sz;
570 	}
571 	return (0);
572 }
573 
574 int
575 image_copyout_region(int fd, lba_t blk, lba_t size)
576 {
577 	struct chunk *ch;
578 	size_t ofs, sz;
579 	int error;
580 
581 	size *= secsz;
582 
583 	while (size > 0) {
584 		ch = image_chunk_find(blk);
585 		if (ch == NULL)
586 			return (EINVAL);
587 		ofs = (blk - ch->ch_block) * secsz;
588 		sz = ch->ch_size - ofs;
589 		sz = ((lba_t)sz < size) ? sz : (size_t)size;
590 		switch (ch->ch_type) {
591 		case CH_TYPE_ZEROES:
592 			error = image_copyout_zeroes(fd, sz);
593 			break;
594 		case CH_TYPE_FILE:
595 			error = image_copyout_file(fd, sz, ch->ch_u.file.fd,
596 			    ch->ch_u.file.ofs + ofs);
597 			break;
598 		case CH_TYPE_MEMORY:
599 			error = image_copyout_memory(fd, sz, ch->ch_u.mem.ptr);
600 			break;
601 		default:
602 			return (EDOOFUS);
603 		}
604 		size -= sz;
605 		blk += sz / secsz;
606 	}
607 	return (0);
608 }
609 
610 int
611 image_data(lba_t blk, lba_t size)
612 {
613 	struct chunk *ch;
614 	lba_t lim;
615 
616 	while (1) {
617 		ch = image_chunk_find(blk);
618 		if (ch == NULL)
619 			return (0);
620 		if (ch->ch_type != CH_TYPE_ZEROES)
621 			return (1);
622 		lim = ch->ch_block + (ch->ch_size / secsz);
623 		if (lim >= blk + size)
624 			return (0);
625 		size -= lim - blk;
626 		blk = lim;
627 	}
628 	/*NOTREACHED*/
629 }
630 
631 lba_t
632 image_get_size(void)
633 {
634 
635 	return (image_size);
636 }
637 
638 int
639 image_set_size(lba_t blk)
640 {
641 	int error;
642 
643 	error = image_chunk_skipto(blk);
644 	if (!error)
645 		image_size = blk;
646 	return (error);
647 }
648 
649 int
650 image_write(lba_t blk, void *buf, ssize_t len)
651 {
652 	struct chunk *ch;
653 
654 	while (len > 0) {
655 		if (!is_empty_sector(buf)) {
656 			ch = image_chunk_find(blk);
657 			if (ch == NULL)
658 				return (ENXIO);
659 			/* We may not be able to write to files. */
660 			if (ch->ch_type == CH_TYPE_FILE)
661 				return (EINVAL);
662 			if (ch->ch_type == CH_TYPE_ZEROES) {
663 				ch = image_chunk_memory(ch, blk);
664 				if (ch == NULL)
665 					return (ENOMEM);
666 			}
667 			assert(ch->ch_type == CH_TYPE_MEMORY);
668 			memcpy(ch->ch_u.mem.ptr, buf, secsz);
669 		}
670 		blk++;
671 		buf = (char *)buf + secsz;
672 		len--;
673 	}
674 	return (0);
675 }
676 
677 static void
678 image_cleanup(void)
679 {
680 	struct chunk *ch;
681 
682 	while ((ch = STAILQ_FIRST(&image_chunks)) != NULL) {
683 		switch (ch->ch_type) {
684 		case CH_TYPE_FILE:
685 			/* We may be closing the same file multiple times. */
686 			if (ch->ch_u.file.fd != -1)
687 				close(ch->ch_u.file.fd);
688 			break;
689 		case CH_TYPE_MEMORY:
690 			free(ch->ch_u.mem.ptr);
691 			break;
692 		default:
693 			break;
694 		}
695 		STAILQ_REMOVE_HEAD(&image_chunks, ch_list);
696 		free(ch);
697 	}
698 	if (image_swap_fd != -1)
699 		close(image_swap_fd);
700 	unlink(image_swap_file);
701 }
702 
703 int
704 image_init(void)
705 {
706 	const char *tmpdir;
707 
708 	STAILQ_INIT(&image_chunks);
709 	image_nchunks = 0;
710 
711 	image_swap_size = 0;
712 	image_swap_pgsz = getpagesize();
713 
714 	if (atexit(image_cleanup) == -1)
715 		return (errno);
716 	if ((tmpdir = getenv("TMPDIR")) == NULL || *tmpdir == '\0')
717 		tmpdir = _PATH_TMP;
718 	snprintf(image_swap_file, sizeof(image_swap_file), "%s/mkimg-XXXXXX",
719 	    tmpdir);
720 	image_swap_fd = mkstemp(image_swap_file);
721 	if (image_swap_fd == -1)
722 		return (errno);
723 	return (0);
724 }
725