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: stable/11/sbin/recoverdisk/recoverdisk.c 344091 2019-02-13 09:28:48Z avos $
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	fclose(file);
127	(void)fprintf(stderr, " done.\n");
128	/*
129	 * Return the number of bytes already read
130	 * (at least not in worklist).
131	 */
132	return (d);
133}
134
135static void
136usage(void)
137{
138	(void)fprintf(stderr, "usage: recoverdisk [-b bigsize] [-r readlist] "
139	    "[-s interval] [-w writelist] source [destination]\n");
140	exit(1);
141}
142
143static void
144sighandler(__unused int sig)
145{
146
147	aborting = 1;
148}
149
150int
151main(int argc, char * const argv[])
152{
153	int ch;
154	int fdr, fdw;
155	off_t t, d, start, len;
156	size_t i, j;
157	int error, state;
158	u_char *buf;
159	u_int sectorsize;
160	off_t stripesize;
161	time_t t1, t2;
162	struct stat sb;
163	u_int n, snapshot = 60;
164
165	while ((ch = getopt(argc, argv, "b:r:w:s:")) != -1) {
166		switch (ch) {
167		case 'b':
168			bigsize = strtoul(optarg, NULL, 0);
169			break;
170		case 'r':
171			rworklist = strdup(optarg);
172			if (rworklist == NULL)
173				err(1, "Cannot allocate enough memory");
174			break;
175		case 's':
176			snapshot = strtoul(optarg, NULL, 0);
177			break;
178		case 'w':
179			wworklist = strdup(optarg);
180			if (wworklist == NULL)
181				err(1, "Cannot allocate enough memory");
182			break;
183		default:
184			usage();
185			/* NOTREACHED */
186		}
187	}
188	argc -= optind;
189	argv += optind;
190
191	if (argc < 1 || argc > 2)
192		usage();
193
194	fdr = open(argv[0], O_RDONLY);
195	if (fdr < 0)
196		err(1, "Cannot open read descriptor %s", argv[0]);
197
198	error = fstat(fdr, &sb);
199	if (error < 0)
200		err(1, "fstat failed");
201	if (S_ISBLK(sb.st_mode) || S_ISCHR(sb.st_mode)) {
202		error = ioctl(fdr, DIOCGSECTORSIZE, &sectorsize);
203		if (error < 0)
204			err(1, "DIOCGSECTORSIZE failed");
205
206		error = ioctl(fdr, DIOCGSTRIPESIZE, &stripesize);
207		if (error == 0 && stripesize > sectorsize)
208			sectorsize = stripesize;
209
210		minsize = sectorsize;
211		bigsize = rounddown(bigsize, sectorsize);
212
213		error = ioctl(fdr, DIOCGMEDIASIZE, &t);
214		if (error < 0)
215			err(1, "DIOCGMEDIASIZE failed");
216	} else {
217		t = sb.st_size;
218	}
219
220	if (bigsize < minsize)
221		bigsize = minsize;
222
223	for (ch = 0; (bigsize >> ch) > minsize; ch++)
224		continue;
225	medsize = bigsize >> (ch / 2);
226	medsize = rounddown(medsize, minsize);
227
228	fprintf(stderr, "Bigsize = %zu, medsize = %zu, minsize = %zu\n",
229	    bigsize, medsize, minsize);
230
231	buf = malloc(bigsize);
232	if (buf == NULL)
233		err(1, "Cannot allocate %zu bytes buffer", bigsize);
234
235	if (argc > 1) {
236		fdw = open(argv[1], O_WRONLY | O_CREAT, DEFFILEMODE);
237		if (fdw < 0)
238			err(1, "Cannot open write descriptor %s", argv[1]);
239		if (ftruncate(fdw, t) < 0)
240			err(1, "Cannot truncate output %s to %jd bytes",
241			    argv[1], (intmax_t)t);
242	} else
243		fdw = -1;
244
245	if (rworklist != NULL) {
246		d = read_worklist(t);
247	} else {
248		new_lump(0, t, 0);
249		d = 0;
250	}
251	if (wworklist != NULL)
252		signal(SIGINT, sighandler);
253
254	t1 = 0;
255	start = len = i = state = 0;
256	PRINT_HEADER;
257	n = 0;
258	for (;;) {
259		lp = TAILQ_FIRST(&lumps);
260		if (lp == NULL)
261			break;
262		while (lp->len > 0 && !aborting) {
263			/* These are only copied for printing stats */
264			start = lp->start;
265			len = lp->len;
266			state = lp->state;
267
268			i = MIN(lp->len, (off_t)bigsize);
269			if (lp->state == 1)
270				i = MIN(lp->len, (off_t)medsize);
271			if (lp->state > 1)
272				i = MIN(lp->len, (off_t)minsize);
273			time(&t2);
274			if (t1 != t2 || lp->len < (off_t)bigsize) {
275				PRINT_STATUS(start, i, len, state, d, t);
276				t1 = t2;
277				if (++n == snapshot) {
278					save_worklist();
279					n = 0;
280				}
281			}
282			if (i == 0) {
283				errx(1, "BOGUS i %10jd", (intmax_t)i);
284			}
285			fflush(stdout);
286			j = pread(fdr, buf, i, lp->start);
287			if (j == i) {
288				d += i;
289				if (fdw >= 0)
290					j = pwrite(fdw, buf, i, lp->start);
291				else
292					j = i;
293				if (j != i)
294					printf("\nWrite error at %jd/%zu\n",
295					    lp->start, i);
296				lp->start += i;
297				lp->len -= i;
298				continue;
299			}
300			printf("\n%jd %zu failed (%s)\n",
301			    lp->start, i, strerror(errno));
302			if (errno == EINVAL) {
303				printf("read() size too big? Try with -b 131072");
304				aborting = 1;
305			}
306			if (errno == ENXIO)
307				aborting = 1;
308			new_lump(lp->start, i, lp->state + 1);
309			lp->start += i;
310			lp->len -= i;
311		}
312		if (aborting) {
313			save_worklist();
314			return (0);
315		}
316		TAILQ_REMOVE(&lumps, lp, list);
317		free(lp);
318	}
319	PRINT_STATUS(start, i, len, state, d, t);
320	save_worklist();
321	printf("\nCompleted\n");
322	return (0);
323}
324