3 files changed, 163 insertions, 71 deletions

diff --git a/tools/tools/recoverdisk/Makefile b/tools/tools/recoverdisk/Makefile
index e86fbb2b96a9..6733a4f8fa20 100644
--- a/tools/tools/recoverdisk/Makefile
+++ b/tools/tools/recoverdisk/Makefile
@@ -1,10 +1,7 @@
 # $FreeBSD$
 
 PROG=	recoverdisk
-
-NO_MAN=
-
-WARNS?=	5
+WARNS?=	6
 
 .include <bsd.prog.mk>
 
diff --git a/tools/tools/recoverdisk/README b/tools/tools/recoverdisk/README
deleted file mode 100644
index 4f7624a88cb4..000000000000
--- a/tools/tools/recoverdisk/README
+++ /dev/null
@@ -1,25 +0,0 @@
-$FreeBSD$
-
-This is a small tool which will read an entire disk(partition) or file
-using 1M blocks and optionally write the read data to a file or disk.
- 
-If a read error happens, the 1M block gets put on the end of the worklist
-and will be retried with 64k blocksize.
- 
-If a read error happens again, the 64k block gets put at the end of the
-worklist and will be retried with single sector reads.
- 
-The program keeps trying until you stop it.
-
-You can refresh a disk:
-
-	recoverdisk /dev/ad1 /dev/ad1
-
-or salvage a floppy:
-
-	touch myfloppy.flp
-	recoverdisk /dev/fd0 myfloppy.flp
-
-or recover a single file from the unreadable media:
-
-	recoverdisk /cdrom/file.avi file.avi
diff --git a/tools/tools/recoverdisk/recoverdisk.c b/tools/tools/recoverdisk/recoverdisk.c
index d30055a2f1b5..2caa1e957370 100644
--- a/tools/tools/recoverdisk/recoverdisk.c
+++ b/tools/tools/recoverdisk/recoverdisk.c
@@ -1,4 +1,4 @@
-/*
+/*-
  * ----------------------------------------------------------------------------
  * "THE BEER-WARE LICENSE" (Revision 42):
  * <phk@FreeBSD.ORG> wrote this file.  As long as you retain this notice you
@@ -8,21 +8,26 @@
  *
  * $FreeBSD$
  */
-#include <stdio.h>
-#include <stdint.h>
-#include <stdlib.h>
+#include <sys/param.h>
+#include <sys/queue.h>
+#include <sys/disk.h>
+#include <sys/stat.h>
+
 #include <err.h>
 #include <errno.h>
 #include <fcntl.h>
+#include <signal.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
 #include <time.h>
 #include <unistd.h>
-#include <sys/queue.h>
-#include <sys/disk.h>
-#include <sys/stat.h>
 
-#define BIGSIZE		(1024 * 1024)
-#define MEDIUMSIZE	(64 * 1024)
-#define MINSIZE		(512)
+volatile sig_atomic_t aborting = 0;
+static size_t bigsize = 1024 * 1024;
+static size_t medsize = 64 * 1024;
+static size_t minsize = 512;
 
 struct lump {
 	off_t			start;
@@ -33,7 +38,6 @@ struct lump {
 
 static TAILQ_HEAD(, lump) lumps = TAILQ_HEAD_INITIALIZER(lumps);
 
-
 static void
 new_lump(off_t start, off_t len, int state)
 {
@@ -48,29 +52,126 @@ new_lump(off_t start, off_t len, int state)
 	TAILQ_INSERT_TAIL(&lumps, lp, list);
 }
 
+static struct lump *lp;
+static char *wworklist = NULL;
+static char *rworklist = NULL;
+
+/* Save the worklist if -w was given */
+static void
+save_worklist(void)
+{
+	FILE *file;
+
+	if (wworklist != NULL) {
+		(void)fprintf(stderr, "\nSaving worklist ...");
+		fflush(stderr);
+
+		file = fopen(wworklist, "w");
+		if (file == NULL)
+			err(1, "Error opening file %s", wworklist);
+
+		for (;;) {
+			lp = TAILQ_FIRST(&lumps);
+			if (lp == NULL)
+				break;
+			fprintf(file, "%jd %jd %d\n",
+			    (intmax_t)lp->start, (intmax_t)lp->len, lp->state);
+			TAILQ_REMOVE(&lumps, lp, list);
+		}
+		(void)fprintf(stderr, " done.\n");
+	}
+}
+
+/* Read the worklist if -r was given */
+static off_t
+read_worklist(off_t t)
+{
+	off_t s, l, d;
+	int state, lines;
+	FILE *file;
+
+	(void)fprintf(stderr, "Reading worklist ...");
+	fflush(stderr);
+	file = fopen(rworklist, "r");
+	if (file == NULL)
+		err(1, "Error opening file %s", rworklist);
+
+	lines = 0;
+	d = t;
+	for (;;) {
+		++lines;
+		if (3 != fscanf(file, "%jd %jd %d\n", &s, &l, &state)) {
+			if (!feof(file))
+				err(1, "Error parsing file %s at line %d",
+				    rworklist, lines);
+			else
+				break;
+		}
+		new_lump(s, l, state);
+		d -= l;
+	}
+	(void)fprintf(stderr, " done.\n");
+	/*
+	 * Return the number of bytes already read
+	 * (at least not in worklist).
+	 */
+	return (d);
+}
+
+static void
+usage(void)
+{
+	(void)fprintf(stderr,
+    "usage: recoverdisk [-r worklist] [-w worklist] source-drive [destination]\n");
+	exit(1);
+}
+
+static void
+sighandler(__unused int sig)
+{
+
+	aborting = 1;
+}
+
 int
-main(int argc, const char **argv)
+main(int argc, char * const argv[])
 {
+	int ch;
 	int fdr, fdw;
-	struct lump *lp;
-	off_t 	t, d;
+	off_t t, d;
 	size_t i, j;
 	int error, flags;
 	u_char *buf;
-	u_int sectorsize, minsize;
+	u_int sectorsize;
 	time_t t1, t2;
 	struct stat sb;
 
+	while ((ch = getopt(argc, argv, "r:w:")) != -1) {
+		switch (ch) {
+		case 'r':
+			rworklist = strdup(optarg);
+			if (rworklist == NULL)
+				err(1, "Cannot allocate enough memory");
+			break;
+		case 'w':
+			wworklist = strdup(optarg);
+			if (wworklist == NULL)
+				err(1, "Cannot allocate enough memory");
+			break;
+		default:
+			usage();
+			/* NOTREACHED */
+		}
+	}
+	argc -= optind;
+	argv += optind;
 
-	if (argc < 2)
-		errx(1, "Usage: %s source-drive [destination]", argv[0]);
+	if (argc < 1 || argc > 2)
+		usage();
 
-	buf = malloc(BIGSIZE);
-	if (buf == NULL)
-		err(1, "Cannot allocate %d bytes buffer", BIGSIZE);
-	fdr = open(argv[1], O_RDONLY);
+	fdr = open(argv[0], O_RDONLY);
 	if (fdr < 0)
-		err(1, "Cannot open read descriptor %s", argv[1]);
+		err(1, "Cannot open read descriptor %s", argv[0]);
 
 	error = fstat(fdr, &sb);
 	if (error < 0)
@@ -80,46 +181,61 @@ main(int argc, const char **argv)
 		error = ioctl(fdr, DIOCGSECTORSIZE, &sectorsize);
 		if (error < 0)
 			err(1, "DIOCGSECTORSIZE failed");
+
+		/*
+		 * Make medsize roughly 64kB, depending on native sector
+		 * size. bigsize has to be a multiple of medsize.
+		 * For media with 2352 sectors, this will
+		 * result in 2352, 63504, and 1016064 bytes.
+		 */
 		minsize = sectorsize;
+		medsize = (medsize / sectorsize) * sectorsize;
+		bigsize = medsize * 16;
 
 		error = ioctl(fdr, DIOCGMEDIASIZE, &t);
 		if (error < 0)
 			err(1, "DIOCGMEDIASIZE failed");
 	} else {
-		sectorsize = 1;
 		t = sb.st_size;
-		minsize = MINSIZE;
 		flags |= O_CREAT | O_TRUNC;
 	}
 
-	if (argc > 2) {
-		fdw = open(argv[2], flags, DEFFILEMODE);
+	buf = malloc(bigsize);
+	if (buf == NULL)
+		err(1, "Cannot allocate %jd bytes buffer", (intmax_t)bigsize);
+
+	if (argc > 1) {
+		fdw = open(argv[1], flags, DEFFILEMODE);
 		if (fdw < 0)
-			err(1, "Cannot open write descriptor %s", argv[2]);
-	} else {
+			err(1, "Cannot open write descriptor %s", argv[1]);
+	} else
 		fdw = -1;
-	}
 
-	new_lump(0, t, 0);
-	d = 0;
+	if (rworklist != NULL) {
+		d = read_worklist(t);
+	} else {
+		new_lump(0, t, 0);
+		d = 0;
+	}
+	if (wworklist != NULL)
+		signal(SIGINT, sighandler);
 
 	t1 = 0;
+	printf("%13s %7s %13s %5s %13s %13s %9s\n",
+	    "start", "size", "len", "state", "done", "remaining", "% done");
 	for (;;) {
 		lp = TAILQ_FIRST(&lumps);
 		if (lp == NULL)
 			break;
-		TAILQ_REMOVE(&lumps, lp, list);
-		while (lp->len > 0) {
-			i = BIGSIZE;
-			if (lp->len < BIGSIZE)
-				i = lp->len;
+		while (lp->len > 0 && !aborting) {
+			i = MIN(lp->len, (off_t)bigsize);
 			if (lp->state == 1)
-				i = MEDIUMSIZE;
+				i = MIN(lp->len, (off_t)medsize);
 			if (lp->state > 1)
-				i = minsize;
+				i = MIN(lp->len, (off_t)minsize);
 			time(&t2);
-			if (t1 != t2 || lp->len < BIGSIZE) {
-				printf("\r%13jd %7zu %13jd %3d %13jd %13jd %.8f",
+			if (t1 != t2 || lp->len < (off_t)bigsize) {
+				printf("\r%13jd %7zu %13jd %5d %13jd %13jd %.7f",
 				    (intmax_t)lp->start,
 				    i, 
 				    (intmax_t)lp->len,
@@ -152,9 +268,13 @@ main(int argc, const char **argv)
 			lp->start += i;
 			lp->len -= i;
 		}
+		if (aborting) {
+			save_worklist();
+			return (0);
+		}
+		TAILQ_REMOVE(&lumps, lp, list);
 		free(lp);
 	}
 	printf("\nCompleted\n");
-	exit (0);
+	return (0);
 }
-