1/*-
2 * SPDX-License-Identifier: Beerware
3 *
4 * ----------------------------------------------------------------------------
5 * "THE BEER-WARE LICENSE" (Revision 42):
6 * <phk@FreeBSD.ORG> wrote this file.  As long as you retain this notice you
7 * can do whatever you want with this stuff. If we meet some day, and you think
8 * this stuff is worth it, you can buy me a beer in return.   Poul-Henning Kamp
9 * ----------------------------------------------------------------------------
10 */
11#include <sys/param.h>
12#include <sys/queue.h>
13#include <sys/disk.h>
14#include <sys/stat.h>
15
16#include <assert.h>
17#include <err.h>
18#include <errno.h>
19#include <math.h>
20#include <fcntl.h>
21#include <signal.h>
22#include <stdint.h>
23#include <stdio.h>
24#include <stdlib.h>
25#include <string.h>
26#include <termios.h>
27#include <time.h>
28#include <unistd.h>
29
30/* Safe printf into a fixed-size buffer */
31#define bprintf(buf, fmt, ...)                                          \
32	do {                                                            \
33		int ibprintf;                                           \
34		ibprintf = snprintf(buf, sizeof buf, fmt, __VA_ARGS__); \
35		assert(ibprintf >= 0 && ibprintf < (int)sizeof buf);    \
36	} while (0)
37
38struct lump {
39	off_t			start;
40	off_t			len;
41	int			state;
42	TAILQ_ENTRY(lump)	list;
43};
44
45struct period {
46	time_t			t0;
47	time_t			t1;
48	char			str[20];
49	off_t			bytes_read;
50	TAILQ_ENTRY(period)	list;
51};
52TAILQ_HEAD(period_head, period);
53
54static volatile sig_atomic_t aborting = 0;
55static int verbose = 0;
56static size_t bigsize = 1024 * 1024;
57static size_t medsize;
58static size_t minsize = 512;
59static off_t tot_size;
60static off_t done_size;
61static char *input;
62static char *wworklist = NULL;
63static char *rworklist = NULL;
64static const char *unreadable_pattern = "_UNREAD_";
65static const int write_errors_are_fatal = 1;
66static int fdr, fdw;
67
68static TAILQ_HEAD(, lump) lumps = TAILQ_HEAD_INITIALIZER(lumps);
69static struct period_head minute = TAILQ_HEAD_INITIALIZER(minute);
70static struct period_head quarter = TAILQ_HEAD_INITIALIZER(quarter);
71static struct period_head hour = TAILQ_HEAD_INITIALIZER(quarter);
72static struct period_head day = TAILQ_HEAD_INITIALIZER(quarter);
73
74/**********************************************************************/
75
76static void
77report_good_read2(time_t now, size_t bytes, struct period_head *ph, time_t dt)
78{
79	struct period *pp;
80	const char *fmt;
81	struct tm tm1;
82
83	pp = TAILQ_FIRST(ph);
84	if (pp == NULL || pp->t1 < now) {
85		pp = calloc(sizeof *pp, 1L);
86		assert(pp != NULL);
87		pp->t0 = (now / dt) * dt;
88		pp->t1 = (now / dt + 1) * dt;
89		assert(localtime_r(&pp->t0, &tm1) != NULL);
90		if (dt < 86400)
91			fmt = "%H:%M";
92		else
93			fmt = "%d%b";
94		assert(strftime(pp->str, sizeof pp->str, fmt, &tm1) != 0);
95		TAILQ_INSERT_HEAD(ph, pp, list);
96	}
97	pp->bytes_read += bytes;
98}
99
100static void
101report_good_read(time_t now, size_t bytes)
102{
103
104	report_good_read2(now, bytes, &minute, 60L);
105	report_good_read2(now, bytes, &quarter, 900L);
106	report_good_read2(now, bytes, &hour, 3600L);
107	report_good_read2(now, bytes, &day, 86400L);
108}
109
110static void
111report_one_period(const char *period, struct period_head *ph)
112{
113	struct period *pp;
114	int n;
115
116	n = 0;
117	printf("%s \xe2\x94\x82", period);
118	TAILQ_FOREACH(pp, ph, list) {
119		if (n == 3) {
120			TAILQ_REMOVE(ph, pp, list);
121			free(pp);
122			break;
123		}
124		if (n++)
125			printf("  \xe2\x94\x82");
126		printf("  %s %14jd", pp->str, pp->bytes_read);
127	}
128	for (; n < 3; n++) {
129		printf("  \xe2\x94\x82");
130		printf("  %5s %14s", "", "");
131	}
132	printf("\x1b[K\n");
133}
134
135static void
136report_periods(void)
137{
138	report_one_period("1m ", &minute);
139	report_one_period("15m", &quarter);
140	report_one_period("1h ", &hour);
141	report_one_period("1d ", &day);
142}
143
144/**********************************************************************/
145
146static void
147set_verbose(void)
148{
149	struct winsize wsz;
150
151	if (!isatty(STDIN_FILENO) || ioctl(STDIN_FILENO, TIOCGWINSZ, &wsz))
152		return;
153	verbose = 1;
154}
155
156static void
157report_header(int eol)
158{
159	printf("%13s %7s %13s %5s %13s %13s %9s",
160	    "start",
161	    "size",
162	    "block-len",
163	    "pass",
164	    "done",
165	    "remaining",
166	    "% done");
167	if (eol)
168		printf("\x1b[K");
169	putchar('\n');
170}
171
172#define REPORTWID 79
173
174static void
175report_hline(const char *how)
176{
177	int j;
178
179	for (j = 0; j < REPORTWID; j++) {
180		if (how && (j == 4 || j == 29 || j == 54)) {
181			printf("%s", how);
182		} else {
183			printf("\xe2\x94\x80");
184		}
185	}
186	printf("\x1b[K\n");
187}
188
189static off_t hist[REPORTWID];
190static off_t last_done = -1;
191
192static void
193report_histogram(const struct lump *lp)
194{
195	off_t j, bucket, fp, fe, k, now;
196	double a;
197	struct lump *lp2;
198
199	bucket = tot_size / REPORTWID;
200	if (tot_size > bucket * REPORTWID)
201		bucket += 1;
202	if (done_size != last_done) {
203		memset(hist, 0, sizeof hist);
204		TAILQ_FOREACH(lp2, &lumps, list) {
205			fp = lp2->start;
206			fe = lp2->start + lp2->len;
207			for (j = fp / bucket; fp < fe; j++) {
208				k = (j + 1) * bucket;
209				if (k > fe)
210					k = fe;
211				k -= fp;
212				hist[j] += k;
213				fp += k;
214			}
215		}
216		last_done = done_size;
217	}
218	now = lp->start / bucket;
219	for (j = 0; j < REPORTWID; j++) {
220		a = round(8 * (double)hist[j] / bucket);
221		assert (a >= 0 && a < 9);
222		if (a == 0 && hist[j])
223			a = 1;
224		if (j == now)
225			printf("\x1b[31m");
226		if (a == 0) {
227			putchar(' ');
228		} else {
229			putchar(0xe2);
230			putchar(0x96);
231			putchar(0x80 + (int)a);
232		}
233		if (j == now)
234			printf("\x1b[0m");
235	}
236	putchar('\n');
237}
238
239static void
240report(const struct lump *lp, size_t sz)
241{
242	struct winsize wsz;
243	int j;
244
245	assert(lp != NULL);
246
247	if (verbose) {
248		printf("\x1b[H%s\x1b[K\n", input);
249		report_header(1);
250	} else {
251		putchar('\r');
252	}
253
254	printf("%13jd %7zu %13jd %5d %13jd %13jd %9.4f",
255	    (intmax_t)lp->start,
256	    sz,
257	    (intmax_t)lp->len,
258	    lp->state,
259	    (intmax_t)done_size,
260	    (intmax_t)(tot_size - done_size),
261	    100*(double)done_size/(double)tot_size
262	);
263
264	if (verbose) {
265		printf("\x1b[K\n");
266		report_hline(NULL);
267		report_histogram(lp);
268		if (TAILQ_EMPTY(&minute)) {
269			report_hline(NULL);
270		} else {
271			report_hline("\xe2\x94\xac");
272			report_periods();
273			report_hline("\xe2\x94\xb4");
274		}
275		j = ioctl(STDIN_FILENO, TIOCGWINSZ, &wsz);
276		if (!j)
277			printf("\x1b[%d;1H", wsz.ws_row);
278	}
279	fflush(stdout);
280}
281
282/**********************************************************************/
283
284static void
285new_lump(off_t start, off_t len, int state)
286{
287	struct lump *lp;
288
289	lp = malloc(sizeof *lp);
290	if (lp == NULL)
291		err(1, "Malloc failed");
292	lp->start = start;
293	lp->len = len;
294	lp->state = state;
295	TAILQ_INSERT_TAIL(&lumps, lp, list);
296}
297
298/**********************************************************************
299 * Save the worklist if -w was given
300 */
301
302static void
303save_worklist(void)
304{
305	FILE *file;
306	struct lump *llp;
307	char buf[PATH_MAX];
308
309	if (fdw >= 0 && fdatasync(fdw))
310		err(1, "Write error, probably disk full");
311
312	if (wworklist != NULL) {
313		bprintf(buf, "%s.tmp", wworklist);
314		(void)fprintf(stderr, "\nSaving worklist ...");
315		(void)fflush(stderr);
316
317		file = fopen(buf, "w");
318		if (file == NULL)
319			err(1, "Error opening file %s", buf);
320
321		TAILQ_FOREACH(llp, &lumps, list)
322			fprintf(file, "%jd %jd %d\n",
323			    (intmax_t)llp->start, (intmax_t)llp->len,
324			    llp->state);
325		(void)fflush(file);
326		if (ferror(file) || fdatasync(fileno(file)) || fclose(file))
327			err(1, "Error writing file %s", buf);
328		if (rename(buf, wworklist))
329			err(1, "Error renaming %s to %s", buf, wworklist);
330		(void)fprintf(stderr, " done.\n");
331	}
332}
333
334/* Read the worklist if -r was given */
335static off_t
336read_worklist(off_t t)
337{
338	off_t s, l, d;
339	int state, lines;
340	FILE *file;
341
342	(void)fprintf(stderr, "Reading worklist ...");
343	(void)fflush(stderr);
344	file = fopen(rworklist, "r");
345	if (file == NULL)
346		err(1, "Error opening file %s", rworklist);
347
348	lines = 0;
349	d = t;
350	for (;;) {
351		++lines;
352		if (3 != fscanf(file, "%jd %jd %d\n", &s, &l, &state)) {
353			if (!feof(file))
354				err(1, "Error parsing file %s at line %d",
355				    rworklist, lines);
356			else
357				break;
358		}
359		new_lump(s, l, state);
360		d -= l;
361	}
362	if (fclose(file))
363		err(1, "Error closing file %s", rworklist);
364	(void)fprintf(stderr, " done.\n");
365	/*
366	 * Return the number of bytes already read
367	 * (at least not in worklist).
368	 */
369	return (d);
370}
371
372/**********************************************************************/
373
374static void
375write_buf(int fd, const void *buf, ssize_t len, off_t where)
376{
377	ssize_t i;
378
379	i = pwrite(fd, buf, len, where);
380	if (i == len)
381		return;
382
383	printf("\nWrite error at %jd/%zu\n\t%s\n",
384	    where, i, strerror(errno));
385	save_worklist();
386	if (write_errors_are_fatal)
387		exit(3);
388}
389
390static void
391fill_buf(char *buf, ssize_t len, const char *pattern)
392{
393	ssize_t sz = strlen(pattern);
394	ssize_t i, j;
395
396	for (i = 0; i < len; i += sz) {
397		j = len - i;
398		if (j > sz)
399			j = sz;
400		memcpy(buf + i, pattern, j);
401	}
402}
403
404/**********************************************************************/
405
406static void
407usage(void)
408{
409	(void)fprintf(stderr, "usage: recoverdisk [-b bigsize] [-r readlist] "
410	    "[-s interval] [-w writelist] source [destination]\n");
411	/* XXX update */
412	exit(1);
413}
414
415static void
416sighandler(__unused int sig)
417{
418
419	aborting = 1;
420}
421
422int
423main(int argc, char * const argv[])
424{
425	int ch;
426	size_t sz, j;
427	int error;
428	char *buf;
429	u_int sectorsize;
430	off_t stripesize;
431	time_t t1, t2;
432	struct stat sb;
433	u_int n, snapshot = 60;
434	static struct lump *lp;
435
436	while ((ch = getopt(argc, argv, "b:r:w:s:u:v")) != -1) {
437		switch (ch) {
438		case 'b':
439			bigsize = strtoul(optarg, NULL, 0);
440			break;
441		case 'r':
442			rworklist = strdup(optarg);
443			if (rworklist == NULL)
444				err(1, "Cannot allocate enough memory");
445			break;
446		case 's':
447			snapshot = strtoul(optarg, NULL, 0);
448			break;
449		case 'u':
450			unreadable_pattern = optarg;
451			break;
452		case 'v':
453			set_verbose();
454			break;
455		case 'w':
456			wworklist = strdup(optarg);
457			if (wworklist == NULL)
458				err(1, "Cannot allocate enough memory");
459			break;
460		default:
461			usage();
462			/* NOTREACHED */
463		}
464	}
465	argc -= optind;
466	argv += optind;
467
468	if (argc < 1 || argc > 2)
469		usage();
470
471	input = argv[0];
472	fdr = open(argv[0], O_RDONLY);
473	if (fdr < 0)
474		err(1, "Cannot open read descriptor %s", argv[0]);
475
476	error = fstat(fdr, &sb);
477	if (error < 0)
478		err(1, "fstat failed");
479	if (S_ISBLK(sb.st_mode) || S_ISCHR(sb.st_mode)) {
480		error = ioctl(fdr, DIOCGSECTORSIZE, &sectorsize);
481		if (error < 0)
482			err(1, "DIOCGSECTORSIZE failed");
483
484		error = ioctl(fdr, DIOCGSTRIPESIZE, &stripesize);
485		if (error == 0 && stripesize > sectorsize)
486			sectorsize = stripesize;
487
488		minsize = sectorsize;
489		bigsize = rounddown(bigsize, sectorsize);
490
491		error = ioctl(fdr, DIOCGMEDIASIZE, &tot_size);
492		if (error < 0)
493			err(1, "DIOCGMEDIASIZE failed");
494	} else {
495		tot_size = sb.st_size;
496	}
497
498	if (bigsize < minsize)
499		bigsize = minsize;
500
501	for (ch = 0; (bigsize >> ch) > minsize; ch++)
502		continue;
503	medsize = bigsize >> (ch / 2);
504	medsize = rounddown(medsize, minsize);
505
506	fprintf(stderr, "Bigsize = %zu, medsize = %zu, minsize = %zu\n",
507	    bigsize, medsize, minsize);
508
509	buf = malloc(bigsize);
510	if (buf == NULL)
511		err(1, "Cannot allocate %zu bytes buffer", bigsize);
512
513	if (argc > 1) {
514		fdw = open(argv[1], O_WRONLY | O_CREAT, DEFFILEMODE);
515		if (fdw < 0)
516			err(1, "Cannot open write descriptor %s", argv[1]);
517		if (ftruncate(fdw, tot_size) < 0)
518			err(1, "Cannot truncate output %s to %jd bytes",
519			    argv[1], (intmax_t)tot_size);
520	} else
521		fdw = -1;
522
523	if (rworklist != NULL) {
524		done_size = read_worklist(tot_size);
525	} else {
526		new_lump(0, tot_size, 0);
527		done_size = 0;
528	}
529	if (wworklist != NULL)
530		signal(SIGINT, sighandler);
531
532	t1 = time(NULL);
533	sz = 0;
534	if (!verbose)
535		report_header(0);
536	else
537		printf("\x1b[2J");
538	n = 0;
539	for (;;) {
540		lp = TAILQ_FIRST(&lumps);
541		if (lp == NULL)
542			break;
543		while (lp->len > 0) {
544
545			if (lp->state == 0)
546				sz = MIN(lp->len, (off_t)bigsize);
547			else if (lp->state == 1)
548				sz = MIN(lp->len, (off_t)medsize);
549			else
550				sz = MIN(lp->len, (off_t)minsize);
551			assert(sz != 0);
552
553			t2 = time(NULL);
554			if (t1 != t2 || lp->len < (off_t)bigsize) {
555				t1 = t2;
556				if (++n == snapshot) {
557					save_worklist();
558					n = 0;
559				}
560				report(lp, sz);
561			}
562
563			j = pread(fdr, buf, sz, lp->start);
564#if 0
565if (!(random() & 0xf)) {
566	j = -1;
567	errno = EIO;
568}
569#endif
570			if (j == sz) {
571				done_size += sz;
572				if (fdw >= 0)
573					write_buf(fdw, buf, sz, lp->start);
574				lp->start += sz;
575				lp->len -= sz;
576				if (verbose && lp->state > 2)
577					report_good_read(t2, sz);
578				continue;
579			}
580			error = errno;
581
582			printf("%jd %zu %d read error (%s)\n",
583			    lp->start, sz, lp->state, strerror(error));
584			if (verbose)
585				report(lp, sz);
586			if (fdw >= 0 && strlen(unreadable_pattern)) {
587				fill_buf(buf, sz, unreadable_pattern);
588				write_buf(fdw, buf, sz, lp->start);
589			}
590			new_lump(lp->start, sz, lp->state + 1);
591			lp->start += sz;
592			lp->len -= sz;
593			if (error == EINVAL) {
594				printf("Try with -b 131072 or lower ?\n");
595				aborting = 1;
596				break;
597			}
598			if (error == ENXIO) {
599				printf("Input device probably detached...\n");
600				aborting = 1;
601				break;
602			}
603		}
604		if (aborting)
605			save_worklist();
606		if (aborting || !TAILQ_NEXT(lp, list))
607			report(lp, sz);
608		if (aborting)
609			break;
610		assert(lp->len == 0);
611		TAILQ_REMOVE(&lumps, lp, list);
612		free(lp);
613	}
614	printf("%s", aborting ? "Aborted\n" : "Completed\n");
615	free(buf);
616	return (0);
617}
618