/* $Id$ */

/*
 * badblock-guess: Quickly recover most of the data from a damaged disk
 * Copyright (C) 2002  Jan Kratochvil <short@ucw.cz>
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; exactly version 2 of the License required
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */


#define _LARGEFILE_SOURCE
#define _LARGEFILE64_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <time.h>
#include <sys/stat.h>
#include <string.h>
#include <limits.h>
#include <ctype.h>
#include <glib.h>
#include <ext2fs/ext2fs.h>


/* configuation */

#define BUF_ZERO_BLOCKS	0x100	/* *BLOCK, =128KB, in BSS section */
#define xDEBUG 1


/* implementation */

#ifndef O_BINARY
#define O_BINARY 0
#endif
#ifndef LINE_MAX
#define LINE_MAX 0x1000
#endif

#define BLOCK (0x200)	/* 512 */

static void finish(void) G_GNUC_NORETURN;

static int src_fd=-1,dst_fd=-1;
static const char *src_name,*dst_name;
static ext2_loff_t src_len,dst_len;	/* in blocks! */
static blk_t src_lenblk,dst_lenblk;	/* ==src_len,==dst_len, just a different type */

static ext2_loff_t stat_lastread,stat_todo,stat_bads,stat_largest,stat_todo_hunks;

struct range {
		ext2_loff_t start,end;
		};
static GTree *todo_tree;	/* struct range *key, NULL value */
static GTree *bad_tree;	/* struct range *key, NULL value */
static GMemChunk *node_memchunk;	/* struct range */

static gint tree_linearize_traverse(struct range *key,gpointer value/*unused*/,struct range **r_fillp/*data*/)
{
	*(*r_fillp)++=*key;
	g_mem_chunk_free(node_memchunk,key);
	return FALSE;
}

/* mallocated result, data-terminated, input "tree" destroyed */
static struct range *tree_linearize(GTree *tree)
{
struct range *r=g_new(struct range,g_tree_nnodes(tree)+1);
struct range *r_fill=r;

	g_tree_traverse(tree,(GTraverseFunc)tree_linearize_traverse,G_IN_ORDER,&r_fill/*data*/);
	g_assert(r_fill==r+g_tree_nnodes(tree));
	r_fill->start=r_fill->end=src_len+1;	/* terminator */
	g_tree_destroy(tree);
	return r;
}

/* returns terminator! */
static struct range *linear_join(struct range *linear,GCompareFunc comparefunc)
{
struct range *linear_term,*dst,*src;

	for (linear_term=linear;linear_term->start<=src_len;linear_term++);
	qsort(linear,linear_term-linear,sizeof(*linear),(int (*)(const void *, const void *))comparefunc);
	for (dst=src=linear;src->start<=src_len;src++) {
		if (dst>linear && dst[-1].end==src->start)
			dst[-1].end=src->end;
		else
			*dst++=*src;
		}
	*dst=*src;	/* terminator */
	return dst;
}

static gint range_compare(const struct range *range1,const struct range *range2)
{
ext2_loff_t len1,len2;
int r;

	g_return_val_if_fail(range1->end>range1->start,0);
	g_return_val_if_fail(range2->end>range2->start,0);

	len1=range1->end-range1->start;
	len2=range2->end-range2->start;

	/* big..small (start with large blocks) */
	if ((r=(len2 > len1) - (len2 < len1)))
		return r;
	/* big..small (reversed processing order) */
	if ((r=(range2->start > range1->start) - (range2->start < range1->start)))
		return r;
	return 0;
}

static gint range_compare_diskorder(const struct range *range1,const struct range *range2)
{
int r;

	g_return_val_if_fail(range1->end>range1->start,0);
	g_return_val_if_fail(range2->end>range2->start,0);

	/* small..big (linear disk processing order) */
	if ((r=(range2->start < range1->start) - (range2->start > range1->start)))
		return r;
	/* small..big (shouldn't have effect, just to get valid compare for ==) */
	if ((r=(range2->end   < range1->end  ) - (range2->end   > range1->end  )))
		return r;
	return 0;
}

static struct range *node_build(ext2_loff_t start,ext2_loff_t end)
{
struct range *r=g_mem_chunk_alloc(node_memchunk);

	r->start=start;
	r->end=end;
	return r;
}

static gint search_end(struct range *key,ext2_loff_t *endp/*data*/)
{
int r;

	if ((r=(*endp < key->end) - (*endp > key->end)))
		return r;
	return 0;
}

static gint search_start(struct range *key,ext2_loff_t *startp/*data*/)
{
int r;

	if ((r=(*startp < key->start) - (*startp > key->start)))
		return r;
	return 0;
}

static void node_add_bad(ext2_loff_t start,ext2_loff_t end)
{
struct range *key,*neigh_left,*neigh_right;

	if (start>=end)
		return;
	key=node_build(start,end);
	if ((neigh_left =g_tree_search(bad_tree,(GCompareFunc)search_end  ,&start))) {
		g_assert(neigh_left->end==key->start);
		key->start=neigh_left->start;
		g_tree_remove(bad_tree,neigh_left);
		g_mem_chunk_free(node_memchunk,neigh_left);
		}
	if ((neigh_right=g_tree_search(bad_tree,(GCompareFunc)search_start,&end  ))) {
		g_assert(neigh_right->start==key->end);
		key->end=neigh_right->end;
		g_tree_remove(bad_tree,neigh_right);
		g_mem_chunk_free(node_memchunk,neigh_right);
		}
	g_tree_insert(bad_tree,key/*key*/,NULL/*value*/);
	stat_bads+=(end-start);
}

static void node_add_todo(ext2_loff_t start,ext2_loff_t end)
{
	if (start>=end)
		return;
	g_tree_insert(todo_tree,node_build(start,end)/*key*/,NULL/*value*/);
	stat_todo_hunks++;
}

static gint traverse_fetchleftkey(gpointer key,gpointer value/*unused*/,gpointer *keyp/*data*/)
{
	*keyp=key;
	return TRUE;
}

static time_t stat_timelast;

static void stat_print(void)
{
time_t now;

	if (stat_timelast==(now=time(NULL)))
		return;
	stat_timelast=now;
	fprintf(stderr,"@%llu/%llu,TODO=%llu,bad=%llu,largest=%llu,hunks=%llu %10s\r",
			(unsigned long long)stat_lastread,
			(unsigned long long)src_len,
			(unsigned long long)stat_todo,
			(unsigned long long)stat_bads,
			(unsigned long long)stat_largest,
			(unsigned long long)stat_todo_hunks,
			"");
	fflush(stderr);
}

static ext2_loff_t simulate_bads=0;

static void user_input(ext2_loff_t pending_start,ext2_loff_t pending_end)
{
fd_set rfds;
struct timeval tv;
char buf[LINE_MAX];

	FD_ZERO(&rfds);
	FD_SET(0,&rfds);
	tv.tv_sec=0;
	tv.tv_usec=0;
	if (select(1,&rfds,NULL,NULL,&tv)<=0)
		return;
	if (buf!=fgets(buf,sizeof(buf),stdin)) {
		fprintf(stderr,"Weird stdin in user_input, assumed line \"f\"\n");
		strcpy(buf,"f");
		}
	stat_timelast=0;	/* reprint statistics line immeadiately */
	switch (*buf) {
		case 'f':
			node_add_todo(pending_start,pending_end);
			finish();	/* finish(): G_GNUC_NORETURN */
			/* NOTREACHED */
		case 'b': {
int bads=atoi(buf+1);

			simulate_bads+=(bads>=1 ? bads : 1);
			} break;
		default:
			fprintf(stderr,"Unrecognized key '%c': 'f'=finish(), 'b[<num>]'=simulate bads\n",
					(isprint(*buf) ? *buf : '?'));
		}
}

static void write_dst(ext2_loff_t start,ext2_loff_t end,const void *buf)
{
size_t len=(end-start)*BLOCK;
ssize_t gotssize;
ext2_loff_t gotext2_loff;

	if (dst_fd==-1)
		return;
	if (start*BLOCK!=(gotext2_loff=ext2fs_llseek(dst_fd,start*BLOCK,SEEK_SET))) {
		fprintf(stderr,"ext2fs_llseek(\"%s\",%llu)=%lld: %m\n",dst_name,(unsigned long long)(start*BLOCK),
				(signed long long)gotext2_loff);
		exit(EXIT_FAILURE);
		}
	if (len!=(gotssize=write(dst_fd,buf,len))) {
		/* see README/"Linux kernel flaw" */
		/* no "start==src_len-1" checking here, we may write larger blocks! */
		if (gotssize==((src_len&~1)-start)*BLOCK
				/* -1 instead of 0 is given when no bytes were written */
				|| (gotssize==-1 && start==src_len-1 && end==src_len))
			return;	/* suppress error message */

		fprintf(stderr,"write(\"%s\",@%llu-@%llu=%llu=%lluB)=%lld: %m\n",dst_name,
				(unsigned long long)start,
				(unsigned long long)end,
				(unsigned long long)(end-start),
				(unsigned long long)((end-start)*BLOCK),
				(  signed long long)gotssize);
		exit(EXIT_FAILURE);
		}
}

static void range_zero(ext2_loff_t start,ext2_loff_t end)
{
static const unsigned char buf_zero[BUF_ZERO_BLOCKS*BLOCK];

	g_return_if_fail(start<end);

	while (start<end) {
		ext2_loff_t mid=((end-start)>BUF_ZERO_BLOCKS ? start+BUF_ZERO_BLOCKS : end);
		g_assert(mid>=start);
		g_assert(mid<=end);

		write_dst(start,mid,buf_zero);

		start=mid;
		}
}

static void process(ext2_loff_t start,ext2_loff_t end)
{
unsigned char block_buf[BLOCK];
ext2_loff_t bads;
ext2_loff_t gotext2_loff;
ssize_t gotssize;

	g_return_if_fail(start<end);
	g_return_if_fail(end<=src_len);

#ifdef DEBUG
	fprintf(stderr,"process([%llu,%llu])\n",(unsigned long long)start,(unsigned long long)end);
#endif

restart:	/* continues when the block was successfuly read */
	stat_print();
	user_input(start,end);

	stat_lastread=start;
	if (start*BLOCK!=(gotext2_loff=ext2fs_llseek(src_fd,start*BLOCK,SEEK_SET))) {
		fprintf(stderr,"ext2fs_llseek(\"%s\",%llu)=%lld: %m\n",src_name,(unsigned long long)(start*BLOCK),
				(signed long long)gotext2_loff);
		exit(EXIT_FAILURE);
		}
	stat_todo--;	/* for the forthcoming read() */
	if (!simulate_bads && (sizeof(block_buf)==(gotssize=read(src_fd,block_buf,sizeof(block_buf)))
					/* see README/"Linux kernel flaw" */
					|| (start==src_len-1 && gotssize==((src_len&~1)-start)*BLOCK))
			) {
		/* success read */
		if (gotssize==sizeof(block_buf))
			write_dst(start,start+1,block_buf);
		else {
			/* handle README/"Linux kernel flaw" */
			fprintf(stderr,"WARNING: Ignored disk-last sector read failure; see README/\"Linux kernel flaw\"\n");
			range_zero(start,start+1);
			}
		start++;
		if (start>=end)	/* finished */
			return;
		goto restart;
		}

	/* failed read */
	if (!simulate_bads)
		bads=1; /* one block was detected */
	else {
		bads=simulate_bads;
		simulate_bads=0;
		}
	if (start+bads>src_len)
		bads=src_len-start;
	node_add_bad(start,start+bads);
	start+=bads;
	while (start<end) {
ext2_loff_t mid=(start+end)/2;

		g_assert(mid<end);
		node_add_todo(mid,end);
		end=mid;
		}
}

static void finish(void)
{
struct range *todo_linear=tree_linearize(todo_tree),*todol;
struct range *bad_linear =tree_linearize(bad_tree ),*badl ;

	linear_join(todo_linear,(GCompareFunc)range_compare_diskorder);

	fputc('\n',stderr);
	for (todol=todo_linear,badl=bad_linear;
			todol->start<=src_len || badl->start<=src_len;) {
ext2_loff_t start,end;

		g_assert(badl->start!=todol->start);
		if (badl->start<todol->start) {
			g_assert(todol->start>=badl->end);
			if (todol->start==badl->end) {
				printf("%llu-%llu\t;BAD=%llu, TODO(@%llu-@%llu)=%llu",
						(unsigned long long)(start=badl->start),
						(unsigned long long)(end=todol->end),
						(unsigned long long)(badl->end-badl->start),
						(unsigned long long)todol->start,
						(unsigned long long)todol->end,
						(unsigned long long)(todol->end-todol->start));
				todol++;
				}
			else
				printf("%llu-%llu\t;BAD=%llu",
						(unsigned long long)(start=badl->start),
						(unsigned long long)(end=badl->end),
						(unsigned long long)(badl->end-badl->start));
			badl++;
			}
		else {
			g_assert(badl->start>todol->start);
			printf("%llu-%llu\t;TODO=%llu",
					(unsigned long long)(start=todol->start),
					(unsigned long long)(end=todol->end),
					(unsigned long long)(todol->end-todol->start));
			todol++;
			}
		fflush(stdout);
		range_zero(start,end);
		putchar('\n');
		}
	exit(EXIT_SUCCESS);
}

int main(int argc,char **argv)
{
	/* prevent block-buffering such as: badblock-guess.c|tee x.log */
	setlinebuf(stdout);
	setlinebuf(stderr);

	if (argc!=2 && argc!=3) {
		fprintf(stderr,"\
badblock-guess, Copyright (C) 2002 Jan Kratochvil <short@ucw.cz>\n\
$Id$\n\
badblock-guess comes with ABSOLUTELY NO WARRANTY.\n\
This is free software, and you are welcome to redistribute it\n\
under certain conditions.\n\
\n\
Syntax: badblock-guess <src_dev> [<dst_dev (OVERWRITTEN & DESTROYED!!!)>]\n\
\n");
		exit(EXIT_FAILURE);
		}
	if (-1==(src_fd=open64((src_name=argv[1]),O_RDONLY|O_BINARY))) {
		fprintf(stderr,"open64(\"%s\",O_RDONLY|O_BINARY): %m\n",src_name);
		exit(EXIT_FAILURE);
		}
	if (argv[2] && -1==(dst_fd=open64((dst_name=argv[2]),O_WRONLY|O_BINARY))) {
		fprintf(stderr,"open64(\"%s\",O_WRONLY|O_BINARY): %m\n",dst_name);
		exit(EXIT_FAILURE);
		}
	if (dst_fd!=-1) {
struct stat src_stat,dst_stat;

			if (fstat(src_fd,&src_stat)) {
				fprintf(stderr,"fstat(\"%s\"): %m\n",src_name);
				exit(EXIT_FAILURE);
				}
			if (fstat(dst_fd,&dst_stat)) {
				fprintf(stderr,"fstat(\"%s\"): %m\n",dst_name);
				exit(EXIT_FAILURE);
				}
      if ((S_ISBLK(src_stat.st_mode) && S_ISBLK(dst_stat.st_mode) && src_stat.st_rdev==dst_stat.st_rdev)
					|| (src_stat.st_dev==dst_stat.st_dev && src_stat.st_ino==dst_stat.st_ino)) {
				fprintf(stderr,"Refusing to operate on two identical devs (src=\"%s\" and dst=\"%s\")\n",src_name,dst_name);
				exit(EXIT_FAILURE);
				}
		}
	if (ext2fs_get_device_size(src_name,BLOCK,&src_lenblk)) {
		fprintf(stderr,"ext2fs_get_device_size(\"%s\",%d,...): %m\n",src_name,BLOCK);
		exit(EXIT_FAILURE);
		}
	src_len=src_lenblk;
	if (src_len<=0) {
		fprintf(stderr,"\"%s\" length %llu <=0\n",src_name,(unsigned long long)src_len);
		exit(EXIT_FAILURE);
		}
	if (dst_fd!=-1) {
		if (ext2fs_get_device_size(dst_name,BLOCK,&dst_lenblk)) {
			fprintf(stderr,"ext2fs_get_device_size(\"%s\",%d,...): %m\n",dst_name,BLOCK);
			exit(EXIT_FAILURE);
			}
		dst_len=dst_lenblk;
		if (dst_len<=0) {
			fprintf(stderr,"\"%s\" length %llu <=0\n",dst_name,(unsigned long long)dst_len);
			exit(EXIT_FAILURE);
			}
		if (dst_len<src_len) {
			fprintf(stderr,"WARNING: Ignoring exceeding trailing data of \"%s\" from %llu up to its end %llu!\n",
					src_name,(unsigned long long)dst_len,(unsigned long long)src_len);
			src_len=dst_len;
			src_lenblk=dst_lenblk;
			}
		}

	node_memchunk=g_mem_chunk_new("node_memchunk",sizeof(struct range)/*atom_size*/,
			0x1000 /*are_size: 64KB*/,G_ALLOC_AND_FREE);
	todo_tree=g_tree_new((GCompareFunc)range_compare);
	bad_tree=g_tree_new((GCompareFunc)range_compare_diskorder);
	node_add_todo(0,src_len);
	stat_todo=src_len;

	for (;;) {
struct range *left=NULL;	/* key */

		g_tree_traverse(todo_tree,(GTraverseFunc)traverse_fetchleftkey,G_IN_ORDER,&left/*data*/);
		if (!left)	/* done */
				break;
#ifndef G_DISABLE_ASSERT
		{
			gint nnodes_preremove=g_tree_nnodes(todo_tree);
#endif
		g_tree_remove(todo_tree,left/*key*/);
#ifndef G_DISABLE_ASSERT
			g_assert(g_tree_nnodes(todo_tree)==nnodes_preremove-1);
			}
#endif
		stat_largest=(left->end-left->start);
		process(left->start,left->end);
		g_mem_chunk_free(node_memchunk,left);
		stat_todo_hunks--;
		}
	finish();	/* finish(): G_GNUC_NORETURN */
	/* NOTREACHED */
}