imagezip.c 37.6 KB
Newer Older
Leigh B. Stoller's avatar
Leigh B. Stoller committed
1
2
/*
 * EMULAB-COPYRIGHT
Mike Hibler's avatar
Mike Hibler committed
3
 * Copyright (c) 2000-2004 University of Utah and the Flux Group.
Leigh B. Stoller's avatar
Leigh B. Stoller committed
4
5
6
 * All rights reserved.
 */

7
8
/*
 * An image zipper.
Mike Hibler's avatar
all:    
Mike Hibler committed
9
10
11
12
13
 *
 * TODO:
 *	Multithread so that we can be reading ahead on the input device
 *	and overlapping IO with compression.  Maybe a third thread for
 *	doing output.
14
 */
15
#include <ctype.h>
16
#include <err.h>
Mac Newbold's avatar
Mac Newbold committed
17
#include <fcntl.h>
18
19
20
#include <fstab.h>
#include <stdio.h>
#include <stdlib.h>
Mac Newbold's avatar
Mac Newbold committed
21
#include <unistd.h>
Mike Hibler's avatar
all:    
Mike Hibler committed
22
#include <string.h>
23
#include <assert.h>
Mike Hibler's avatar
all:    
Mike Hibler committed
24
#include <errno.h>
25
26
27
#include <sys/param.h>
#include <sys/time.h>
#include <sys/stat.h>
Mac Newbold's avatar
Mac Newbold committed
28
#include <zlib.h>
Mike Hibler's avatar
Mike Hibler committed
29

30
#include "imagehdr.h"
Mike Hibler's avatar
Mike Hibler committed
31
32
#include "sliceinfo.h"
#include "global.h"
33
34

#define min(a,b) ((a) <= (b) ? (a) : (b))
Mac Newbold's avatar
Mac Newbold committed
35

36
char	*infilename;
37
int	infd, outfd, outcanseek;
38
int	secsize	  = 512;	/* XXX bytes. */
39
int	debug	  = 0;
Mike Hibler's avatar
all:    
Mike Hibler committed
40
int	dots	  = 0;
41
int     info      = 0;
42
int     version   = 0;
43
int     slicemode = 0;
44
int     maxmode   = 0;
45
int     slice     = 0;
46
int	level	  = 4;
47
long	dev_bsize = 1;
48
int	oldstyle  = 0;
49
int	frangesize= 64;	/* 32k */
50
51
int	forcereads= 0;
int	retrywrites= 1;
Mike Hibler's avatar
Mike Hibler committed
52
int	dorelocs  = 1;
Mike Hibler's avatar
Mike Hibler committed
53
off_t	datawritten;
54
partmap_t ignore, forceraw;
Mike Hibler's avatar
all:    
Mike Hibler committed
55

56
57
58
59
#define HDRUSED(reg, rel) \
    (sizeof(blockhdr_t) + \
    (reg) * sizeof(struct region) + (rel) * sizeof(struct blockreloc))

60
61
62
63
64
65
66
67
68
69
70
/*
 * We want to be able to compress slices by themselves, so we need
 * to know where the slice starts when reading the input file for
 * compression. 
 *
 * These numbers are in sectors.
 */
long	inputminsec	= 0;
long    inputmaxsec	= 0;	/* 0 means the entire input image */

/*
71
 * A list of data ranges. 
72
 */
73
struct range {
Mike Hibler's avatar
all:    
Mike Hibler committed
74
75
76
	uint32_t	start;		/* In sectors */
	uint32_t	size;		/* In sectors */
	void		*data;
77
	struct range	*next;
78
};
Mike Hibler's avatar
all:    
Mike Hibler committed
79
struct range	*ranges, *skips, *fixups;
80
int		numranges, numskips;
Mike Hibler's avatar
all:    
Mike Hibler committed
81
struct blockreloc	*relocs;
82
int			numregions, numrelocs;
Mike Hibler's avatar
all:    
Mike Hibler committed
83
84

void	addskip(uint32_t start, uint32_t size);
85
void	dumpskips(int verbose);
Mike Hibler's avatar
Mike Hibler committed
86
87
void	sortrange(struct range *head, int domerge,
		  int (*rangecmp)(struct range *, struct range *));
88
void    makeranges(void);
89
void	dumpranges(int verbose);
Mike Hibler's avatar
all:    
Mike Hibler committed
90
91
92
void	addfixup(off_t offset, off_t poffset, off_t size, void *data,
		 int reloctype);
void	addreloc(off_t offset, off_t size, int reloctype);
Mike Hibler's avatar
Mike Hibler committed
93
static int cmpfixups(struct range *r1, struct range *r2);
94
95

/* Forward decls */
96
int	read_image(u_int32_t start, int pstart, u_int32_t extstart);
Mike Hibler's avatar
all:    
Mike Hibler committed
97
int	read_raw(void);
98
99
100
int	compress_image(void);
void	usage(void);

Mike Hibler's avatar
Mike Hibler committed
101
102
103
104
105
static SLICEMAP_PROCESS_PROTO(read_slice);

struct slicemap fsmap[] = {
	{ DOSPTYP_UNUSED,	"UNUSED",	0 },
#ifdef WITH_FFS
Mike Hibler's avatar
Mike Hibler committed
106
107
	{ DOSPTYP_386BSD,	"FreeBSD FFS",	read_bsdslice },
	{ DOSPTYP_OPENBSD,	"OpenBSD FFS",	read_bsdslice },
Mike Hibler's avatar
Mike Hibler committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
#endif
#ifdef WITH_EXTFS
	{ DOSPTYP_LINUX,	"Linux EXT",	read_linuxslice },
	{ DOSPTYP_LINSWP,	"Linux SWP",	read_linuxswap },
#endif
#ifdef WITH_NTFS
	{ DOSPTYP_NTFS,		"NTFS",		read_ntfsslice },
#endif
#ifdef WITH_FAT
	{ DOSPTYP_FAT12,	"FAT12",	read_fatslice },
	{ DOSPTYP_FAT16,	"FAT16",	read_fatslice },
	{ DOSPTYP_FAT16L,	"FAT16L",	read_fatslice },
	{ DOSPTYP_FAT16L_LBA,	"FAT16 LBA",	read_fatslice },
	{ DOSPTYP_FAT32,	"FAT32",	read_fatslice },
	{ DOSPTYP_FAT32_LBA,	"FAT32 LBA",	read_fatslice },
#endif
	{ DOSPTYP_EXT,		"DOSEXT",	0 },
	{ DOSPTYP_EXT_LBA,	"DOSEXT LBA",	0 },
	{ -1,			"",		0 },
};

static inline struct slicemap *
getslicemap(int stype)
{
	struct slicemap *smap;

	for (smap = fsmap; smap->type != -1; smap++)
		if (smap->type == stype)
			return smap;
	return 0;
}
139

140
141
#define IORETRIES	10

Mike Hibler's avatar
Mike Hibler committed
142
143
144
145
146
/*
 * Assert the hell out of it...
 */
off_t
devlseek(int fd, off_t off, int whence)
147
148
149
150
151
152
153
154
{
	off_t noff;
	assert((off & (DEV_BSIZE-1)) == 0);
	noff = lseek(fd, off, whence);
	assert(noff == (off_t)-1 || (noff & (DEV_BSIZE-1)) == 0);
	return noff;
}

155
156
157
158
159
160
/*
 * Wrap up read in a retry mechanism to persist in the face of IO errors,
 * even faking data if requested.
 */
ssize_t
devread(int fd, void *buf, size_t nbytes)
161
{
162
163
164
165
166
	int		cc, i, count;
	off_t		startoffset;

#ifndef linux
	assert((nbytes & (DEV_BSIZE-1)) == 0);
167
#endif
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
	if (!forcereads)
		return read(fd, buf, nbytes);

	if ((startoffset = lseek(fd, (off_t) 0, SEEK_CUR)) < 0) {
		perror("devread: seeking to get input file ptr");
		exit(1);
	}

	count = 0;
	for (i = 0; i < IORETRIES; i++) {
		while (nbytes) {
			cc = read(fd, buf, nbytes);
			if (cc == 0)
				break;

			if (cc > 0) {
				nbytes -= cc;
				buf    += cc;
				count  += cc;
				continue;
			}

			if (i == 0) 
				fprintf(stderr, "read failed: %s, "
					"will retry %d more times\n",
					strerror(errno), IORETRIES-1);
	
			nbytes += count;
			buf    -= count;
			count   = 0;
			goto again;
		}
		return count;

	again:
		if (lseek(fd, startoffset, SEEK_SET) < 0) {
			perror("devread: seeking to set file ptr");
			exit(1);
		}
	}

	fprintf(stderr, "devread: read failed in sector range [%u-%u], "
		"returning zeros\n",
		bytestosec(startoffset), bytestosec(startoffset+nbytes));
	memset(buf, 0, nbytes);
	return nbytes;
}
215

216
217
218
219
220
/*
 * Wrap up write in a retry mechanism to protect against transient NFS
 * errors causing a fatal error. 
 */
ssize_t
221
devwrite(int fd, const void *buf, size_t nbytes)
222
223
{
	int		cc, i, count = 0;
Mike Hibler's avatar
all:    
Mike Hibler committed
224
	off_t		startoffset = 0;
225

226
	if (retrywrites && outcanseek &&
227
	    ((startoffset = lseek(fd, (off_t) 0, SEEK_CUR)) < 0)) {
228
		perror("devwrite: seeking to get output file ptr");
229
230
		exit(1);
	}
231

232
	for (i = 0; i < IORETRIES; i++) {
233
		while (nbytes) {
234
235
236
237
238
239
			cc = write(fd, buf, nbytes);

			if (cc > 0) {
				nbytes -= cc;
				buf    += cc;
				count  += cc;
240
				continue;
241
			}
242
243
244

			if (!retrywrites)
				return cc;
245

246
			if (i == 0) 
247
				perror("write error: will retry");
248
	
249
			sleep(1);
250
251
252
253
			nbytes += count;
			buf    -= count;
			count   = 0;
			goto again;
254
		}
255
		if (retrywrites && fsync(fd) < 0) {
256
257
258
259
260
261
262
			perror("fsync error: will retry");
			sleep(1);
			nbytes += count;
			buf    -= count;
			count   = 0;
			goto again;
		}
Mike Hibler's avatar
Mike Hibler committed
263
		datawritten += count;
264
		return count;
265
	again:
266
		if (lseek(fd, startoffset, SEEK_SET) < 0) {
267
			perror("devwrite: seeking to set file ptr");
268
269
			exit(1);
		}
270
	}
271
272
273
	perror("write error: busted for too long");
	fflush(stderr);
	exit(1);
274
275
}

276
277
278
279
280
281
static int
setpartition(partmap_t map, char *str)
{
	int dospart;
	char bsdpart;

282
283
284
285
286
287
288
289
290
	if (isdigit(str[1])) {
		bsdpart = str[2];
		str[2] = '\0';
	} else {
		bsdpart = str[1];
		str[1] = '\0';
	}
	dospart = atoi(str);
	if (dospart < 1 || dospart > MAXSLICES)
291
292
293
294
295
296
297
298
299
300
301
302
303
304
		return EINVAL;

	/* common case: apply to complete DOS partition */
	if (bsdpart == '\0') {
		map[dospart-1] = ~0;
		return 0;
	}

	if (bsdpart < 'a' || bsdpart > 'p')
		return EINVAL;

	map[dospart-1] |= (1 << (bsdpart - 'a'));
	return 0;
}
Mac Newbold's avatar
Mac Newbold committed
305
306

int
307
308
309
main(argc, argv)
	int argc;
	char *argv[];
Mac Newbold's avatar
Mac Newbold committed
310
{
311
312
313
	int	ch, rval;
	char	*outfilename = 0;
	int	rawmode	  = 0;
Mike Hibler's avatar
Mike Hibler committed
314
	int	slicetype = 0;
315
	extern char build_info[];
Mac Newbold's avatar
Mac Newbold committed
316

Mike Hibler's avatar
Mike Hibler committed
317
	while ((ch = getopt(argc, argv, "vlbnNdihrs:c:z:oI:1F:DR:S:X")) != -1)
318
		switch(ch) {
319
320
321
		case 'v':
			version++;
			break;
322
323
324
		case 'i':
			info++;
			break;
Mike Hibler's avatar
Mike Hibler committed
325
		case 'D':
326
			retrywrites = 0;
Mike Hibler's avatar
Mike Hibler committed
327
			break;
328
329
330
331
		case 'd':
			debug++;
			break;
		case 'l':
Mike Hibler's avatar
Mike Hibler committed
332
			slicetype = DOSPTYP_LINUX;
333
334
			break;
		case 'b':
Mike Hibler's avatar
Mike Hibler committed
335
			slicetype = DOSPTYP_386BSD;
336
			break;
Mike Hibler's avatar
Mike Hibler committed
337
338
339
		case 'N':
			dorelocs = 0;
			break;
340
		case 'n':
Mike Hibler's avatar
Mike Hibler committed
341
			slicetype = DOSPTYP_NTFS;
342
			break;
Mike Hibler's avatar
all:    
Mike Hibler committed
343
344
345
		case 'o':
			dots++;
			break;
346
347
348
		case 'r':
			rawmode++;
			break;
Mike Hibler's avatar
Mike Hibler committed
349
350
351
		case 'S':
			slicetype = atoi(optarg);
			break;
352
353
354
355
		case 's':
			slicemode = 1;
			slice = atoi(optarg);
			break;
356
357
		case 'z':
			level = atoi(optarg);
Mike Hibler's avatar
all:    
Mike Hibler committed
358
			if (level < 0 || level > 9)
359
360
				usage();
			break;
361
		case 'c':
362
			maxmode     = 1;
363
364
			inputmaxsec = atoi(optarg);
			break;
Mike Hibler's avatar
all:    
Mike Hibler committed
365
		case 'I':
366
			if (setpartition(ignore, optarg))
Mike Hibler's avatar
all:    
Mike Hibler committed
367
368
				usage();
			break;
Mike Hibler's avatar
Mike Hibler committed
369
		case 'R':
370
			if (setpartition(forceraw, optarg))
Mike Hibler's avatar
Mike Hibler committed
371
372
				usage();
			break;
373
374
375
		case '1':
			oldstyle = 1;
			break;
376
377
378
379
380
		case 'F':
			frangesize = atoi(optarg);
			if (frangesize < 0)
				usage();
			break;
381
382
383
		case 'X':
			forcereads++;
			break;
384
385
386
387
388
389
390
		case 'h':
		case '?':
		default:
			usage();
		}
	argc -= optind;
	argv += optind;
Mac Newbold's avatar
Mac Newbold committed
391

392
393
	if (version || info || debug) {
		fprintf(stderr, "%s\n", build_info);
Mike Hibler's avatar
Mike Hibler committed
394
395
396
397
398
399
400
401
		if (version) {
			fprintf(stderr, "Supports");
			for (ch = 1; fsmap[ch].type != -1; ch++)
				if (fsmap[ch].process != 0)
					fprintf(stderr, "%c %s",
						ch > 1 ? ',' : ':',
						fsmap[ch].desc);
			fprintf(stderr, "\n");
402
			exit(0);
Mike Hibler's avatar
Mike Hibler committed
403
		}
404
405
	}

406
407
	if (argc < 1 || argc > 2)
		usage();
Mac Newbold's avatar
Mac Newbold committed
408

409
410
411
	if (slicemode && (slice < 1 || slice > MAXSLICES)) {
		fprintf(stderr, "Slice must be a DOS partition (1-4) "
			"or extended DOS partition (5-%d)\n\n", MAXSLICES);
412
413
		usage();
	}
414
415
416
	if (maxmode && slicemode) {
		fprintf(stderr, "Count option (-c) cannot be used with "
			"the slice (-s) option\n\n");
417
418
419
420
421
		usage();
	}
	if (!info && argc != 2) {
		fprintf(stderr, "Must specify an output filename!\n\n");
		usage();
Mac Newbold's avatar
Mac Newbold committed
422
	}
423
424
	else
		outfilename = argv[1];
Mac Newbold's avatar
Mac Newbold committed
425

426
427
428
	if (info && !debug)
		debug++;

Mike Hibler's avatar
Mike Hibler committed
429
430
431
	if (!slicemode && dorelocs)
		dorelocs = 0;

432
433
434
	infilename = argv[0];
	if ((infd = open(infilename, O_RDONLY, 0)) < 0) {
		perror(infilename);
Mac Newbold's avatar
Mac Newbold committed
435
436
		exit(1);
	}
437

Mike Hibler's avatar
Mike Hibler committed
438
439
440
441
442
	if (slicetype != 0) {
		rval = read_slice(-1, slicetype, 0, 0, infilename, infd);
		if (rval == -1)
			fprintf(stderr, ", cannot process\n");
	} else if (rawmode)
443
444
		rval = read_raw();
	else
445
		rval = read_image(DOSBBSECTOR, 0, 0);
446
447
448
449
450
	if (rval) {
		fprintf(stderr, "* * * Aborting * * *\n");
		exit(1);
	}

Mike Hibler's avatar
Mike Hibler committed
451
	sortrange(skips, 1, 0);
452
	if (debug)
453
		dumpskips(info || debug > 2);
454
	makeranges();
455
456
	if (debug)
		dumpranges(info || debug > 2);
Mike Hibler's avatar
Mike Hibler committed
457
	sortrange(fixups, 0, cmpfixups);
458
	fflush(stderr);
459
460
461
462
463
464

	if (info) {
		close(infd);
		exit(0);
	}

465
466
467
468
469
470
471
472
473
474
475
	if (strcmp(outfilename, "-")) {
		if ((outfd = open(outfilename, O_RDWR|O_CREAT|O_TRUNC, 0666))
		    < 0) {
			perror("opening output file");
			exit(1);
		}
		outcanseek = 1;
	}
	else {
		outfd = fileno(stdout);
		outcanseek = 0;
476
		retrywrites = 0;
Mac Newbold's avatar
Mac Newbold committed
477
	}
478
479
	compress_image();
	
480
	fflush(stderr);
481
	close(infd);
482
483
	if (outcanseek)
		close(outfd);
484
	exit(0);
485
486
}

Mike Hibler's avatar
Mike Hibler committed
487
488
489
490
491
492
493
494
495
496
497
498
499
500
static int
read_slice(int snum, int stype, u_int32_t start, u_int32_t size,
	   char *sname, int sfd)
{
	struct slicemap *smap = getslicemap(stype);

	if (smap && smap->process)
		return (*smap->process)(snum, stype, start, size, sname, sfd);
	
	fprintf(stderr, "Slice %d is an unknown type %#x (%s)",
		snum+1, stype, smap ? smap->desc : "??");
	return -1;
}

501
502
503
504
/*
 * Parse the DOS partition table and dispatch to the individual readers.
 */
int
505
read_image(u_int32_t bbstart, int pstart, u_int32_t extstart)
506
507
{
	int		i, cc, rval = 0;
Mike Hibler's avatar
Mike Hibler committed
508
509
	struct slicemap	*smap;
	struct doslabel doslabel;
510

511
512
513
514
	if (devlseek(infd, sectobytes(bbstart), SEEK_SET) < 0) {
		warn("Could not seek to DOS label at sector %u", bbstart);
		return 1;
	}
515
	if ((cc = devread(infd, doslabel.pad2, DOSPARTSIZE)) < 0) {
516
		warn("Could not read DOS label at sector %u", bbstart);
517
518
519
		return 1;
	}
	if (cc != DOSPARTSIZE) {
520
521
		warnx("Could not get the entire DOS label at sector %u",
		      bbstart);
522
523
524
 		return 1;
	}
	if (doslabel.magic != BOOT_MAGIC) {
525
526
		warnx("Wrong magic number in DOS partition table at sector %u",
		      bbstart);
527
528
529
530
 		return 1;
	}

	if (debug) {
531
532
533
534
535
536
		if (bbstart == 0)
			fprintf(stderr, "DOS Partitions:\n");
		else
			fprintf(stderr,
				"DOS Partitions in Extended table at %u\n",
				bbstart);
537
		for (i = 0; i < NDOSPART; i++) {
538
539
540
541
			u_int32_t start;
			int bsdix = pstart + i;

			fprintf(stderr, "  P%d: ", bsdix + 1);
Mike Hibler's avatar
Mike Hibler committed
542
543
544
545
546
547
			smap = getslicemap(doslabel.parts[i].dp_typ);
			if (smap == 0)
				fprintf(stderr, "%-10s", "UNKNOWN");
			else
				fprintf(stderr, "%-10s", smap->desc);

548
549
			start = doslabel.parts[i].dp_start;
#if 0
Mike Hibler's avatar
Mike Hibler committed
550
			/* Make start sector absolute */
551
552
553
554
555
			if (ISEXT(doslabel.parts[i].dp_typ))
				start += extstart;
			else
				start += bbstart;
#endif
556
			fprintf(stderr, "  start %9d, size %9d\n",
557
				start, doslabel.parts[i].dp_size);
558
		}
559
		fprintf(stderr, "\n");
560
561
562
563
564
565
566
	}

	/*
	 * Now operate on individual slices. 
	 */
	for (i = 0; i < NDOSPART; i++) {
		unsigned char	type  = doslabel.parts[i].dp_typ;
567
		u_int32_t	start = bbstart + doslabel.parts[i].dp_start;
568
		u_int32_t	size  = doslabel.parts[i].dp_size;
569
		int		bsdix = pstart + i;
570

571
		if (slicemode && bsdix + 1 != slice && !ISEXT(type))
572
573
			continue;
		
574
575
		if (ignore[bsdix]) {
			if (!ISBSD(type) || ignore[bsdix] == ~0)
576
				type = DOSPTYP_UNUSED;
577
578
		} else if (forceraw[bsdix]) {
			if (!ISBSD(type) || forceraw[bsdix] == ~0) {
579
580
				fprintf(stderr,
					"  Slice %d, forcing raw compression\n",
581
					bsdix + 1);
582
583
				goto skipcheck;
			}
Mike Hibler's avatar
Mike Hibler committed
584
		}
Mike Hibler's avatar
all:    
Mike Hibler committed
585

Mike Hibler's avatar
Mike Hibler committed
586
		smap = getslicemap(type);
587
		switch (type) {
588
589
590
591
592
593
594
595
596
		case DOSPTYP_EXT:
		case DOSPTYP_EXT_LBA:
			/*
			 * XXX extended partition start sectors are
			 * relative to the first extended partition found
			 */
			rval = read_image(extstart + doslabel.parts[i].dp_start,
					  pstart + NDOSPART,
					  extstart ?: start);
597
598
			/* XXX for inputmaxsec calculation below */
			start = extstart + doslabel.parts[i].dp_start;
599
600
			break;

Mike Hibler's avatar
all:    
Mike Hibler committed
601
		case DOSPTYP_UNUSED:
602
			fprintf(stderr,
603
604
				"  Slice %d %s, NOT SAVING.\n", bsdix + 1,
				ignore[bsdix] ? "ignored" : "is unused");
Mike Hibler's avatar
Mike Hibler committed
605
			if (size > 0)
606
				addskip(start, size);
Mike Hibler's avatar
all:    
Mike Hibler committed
607
			break;
Mike Hibler's avatar
Mike Hibler committed
608

Mike Hibler's avatar
all:    
Mike Hibler committed
609
		default:
Mike Hibler's avatar
Mike Hibler committed
610
611
612
613
614
615
			rval = read_slice(bsdix, type, start, size,
					  infilename, infd);
			if (rval == -1) {
				fprintf(stderr, ", forcing raw compression\n");
				rval = 0;
			}
Mike Hibler's avatar
all:    
Mike Hibler committed
616
			break;
617
		}
618
		if (rval) {
619
620
621
622
623
624
			if (!ISEXT(type))
				fprintf(stderr,
					"  Filesystem specific error "
					"in Slice %d, "
					"use -R%d to force raw compression.\n",
					bsdix + 1, bsdix + 1);
625
			break;
626
		}
627
		
Mike Hibler's avatar
Mike Hibler committed
628
	skipcheck:
629
630
631
632
633
634
635
636
637
638
639
640
641
642
		/*
		 * In slicemode, we need to set the bounds of compression.
		 * Slice is a DOS partition number (1-4). If not in slicemode,
		 * we cannot set the bounds according to the doslabel since its
		 * possible that someone will create a disk with empty space
		 * before the first partition (typical, to start partition 1
		 * at the second cylinder) or after the last partition (Mike!).
		 * However, do not set the inputminsec since we usually want the
		 * stuff before the first partition, which is the boot stuff.
		 */
		if (slicemode && slice == bsdix + 1) {
			inputminsec = start;
			inputmaxsec = start + size;
		} else if (!slicemode && !maxmode) {
643
644
645
			if (start + size > inputmaxsec)
				inputmaxsec = start + size;
		}
646
647
648
649
650
	}

	return rval;
}

651

652
653
/*
 * For a raw image (something we know nothing about), we report the size
654
 * and compress the entire thing (that is, there are no skip ranges).
655
656
657
658
659
660
 */
int
read_raw(void)
{
	off_t	size;

661
	if ((size = devlseek(infd, (off_t) 0, SEEK_END)) < 0) {
662
663
664
665
666
		warn("lseeking to end of raw image");
		return 1;
	}

	if (debug) {
667
668
		fprintf(stderr, "  Raw Image\n");
		fprintf(stderr, "        start %12d, size %12qd\n", 0, size);
669
670
671
672
673
	}
	return 0;
}

char *usagestr = 
674
 "usage: imagezip [-vihor] [-s #] <image | device> [outputfilename]\n"
Mike Hibler's avatar
all:    
Mike Hibler committed
675
 " -v             Print version info and exit\n"
676
 " -i             Info mode only.  Do not write an output file\n"
Mike Hibler's avatar
all:    
Mike Hibler committed
677
 " -h             Print this help message\n"
678
679
680
681
682
683
684
685
686
 " -o             Print progress indicating dots\n"
 " -r             Generate a `raw' image.  No FS compression is attempted\n"
 " -s slice       Compress a particular slice (DOS numbering 1-4)\n"
 " image | device The input image or a device special file (ie: /dev/ad0)\n"
 " outputfilename The output file ('-' for stdout)\n"
 "\n"
 " Advanced options\n"
 " -z level       Set the compression level.  Range 0-9 (0==none, default==4)\n"
 " -I slice       Ignore (skip) the indicated slice (not with slice mode)\n"
Mike Hibler's avatar
Mike Hibler committed
687
 " -R slice       Force raw compression of the indicated slice (not with slice mode)\n"
688
 " -c count       Compress <count> number of sectors (not with slice mode)\n"
Mike Hibler's avatar
Mike Hibler committed
689
 " -D             Do `dangerous' writes (don't check for async errors)\n"
690
 " -1             Output a version one image file\n"
691
692
 "\n"
 " Debugging options (not to be used by mere mortals!)\n"
693
 " -d             Turn on debugging.  Multiple -d options increase output\n"
694
695
 " -b             FreeBSD slice only.  Input must be a FreeBSD FFS slice\n"
 " -l             Linux slice only.  Input must be a Linux EXT2FS slice\n"
Mike Hibler's avatar
Mike Hibler committed
696
697
 " -n             NTFS slice only.  Input must be an NTFS slice\n"
 " -S DOS-ptype   Treat the input device as containing a slice of the given type\n";
698
699
700
701
702
703
704
705
706

void
usage()
{
	fprintf(stderr, usagestr);
	exit(1);
}

void
Mike Hibler's avatar
all:    
Mike Hibler committed
707
addskip(uint32_t start, uint32_t size)
708
{
709
	struct range	   *skip;
710

Mike Hibler's avatar
Mike Hibler committed
711
	if (size == 0 || size < frangesize)
712
713
		return;

714
	if ((skip = (struct range *) malloc(sizeof(*skip))) == NULL) {
715
716
717
718
		fprintf(stderr, "No memory for skip range, "
			"try again with '-F <numsect>'\n"
			"where <numsect> is greater than the current %d\n",
			frangesize);
719
720
721
		exit(1);
	}
	
722
723
724
725
726
	skip->start = start;
	skip->size  = size;
	skip->next  = skips;
	skips       = skip;
	numskips++;
727
728
}

Mike Hibler's avatar
all:    
Mike Hibler committed
729
void
730
dumpskips(int verbose)
Mike Hibler's avatar
all:    
Mike Hibler committed
731
732
{
	struct range	*pskip;
733
	uint32_t	offset = 0, total = 0;
Mike Hibler's avatar
all:    
Mike Hibler committed
734
735
736
737

	if (!skips)
		return;

738
739
740
741
742
	if (verbose) {
		fprintf(stderr, "\nMin sector %lu, Max sector %lu\n",
			inputminsec, inputmaxsec);
		fprintf(stderr, "Skip ranges (start/size) in sectors:\n");
	}
Mike Hibler's avatar
Mike Hibler committed
743

Mike Hibler's avatar
all:    
Mike Hibler committed
744
745
	pskip = skips;
	while (pskip) {
746
		if (verbose)
Mike Hibler's avatar
all:    
Mike Hibler committed
747
748
			fprintf(stderr,
				"  %12d    %9d\n", pskip->start, pskip->size);
Mike Hibler's avatar
Mike Hibler committed
749
		assert(pskip->start >= offset);
750
		offset = pskip->start + pskip->size;
Mike Hibler's avatar
all:    
Mike Hibler committed
751
752
753
754
755
756
757
758
		total += pskip->size;
		pskip  = pskip->next;
	}
	
	fprintf(stderr, "Total Number of Free Sectors: %d (bytes %qd)\n",
		total, sectobytes(total));
}

759
760
761
762
/*
 * A very dumb bubblesort!
 */
void
Mike Hibler's avatar
Mike Hibler committed
763
764
sortrange(struct range *head, int domerge,
	  int (*rangecmp)(struct range *, struct range *))
765
{
Mike Hibler's avatar
all:    
Mike Hibler committed
766
	struct range	*prange, tmp, *ptmp;
767
768
	int		changed = 1;

Mike Hibler's avatar
all:    
Mike Hibler committed
769
	if (head == NULL)
770
771
772
773
774
		return;
	
	while (changed) {
		changed = 0;

Mike Hibler's avatar
all:    
Mike Hibler committed
775
776
777
		prange = head;
		while (prange) {
			if (prange->next &&
Mike Hibler's avatar
Mike Hibler committed
778
779
			    (prange->start > prange->next->start ||
			     (rangecmp && (*rangecmp)(prange, prange->next)))) {
Mike Hibler's avatar
all:    
Mike Hibler committed
780
781
				tmp.start = prange->start;
				tmp.size  = prange->size;
Mike Hibler's avatar
Mike Hibler committed
782
				tmp.data  = prange->data;
783

Mike Hibler's avatar
all:    
Mike Hibler committed
784
785
				prange->start = prange->next->start;
				prange->size  = prange->next->size;
Mike Hibler's avatar
Mike Hibler committed
786
				prange->data  = prange->next->data;
Mike Hibler's avatar
all:    
Mike Hibler committed
787
788
				prange->next->start = tmp.start;
				prange->next->size  = tmp.size;
Mike Hibler's avatar
Mike Hibler committed
789
				prange->next->data  = tmp.data;
790
791
792

				changed = 1;
			}
Mike Hibler's avatar
all:    
Mike Hibler committed
793
			prange  = prange->next;
794
795
796
		}
	}

Mike Hibler's avatar
all:    
Mike Hibler committed
797
798
799
	if (!domerge)
		return;

800
	/*
Mike Hibler's avatar
all:    
Mike Hibler committed
801
	 * Now look for contiguous free regions and combine them.
802
	 */
Mike Hibler's avatar
all:    
Mike Hibler committed
803
804
	prange = head;
	while (prange) {
805
	again:
Mike Hibler's avatar
all:    
Mike Hibler committed
806
807
808
		if (prange->next &&
		    prange->start + prange->size == prange->next->start) {
			prange->size += prange->next->size;
809
			
Mike Hibler's avatar
all:    
Mike Hibler committed
810
811
			ptmp        = prange->next;
			prange->next = prange->next->next;
812
813
814
			free(ptmp);
			goto again;
		}
Mike Hibler's avatar
all:    
Mike Hibler committed
815
		prange  = prange->next;
816
817
818
819
	}
}

/*
820
 * Life is easier if I think in terms of the valid ranges instead of
Mike Hibler's avatar
all:    
Mike Hibler committed
821
822
 * the free ranges. So, convert them.  Note that if there were no skips,
 * we create a single range covering the entire partition.
823
 */
824
825
void
makeranges(void)
826
{
Mike Hibler's avatar
all:    
Mike Hibler committed
827
828
	struct range	*pskip, *ptmp, *range, **lastrange;
	uint32_t	offset;
829
	
Mike Hibler's avatar
all:    
Mike Hibler committed
830
831
	offset = inputminsec;
	lastrange = &ranges;
832

833
834
835
836
837
838
839
840
841
842
843
844
	pskip = skips;
	while (pskip) {
		if ((range = (struct range *)
		             malloc(sizeof(*range))) == NULL) {
			fprintf(stderr, "Out of memory!\n");
			exit(1);
		}
		range->start = offset;
		range->size  = pskip->start - offset;
		range->next  = 0;
		offset       = pskip->start + pskip->size;
		
Mike Hibler's avatar
all:    
Mike Hibler committed
845
846
		*lastrange = range;
		lastrange = &range->next;
847
848
849
850
851
852
		numranges++;

		ptmp  = pskip;
		pskip = pskip->next;
		free(ptmp);
	}
853
	/*
854
	 * Last piece, but only if there is something to compress.
855
	 */
Mike Hibler's avatar
all:    
Mike Hibler committed
856
	if (inputmaxsec == 0 || (inputmaxsec - offset) != 0) {
857
		assert(inputmaxsec == 0 || inputmaxsec > offset);
Mike Hibler's avatar
all:    
Mike Hibler committed
858
		if ((range = (struct range *)malloc(sizeof(*range))) == NULL) {
859
860
861
862
			fprintf(stderr, "Out of memory!\n");
			exit(1);
		}
		range->start = offset;
863
	
864
865
866
867
868
869
870
871
872
873
		/*
		 * A bug in FreeBSD causes lseek on a device special file to
		 * return 0 all the time! Well we want to be able to read
		 * directly out of a raw disk (/dev/rad0), so we need to
		 * use the compressor to figure out the actual size when it
		 * isn't known beforehand.
		 *
		 * Mark the last range with 0 so compression goes to end
		 * if we don't know where it is.
		 */
Mike Hibler's avatar
all:    
Mike Hibler committed
874
		if (inputmaxsec)
875
876
877
			range->size = inputmaxsec - offset;
		else
			range->size = 0;
878
879
		range->next = 0;

Mike Hibler's avatar
all:    
Mike Hibler committed
880
		*lastrange = range;
881
882
		numranges++;
	}
883
}
884

885
886
887
888
889
890
void
dumpranges(int verbose)
{
	struct range *range;
	uint32_t total = 0;

Mike Hibler's avatar
Mike Hibler committed
891
	if (verbose)
892
		fprintf(stderr, "\nAllocated ranges (start/size) in sectors:\n");
Mike Hibler's avatar
Mike Hibler committed
893
894
895
	range = ranges;
	while (range) {
		if (verbose)
896
897
			fprintf(stderr, "  %12d    %9d\n",
				range->start, range->size);
Mike Hibler's avatar
Mike Hibler committed
898
899
		total += range->size;
		range = range->next;
900
	}
901
902
903
	fprintf(stderr,
		"Total Number of Valid Sectors: %d (bytes %qd)\n",
		total, sectobytes(total));
904
905
906
}

/*
Mike Hibler's avatar
all:    
Mike Hibler committed
907
908
909
 * Fixup descriptor handling.
 *
 * Fixups are modifications that need to be made to file data prior
Mike Hibler's avatar
Mike Hibler committed
910
 * to compressing.
911
 */
Mike Hibler's avatar
all:    
Mike Hibler committed
912
913
914
915
916
917
918
struct fixup {
	off_t offset;	/* disk offset */
	off_t poffset;	/* partition offset */
	off_t size;
	int reloctype;
	char data[0];
};
919

Mike Hibler's avatar
all:    
Mike Hibler committed
920
921
void
addfixup(off_t offset, off_t poffset, off_t size, void *data, int reloctype)
922
{
Mike Hibler's avatar
all:    
Mike Hibler committed
923
924
	struct range *entry;
	struct fixup *fixup;
925

926
927
928
929
930
931
932
933
934
935
	if (oldstyle) {
		static int warned;

		if (!warned) {
			fprintf(stderr, "WARNING: no fixups in V1 images\n");
			warned = 1;
		}
		return;
	}

Mike Hibler's avatar
all:    
Mike Hibler committed
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
	if ((entry = malloc(sizeof(*entry))) == NULL ||
	    (fixup = malloc(sizeof(*fixup) + (int)size)) == NULL) {
		fprintf(stderr, "Out of memory!\n");
		exit(1);
	}
	
	entry->start = bytestosec(offset);
	entry->size  = bytestosec(size + secsize - 1);
	entry->data  = fixup;
	
	fixup->offset    = offset;
	fixup->poffset   = poffset;
	fixup->size      = size;
	fixup->reloctype = reloctype;
	memcpy(fixup->data, data, size);

	entry->next  = fixups;
	fixups       = entry;
}
955

Mike Hibler's avatar
Mike Hibler committed
956
957
958
959
960
961
962
963
964
965
966
967
968
969
/*
 * Return 1 if r1 > r2
 */
static int
cmpfixups(struct range *r1, struct range *r2)
{
	if (r1->start > r2->start ||
	    (r1->start == r2->start &&
	     ((struct fixup *)r1->data)->offset >
	     ((struct fixup *)r2->data)->offset))
		return 1;
	return 0;
}

Mike Hibler's avatar
all:    
Mike Hibler committed
970
971
972
973
974
975
976
void
applyfixups(off_t offset, off_t size, void *data)
{
	struct range **prev, *entry;
	struct fixup *fp;
	uint32_t coff, clen;

Mike Hibler's avatar
Mike Hibler committed
977
978
	prev = &fixups;
	while ((entry = *prev) != NULL) {
Mike Hibler's avatar
all:    
Mike Hibler committed
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
		fp = entry->data;

		if (offset < fp->offset+fp->size && offset+size > fp->offset) {
			/* XXX lazy: fixup must be totally contained */
			assert(offset <= fp->offset);
			assert(fp->offset+fp->size <= offset+size);

			coff = (u_int32_t)(fp->offset - offset);
			clen = (u_int32_t)fp->size;
			if (debug > 1)
				fprintf(stderr,
					"Applying fixup [%qu-%qu] "
					"to [%qu-%qu]\n",
					fp->offset, fp->offset+fp->size,
					offset, offset+size);
			memcpy(data+coff, fp->data, clen);

			/* create a reloc if necessary */
			if (fp->reloctype != RELOC_NONE)
				addreloc(fp->offset - fp->poffset,
					 fp->size, fp->reloctype);

			*prev = entry->next;
			free(fp);
			free(entry);
Mike Hibler's avatar
Mike Hibler committed
1004
1005
		} else
			prev = &entry->next;
Mike Hibler's avatar
all:    
Mike Hibler committed
1006
1007
	}
}
Mac Newbold's avatar
Mac Newbold committed
1008

Mike Hibler's avatar
all:    
Mike Hibler committed
1009
1010
1011
1012
void
addreloc(off_t offset, off_t size, int reloctype)
{
	struct blockreloc *reloc;
1013

1014
1015
	assert(!oldstyle);

Mike Hibler's avatar
all:    
Mike Hibler committed
1016
	numrelocs++;
1017
1018
1019
1020
1021
1022
	if (HDRUSED(numregions, numrelocs) > DEFAULTREGIONSIZE) {
		fprintf(stderr, "Over filled region/reloc table (%d/%d)\n",
			numregions, numrelocs);
		exit(1);
	}

Mike Hibler's avatar
all:    
Mike Hibler committed
1023
1024
1025
1026
1027
	relocs = realloc(relocs, numrelocs * sizeof(struct blockreloc));
	if (relocs == NULL) {
		fprintf(stderr, "Out of memory!\n");
		exit(1);
	}
Mac Newbold's avatar
Mac Newbold committed
1028

Mike Hibler's avatar
all:    
Mike Hibler committed
1029
1030
1031
1032
1033
1034
	reloc = &relocs[numrelocs-1];
	reloc->type = reloctype;
	reloc->sector = bytestosec(offset);
	reloc->sectoff = offset - sectobytes(reloc->sector);
	reloc->size = size;
}
1035

Mike Hibler's avatar
all:    
Mike Hibler committed
1036
1037
1038
1039
1040
1041
1042
void
freerelocs(void)
{
	numrelocs = 0;
	free(relocs);
	relocs = NULL;
}
Mac Newbold's avatar
Mac Newbold committed
1043

Mike Hibler's avatar
all:    
Mike Hibler committed
1044
1045
1046
1047
1048
1049
1050
/*
 * Compress the image.
 */
static u_char   output_buffer[SUBBLOCKSIZE];
static int	buffer_offset;
static off_t	inputoffset;
static struct timeval cstamp;
Mike Hibler's avatar
Mike Hibler committed
1051
static long long bytescompressed;
1052

Mike Hibler's avatar
all:    
Mike Hibler committed
1053
1054
static off_t	compress_chunk(off_t, off_t, int *, uint32_t *);
static int	compress_finish(uint32_t *subblksize);
Mike Hibler's avatar
Mike Hibler committed
1055
static void	compress_status(int sig);
1056

Mike Hibler's avatar
all:    
Mike Hibler committed
1057
1058
1059
1060
1061
1062
/*
 * Loop through the image, compressing the allocated ranges.
 */
int
compress_image(void)
{
1063
	int		cc, full, i, count, chunkno;
Mike Hibler's avatar
all:    
Mike Hibler committed
1064
1065
1066
1067
1068
1069
	off_t		size = 0, outputoffset;
	off_t		tmpoffset, rangesize;
	struct range	*prange;
	blockhdr_t	*blkhdr;
	struct region	*curregion, *regions;
	struct timeval  estamp;
Mike Hibler's avatar
Mike Hibler committed
1070
	char		*buf;
Mike Hibler's avatar
all:    
Mike Hibler committed
1071
1072
	uint32_t	cursect = 0;
	struct region	*lreg;
1073

Mike Hibler's avatar
all:    
Mike Hibler committed
1074
1075
1076
1077
1078
	gettimeofday(&cstamp, 0);
	inputoffset = 0;
#ifdef SIGINFO
	signal(SIGINFO, compress_status);
#endif
1079

Mike Hibler's avatar
Mike Hibler committed
1080
1081
	buf = output_buffer;
	memset(buf, 0, DEFAULTREGIONSIZE);
1082
1083
1084
1085
1086
	blkhdr = (blockhdr_t *) buf;
	if (oldstyle)
		regions = (struct region *)((struct blockhdr_V1 *)blkhdr + 1);
	else
		regions = (struct region *)(blkhdr + 1);
Mike Hibler's avatar
all: <