xref: /freebsd/usr.bin/mkimg/image.c (revision 413a368c90ea8621329ec702d46501ac7a4ffb4b)
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/types.h>
31 #include <assert.h>
32 #include <errno.h>
33 #include <limits.h>
34 #include <paths.h>
35 #include <stdio.h>
36 #include <stdlib.h>
37 #include <unistd.h>
38 
39 #include "image.h"
40 #include "mkimg.h"
41 
42 #define	BUFFER_SIZE	(1024*1024)
43 
44 static char image_tmpfile[PATH_MAX];
45 static int image_fd = -1;
46 static lba_t image_size;
47 
48 static void
49 cleanup(void)
50 {
51 
52 	if (image_fd != -1)
53 		close(image_fd);
54 	unlink(image_tmpfile);
55 }
56 
57 int
58 image_copyin(lba_t blk, int fd, uint64_t *sizep)
59 {
60 	char *buffer;
61 	uint64_t bytesize;
62 	ssize_t bcnt, rdsz;
63 	int error, partial;
64 
65 	assert(BUFFER_SIZE % secsz == 0);
66 
67 	buffer = malloc(BUFFER_SIZE);
68 	if (buffer == NULL)
69 		return (ENOMEM);
70 	bytesize = 0;
71 	partial = 0;
72 	while (1) {
73 		rdsz = read(fd, buffer, BUFFER_SIZE);
74 		if (rdsz <= 0) {
75 			error = (rdsz < 0) ? errno : 0;
76 			break;
77 		}
78 		if (partial)
79 			abort();
80 		bytesize += rdsz;
81 		bcnt = (rdsz + secsz - 1) / secsz;
82 		error = image_write(blk, buffer, bcnt);
83 		if (error)
84 			break;
85 		blk += bcnt;
86 		partial = ((ssize_t)(bcnt * secsz) != rdsz) ? 1 : 0;
87 	}
88 	free(buffer);
89 	if (sizep != NULL)
90 		*sizep = bytesize;
91 	return (error);
92 }
93 
94 int
95 image_copyout(int fd)
96 {
97 	off_t ofs;
98 	int error;
99 
100 	error = image_copyout_region(fd, 0, image_size);
101 	if (error)
102 		return (error);
103 
104 	ofs = lseek(fd, 0L, SEEK_CUR);
105 	if (ofs == -1)
106 		return (0);
107 	error = (ftruncate(fd, ofs) == -1) ? errno : 0;
108 	return (error);
109 }
110 
111 int
112 image_copyout_region(int fd, lba_t blk, lba_t size)
113 {
114 	char *buffer;
115 	off_t ofs;
116 	size_t sz;
117 	ssize_t rdsz, wrsz;
118 	int error;
119 
120 	ofs = lseek(fd, 0L, SEEK_CUR);
121 
122 	blk *= secsz;
123 	if (lseek(image_fd, blk, SEEK_SET) != blk)
124 		return (errno);
125 	buffer = malloc(BUFFER_SIZE);
126 	if (buffer == NULL)
127 		return (errno);
128 	error = 0;
129 	size *= secsz;
130 	while (size > 0) {
131 		sz = (BUFFER_SIZE < size) ? BUFFER_SIZE : size;
132 		rdsz = read(image_fd, buffer, sz);
133 		if (rdsz <= 0) {
134 			error = (rdsz < 0) ? errno : 0;
135 			break;
136 		}
137 		wrsz = (ofs == -1) ?
138 		    write(fd, buffer, rdsz) :
139 		    sparse_write(fd, buffer, rdsz);
140 		if (wrsz < 0) {
141 			error = errno;
142 			break;
143 		}
144 		assert(wrsz == rdsz);
145 		size -= rdsz;
146 	}
147 	free(buffer);
148 	return (error);
149 }
150 
151 lba_t
152 image_get_size(void)
153 {
154 
155 	return (image_size);
156 }
157 
158 int
159 image_set_size(lba_t blk)
160 {
161 
162 	image_size = blk;
163 	if (ftruncate(image_fd, blk * secsz) == -1)
164 		return (errno);
165 	return (0);
166 }
167 
168 int
169 image_write(lba_t blk, void *buf, ssize_t len)
170 {
171 
172 	blk *= secsz;
173 	if (lseek(image_fd, blk, SEEK_SET) != blk)
174 		return (errno);
175 	len *= secsz;
176 	if (sparse_write(image_fd, buf, len) != len)
177 		return (errno);
178 	return (0);
179 }
180 
181 int
182 image_init(void)
183 {
184 	const char *tmpdir;
185 
186 	if (atexit(cleanup) == -1)
187 		return (errno);
188 	if ((tmpdir = getenv("TMPDIR")) == NULL || *tmpdir == '\0')
189 		tmpdir = _PATH_TMP;
190 	snprintf(image_tmpfile, sizeof(image_tmpfile), "%s/mkimg-XXXXXX",
191 	    tmpdir);
192 	image_fd = mkstemp(image_tmpfile);
193 	if (image_fd == -1)
194 		return (errno);
195 	return (0);
196 }
197