/* This program implements Block ECC using a Forward Error Correction (FEC)
code based on Vandermonde (VDM) matrices in GF(2^8) due to Luigi Rizzo.

Given a blocksize, and the VDM parameters K and N, N blocks are written for
every chunk of K input packets, so that N - K blocks can be lost per chunk,
without error.

This means that large consecutive sections of files can be lost, and the
data recovered.  For example, diskettes typically lose whole sectors at
once, or related groups of sectors, or even entire tracks.  Data written to
a diskette with this program may be recovered even with many read errors.

This program is distributed under the GNU GENERAL PUBLIC LICENSE, version 2. */

#include "defs.h"
#include "util.h"
#include "md5.h"
#include "fec.h"
#include "version.h"

char Usage[] =
	"usage: %s [-v] [-d] [-b<blocksize>] [-n<n>] [-k<k>] [file]\n";
#define USAGE() msg(Usage, Progname)

/* Default parameters, suitable for standard diskettes. */

int Blocksize = 1024;
int N = 18;
int K = 14;

int Decode = FALSE;             /* -d means decode, else encode */
int Verbose = 0;                /* -v or -vv print informative stuff */

int Pipe = FALSE;               /* set automatically if the input is a pipe */

/* Various parameters used by both the encoder and the decoder. */

void *Fec;                      /* opaque FEC code tables & stuff */
u_char *Chunk;                  /* space for K buffers of size Blocksize */
u_char **Packet;                /* K pointers into Chunk */
u_char *Block;                  /* another buffer of size Blocksize */
int Packlen;                    /* size of data packets within blocks */
int Chunksize;                  /* length of encoding unit == K packets */
int *Idx;                       /* index array of size K, for FEC decoder */

/* Block footer. */

#define BLOCK_FOOT_SIZE (MD5_DIGEST_SIZE + 2)

/* Chunk header. */

#define CHUNK_HEAD_SIZE 8       /* clen and chkid */
#define CHUNK_FINAL 0xFFFFFFFF  /* flag last chunk */

/* Stats. */

long Read_count;
long Write_count;

/* Forward decls. */

void encode(FILE *);
int read_chunk(u_int32_t *chkid, FILE *infile);
void write_block(int id, u_char *packet);

void decode(FILE *);
u_int32_t decode_chunk(long offset, FILE *infile, u_int32_t *clenp);
int read_block(long pos, u_char *block, FILE *infile);
int check_block(u_char *block);

int main(int argc, char **argv)
{
	int i, c;
	FILE *infile;

	Progname = *argv;

	/* Parse option arguments. */

	while ((c = getopt(argc, argv, "vdb:n:k:")) != EOF) {
		switch (c) {

		case 'v':
			Verbose++;
			break;

		case 'd':
			Decode = TRUE;
			break;

		case 'b':
			Blocksize = atoik(optarg);
			break;

		case 'n':
			N = atoi(optarg);
			break;

		case 'k':
			K = atoi(optarg);
			break;

		default:
			USAGE();
			exit(1);
		}
	}
	argc -= optind;
	argv += optind;

	/* Error check arguments. */

	if (N < 2 || N > GF_SIZE + 1 || K < 1 || K > N) {
		fatalerr("invalid (k, n)");
	}
	if (Blocksize < BLOCK_FOOT_SIZE + CHUNK_HEAD_SIZE) {
		fatalerr("invalid blocksize");
	}

	if (argc > 1) {
		USAGE();
		exit(1);
	}

	/* Read from stdin or the given file. */

	infile = stdin;
	if (argc == 1) {
		infile = fileopen(argv[0], "r");
	}

	/* Initialize.  Allocate enough memory for these arrays to be used
	by both the encoder and the decoder. */

	Fec = fec_new(K, N);
	Chunk = new_array(u_char, K * Blocksize);
	Packet = new_array(u_char *, K);
	for (i = 0; i < K; i++) {
		Packet[i] = Chunk + i * Blocksize;
	}
	Block = new_array(u_char, Blocksize);
	Packlen = Blocksize - BLOCK_FOOT_SIZE;
	Chunksize = Packlen * K;
	Idx = new_array(int, K);

	/* If this program is called *decode, or if -d was passed,
	we're a decoder, otherwise we're an encoder. */

	i = strlen(Progname);
	if (i >= 6 && strcmp(&Progname[i - 6], "decode") == 0) {
		Decode = TRUE;
	}

	/* In decode mode, we expect read errors, which leave the stream
	position undefined.  So we use seek.  However, it is sometimes
	useful to do the reading with a different program, e.g. dd, and
	pipe the output to the decoder.  So we also want to allow reading
	from non-seekable streams.  We detect whether or not the stream is
	seekable by trying to seek to the start and checking for errors. */

	if (fseek(infile, 0L, SEEK_SET) < 0) {
		if (errno == ESPIPE) {
			Pipe = TRUE;
		}
	}

	if (Verbose) {
		msg("Vandermonde Forward Error Correction v%s\n", VERSION);
		msg("(K, N) = (%d, %d)\n", K, N);
		msg("Blocksize = %d\n", Blocksize);
		if (Decode) {
			msg("Decode mode.\n");
			if (Pipe) {
				msg("Input is a pipe.\n");
			}
			if (Verbose > 1) {
				msg("Reading data in chunks of %d bytes.\n",
					N * Blocksize);
				msg("Writing data in chunks of %d bytes.\n",
					K * Packlen - CHUNK_HEAD_SIZE);
			}
		} else {
			msg("Encode mode.\n");
			if (Verbose > 1) {
				msg("Reading data in chunks of %d bytes.\n",
					K * Packlen - CHUNK_HEAD_SIZE);
				msg("Writing data in chunks of %d bytes.\n",
					N * Blocksize);
			}
		}
	}

	Read_count = 0;
	Write_count = 0;

	if (Decode) {
		decode(infile);
	} else {
		encode(infile);
	}

	if (!Decode && Verbose) {
		msg("Read %ld bytes, wrote %ld bytes.\n",
			Read_count, Write_count);
		msg("Expansion factor = %g.\n",
			(double)Write_count / Read_count);
	}

	return 0;
}

/*** The encoder. ***/

/* Read chunks of data, encode them, and write them out.  Each input
chunk yields N output blocks. */

void encode(FILE *infile)
{
	int i;
	u_int32_t chkid;

	chkid = 0;
	while (read_chunk(&chkid, infile)) {
		for (i = 0; i < N; i++) {

			/* Get the i'th code block for this set of packets,
			and write it out. */

			fec_encode(Fec, Packet, Block, i, Packlen);
			write_block(i, Block);
		}
		if (chkid == CHUNK_FINAL) {
			break;
		}
		chkid++;
	}
}

/* Read a chunk from the input stream and return the amount read.  On EOF,
the chunk is partially filled but padded with zeros, and the chkid is set
to CHUNK_FINAL.  In any case the chunk header (clen and chkid) is set. */

int read_chunk(u_int32_t *chkid, FILE *infile)
{
	int i, len, clen, rlen;
	u_char *data;
	u_int32_t *p;

	/* Initialize and zero pad. */

	clen = 0;
	memset(Chunk, 0, K * Blocksize);

	/* Try to read a whole chunk. */

	for (i = 0; i < K; i++) {
		len = Packlen;
		data = Packet[i];

		/* The first packet in a chunk has a header. */

		if (i == 0) {
			len -= CHUNK_HEAD_SIZE;
			data += CHUNK_HEAD_SIZE;
		}

		/* Read the packet.  A short read ends the chunk. */

		rlen = fread(data, 1, len, infile);
		clen += rlen;
		Read_count += rlen;
		if (rlen < len) {
			break;
		}
	}

	/* Get EOF status by trying to read a byte and then pushing it back. */

	i = getc(infile);
	if (i == EOF) {
		*chkid = CHUNK_FINAL;
	} else {
		ungetc(i, infile);
	}

	/* Set the header. */

	p = (u_int32_t *)Packet[0];
	*p++ = htonl(clen);
	*p = htonl(*chkid);

	if (Verbose > 1) {
		msg("read a chunk of size %d\n", clen);
	}

	return clen;
}

/* Write a block consisting of an output packet and a hash. */

void write_block(int id, u_char *packet)
{
	u_char footer[BLOCK_FOOT_SIZE];
	struct MD5Context ctx;

	/* First, write the output packet. */

	if (fwrite(packet, 1, Packlen, stdout) != Packlen) {
		fatalerr("write error");
	}

	/* Put the block id and K into the first two bytes of the footer. */

	footer[0] = id;
	footer[1] = K;

	/* Then, make an MD5 hash and write the block footer. */

	MD5Init(&ctx);
	MD5Update(&ctx, packet, Packlen);
	MD5Update(&ctx, footer, 2);
	MD5Final(footer + 2, &ctx);

	if (fwrite(footer, 1, BLOCK_FOOT_SIZE, stdout) != BLOCK_FOOT_SIZE) {
		fatalerr("write error");
	}

	Write_count += Packlen + BLOCK_FOOT_SIZE;
}

/*** The decoder. ***/

/* Read blocks of data N at a time, decode them, and write them out.  Each N
input blocks yield one output chunk of K packets. */

void decode(FILE *infile)
{
	int i, j;
	u_char *data;
	long offset;
	u_int32_t clen, chkid, id;

	offset = 0;
	id = 0;
	do {
		chkid = decode_chunk(offset, infile, &clen);
		if (chkid != CHUNK_FINAL && chkid != id) {
			fatalerr("chunk out of order at %ld id %ld expecting %ld",
				offset, (long)chkid, (long)id);
		}
		id++;

		/* We've got a good chunk.  Write out K packets,
		or clen bytes, whichever is first. */

		i = 0;
		while (clen > 0 && i < K) {
			data = Packet[i];
			j = Packlen;

			/* The first packet in a chunk has a header. */

			if (i == 0) {
				j -= CHUNK_HEAD_SIZE;
				data += CHUNK_HEAD_SIZE;
			}

			/* The last packet may be short (clen < j), so
			only write what's there. */

			if (clen < j) {
				j = clen;
			}
			fwrite(data, 1, j, stdout);
			i++;
			clen -= j;
		}
		offset += N * Blocksize;
	} while (chkid != CHUNK_FINAL);
}

/* From the next N input blocks starting at offset, we need K good ones.
If this is a pipe, we'll read all N, but otherwise we'll stop after K.
Return the chkid and clen from a good chunk. */

u_int32_t decode_chunk(long offset, FILE *infile, u_int32_t *clenp)
{
	int i, id, blkid;
	long pos;
	u_int32_t *p, chkid;

	/* Read K good blocks.  id is the FEC block number relative to
	this chunk, i counts how many good ones we've read. */

	i = id = 0;
	while (i < K && id < N) {

		/* pos is where we think the block is.  It's ignored for
		pipes, so in that case we'd better be in the right place
		(if we're not, the hash will probably fail). */

		pos = offset + id * Blocksize;
		blkid = read_block(pos, Packet[i], infile);
		if (blkid == id) {

			/* Got a good one, save it. */

			Idx[i++] = id;
		} else if (Verbose > 1) {
			msg("bad block at %ld id %d expecting %d\n",
				pos, blkid, id);
		}
		id++;
	}
	if (i < K) {
		fatalerr("unrecoverable error -- too many bad blocks");
	}

	/* If this is a pipe, we want to read the next N - id blocks.
	If there are read errors, we'll be out of sync and likely fail. */

	if (Pipe) {
		while (id < N) {
			fread(Block, 1, Blocksize, infile);
			id++;
		}
	}

	/* We have K good packets.  Decode them in place. */

	fec_decode(Fec, Packet, Idx, Packlen);

	/* Get clen and chkid. */

	p = (u_int32_t *)Packet[0];
	*clenp = ntohl(*p);
	p++;
	chkid = ntohl(*p);

	if (Verbose > 1) {
		msg("read %d blocks, clen = %d\n", K, *clenp);
	}
	return chkid;
}

/* Read a block and check the hash.  If a seek or read error occurs the
block is presumed bad.  Return the block id for a good block, a negative
code for a bad one.  Since a failed read() may or may not change the file
position, we explicitly seek to the correct position before every read,
unless this is a pipe. */

int read_block(long pos, u_char *block, FILE *infile)
{
	u_char *p;

	if (!Pipe && fseek(infile, pos, SEEK_SET) < 0) {
		return -1;
	}
	if (fread(block, 1, Blocksize, infile) != Blocksize) {
		return -2;
	}
	if (check_block(block)) {
		p = block + Packlen;
		if (p[1] != K) {
			fatalerr("incorrect K value");
		}
		return p[0];
	}
	return -3;
}

/* Check the hash on a block.  Return TRUE if it's OK, FALSE otherwise. */

int check_block(u_char *block)
{
	u_char digest[MD5_DIGEST_SIZE];
	struct MD5Context ctx;

	/* Make an MD5 hash of the packet data and the first two bytes of
	the footer. */

	MD5Init(&ctx);
	MD5Update(&ctx, block, Packlen + 2);
	MD5Final(digest, &ctx);
	if (memcmp(digest, block + Packlen + 2, MD5_DIGEST_SIZE) == 0) {
		return TRUE;
	}
	return FALSE;
}
