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/param.h>
28#include <sys/stat.h>
29#include <errno.h>
30#include <err.h>
31#include <fcntl.h>
32#include <getopt.h>
33#include <libutil.h>
34#include <limits.h>
35#include <stdbool.h>
36#include <stdint.h>
37#include <stdio.h>
38#include <stdlib.h>
39#include <string.h>
40#include <sysexits.h>
41#include <unistd.h>
42
43#include "image.h"
44#include "format.h"
45#include "mkimg.h"
46#include "scheme.h"
47
48#define	LONGOPT_FORMATS		0x01000001
49#define	LONGOPT_SCHEMES		0x01000002
50#define	LONGOPT_VERSION		0x01000003
51#define	LONGOPT_CAPACITY	0x01000004
52
53static struct option longopts[] = {
54	{ "formats", no_argument, NULL, LONGOPT_FORMATS },
55	{ "schemes", no_argument, NULL, LONGOPT_SCHEMES },
56	{ "version", no_argument, NULL, LONGOPT_VERSION },
57	{ "capacity", required_argument, NULL, LONGOPT_CAPACITY },
58	{ NULL, 0, NULL, 0 }
59};
60
61static uint64_t min_capacity = 0;
62static uint64_t max_capacity = 0;
63
64struct partlisthead partlist = TAILQ_HEAD_INITIALIZER(partlist);
65u_int nparts = 0;
66
67u_int unit_testing;
68u_int verbose;
69
70u_int ncyls = 0;
71u_int nheads = 1;
72u_int nsecs = 1;
73u_int secsz = 512;
74u_int blksz = 0;
75uint32_t active_partition = 0;
76
77static void
78print_formats(int usage)
79{
80	struct mkimg_format *f;
81	const char *sep;
82
83	if (usage) {
84		fprintf(stderr, "    formats:\n");
85		f = NULL;
86		while ((f = format_iterate(f)) != NULL) {
87			fprintf(stderr, "\t%s\t-  %s\n", f->name,
88			    f->description);
89		}
90	} else {
91		sep = "";
92		f = NULL;
93		while ((f = format_iterate(f)) != NULL) {
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;
105	const char *sep;
106
107	if (usage) {
108		fprintf(stderr, "    schemes:\n");
109		s = NULL;
110		while ((s = scheme_iterate(s)) != NULL) {
111			fprintf(stderr, "\t%s\t-  %s\n", s->name,
112			    s->description);
113		}
114	} else {
115		sep = "";
116		s = NULL;
117		while ((s = scheme_iterate(s)) != NULL) {
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-a <num>\t-  mark num'th partition as active\n");
152	fprintf(stderr, "\t-b <file>\t-  file containing boot code\n");
153	fprintf(stderr, "\t-c <num>\t-  minimum capacity (in bytes) of the disk\n");
154	fprintf(stderr, "\t-C <num>\t-  maximum capacity (in bytes) of the disk\n");
155	fprintf(stderr, "\t-f <format>\n");
156	fprintf(stderr, "\t-o <file>\t-  file to write image into\n");
157	fprintf(stderr, "\t-p <partition>\n");
158	fprintf(stderr, "\t-s <scheme>\n");
159	fprintf(stderr, "\t-v\t\t-  increase verbosity\n");
160	fprintf(stderr, "\t-y\t\t-  [developers] enable unit test\n");
161	fprintf(stderr, "\t-H <num>\t-  number of heads to simulate\n");
162	fprintf(stderr, "\t-P <num>\t-  physical sector size\n");
163	fprintf(stderr, "\t-S <num>\t-  logical sector size\n");
164	fprintf(stderr, "\t-T <num>\t-  number of tracks to simulate\n");
165	fputc('\n', stderr);
166	print_formats(1);
167	fputc('\n', stderr);
168	print_schemes(1);
169	fputc('\n', stderr);
170	fprintf(stderr, "    partition specification:\n");
171	fprintf(stderr, "\t<t>[/<l>]::<size>[:[+]<offset>]\t-  "
172	    "empty partition of given size and\n\t\t\t\t\t"
173	    "   optional relative or absolute offset\n");
174	fprintf(stderr, "\t<t>[/<l>]:=<file>[:[+]offset]\t-  partition "
175	    "content and size are\n\t\t\t\t\t"
176	    "   determined by the named file and\n"
177	    "\t\t\t\t\t   optional relative or absolute offset\n");
178	fprintf(stderr, "\t<t>[/<l>]:-<cmd>\t\t-  partition content and size "
179	    "are taken\n\t\t\t\t\t   from the output of the command to run\n");
180	fprintf(stderr, "\t-\t\t\t\t-  unused partition entry\n");
181	fprintf(stderr, "\t    where:\n");
182	fprintf(stderr, "\t\t<t>\t-  scheme neutral partition type\n");
183	fprintf(stderr, "\t\t<l>\t-  optional scheme-dependent partition "
184	    "label\n");
185
186	exit(EX_USAGE);
187}
188
189static int
190parse_uint32(uint32_t *valp, uint32_t min, uint32_t max, const char *arg)
191{
192	uint64_t val;
193
194	if (expand_number(arg, &val) == -1)
195		return (errno);
196	if (val > UINT_MAX || val < (uint64_t)min || val > (uint64_t)max)
197		return (EINVAL);
198	*valp = (uint32_t)val;
199	return (0);
200}
201
202static int
203parse_uint64(uint64_t *valp, uint64_t min, uint64_t max, const char *arg)
204{
205	uint64_t val;
206
207	if (expand_number(arg, &val) == -1)
208		return (errno);
209	if (val < min || val > max)
210		return (EINVAL);
211	*valp = val;
212	return (0);
213}
214
215static int
216pwr_of_two(u_int nr)
217{
218
219	return (((nr & (nr - 1)) == 0) ? 1 : 0);
220}
221
222/*
223 * A partition specification has the following format:
224 *	<type> ':' <kind> <contents>
225 * where:
226 *	type	  the partition type alias
227 *	kind	  the interpretation of the contents specification
228 *		  ':'   contents holds the size of an empty partition
229 *		  '='   contents holds the name of a file to read
230 *		  '-'   contents holds a command to run; the output of
231 *			which is the contents of the partition.
232 *	contents  the specification of a partition's contents
233 *
234 * A specification that is a single dash indicates an unused partition
235 * entry.
236 */
237static int
238parse_part(const char *spec)
239{
240	struct part *part;
241	char *sep;
242	size_t len;
243	int error;
244
245	if (strcmp(spec, "-") == 0) {
246		nparts++;
247		return (0);
248	}
249
250	part = calloc(1, sizeof(struct part));
251	if (part == NULL)
252		return (ENOMEM);
253
254	sep = strchr(spec, ':');
255	if (sep == NULL) {
256		error = EINVAL;
257		goto errout;
258	}
259	len = sep - spec + 1;
260	if (len < 2) {
261		error = EINVAL;
262		goto errout;
263	}
264	part->alias = malloc(len);
265	if (part->alias == NULL) {
266		error = ENOMEM;
267		goto errout;
268	}
269	strlcpy(part->alias, spec, len);
270	spec = sep + 1;
271
272	switch (*spec) {
273	case ':':
274		part->kind = PART_KIND_SIZE;
275		break;
276	case '=':
277		part->kind = PART_KIND_FILE;
278		break;
279	case '-':
280		part->kind = PART_KIND_PIPE;
281		break;
282	default:
283		error = EINVAL;
284		goto errout;
285	}
286	spec++;
287
288	part->contents = strdup(spec);
289	if (part->contents == NULL) {
290		error = ENOMEM;
291		goto errout;
292	}
293
294	spec = part->alias;
295	sep = strchr(spec, '/');
296	if (sep != NULL) {
297		*sep++ = '\0';
298		if (strlen(part->alias) == 0 || strlen(sep) == 0) {
299			error = EINVAL;
300			goto errout;
301		}
302		part->label = strdup(sep);
303		if (part->label == NULL) {
304			error = ENOMEM;
305			goto errout;
306		}
307	}
308
309	part->index = nparts;
310	TAILQ_INSERT_TAIL(&partlist, part, link);
311	nparts++;
312	return (0);
313
314 errout:
315	if (part->alias != NULL)
316		free(part->alias);
317	free(part);
318	return (error);
319}
320
321#if defined(SPARSE_WRITE)
322ssize_t
323sparse_write(int fd, const void *ptr, size_t sz)
324{
325	const char *buf, *p;
326	off_t ofs;
327	size_t len;
328	ssize_t wr, wrsz;
329
330	buf = ptr;
331	wrsz = 0;
332	p = memchr(buf, 0, sz);
333	while (sz > 0) {
334		len = (p != NULL) ? (size_t)(p - buf) : sz;
335		if (len > 0) {
336			len = (len + secsz - 1) & ~(secsz - 1);
337			if (len > sz)
338				len = sz;
339			wr = write(fd, buf, len);
340			if (wr < 0)
341				return (-1);
342		} else {
343			while (len < sz && *p++ == '\0')
344				len++;
345			if (len < sz)
346				len &= ~(secsz - 1);
347			if (len == 0)
348				continue;
349			ofs = lseek(fd, len, SEEK_CUR);
350			if (ofs < 0)
351				return (-1);
352			wr = len;
353		}
354		buf += wr;
355		sz -= wr;
356		wrsz += wr;
357		p = memchr(buf, 0, sz);
358	}
359	return (wrsz);
360}
361#endif /* SPARSE_WRITE */
362
363void
364mkimg_chs(lba_t lba, u_int maxcyl, u_int *cylp, u_int *hdp, u_int *secp)
365{
366	u_int hd, sec;
367
368	*cylp = *hdp = *secp = ~0U;
369	if (nsecs == 1 || nheads == 1)
370		return;
371
372	sec = lba % nsecs + 1;
373	lba /= nsecs;
374	hd = lba % nheads;
375	lba /= nheads;
376	if (lba > maxcyl)
377		return;
378
379	*cylp = lba;
380	*hdp = hd;
381	*secp = sec;
382}
383
384static int
385capacity_resize(lba_t end)
386{
387	lba_t min_capsz, max_capsz;
388
389	min_capsz = (min_capacity + secsz - 1) / secsz;
390	max_capsz = (max_capacity + secsz - 1) / secsz;
391
392	if (max_capsz != 0 && end > max_capsz)
393		return (ENOSPC);
394	if (end >= min_capsz)
395		return (0);
396
397	return (image_set_size(min_capsz));
398}
399
400static void
401mkimg_validate(void)
402{
403	struct part *part, *part2;
404	lba_t start, end, start2, end2;
405	int i, j;
406
407	i = 0;
408
409	TAILQ_FOREACH(part, &partlist, link) {
410		start = part->block;
411		end = part->block + part->size;
412		j = i + 1;
413		part2 = TAILQ_NEXT(part, link);
414		if (part2 == NULL)
415			break;
416
417		TAILQ_FOREACH_FROM(part2, &partlist, link) {
418			start2 = part2->block;
419			end2 = part2->block + part2->size;
420
421			if ((start >= start2 && start < end2) ||
422			    (end > start2 && end <= end2)) {
423				errx(1, "partition %d overlaps partition %d",
424				    i, j);
425			}
426
427			j++;
428		}
429
430		i++;
431	}
432}
433
434static void
435mkimg(void)
436{
437	FILE *fp;
438	struct part *part;
439	lba_t block, blkoffset;
440	uint64_t bytesize, byteoffset;
441	char *size, *offset;
442	bool abs_offset;
443	int error, fd;
444
445	/* First check partition information */
446	TAILQ_FOREACH(part, &partlist, link) {
447		error = scheme_check_part(part);
448		if (error)
449			errc(EX_DATAERR, error, "partition %d", part->index+1);
450	}
451
452	block = scheme_metadata(SCHEME_META_IMG_START, 0);
453	abs_offset = false;
454	TAILQ_FOREACH(part, &partlist, link) {
455		byteoffset = blkoffset = 0;
456		abs_offset = false;
457
458		/* Look for an offset. Set size too if we can. */
459		switch (part->kind) {
460		case PART_KIND_SIZE:
461		case PART_KIND_FILE:
462			offset = part->contents;
463			size = strsep(&offset, ":");
464			if (part->kind == PART_KIND_SIZE &&
465			    expand_number(size, &bytesize) == -1)
466				error = errno;
467			if (offset != NULL) {
468				if (*offset != '+')
469					abs_offset = true;
470				else
471					offset++;
472				if (expand_number(offset, &byteoffset) == -1)
473					error = errno;
474			}
475			break;
476		}
477
478		/* Work out exactly where the partition starts. */
479		blkoffset = (byteoffset + secsz - 1) / secsz;
480		if (abs_offset)
481			block = scheme_metadata(SCHEME_META_PART_ABSOLUTE,
482			    blkoffset);
483		else
484			block = scheme_metadata(SCHEME_META_PART_BEFORE,
485			    block + blkoffset);
486		part->block = block;
487
488		if (verbose)
489			fprintf(stderr, "partition %d: starting block %llu "
490			    "... ", part->index + 1, (long long)part->block);
491
492		/* Pull in partition contents, set size if we haven't yet. */
493		switch (part->kind) {
494		case PART_KIND_FILE:
495			fd = open(part->contents, O_RDONLY, 0);
496			if (fd != -1) {
497				error = image_copyin(block, fd, &bytesize);
498				close(fd);
499			} else
500				error = errno;
501			break;
502		case PART_KIND_PIPE:
503			fp = popen(part->contents, "r");
504			if (fp != NULL) {
505				fd = fileno(fp);
506				error = image_copyin(block, fd, &bytesize);
507				pclose(fp);
508			} else
509				error = errno;
510			break;
511		}
512		if (error)
513			errc(EX_IOERR, error, "partition %d", part->index + 1);
514		part->size = (bytesize + secsz - 1) / secsz;
515		if (verbose) {
516			bytesize = part->size * secsz;
517			fprintf(stderr, "size %llu bytes (%llu blocks)\n",
518			     (long long)bytesize, (long long)part->size);
519			if (abs_offset) {
520				fprintf(stderr,
521				    "    location %llu bytes (%llu blocks)\n",
522				    (long long)byteoffset,
523				    (long long)blkoffset);
524			} else if (blkoffset > 0) {
525				fprintf(stderr,
526				    "    offset %llu bytes (%llu blocks)\n",
527				    (long long)byteoffset,
528				    (long long)blkoffset);
529			}
530		}
531		block = scheme_metadata(SCHEME_META_PART_AFTER,
532		    part->block + part->size);
533	}
534
535	mkimg_validate();
536
537	block = scheme_metadata(SCHEME_META_IMG_END, block);
538	error = image_set_size(block);
539	if (!error) {
540		error = capacity_resize(block);
541		block = image_get_size();
542	}
543	if (!error) {
544		error = format_resize(block);
545		block = image_get_size();
546	}
547	if (error)
548		errc(EX_IOERR, error, "image sizing");
549	ncyls = block / (nsecs * nheads);
550	error = scheme_write(block);
551	if (error)
552		errc(EX_IOERR, error, "writing metadata");
553}
554
555int
556main(int argc, char *argv[])
557{
558	const char *format_name;
559	int bcfd, outfd;
560	int c, error;
561
562	bcfd = -1;
563	outfd = 1;	/* Write to stdout by default */
564	while ((c = getopt_long(argc, argv, "a:b:c:C:f:o:p:s:vyH:P:S:T:",
565	    longopts, NULL)) != -1) {
566		switch (c) {
567		case 'a':	/* ACTIVE PARTITION, if supported */
568			error = parse_uint32(&active_partition, 1, 100, optarg);
569			if (error)
570				errc(EX_DATAERR, error, "Partition ordinal");
571			break;
572		case 'b':	/* BOOT CODE */
573			if (bcfd != -1)
574				usage("multiple bootcode given");
575			bcfd = open(optarg, O_RDONLY, 0);
576			if (bcfd == -1)
577				err(EX_UNAVAILABLE, "%s", optarg);
578			break;
579		case 'c':	/* MINIMUM CAPACITY */
580			error = parse_uint64(&min_capacity, 1, INT64_MAX, optarg);
581			if (error)
582				errc(EX_DATAERR, error, "minimum capacity in bytes");
583			break;
584		case 'C':	/* MAXIMUM CAPACITY */
585			error = parse_uint64(&max_capacity, 1, INT64_MAX, optarg);
586			if (error)
587				errc(EX_DATAERR, error, "maximum capacity in bytes");
588			break;
589		case 'f':	/* OUTPUT FORMAT */
590			if (format_selected() != NULL)
591				usage("multiple formats given");
592			error = format_select(optarg);
593			if (error)
594				errc(EX_DATAERR, error, "format");
595			break;
596		case 'o':	/* OUTPUT FILE */
597			if (outfd != 1)
598				usage("multiple output files given");
599			outfd = open(optarg, O_WRONLY | O_CREAT | O_TRUNC,
600			    S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
601			if (outfd == -1)
602				err(EX_CANTCREAT, "%s", optarg);
603			break;
604		case 'p':	/* PARTITION */
605			error = parse_part(optarg);
606			if (error)
607				errc(EX_DATAERR, error, "partition");
608			break;
609		case 's':	/* SCHEME */
610			if (scheme_selected() != NULL)
611				usage("multiple schemes given");
612			error = scheme_select(optarg);
613			if (error)
614				errc(EX_DATAERR, error, "scheme");
615			break;
616		case 'y':
617			unit_testing++;
618			break;
619		case 'v':
620			verbose++;
621			break;
622		case 'H':	/* GEOMETRY: HEADS */
623			error = parse_uint32(&nheads, 1, 255, optarg);
624			if (error)
625				errc(EX_DATAERR, error, "number of heads");
626			break;
627		case 'P':	/* GEOMETRY: PHYSICAL SECTOR SIZE */
628			error = parse_uint32(&blksz, 512, INT_MAX+1U, optarg);
629			if (error == 0 && !pwr_of_two(blksz))
630				error = EINVAL;
631			if (error)
632				errc(EX_DATAERR, error, "physical sector size");
633			break;
634		case 'S':	/* GEOMETRY: LOGICAL SECTOR SIZE */
635			error = parse_uint32(&secsz, 512, INT_MAX+1U, optarg);
636			if (error == 0 && !pwr_of_two(secsz))
637				error = EINVAL;
638			if (error)
639				errc(EX_DATAERR, error, "logical sector size");
640			break;
641		case 'T':	/* GEOMETRY: TRACK SIZE */
642			error = parse_uint32(&nsecs, 1, 63, optarg);
643			if (error)
644				errc(EX_DATAERR, error, "track size");
645			break;
646		case LONGOPT_FORMATS:
647			print_formats(0);
648			exit(EX_OK);
649			/*NOTREACHED*/
650		case LONGOPT_SCHEMES:
651			print_schemes(0);
652			exit(EX_OK);
653			/*NOTREACHED*/
654		case LONGOPT_VERSION:
655			print_version();
656			exit(EX_OK);
657			/*NOTREACHED*/
658		case LONGOPT_CAPACITY:
659			error = parse_uint64(&min_capacity, 1, INT64_MAX, optarg);
660			if (error)
661				errc(EX_DATAERR, error, "capacity in bytes");
662			max_capacity = min_capacity;
663			break;
664		default:
665			usage("unknown option");
666		}
667	}
668
669	if (argc > optind)
670		usage("trailing arguments");
671	if (scheme_selected() == NULL && nparts > 0)
672		usage("no scheme");
673	if (nparts == 0 && min_capacity == 0)
674		usage("no partitions");
675	if (max_capacity != 0 && min_capacity > max_capacity)
676		usage("minimum capacity cannot be larger than the maximum one");
677
678	if (secsz > blksz) {
679		if (blksz != 0)
680			errx(EX_DATAERR, "the physical block size cannot "
681			    "be smaller than the sector size");
682		blksz = secsz;
683	}
684
685	if (secsz > scheme_max_secsz())
686		errx(EX_DATAERR, "maximum sector size supported is %u; "
687		    "size specified is %u", scheme_max_secsz(), secsz);
688
689	if (nparts > scheme_max_parts())
690		errx(EX_DATAERR, "%d partitions supported; %d given",
691		    scheme_max_parts(), nparts);
692
693	if (format_selected() == NULL)
694		format_select("raw");
695
696	if (bcfd != -1) {
697		error = scheme_bootcode(bcfd);
698		close(bcfd);
699		if (error)
700			errc(EX_DATAERR, error, "boot code");
701	}
702
703	format_name = format_selected()->name;
704	if (verbose) {
705		fprintf(stderr, "Logical sector size: %u\n", secsz);
706		fprintf(stderr, "Physical block size: %u\n", blksz);
707		fprintf(stderr, "Sectors per track:   %u\n", nsecs);
708		fprintf(stderr, "Number of heads:     %u\n", nheads);
709		fputc('\n', stderr);
710		if (scheme_selected())
711			fprintf(stderr, "Partitioning scheme: %s\n",
712			    scheme_selected()->name);
713		fprintf(stderr, "Output file format:  %s\n",
714		    format_name);
715		fputc('\n', stderr);
716	}
717
718#if defined(SPARSE_WRITE)
719	/*
720	 * sparse_write() fails if output is not seekable so fail early
721	 * not wasting some load unless output format is raw
722	 */
723	if (strcmp("raw", format_name) &&
724	    lseek(outfd, (off_t)0, SEEK_CUR) == -1 && errno == ESPIPE)
725		errx(EX_USAGE, "%s: output must be seekable", format_name);
726#endif
727
728	error = image_init();
729	if (error)
730		errc(EX_OSERR, error, "cannot initialize");
731
732	mkimg();
733
734	if (verbose) {
735		fputc('\n', stderr);
736		fprintf(stderr, "Number of cylinders: %u\n", ncyls);
737	}
738
739	error = format_write(outfd);
740	if (error)
741		errc(EX_IOERR, error, "writing image");
742
743	return (0);
744}
745