recoverdisk.c revision 227081 
1/*-
2 * ----------------------------------------------------------------------------
3 * "THE BEER-WARE LICENSE" (Revision 42):
4 * <phk@FreeBSD.ORG> wrote this file.  As long as you retain this notice you
5 * can do whatever you want with this stuff. If we meet some day, and you think
6 * this stuff is worth it, you can buy me a beer in return.   Poul-Henning Kamp
7 * ----------------------------------------------------------------------------
8 *
9 * $FreeBSD: head/sbin/recoverdisk/recoverdisk.c 227081 2011-11-04 13:36:02Z ed $
10 */
11#include <sys/param.h>
12#include <sys/queue.h>
13#include <sys/disk.h>
14#include <sys/stat.h>
15
16#include <err.h>
17#include <errno.h>
18#include <fcntl.h>
19#include <signal.h>
20#include <stdint.h>
21#include <stdio.h>
22#include <stdlib.h>
23#include <string.h>
24#include <time.h>
25#include <unistd.h>
26
27static volatile sig_atomic_t aborting = 0;
28static size_t bigsize = 1024 * 1024;
29static size_t medsize;
30static size_t minsize = 512;
31
32struct lump {
33	off_t			start;
34	off_t			len;
35	int			state;
36	TAILQ_ENTRY(lump)	list;
37};
38
39static TAILQ_HEAD(, lump) lumps = TAILQ_HEAD_INITIALIZER(lumps);
40
41static void
42new_lump(off_t start, off_t len, int state)
43{
44	struct lump *lp;
45
46	lp = malloc(sizeof *lp);
47	if (lp == NULL)
48		err(1, "Malloc failed");
49	lp->start = start;
50	lp->len = len;
51	lp->state = state;
52	TAILQ_INSERT_TAIL(&lumps, lp, list);
53}
54
55static struct lump *lp;
56static char *wworklist = NULL;
57static char *rworklist = NULL;
58
59
60#define PRINT_HEADER \
61	printf("%13s %7s %13s %5s %13s %13s %9s\n", \
62		"start", "size", "block-len", "state", "done", "remaining", "% done")
63
64#define PRINT_STATUS(start, i, len, state, d, t) \
65	printf("\r%13jd %7zu %13jd %5d %13jd %13jd %9.5f", \
66		(intmax_t)start, \
67		i,  \
68		(intmax_t)len, \
69		state, \
70		(intmax_t)d, \
71		(intmax_t)(t - d), \
72		100*(double)d/(double)t)
73
74/* Save the worklist if -w was given */
75static void
76save_worklist(void)
77{
78	FILE *file;
79	struct lump *llp;
80
81	if (wworklist != NULL) {
82		(void)fprintf(stderr, "\nSaving worklist ...");
83		fflush(stderr);
84
85		file = fopen(wworklist, "w");
86		if (file == NULL)
87			err(1, "Error opening file %s", wworklist);
88
89		TAILQ_FOREACH(llp, &lumps, list)
90			fprintf(file, "%jd %jd %d\n",
91			    (intmax_t)llp->start, (intmax_t)llp->len,
92			    llp->state);
93		fclose(file);
94		(void)fprintf(stderr, " done.\n");
95	}
96}
97
98/* Read the worklist if -r was given */
99static off_t
100read_worklist(off_t t)
101{
102	off_t s, l, d;
103	int state, lines;
104	FILE *file;
105
106	(void)fprintf(stderr, "Reading worklist ...");
107	fflush(stderr);
108	file = fopen(rworklist, "r");
109	if (file == NULL)
110		err(1, "Error opening file %s", rworklist);
111
112	lines = 0;
113	d = t;
114	for (;;) {
115		++lines;
116		if (3 != fscanf(file, "%jd %jd %d\n", &s, &l, &state)) {
117			if (!feof(file))
118				err(1, "Error parsing file %s at line %d",
119				    rworklist, lines);
120			else
121				break;
122		}
123		new_lump(s, l, state);
124		d -= l;
125	}
126	(void)fprintf(stderr, " done.\n");
127	/*
128	 * Return the number of bytes already read
129	 * (at least not in worklist).
130	 */
131	return (d);
132}
133
134static void
135usage(void)
136{
137	(void)fprintf(stderr, "usage: recoverdisk [-b bigsize] [-r readlist] "
138	    "[-s interval] [-w writelist] source [destination]\n");
139	exit(1);
140}
141
142static void
143sighandler(__unused int sig)
144{
145
146	aborting = 1;
147}
148
149int
150main(int argc, char * const argv[])
151{
152	int ch;
153	int fdr, fdw;
154	off_t t, d, start, len;
155	size_t i, j;
156	int error, state;
157	u_char *buf;
158	u_int sectorsize;
159	time_t t1, t2;
160	struct stat sb;
161	u_int n, snapshot = 60;
162
163	while ((ch = getopt(argc, argv, "b:r:w:s:")) != -1) {
164		switch (ch) {
165		case 'b':
166			bigsize = strtoul(optarg, NULL, 0);
167			break;
168		case 'r':
169			rworklist = strdup(optarg);
170			if (rworklist == NULL)
171				err(1, "Cannot allocate enough memory");
172			break;
173		case 's':
174			snapshot = strtoul(optarg, NULL, 0);
175			break;
176		case 'w':
177			wworklist = strdup(optarg);
178			if (wworklist == NULL)
179				err(1, "Cannot allocate enough memory");
180			break;
181		default:
182			usage();
183			/* NOTREACHED */
184		}
185	}
186	argc -= optind;
187	argv += optind;
188
189	if (argc < 1 || argc > 2)
190		usage();
191
192	fdr = open(argv[0], O_RDONLY);
193	if (fdr < 0)
194		err(1, "Cannot open read descriptor %s", argv[0]);
195
196	error = fstat(fdr, &sb);
197	if (error < 0)
198		err(1, "fstat failed");
199	if (S_ISBLK(sb.st_mode) || S_ISCHR(sb.st_mode)) {
200		error = ioctl(fdr, DIOCGSECTORSIZE, &sectorsize);
201		if (error < 0)
202			err(1, "DIOCGSECTORSIZE failed");
203
204		minsize = sectorsize;
205		bigsize = (bigsize / sectorsize) * sectorsize;
206
207		error = ioctl(fdr, DIOCGMEDIASIZE, &t);
208		if (error < 0)
209			err(1, "DIOCGMEDIASIZE failed");
210	} else {
211		t = sb.st_size;
212	}
213
214	if (bigsize < minsize)
215		bigsize = minsize;
216
217	for (ch = 0; (bigsize >> ch) > minsize; ch++)
218		continue;
219	medsize = bigsize >> (ch / 2);
220	medsize = (medsize / minsize) * minsize;
221
222	fprintf(stderr, "Bigsize = %zu, medsize = %zu, minsize = %zu\n",
223	    bigsize, medsize, minsize);
224
225	buf = malloc(bigsize);
226	if (buf == NULL)
227		err(1, "Cannot allocate %zu bytes buffer", bigsize);
228
229	if (argc > 1) {
230		fdw = open(argv[1], O_WRONLY | O_CREAT, DEFFILEMODE);
231		if (fdw < 0)
232			err(1, "Cannot open write descriptor %s", argv[1]);
233		if (ftruncate(fdw, t) < 0)
234			err(1, "Cannot truncate output %s to %jd bytes",
235			    argv[1], (intmax_t)t);
236	} else
237		fdw = -1;
238
239	if (rworklist != NULL) {
240		d = read_worklist(t);
241	} else {
242		new_lump(0, t, 0);
243		d = 0;
244	}
245	if (wworklist != NULL)
246		signal(SIGINT, sighandler);
247
248	t1 = 0;
249	start = len = i = state = 0;
250	PRINT_HEADER;
251	n = 0;
252	for (;;) {
253		lp = TAILQ_FIRST(&lumps);
254		if (lp == NULL)
255			break;
256		while (lp->len > 0 && !aborting) {
257			/* These are only copied for printing stats */
258			start = lp->start;
259			len = lp->len;
260			state = lp->state;
261
262			i = MIN(lp->len, (off_t)bigsize);
263			if (lp->state == 1)
264				i = MIN(lp->len, (off_t)medsize);
265			if (lp->state > 1)
266				i = MIN(lp->len, (off_t)minsize);
267			time(&t2);
268			if (t1 != t2 || lp->len < (off_t)bigsize) {
269				PRINT_STATUS(start, i, len, state, d, t);
270				t1 = t2;
271				if (++n == snapshot) {
272					save_worklist();
273					n = 0;
274				}
275			}
276			if (i == 0) {
277				errx(1, "BOGUS i %10jd", (intmax_t)i);
278			}
279			fflush(stdout);
280			j = pread(fdr, buf, i, lp->start);
281			if (j == i) {
282				d += i;
283				if (fdw >= 0)
284					j = pwrite(fdw, buf, i, lp->start);
285				else
286					j = i;
287				if (j != i)
288					printf("\nWrite error at %jd/%zu\n",
289					    lp->start, i);
290				lp->start += i;
291				lp->len -= i;
292				continue;
293			}
294			printf("\n%jd %zu failed (%s)\n",
295			    lp->start, i, strerror(errno));
296			if (errno == EINVAL) {
297				printf("read() size too big? Try with -b 131072");
298				aborting = 1;
299			}
300			if (errno == ENXIO)
301				aborting = 1;
302			new_lump(lp->start, i, lp->state + 1);
303			lp->start += i;
304			lp->len -= i;
305		}
306		if (aborting) {
307			save_worklist();
308			return (0);
309		}
310		TAILQ_REMOVE(&lumps, lp, list);
311		free(lp);
312	}
313	PRINT_STATUS(start, i, len, state, d, t);
314	save_worklist();
315	printf("\nCompleted\n");
316	return (0);
317}
318