xref: /freebsd/usr.bin/mkimg/mkimg.c (revision cc3f4b99653c34ae64f8a1fddea370abefef680e)
1 /*-
2  * Copyright (c) 2013,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/linker_set.h>
31 #include <sys/queue.h>
32 #include <sys/stat.h>
33 #include <sys/types.h>
34 #include <errno.h>
35 #include <err.h>
36 #include <fcntl.h>
37 #include <libutil.h>
38 #include <limits.h>
39 #include <stdio.h>
40 #include <stdlib.h>
41 #include <string.h>
42 #include <sysexits.h>
43 #include <unistd.h>
44 
45 #include "mkimg.h"
46 #include "scheme.h"
47 
48 #if !defined(SPARSE_WRITE)
49 #define	sparse_write	write
50 #endif
51 
52 #define	BUFFER_SIZE	(1024*1024)
53 
54 struct partlisthead partlist = STAILQ_HEAD_INITIALIZER(partlist);
55 u_int nparts = 0;
56 
57 u_int verbose;
58 
59 u_int ncyls = 0;
60 u_int nheads = 1;
61 u_int nsecs = 1;
62 u_int secsz = 512;
63 u_int blksz = 0;
64 
65 static int bcfd = -1;
66 static int outfd = 0;
67 static int tmpfd = -1;
68 
69 static char tmpfname[] = "/tmp/mkimg-XXXXXX";
70 
71 static void
72 cleanup(void)
73 {
74 
75 	if (tmpfd != -1)
76 		close(tmpfd);
77 	unlink(tmpfname);
78 }
79 
80 static void
81 usage(const char *why)
82 {
83 	struct mkimg_scheme *s, **iter;
84 
85 	warnx("error: %s", why);
86 	fprintf(stderr, "\nusage: %s <options>\n", getprogname());
87 
88 	fprintf(stderr, "    options:\n");
89 	fprintf(stderr, "\t-b <file>\t-  file containing boot code\n");
90 	fprintf(stderr, "\t-o <file>\t-  file to write image into\n");
91 	fprintf(stderr, "\t-p <partition>\n");
92 	fprintf(stderr, "\t-s <scheme>\n");
93 	fprintf(stderr, "\t-H <num>\t-  number of heads to simulate\n");
94 	fprintf(stderr, "\t-P <num>\t-  physical sector size\n");
95 	fprintf(stderr, "\t-S <num>\t-  logical sector size\n");
96 	fprintf(stderr, "\t-T <num>\t-  number of tracks to simulate\n");
97 
98 	fprintf(stderr, "    schemes:\n");
99 	SET_FOREACH(iter, schemes) {
100 		s = *iter;
101 		fprintf(stderr, "\t%s\t-  %s\n", s->name, s->description);
102 	}
103 
104 	fprintf(stderr, "    partition specification:\n");
105 	fprintf(stderr, "\t<t>[/<l>]::<size>\t-  empty partition of given "
106 	    "size\n");
107 	fprintf(stderr, "\t<t>[/<l>]:=<file>\t-  partition content and size "
108 	    "are determined\n\t\t\t\t   by the named file\n");
109 	fprintf(stderr, "\t<t>[/<l>]:-<cmd>\t-  partition content and size "
110 	    "are taken from\n\t\t\t\t   the output of the command to run\n");
111 	fprintf(stderr, "\t    where:\n");
112 	fprintf(stderr, "\t\t<t>\t-  scheme neutral partition type\n");
113 	fprintf(stderr, "\t\t<l>\t-  optional scheme-dependent partition "
114 	    "label\n");
115 
116 	exit(EX_USAGE);
117 }
118 
119 static int
120 parse_number(u_int *valp, u_int min, u_int max, const char *arg)
121 {
122 	uint64_t val;
123 
124 	if (expand_number(arg, &val) == -1)
125 		return (errno);
126 	if (val > UINT_MAX || val < (uint64_t)min || val > (uint64_t)max)
127 		return (EINVAL);
128 	*valp = (u_int)val;
129 	return (0);
130 }
131 
132 static int
133 pwr_of_two(u_int nr)
134 {
135 
136 	return (((nr & (nr - 1)) == 0) ? 1 : 0);
137 }
138 
139 /*
140  * A partition specification has the following format:
141  *	<type> ':' <kind> <contents>
142  * where:
143  *	type	  the partition type alias
144  *	kind	  the interpretation of the contents specification
145  *		  ':'   contents holds the size of an empty partition
146  *		  '='   contents holds the name of a file to read
147  *		  '-'   contents holds a command to run; the output of
148  *			which is the contents of the partition.
149  *	contents  the specification of a partition's contents
150  */
151 static int
152 parse_part(const char *spec)
153 {
154 	struct part *part;
155 	char *sep;
156 	size_t len;
157 	int error;
158 
159 	part = calloc(1, sizeof(struct part));
160 	if (part == NULL)
161 		return (ENOMEM);
162 
163 	sep = strchr(spec, ':');
164 	if (sep == NULL) {
165 		error = EINVAL;
166 		goto errout;
167 	}
168 	len = sep - spec + 1;
169 	if (len < 2) {
170 		error = EINVAL;
171 		goto errout;
172 	}
173 	part->alias = malloc(len);
174 	if (part->alias == NULL) {
175 		error = ENOMEM;
176 		goto errout;
177 	}
178 	strlcpy(part->alias, spec, len);
179 	spec = sep + 1;
180 
181 	switch (*spec) {
182 	case ':':
183 		part->kind = PART_KIND_SIZE;
184 		break;
185 	case '=':
186 		part->kind = PART_KIND_FILE;
187 		break;
188 	case '-':
189 		part->kind = PART_KIND_PIPE;
190 		break;
191 	default:
192 		error = EINVAL;
193 		goto errout;
194 	}
195 	spec++;
196 
197 	part->contents = strdup(spec);
198 	if (part->contents == NULL) {
199 		error = ENOMEM;
200 		goto errout;
201 	}
202 
203 	spec = part->alias;
204 	sep = strchr(spec, '/');
205 	if (sep != NULL) {
206 		*sep++ = '\0';
207 		if (strlen(part->alias) == 0 || strlen(sep) == 0) {
208 			error = EINVAL;
209 			goto errout;
210 		}
211 		part->label = strdup(sep);
212 		if (part->label == NULL) {
213 			error = ENOMEM;
214 			goto errout;
215 		}
216 	}
217 
218 	part->index = nparts;
219 	STAILQ_INSERT_TAIL(&partlist, part, link);
220 	nparts++;
221 	return (0);
222 
223  errout:
224 	if (part->alias != NULL)
225 		free(part->alias);
226 	free(part);
227 	return (error);
228 }
229 
230 #if defined(SPARSE_WRITE)
231 static ssize_t
232 sparse_write(int fd, const char *buf, size_t sz)
233 {
234 	const char *p;
235 	off_t ofs;
236 	size_t len;
237 	ssize_t wr, wrsz;
238 
239 	wrsz = 0;
240 	p = memchr(buf, 0, sz);
241 	while (sz > 0) {
242 		len = (p != NULL) ? (size_t)(p - buf) : sz;
243 		if (len > 0) {
244 			len = (len + secsz - 1) & ~(secsz - 1);
245 			if (len > sz)
246 				len = sz;
247 			wr = write(fd, buf, len);
248 			if (wr < 0)
249 				return (-1);
250 		} else {
251 			while (len < sz && *p++ == '\0')
252 				len++;
253 			if (len < sz)
254 				len &= ~(secsz - 1);
255 			if (len == 0)
256 				continue;
257 			ofs = lseek(fd, len, SEEK_CUR);
258 			if (ofs < 0)
259 				return (-1);
260 			wr = len;
261 		}
262 		buf += wr;
263 		sz -= wr;
264 		wrsz += wr;
265 		p = memchr(buf, 0, sz);
266 	}
267 	return (wrsz);
268 }
269 #endif /* SPARSE_WRITE */
270 
271 static int
272 fdcopy(int src, int dst, uint64_t *count)
273 {
274 	char *buffer;
275 	off_t ofs;
276 	ssize_t rdsz, wrsz;
277 
278 	/* A return value of -1 means that we can't write a sparse file. */
279 	ofs = lseek(dst, 0L, SEEK_CUR);
280 
281 	if (count != NULL)
282 		*count = 0;
283 
284 	buffer = malloc(BUFFER_SIZE);
285 	if (buffer == NULL)
286 		return (errno);
287 	while (1) {
288 		rdsz = read(src, buffer, BUFFER_SIZE);
289 		if (rdsz <= 0) {
290 			free(buffer);
291 			return ((rdsz < 0) ? errno : 0);
292 		}
293 		if (count != NULL)
294 			*count += rdsz;
295 		wrsz = (ofs == -1) ?
296 		    write(dst, buffer, rdsz) :
297 		    sparse_write(dst, buffer, rdsz);
298 		if (wrsz < 0)
299 			break;
300 	}
301 	free(buffer);
302 	return (errno);
303 }
304 
305 static int
306 mkimg_seek(int fd, lba_t blk)
307 {
308 	off_t off;
309 
310 	off = blk * secsz;
311 	if (lseek(fd, off, SEEK_SET) != off)
312 		return (errno);
313 	return (0);
314 }
315 
316 int
317 mkimg_write(int fd, lba_t blk, void *buf, ssize_t len)
318 {
319 
320 	blk *= secsz;
321 	if (lseek(fd, blk, SEEK_SET) != blk)
322 		return (errno);
323 	len *= secsz;
324 	if (write(fd, buf, len) != len)
325 		return (errno);
326 	return (0);
327 }
328 
329 static void
330 mkimg(int bfd)
331 {
332 	FILE *fp;
333 	struct part *part;
334 	lba_t block;
335 	off_t bytesize;
336 	int error, fd;
337 
338 	error = scheme_bootcode(bfd);
339 	if (error)
340 		errc(EX_DATAERR, error, "boot code");
341 
342 	/* First check partition information */
343 	STAILQ_FOREACH(part, &partlist, link) {
344 		error = scheme_check_part(part);
345 		if (error)
346 			errc(EX_DATAERR, error, "partition %d", part->index+1);
347 	}
348 
349 	block = scheme_metadata(SCHEME_META_IMG_START, 0);
350 	STAILQ_FOREACH(part, &partlist, link) {
351 		block = scheme_metadata(SCHEME_META_PART_BEFORE, block);
352 		if (verbose)
353 			fprintf(stderr, "partition %d: starting block %llu "
354 			    "... ", part->index + 1, (long long)block);
355 		part->block = block;
356 		error = mkimg_seek(tmpfd, block);
357 		switch (part->kind) {
358 		case PART_KIND_SIZE:
359 			if (expand_number(part->contents, &bytesize) == -1)
360 				error = errno;
361 			break;
362 		case PART_KIND_FILE:
363 			fd = open(part->contents, O_RDONLY, 0);
364 			if (fd != -1) {
365 				error = fdcopy(fd, tmpfd, &bytesize);
366 				close(fd);
367 			} else
368 				error = errno;
369 			break;
370 		case PART_KIND_PIPE:
371 			fp = popen(part->contents, "r");
372 			if (fp != NULL) {
373 				error = fdcopy(fileno(fp), tmpfd, &bytesize);
374 				pclose(fp);
375 			} else
376 				error = errno;
377 			break;
378 		}
379 		if (error)
380 			errc(EX_IOERR, error, "partition %d", part->index + 1);
381 		part->size = (bytesize + secsz - 1) / secsz;
382 		if (verbose) {
383 			bytesize = part->size * secsz;
384 			fprintf(stderr, "size %llu bytes (%llu blocks)\n",
385 			     (long long)bytesize, (long long)part->size);
386 		}
387 		block = scheme_metadata(SCHEME_META_PART_AFTER,
388 		    part->block + part->size);
389 	}
390 
391 	block = scheme_metadata(SCHEME_META_IMG_END, block);
392 	error = (scheme_write(tmpfd, block));
393 }
394 
395 int
396 main(int argc, char *argv[])
397 {
398 	int c, error;
399 
400 	while ((c = getopt(argc, argv, "b:o:p:s:vH:P:S:T:")) != -1) {
401 		switch (c) {
402 		case 'b':	/* BOOT CODE */
403 			if (bcfd != -1)
404 				usage("multiple bootcode given");
405 			bcfd = open(optarg, O_RDONLY, 0);
406 			if (bcfd == -1)
407 				err(EX_UNAVAILABLE, "%s", optarg);
408 			break;
409 		case 'o':	/* OUTPUT FILE */
410 			if (outfd != 0)
411 				usage("multiple output files given");
412 			outfd = open(optarg, O_WRONLY | O_CREAT | O_TRUNC,
413 			    S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
414 			if (outfd == -1)
415 				err(EX_CANTCREAT, "%s", optarg);
416 			break;
417 		case 'p':	/* PARTITION */
418 			error = parse_part(optarg);
419 			if (error)
420 				errc(EX_DATAERR, error, "partition");
421 			break;
422 		case 's':	/* SCHEME */
423 			if (scheme_selected() != NULL)
424 				usage("multiple schemes given");
425 			error = scheme_select(optarg);
426 			if (error)
427 				errc(EX_DATAERR, error, "scheme");
428 			break;
429 		case 'v':
430 			verbose++;
431 			break;
432 		case 'H':	/* GEOMETRY: HEADS */
433 			error = parse_number(&nheads, 1, 255, optarg);
434 			if (error)
435 				errc(EX_DATAERR, error, "number of heads");
436 			break;
437 		case 'P':	/* GEOMETRY: PHYSICAL SECTOR SIZE */
438 			error = parse_number(&blksz, 512, INT_MAX+1U, optarg);
439 			if (error == 0 && !pwr_of_two(blksz))
440 				error = EINVAL;
441 			if (error)
442 				errc(EX_DATAERR, error, "physical sector size");
443 			break;
444 		case 'S':	/* GEOMETRY: LOGICAL SECTOR SIZE */
445 			error = parse_number(&secsz, 512, INT_MAX+1U, optarg);
446 			if (error == 0 && !pwr_of_two(secsz))
447 				error = EINVAL;
448 			if (error)
449 				errc(EX_DATAERR, error, "logical sector size");
450 			break;
451 		case 'T':	/* GEOMETRY: TRACK SIZE */
452 			error = parse_number(&nsecs, 1, 63, optarg);
453 			if (error)
454 				errc(EX_DATAERR, error, "track size");
455 			break;
456 		default:
457 			usage("unknown option");
458 		}
459 	}
460 
461 	if (argc > optind)
462 		usage("trailing arguments");
463 	if (scheme_selected() == NULL)
464 		usage("no scheme");
465 	if (nparts == 0)
466 		usage("no partitions");
467 
468 	if (secsz > blksz) {
469 		if (blksz != 0)
470 			errx(EX_DATAERR, "the physical block size cannot "
471 			    "be smaller than the sector size");
472 		blksz = secsz;
473 	}
474 
475 	if (secsz > scheme_max_secsz())
476 		errx(EX_DATAERR, "maximum sector size supported is %u; "
477 		    "size specified is %u", scheme_max_secsz(), secsz);
478 
479 	if (nparts > scheme_max_parts())
480 		errx(EX_DATAERR, "%d partitions supported; %d given",
481 		    scheme_max_parts(), nparts);
482 
483 	if (outfd == 0) {
484 		if (atexit(cleanup) == -1)
485 			err(EX_OSERR, "cannot register cleanup function");
486 		outfd = 1;
487 		tmpfd = mkstemp(tmpfname);
488 		if (tmpfd == -1)
489 			err(EX_OSERR, "cannot create temporary file");
490 	} else
491 		tmpfd = outfd;
492 
493 	if (verbose) {
494 		fprintf(stderr, "Logical sector size: %u\n", secsz);
495 		fprintf(stderr, "Physical block size: %u\n", blksz);
496 		fprintf(stderr, "Sectors per track:   %u\n", nsecs);
497 		fprintf(stderr, "Number of heads:     %u\n", nheads);
498 	}
499 
500 	mkimg(bcfd);
501 
502 	if (verbose)
503 		fprintf(stderr, "Number of cylinders: %u\n", ncyls);
504 
505 	if (tmpfd != outfd) {
506 		error = mkimg_seek(tmpfd, 0);
507 		if (error == 0)
508 			error = fdcopy(tmpfd, outfd, NULL);
509 		if (error)
510 			errc(EX_IOERR, error, "writing to stdout");
511 	}
512 
513 	return (0);
514 }
515