xref: /freebsd/sbin/recoverdisk/recoverdisk.c (revision 39beb93c3f8bdbf72a61fda42300b5ebed7390c8)
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 volatile sig_atomic_t aborting = 0;
28 static size_t bigsize = 1024 * 1024;
29 static size_t medsize = 64 * 1024;
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 
80 	if (wworklist != NULL) {
81 		(void)fprintf(stderr, "\nSaving worklist ...");
82 		fflush(stderr);
83 
84 		file = fopen(wworklist, "w");
85 		if (file == NULL)
86 			err(1, "Error opening file %s", wworklist);
87 
88 		for (;;) {
89 			lp = TAILQ_FIRST(&lumps);
90 			if (lp == NULL)
91 				break;
92 			fprintf(file, "%jd %jd %d\n",
93 			    (intmax_t)lp->start, (intmax_t)lp->len, lp->state);
94 			TAILQ_REMOVE(&lumps, lp, list);
95 		}
96 		(void)fprintf(stderr, " done.\n");
97 	}
98 }
99 
100 /* Read the worklist if -r was given */
101 static off_t
102 read_worklist(off_t t)
103 {
104 	off_t s, l, d;
105 	int state, lines;
106 	FILE *file;
107 
108 	(void)fprintf(stderr, "Reading worklist ...");
109 	fflush(stderr);
110 	file = fopen(rworklist, "r");
111 	if (file == NULL)
112 		err(1, "Error opening file %s", rworklist);
113 
114 	lines = 0;
115 	d = t;
116 	for (;;) {
117 		++lines;
118 		if (3 != fscanf(file, "%jd %jd %d\n", &s, &l, &state)) {
119 			if (!feof(file))
120 				err(1, "Error parsing file %s at line %d",
121 				    rworklist, lines);
122 			else
123 				break;
124 		}
125 		new_lump(s, l, state);
126 		d -= l;
127 	}
128 	(void)fprintf(stderr, " done.\n");
129 	/*
130 	 * Return the number of bytes already read
131 	 * (at least not in worklist).
132 	 */
133 	return (d);
134 }
135 
136 static void
137 usage(void)
138 {
139 	(void)fprintf(stderr,
140     "usage: recoverdisk [-r worklist] [-w worklist] source-drive [destination]\n");
141 	exit(1);
142 }
143 
144 static void
145 sighandler(__unused int sig)
146 {
147 
148 	aborting = 1;
149 }
150 
151 int
152 main(int argc, char * const argv[])
153 {
154 	int ch;
155 	int fdr, fdw;
156 	off_t t, d, start, len;
157 	size_t i, j;
158 	int error, flags, state;
159 	u_char *buf;
160 	u_int sectorsize;
161 	time_t t1, t2;
162 	struct stat sb;
163 
164 	while ((ch = getopt(argc, argv, "r:w:")) != -1) {
165 		switch (ch) {
166 		case 'r':
167 			rworklist = strdup(optarg);
168 			if (rworklist == NULL)
169 				err(1, "Cannot allocate enough memory");
170 			break;
171 		case 'w':
172 			wworklist = strdup(optarg);
173 			if (wworklist == NULL)
174 				err(1, "Cannot allocate enough memory");
175 			break;
176 		default:
177 			usage();
178 			/* NOTREACHED */
179 		}
180 	}
181 	argc -= optind;
182 	argv += optind;
183 
184 	if (argc < 1 || argc > 2)
185 		usage();
186 
187 	fdr = open(argv[0], O_RDONLY);
188 	if (fdr < 0)
189 		err(1, "Cannot open read descriptor %s", argv[0]);
190 
191 	error = fstat(fdr, &sb);
192 	if (error < 0)
193 		err(1, "fstat failed");
194 	flags = O_WRONLY;
195 	if (S_ISBLK(sb.st_mode) || S_ISCHR(sb.st_mode)) {
196 		error = ioctl(fdr, DIOCGSECTORSIZE, &sectorsize);
197 		if (error < 0)
198 			err(1, "DIOCGSECTORSIZE failed");
199 
200 		/*
201 		 * Make medsize roughly 64kB, depending on native sector
202 		 * size. bigsize has to be a multiple of medsize.
203 		 * For media with 2352 sectors, this will
204 		 * result in 2352, 63504, and 1016064 bytes.
205 		 */
206 		minsize = sectorsize;
207 		medsize = (medsize / sectorsize) * sectorsize;
208 		bigsize = medsize * 16;
209 
210 		error = ioctl(fdr, DIOCGMEDIASIZE, &t);
211 		if (error < 0)
212 			err(1, "DIOCGMEDIASIZE failed");
213 	} else {
214 		t = sb.st_size;
215 		flags |= O_CREAT | O_TRUNC;
216 	}
217 
218 	buf = malloc(bigsize);
219 	if (buf == NULL)
220 		err(1, "Cannot allocate %jd bytes buffer", (intmax_t)bigsize);
221 
222 	if (argc > 1) {
223 		fdw = open(argv[1], flags, DEFFILEMODE);
224 		if (fdw < 0)
225 			err(1, "Cannot open write descriptor %s", argv[1]);
226 	} else
227 		fdw = -1;
228 
229 	if (rworklist != NULL) {
230 		d = read_worklist(t);
231 	} else {
232 		new_lump(0, t, 0);
233 		d = 0;
234 	}
235 	if (wworklist != NULL)
236 		signal(SIGINT, sighandler);
237 
238 	t1 = 0;
239 	start = len = i = state = 0;
240 	PRINT_HEADER;
241 	for (;;) {
242 		lp = TAILQ_FIRST(&lumps);
243 		if (lp == NULL)
244 			break;
245 		while (lp->len > 0 && !aborting) {
246 			/* These are only copied for printing stats */
247 			start = lp->start;
248 			len = lp->len;
249 			state = lp->state;
250 
251 			i = MIN(lp->len, (off_t)bigsize);
252 			if (lp->state == 1)
253 				i = MIN(lp->len, (off_t)medsize);
254 			if (lp->state > 1)
255 				i = MIN(lp->len, (off_t)minsize);
256 			time(&t2);
257 			if (t1 != t2 || lp->len < (off_t)bigsize) {
258 				PRINT_STATUS(start, i, len, state, d, t);
259 				t1 = t2;
260 			}
261 			if (i == 0) {
262 				errx(1, "BOGUS i %10jd", (intmax_t)i);
263 			}
264 			fflush(stdout);
265 			j = pread(fdr, buf, i, lp->start);
266 			if (j == i) {
267 				d += i;
268 				if (fdw >= 0)
269 					j = pwrite(fdw, buf, i, lp->start);
270 				else
271 					j = i;
272 				if (j != i)
273 					printf("\nWrite error at %jd/%zu\n",
274 					    lp->start, i);
275 				lp->start += i;
276 				lp->len -= i;
277 				continue;
278 			}
279 			printf("\n%jd %zu failed (%s)\n",
280 			    lp->start, i, strerror(errno));
281 			if (errno == ENXIO)
282 				aborting = 1;
283 			new_lump(lp->start, i, lp->state + 1);
284 			lp->start += i;
285 			lp->len -= i;
286 		}
287 		if (aborting) {
288 			save_worklist();
289 			return (0);
290 		}
291 		TAILQ_REMOVE(&lumps, lp, list);
292 		free(lp);
293 	}
294 	PRINT_STATUS(start, i, len, state, d, t);
295 	printf("\nCompleted\n");
296 	return (0);
297 }
298