image.c revision 274634 
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: stable/10/usr.bin/mkimg/image.c 274634 2014-11-17 17:17:10Z marcel $");
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
47struct 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
66static STAILQ_HEAD(chunk_head, chunk) image_chunks;
67static u_int image_nchunks;
68
69static char image_swap_file[PATH_MAX];
70static int image_swap_fd = -1;
71static u_int image_swap_pgsz;
72static off_t image_swap_size;
73
74static lba_t image_size;
75
76static int
77is_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
96static off_t
97image_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
120static struct chunk *
121image_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
139static size_t
140image_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
156static struct chunk *
157image_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
200static int
201image_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
237static int
238image_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
268static int
269image_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
293static void *
294image_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
311static int
312image_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
326static int
327image_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
374static int
375image_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
467int
468image_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
487int
488image_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
498int
499image_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
511static int
512image_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
520static int
521image_copyout_zeroes(int fd, size_t size)
522{
523	static uint8_t *zeroes = NULL;
524	size_t sz;
525	int error;
526
527	if (lseek(fd, (off_t)size, 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 (size > 0) {
541		sz = (size > secsz) ? secsz : size;
542		error = image_copyout_memory(fd, sz, zeroes);
543		if (error)
544			return (error);
545		size -= sz;
546	}
547	return (0);
548}
549
550static int
551image_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
574int
575image_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
610int
611image_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
631lba_t
632image_get_size(void)
633{
634
635	return (image_size);
636}
637
638int
639image_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
649int
650image_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
677static void
678image_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
703int
704image_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