xref: /freebsd/usr.bin/mkimg/image.c (revision a0e793cbf1951d07fc47a0d9ea389d7dacba5213)
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 	int error;
98 
99 	error = image_copyout_region(fd, 0, image_size);
100 	if (!error)
101 		error = image_copyout_done(fd);
102 	return (error);
103 }
104 
105 int
106 image_copyout_done(int fd)
107 {
108 	off_t ofs;
109 	int error;
110 
111 	ofs = lseek(fd, 0L, SEEK_CUR);
112 	if (ofs == -1)
113 		return (0);
114 	error = (ftruncate(fd, ofs) == -1) ? errno : 0;
115 	return (error);
116 }
117 
118 int
119 image_copyout_region(int fd, lba_t blk, lba_t size)
120 {
121 	char *buffer;
122 	off_t ofs;
123 	size_t sz;
124 	ssize_t rdsz, wrsz;
125 	int error;
126 
127 	ofs = lseek(fd, 0L, SEEK_CUR);
128 
129 	blk *= secsz;
130 	if (lseek(image_fd, blk, SEEK_SET) != blk)
131 		return (errno);
132 	buffer = malloc(BUFFER_SIZE);
133 	if (buffer == NULL)
134 		return (errno);
135 	error = 0;
136 	size *= secsz;
137 	while (size > 0) {
138 		sz = (BUFFER_SIZE < size) ? BUFFER_SIZE : size;
139 		rdsz = read(image_fd, buffer, sz);
140 		if (rdsz <= 0) {
141 			error = (rdsz < 0) ? errno : 0;
142 			break;
143 		}
144 		wrsz = (ofs == -1) ?
145 		    write(fd, buffer, rdsz) :
146 		    sparse_write(fd, buffer, rdsz);
147 		if (wrsz < 0) {
148 			error = errno;
149 			break;
150 		}
151 		assert(wrsz == rdsz);
152 		size -= rdsz;
153 	}
154 	free(buffer);
155 	return (error);
156 }
157 
158 int
159 image_data(lba_t blk, lba_t size)
160 {
161 	char *buffer, *p;
162 
163 	blk *= secsz;
164 	if (lseek(image_fd, blk, SEEK_SET) != blk)
165 		return (1);
166 
167 	size *= secsz;
168 	buffer = malloc(size);
169 	if (buffer == NULL)
170 		return (1);
171 
172 	if (read(image_fd, buffer, size) != (ssize_t)size) {
173 		free(buffer);
174 		return (1);
175 	}
176 
177 	p = buffer;
178 	while (size > 0 && *p == '\0')
179 		size--, p++;
180 
181 	free(buffer);
182 	return ((size == 0) ? 0 : 1);
183 }
184 
185 lba_t
186 image_get_size(void)
187 {
188 
189 	return (image_size);
190 }
191 
192 int
193 image_set_size(lba_t blk)
194 {
195 
196 	image_size = blk;
197 	if (ftruncate(image_fd, blk * secsz) == -1)
198 		return (errno);
199 	return (0);
200 }
201 
202 int
203 image_write(lba_t blk, void *buf, ssize_t len)
204 {
205 
206 	blk *= secsz;
207 	if (lseek(image_fd, blk, SEEK_SET) != blk)
208 		return (errno);
209 	len *= secsz;
210 	if (sparse_write(image_fd, buf, len) != len)
211 		return (errno);
212 	return (0);
213 }
214 
215 int
216 image_init(void)
217 {
218 	const char *tmpdir;
219 
220 	if (atexit(cleanup) == -1)
221 		return (errno);
222 	if ((tmpdir = getenv("TMPDIR")) == NULL || *tmpdir == '\0')
223 		tmpdir = _PATH_TMP;
224 	snprintf(image_tmpfile, sizeof(image_tmpfile), "%s/mkimg-XXXXXX",
225 	    tmpdir);
226 	image_fd = mkstemp(image_tmpfile);
227 	if (image_fd == -1)
228 		return (errno);
229 	return (0);
230 }
231