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