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 <sys/uuid.h>
35#include <errno.h>
36#include <err.h>
37#include <fcntl.h>
38#include <getopt.h>
39#include <libutil.h>
40#include <limits.h>
41#include <stdio.h>
42#include <stdlib.h>
43#include <string.h>
44#include <sysexits.h>
45#include <unistd.h>
46
47#include "image.h"
48#include "format.h"
49#include "mkimg.h"
50#include "scheme.h"
51
52#define	LONGOPT_FORMATS	0x01000001
53#define	LONGOPT_SCHEMES	0x01000002
54#define	LONGOPT_VERSION	0x01000003
55
56static struct option longopts[] = {
57	{ "formats", no_argument, NULL, LONGOPT_FORMATS },
58	{ "schemes", no_argument, NULL, LONGOPT_SCHEMES },
59	{ "version", no_argument, NULL, LONGOPT_VERSION },
60	{ NULL, 0, NULL, 0 }
61};
62
63static uint64_t capacity;
64
65struct partlisthead partlist = STAILQ_HEAD_INITIALIZER(partlist);
66u_int nparts = 0;
67
68u_int unit_testing;
69u_int verbose;
70
71u_int ncyls = 0;
72u_int nheads = 1;
73u_int nsecs = 1;
74u_int secsz = 512;
75u_int blksz = 0;
76
77static void
78print_formats(int usage)
79{
80	struct mkimg_format *f, **f_iter;
81	const char *sep;
82
83	if (usage) {
84		fprintf(stderr, "    formats:\n");
85		SET_FOREACH(f_iter, formats) {
86			f = *f_iter;
87			fprintf(stderr, "\t%s\t-  %s\n", f->name,
88			    f->description);
89		}
90	} else {
91		sep = "";
92		SET_FOREACH(f_iter, formats) {
93			f = *f_iter;
94			printf("%s%s", sep, f->name);
95			sep = " ";
96		}
97		putchar('\n');
98	}
99}
100
101static void
102print_schemes(int usage)
103{
104	struct mkimg_scheme *s, **s_iter;
105	const char *sep;
106
107	if (usage) {
108		fprintf(stderr, "    schemes:\n");
109		SET_FOREACH(s_iter, schemes) {
110			s = *s_iter;
111			fprintf(stderr, "\t%s\t-  %s\n", s->name,
112			    s->description);
113		}
114	} else {
115		sep = "";
116		SET_FOREACH(s_iter, schemes) {
117			s = *s_iter;
118			printf("%s%s", sep, s->name);
119			sep = " ";
120		}
121		putchar('\n');
122	}
123}
124
125static void
126print_version(void)
127{
128	u_int width;
129
130#ifdef __LP64__
131	width = 64;
132#else
133	width = 32;
134#endif
135	printf("mkimg %u (%u-bit)\n", MKIMG_VERSION, width);
136}
137
138static void
139usage(const char *why)
140{
141
142	warnx("error: %s", why);
143	fputc('\n', stderr);
144	fprintf(stderr, "usage: %s <options>\n", getprogname());
145
146	fprintf(stderr, "    options:\n");
147	fprintf(stderr, "\t--formats\t-  list image formats\n");
148	fprintf(stderr, "\t--schemes\t-  list partition schemes\n");
149	fprintf(stderr, "\t--version\t-  show version information\n");
150	fputc('\n', stderr);
151	fprintf(stderr, "\t-b <file>\t-  file containing boot code\n");
152	fprintf(stderr, "\t-c <num>\t-  capacity (in bytes) of the disk\n");
153	fprintf(stderr, "\t-f <format>\n");
154	fprintf(stderr, "\t-o <file>\t-  file to write image into\n");
155	fprintf(stderr, "\t-p <partition>\n");
156	fprintf(stderr, "\t-s <scheme>\n");
157	fprintf(stderr, "\t-v\t\t-  increase verbosity\n");
158	fprintf(stderr, "\t-y\t\t-  [developers] enable unit test\n");
159	fprintf(stderr, "\t-H <num>\t-  number of heads to simulate\n");
160	fprintf(stderr, "\t-P <num>\t-  physical sector size\n");
161	fprintf(stderr, "\t-S <num>\t-  logical sector size\n");
162	fprintf(stderr, "\t-T <num>\t-  number of tracks to simulate\n");
163	fputc('\n', stderr);
164	print_formats(1);
165	fputc('\n', stderr);
166	print_schemes(1);
167	fputc('\n', stderr);
168	fprintf(stderr, "    partition specification:\n");
169	fprintf(stderr, "\t<t>[/<l>]::<size>\t-  empty partition of given "
170	    "size\n");
171	fprintf(stderr, "\t<t>[/<l>]:=<file>\t-  partition content and size "
172	    "are determined\n\t\t\t\t   by the named file\n");
173	fprintf(stderr, "\t<t>[/<l>]:-<cmd>\t-  partition content and size "
174	    "are taken from\n\t\t\t\t   the output of the command to run\n");
175	fprintf(stderr, "\t-\t\t\t-  unused partition entry\n");
176	fprintf(stderr, "\t    where:\n");
177	fprintf(stderr, "\t\t<t>\t-  scheme neutral partition type\n");
178	fprintf(stderr, "\t\t<l>\t-  optional scheme-dependent partition "
179	    "label\n");
180
181	exit(EX_USAGE);
182}
183
184static int
185parse_uint32(uint32_t *valp, uint32_t min, uint32_t max, const char *arg)
186{
187	uint64_t val;
188
189	if (expand_number(arg, &val) == -1)
190		return (errno);
191	if (val > UINT_MAX || val < (uint64_t)min || val > (uint64_t)max)
192		return (EINVAL);
193	*valp = (uint32_t)val;
194	return (0);
195}
196
197static int
198parse_uint64(uint64_t *valp, uint64_t min, uint64_t max, const char *arg)
199{
200	uint64_t val;
201
202	if (expand_number(arg, &val) == -1)
203		return (errno);
204	if (val < min || val > max)
205		return (EINVAL);
206	*valp = val;
207	return (0);
208}
209
210static int
211pwr_of_two(u_int nr)
212{
213
214	return (((nr & (nr - 1)) == 0) ? 1 : 0);
215}
216
217/*
218 * A partition specification has the following format:
219 *	<type> ':' <kind> <contents>
220 * where:
221 *	type	  the partition type alias
222 *	kind	  the interpretation of the contents specification
223 *		  ':'   contents holds the size of an empty partition
224 *		  '='   contents holds the name of a file to read
225 *		  '-'   contents holds a command to run; the output of
226 *			which is the contents of the partition.
227 *	contents  the specification of a partition's contents
228 *
229 * A specification that is a single dash indicates an unused partition
230 * entry.
231 */
232static int
233parse_part(const char *spec)
234{
235	struct part *part;
236	char *sep;
237	size_t len;
238	int error;
239
240	if (strcmp(spec, "-") == 0) {
241		nparts++;
242		return (0);
243	}
244
245	part = calloc(1, sizeof(struct part));
246	if (part == NULL)
247		return (ENOMEM);
248
249	sep = strchr(spec, ':');
250	if (sep == NULL) {
251		error = EINVAL;
252		goto errout;
253	}
254	len = sep - spec + 1;
255	if (len < 2) {
256		error = EINVAL;
257		goto errout;
258	}
259	part->alias = malloc(len);
260	if (part->alias == NULL) {
261		error = ENOMEM;
262		goto errout;
263	}
264	strlcpy(part->alias, spec, len);
265	spec = sep + 1;
266
267	switch (*spec) {
268	case ':':
269		part->kind = PART_KIND_SIZE;
270		break;
271	case '=':
272		part->kind = PART_KIND_FILE;
273		break;
274	case '-':
275		part->kind = PART_KIND_PIPE;
276		break;
277	default:
278		error = EINVAL;
279		goto errout;
280	}
281	spec++;
282
283	part->contents = strdup(spec);
284	if (part->contents == NULL) {
285		error = ENOMEM;
286		goto errout;
287	}
288
289	spec = part->alias;
290	sep = strchr(spec, '/');
291	if (sep != NULL) {
292		*sep++ = '\0';
293		if (strlen(part->alias) == 0 || strlen(sep) == 0) {
294			error = EINVAL;
295			goto errout;
296		}
297		part->label = strdup(sep);
298		if (part->label == NULL) {
299			error = ENOMEM;
300			goto errout;
301		}
302	}
303
304	part->index = nparts;
305	STAILQ_INSERT_TAIL(&partlist, part, link);
306	nparts++;
307	return (0);
308
309 errout:
310	if (part->alias != NULL)
311		free(part->alias);
312	free(part);
313	return (error);
314}
315
316#if defined(SPARSE_WRITE)
317ssize_t
318sparse_write(int fd, const void *ptr, size_t sz)
319{
320	const char *buf, *p;
321	off_t ofs;
322	size_t len;
323	ssize_t wr, wrsz;
324
325	buf = ptr;
326	wrsz = 0;
327	p = memchr(buf, 0, sz);
328	while (sz > 0) {
329		len = (p != NULL) ? (size_t)(p - buf) : sz;
330		if (len > 0) {
331			len = (len + secsz - 1) & ~(secsz - 1);
332			if (len > sz)
333				len = sz;
334			wr = write(fd, buf, len);
335			if (wr < 0)
336				return (-1);
337		} else {
338			while (len < sz && *p++ == '\0')
339				len++;
340			if (len < sz)
341				len &= ~(secsz - 1);
342			if (len == 0)
343				continue;
344			ofs = lseek(fd, len, SEEK_CUR);
345			if (ofs < 0)
346				return (-1);
347			wr = len;
348		}
349		buf += wr;
350		sz -= wr;
351		wrsz += wr;
352		p = memchr(buf, 0, sz);
353	}
354	return (wrsz);
355}
356#endif /* SPARSE_WRITE */
357
358void
359mkimg_chs(lba_t lba, u_int maxcyl, u_int *cylp, u_int *hdp, u_int *secp)
360{
361	u_int hd, sec;
362
363	*cylp = *hdp = *secp = ~0U;
364	if (nsecs == 1 || nheads == 1)
365		return;
366
367	sec = lba % nsecs + 1;
368	lba /= nsecs;
369	hd = lba % nheads;
370	lba /= nheads;
371	if (lba > maxcyl)
372		return;
373
374	*cylp = lba;
375	*hdp = hd;
376	*secp = sec;
377}
378
379void
380mkimg_uuid(struct uuid *uuid)
381{
382	static uint8_t gen[sizeof(struct uuid)];
383	u_int i;
384
385	if (!unit_testing) {
386		uuidgen(uuid, 1);
387		return;
388	}
389
390	for (i = 0; i < sizeof(gen); i++)
391		gen[i]++;
392	memcpy(uuid, gen, sizeof(uuid_t));
393}
394
395static int
396capacity_resize(lba_t end)
397{
398	lba_t capsz;
399
400	capsz = (capacity + secsz - 1) / secsz;
401	if (end >= capsz)
402		return (0);
403	return (image_set_size(capsz));
404}
405
406static void
407mkimg(void)
408{
409	FILE *fp;
410	struct part *part;
411	lba_t block;
412	off_t bytesize;
413	int error, fd;
414
415	/* First check partition information */
416	STAILQ_FOREACH(part, &partlist, link) {
417		error = scheme_check_part(part);
418		if (error)
419			errc(EX_DATAERR, error, "partition %d", part->index+1);
420	}
421
422	block = scheme_metadata(SCHEME_META_IMG_START, 0);
423	STAILQ_FOREACH(part, &partlist, link) {
424		block = scheme_metadata(SCHEME_META_PART_BEFORE, block);
425		if (verbose)
426			fprintf(stderr, "partition %d: starting block %llu "
427			    "... ", part->index + 1, (long long)block);
428		part->block = block;
429		switch (part->kind) {
430		case PART_KIND_SIZE:
431			if (expand_number(part->contents, &bytesize) == -1)
432				error = errno;
433			break;
434		case PART_KIND_FILE:
435			fd = open(part->contents, O_RDONLY, 0);
436			if (fd != -1) {
437				error = image_copyin(block, fd, &bytesize);
438				close(fd);
439			} else
440				error = errno;
441			break;
442		case PART_KIND_PIPE:
443			fp = popen(part->contents, "r");
444			if (fp != NULL) {
445				fd = fileno(fp);
446				error = image_copyin(block, fd, &bytesize);
447				pclose(fp);
448			} else
449				error = errno;
450			break;
451		}
452		if (error)
453			errc(EX_IOERR, error, "partition %d", part->index + 1);
454		part->size = (bytesize + secsz - 1) / secsz;
455		if (verbose) {
456			bytesize = part->size * secsz;
457			fprintf(stderr, "size %llu bytes (%llu blocks)\n",
458			     (long long)bytesize, (long long)part->size);
459		}
460		block = scheme_metadata(SCHEME_META_PART_AFTER,
461		    part->block + part->size);
462	}
463
464	block = scheme_metadata(SCHEME_META_IMG_END, block);
465	error = image_set_size(block);
466	if (!error)
467		error = capacity_resize(block);
468	if (!error)
469		error = format_resize(block);
470	if (error)
471		errc(EX_IOERR, error, "image sizing");
472	block = image_get_size();
473	ncyls = block / (nsecs * nheads);
474	error = scheme_write(block);
475	if (error)
476		errc(EX_IOERR, error, "writing metadata");
477}
478
479int
480main(int argc, char *argv[])
481{
482	int bcfd, outfd;
483	int c, error;
484
485	bcfd = -1;
486	outfd = 1;	/* Write to stdout by default */
487	while ((c = getopt_long(argc, argv, "b:c:f:o:p:s:vyH:P:S:T:",
488	    longopts, NULL)) != -1) {
489		switch (c) {
490		case 'b':	/* BOOT CODE */
491			if (bcfd != -1)
492				usage("multiple bootcode given");
493			bcfd = open(optarg, O_RDONLY, 0);
494			if (bcfd == -1)
495				err(EX_UNAVAILABLE, "%s", optarg);
496			break;
497		case 'c':	/* CAPACITY */
498			error = parse_uint64(&capacity, 1, OFF_MAX, optarg);
499			if (error)
500				errc(EX_DATAERR, error, "capacity in bytes");
501			break;
502		case 'f':	/* OUTPUT FORMAT */
503			if (format_selected() != NULL)
504				usage("multiple formats given");
505			error = format_select(optarg);
506			if (error)
507				errc(EX_DATAERR, error, "format");
508			break;
509		case 'o':	/* OUTPUT FILE */
510			if (outfd != 1)
511				usage("multiple output files given");
512			outfd = open(optarg, O_WRONLY | O_CREAT | O_TRUNC,
513			    S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
514			if (outfd == -1)
515				err(EX_CANTCREAT, "%s", optarg);
516			break;
517		case 'p':	/* PARTITION */
518			error = parse_part(optarg);
519			if (error)
520				errc(EX_DATAERR, error, "partition");
521			break;
522		case 's':	/* SCHEME */
523			if (scheme_selected() != NULL)
524				usage("multiple schemes given");
525			error = scheme_select(optarg);
526			if (error)
527				errc(EX_DATAERR, error, "scheme");
528			break;
529		case 'y':
530			unit_testing++;
531			break;
532		case 'v':
533			verbose++;
534			break;
535		case 'H':	/* GEOMETRY: HEADS */
536			error = parse_uint32(&nheads, 1, 255, optarg);
537			if (error)
538				errc(EX_DATAERR, error, "number of heads");
539			break;
540		case 'P':	/* GEOMETRY: PHYSICAL SECTOR SIZE */
541			error = parse_uint32(&blksz, 512, INT_MAX+1U, optarg);
542			if (error == 0 && !pwr_of_two(blksz))
543				error = EINVAL;
544			if (error)
545				errc(EX_DATAERR, error, "physical sector size");
546			break;
547		case 'S':	/* GEOMETRY: LOGICAL SECTOR SIZE */
548			error = parse_uint32(&secsz, 512, INT_MAX+1U, optarg);
549			if (error == 0 && !pwr_of_two(secsz))
550				error = EINVAL;
551			if (error)
552				errc(EX_DATAERR, error, "logical sector size");
553			break;
554		case 'T':	/* GEOMETRY: TRACK SIZE */
555			error = parse_uint32(&nsecs, 1, 63, optarg);
556			if (error)
557				errc(EX_DATAERR, error, "track size");
558			break;
559		case LONGOPT_FORMATS:
560			print_formats(0);
561			exit(EX_OK);
562			/*NOTREACHED*/
563		case LONGOPT_SCHEMES:
564			print_schemes(0);
565			exit(EX_OK);
566			/*NOTREACHED*/
567		case LONGOPT_VERSION:
568			print_version();
569			exit(EX_OK);
570			/*NOTREACHED*/
571		default:
572			usage("unknown option");
573		}
574	}
575
576	if (argc > optind)
577		usage("trailing arguments");
578	if (scheme_selected() == NULL && nparts > 0)
579		usage("no scheme");
580	if (nparts == 0 && capacity == 0)
581		usage("no partitions");
582
583	if (secsz > blksz) {
584		if (blksz != 0)
585			errx(EX_DATAERR, "the physical block size cannot "
586			    "be smaller than the sector size");
587		blksz = secsz;
588	}
589
590	if (secsz > scheme_max_secsz())
591		errx(EX_DATAERR, "maximum sector size supported is %u; "
592		    "size specified is %u", scheme_max_secsz(), secsz);
593
594	if (nparts > scheme_max_parts())
595		errx(EX_DATAERR, "%d partitions supported; %d given",
596		    scheme_max_parts(), nparts);
597
598	if (format_selected() == NULL)
599		format_select("raw");
600
601	if (bcfd != -1) {
602		error = scheme_bootcode(bcfd);
603		close(bcfd);
604		if (error)
605			errc(EX_DATAERR, error, "boot code");
606	}
607
608	if (verbose) {
609		fprintf(stderr, "Logical sector size: %u\n", secsz);
610		fprintf(stderr, "Physical block size: %u\n", blksz);
611		fprintf(stderr, "Sectors per track:   %u\n", nsecs);
612		fprintf(stderr, "Number of heads:     %u\n", nheads);
613		fputc('\n', stderr);
614		if (scheme_selected())
615			fprintf(stderr, "Partitioning scheme: %s\n",
616			    scheme_selected()->name);
617		fprintf(stderr, "Output file format:  %s\n",
618		    format_selected()->name);
619		fputc('\n', stderr);
620	}
621
622	error = image_init();
623	if (error)
624		errc(EX_OSERR, error, "cannot initialize");
625
626	mkimg();
627
628	if (verbose) {
629		fputc('\n', stderr);
630		fprintf(stderr, "Number of cylinders: %u\n", ncyls);
631	}
632
633	error = format_write(outfd);
634	if (error)
635		errc(EX_IOERR, error, "writing image");
636
637	return (0);
638}
639