xref: /freebsd/sbin/recoverdisk/recoverdisk.c (revision 6472ac3d8a86336899b6cfb789a4cd9897e3fab5)
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$
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 
27 static volatile sig_atomic_t aborting = 0;
28 static size_t bigsize = 1024 * 1024;
29 static size_t medsize;
30 static size_t minsize = 512;
31 
32 struct lump {
33 	off_t			start;
34 	off_t			len;
35 	int			state;
36 	TAILQ_ENTRY(lump)	list;
37 };
38 
39 static TAILQ_HEAD(, lump) lumps = TAILQ_HEAD_INITIALIZER(lumps);
40 
41 static void
42 new_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 
55 static struct lump *lp;
56 static char *wworklist = NULL;
57 static 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 */
75 static void
76 save_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 */
99 static off_t
100 read_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 
134 static void
135 usage(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 
142 static void
143 sighandler(__unused int sig)
144 {
145 
146 	aborting = 1;
147 }
148 
149 int
150 main(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