////////////////////////////////////////////////////////////////////////////////
//
//	xfrchk.cpp - a program for verifying I/O bus transfer integrity
//
//	(C)2009, Josh Moyer <jmoyer@nodomain.net>.  All rights reserved.
//
////////////////////////////////////////////////////////////////////////////////

#define DBG

#include <errno.h>
#include <fcntl.h>
#include <iostream>
#include <linux/fs.h>
#include <stdlib.h>
#include <string>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

using namespace std;

class xfrchk
{
public:
	size_t		 bs;	// buffer size

	short		 cc;	// cache control

	int			 fd;	// file descriptor
	string		 fp;	// file path
	int			 ff;	// file open flags
	struct stat	*fs;	// file stats

	off64_t		 ls;	// file start location	
	off64_t		 lr;	// file current read location
	off64_t		 lw;	// file current write location
	off64_t		 le;	// file end location
	
	void		*pc;	// comparison buffer
	void		*pi;	// input buffer
	void		*po;	// output buffer

	long		 ps;	// page size

	int 		 ss;	// sector size
	
	int			  r;	// retval

	// filesystem cache control constants
	static const short cc_p = 0x0;	// POSIX (posix_fadvise)
	static const short cc_m = 0x1;	// BSD/Linux (madvise)
	static const short cc_d = 0x2;	// Linux (open(,,ODIRECT)

	// seek patterns
	static const short sk_asc = 0x0;	// sequential, incrementing
	static const short sk_desc = 0x1;	// sequential, decrementing
	static const short sk_rnd = 0x2;	// random

	xfrchk(string pathname, size_t buffersize, short cachecontrol)
	{
		bs = buffersize;
		cc = cachecontrol;
		fd = -1;			
		ff = O_RDWR | O_LARGEFILE; //| O_SYNC
		fp = pathname;
		fs = new struct stat;
		le = NULL;
		lr = NULL;
		ls = NULL;
		lw = NULL;
		pc = NULL;
		pi = NULL;
		po = NULL;
		ps = sysconf(_SC_PAGESIZE);
	 	r = 0;
		ss = -1;

		posix_memalign(&pc, (size_t)ps, bs);
		posix_memalign(&pi, (size_t)ps, bs);
		posix_memalign(&po, (size_t)ps, bs);
		
		if (pc == NULL || pi == NULL || po == NULL || fs == NULL)
		{
			cerr << "failed to allocate buffers" << endl;
			throw;
		}
		
		#ifdef DBG
		cout <<
			"allocated pc@" << pc <<
			", pi@" << pi <<
			", and po@" << po << endl;
		#endif DBG

		// ensure that data is actually tranferred across the storage bus and 
		// is not just tranferred from the filesystem cache.  O_DIRECT *is*
		// best here -- Linus' position against it is curious considering that
		// posix_[f|m]advise() advice can be ignored by the system.
		if (cc == cc_d)	fd = open(fp.c_str(), O_DIRECT | ff);
		else fd = open(fp.c_str(), ff);
		
	   	if (fd == -1)
		{
			cerr << "opening " << fp <<
			" failed with " << strerror(errno) << endl;
			throw;
		}
		
		// DIOCGSECTORSIZE on BSDs?
		if (ioctl(fd, BLKSSZGET, &ss)) cerr <<
			"unable to obtain device sector size." << endl;

		if (fstat(fd, fs))
		{
			cerr << "fstat returned " << strerror(errno) << endl;
			throw;
		}

		#ifdef DBG
		cout << "st_size: " << (fs->st_size) << endl;
		#endif DBG

		r = 0;
		errno = 0;

		switch (cc)
		{
			case cc_m:
				r |= madvise(pc, bs, MADV_DONTNEED);
				r |= madvise(pi, bs, MADV_DONTNEED);
				r |= madvise(po, bs, MADV_DONTNEED);
				break;
			case cc_p:
				errno = posix_fadvise(
					fd,
					0, 
					(fs->st_size>0) ? fs->st_size : 0,
					POSIX_FADV_DONTNEED | POSIX_FADV_RANDOM);
				break;
			case cc_d:
				break;
			default:
				cerr << "unknown cache control method specified" << endl;
				throw;
				break;
		}
		
		if (r || errno)
		{
			cerr <<
				"system's response to advice was " << strerror(errno) << endl;
			throw; 
		}
	}
	
	// Contents of po may change after write() if there was a transfer error.
	// Because of this, using write() incurs the risk of data loss.

	// Writes the contents of buffer po to disk starting at location lw, 
	// with an optional number of verifies.
	int write(off64_t *lw, unsigned int verifies)
	{
		int r = 0;
		unsigned int v = verifies;
		
		r = pwrite(fd, po, bs, (*lw)*(off64_t)ss);

		if (r != (int)bs) return r;

		if (v > 0) r = this->read(lw, v, po);

		if (r != (int)bs) return r;

		*lw = *lw + (bs/ss);

		return r;
	}
	
	void fillbuf(long p, void *bp, size_t bs)
	{
	  for (size_t i=0; i < (bs/sizeof(p)); i++)
		{
			*(((long*)bp)+i) = p;
		}
	}

	//verify failed (-2)
	//eof 0
	//os reported error -1
	//incomplete read 
	int read(off64_t *lr, int rereads)
	{
		return read(lr, rereads, NULL);
	}

	int read(off64_t *lr, int rereads, void *po)
	{

		long i = 0;			// iteration counter
		int r = 0;			// result code
		int rr = rereads;
		void *pp;			// previous buffer location for compares
	    void *ps;			// saved location of pi when po passed in
		bool v = false;		// verify?

		if (po != NULL)
		{
		  	v = true;
			ps = pi;
		  	pi = po;
		}
		
		#ifdef DBG
		cout << "pre-read: pc@" << pc <<
			", pi@" << pi <<
			", and po@" << po << endl;
		#endif DBG

		do
		{
			cout << *lr << "(" << rr << ")";			

			pp = pi;
			pi = pc;
			pc = pp;
			
			r = pread(fd, pi, bs, (*lr)*(off64_t)ss);
			
			if (r != (long)bs)
			{
				if (r < 1) cout << ":"<< r << "/" << bs << endl;
				return r;
       		}

			if (v)
			{				
				for (i=0 ; i<(long)bs/sizeof(long) ; i++)
				{
					if (*((long*)pi+i) != *((long*)pc+i))
					{
						cout << "!" << endl;
						errno = EIO;
						r = -1;
						goto postread;
					}
				}
			}
			
			v = true;
	
			cout << "\033[0K\r";  // VT100 escape sequence
		} while (rr--);

		*lr = *lr + (bs/ss);
		
		postread:
		#ifdef DBG
		cout <<
			"post-read: pc@" << pc <<
			", pi@" << pi <<
			", and po@" << po << endl;
		#endif DBG

		if (po != NULL) (po == pi) ? pi = ps : pc = ps;

	  	#ifdef DBG
		cout <<
			"returning: pc@" << pc <<
			", pi@" << pi <<
			", and po@" << po << endl;
		#endif DBG

		return r;
	}

	
	virtual ~xfrchk()
	{
		delete fs;
		fd = close(fd);
		
		if (fd != 0) 
		{
			cerr <<
				"closing " << fp <<
				" failed with " << strerror(errno) << endl;
			throw;
		}
		
		#ifdef DBG
		cout <<
			"about to free() pc@" << pc <<
			", pi@" << pi <<
			", and po@" << po << endl;
		#endif DBG

		free(pc);
		free(pi);
		free(po);
	}
};

void test_constructor()
{
 	xfrchk *xc;
	
	// getpagesize(2) & ioctl(getsecsize) 
	string fp[] = {
	 	//"/dev/null",		// fails with O_DIRECT as invalidarg
		"/dev/sdd",
		//"/proc",			// fails because is dir
		//"/dev/foobar",	// fails with file not found
		//"/dev/nvidia0",	// fails with invalidarg on O_DIRECT   
		"/usr/share/dict/american-english"
	};
	int bs[] = { /*4096, 4133,*/ 512,/* 32768, 35000,*/ 23, 1, 0/* -1 */ };
	short cc[] = { xfrchk::cc_p, xfrchk::cc_d, /*xfrchk::cc_m, -1*/ }; 
	
	for (unsigned int d = 1;
		 d != 0;
		 --d)
	 {
		for (unsigned int cci = 0;
			 cci < (sizeof(cc)/sizeof(cc[cci]));
			 cci++)
		{
			for (unsigned int bsi = 0;
				 bsi < (sizeof(bs)/sizeof(bs[bsi]));
				 bsi++)
			{
				for (unsigned int fpi = 0;
					 fpi < (sizeof(fp)/sizeof(fp[fpi]));
					 fpi++)
				{
					cout <<
						fp[fpi] << ", " <<
						bs[bsi] << ", " <<
						cc[cci] << endl;
					xc = new xfrchk(fp[fpi], bs[bsi], cc[cci]);
					if (d) delete xc;
				}
			}
		}
	}
}

void test_read(xfrchk *xc)
{
	xc->lr = 0;
	xc->read(&xc->lr,NULL);
	xc->lr = 0;
	xc->read(&xc->lr,0);
	//xc->lr = 0;
	//xc->read(xc->lr,-1); // infinite loop
	xc->lr = 0;
	xc->read(&xc->lr,256);
	//xc->lr = 0;
	//xc->read(0,1024*1024);

	//for (size_t i=0;i < 0x4000;i++)
	//xc->bs = 512;
	xc->bs = 4096;
	//xc->bs = 65536;
	for (xc->lr = 0; xc->lr < 625142448; xc->lr++)
	{
		xc->read(&xc->lr, 0);
	}
}

void test_write(xfrchk *xc)
{
	int r = 0;
	xc->lw = 0 ; xc->write(&xc->lw,0);
	xc->lw = 0 ; xc->write(&xc->lw,1);
	xc->lw = 0 ; xc->write(&xc->lw,2);
	
	xc->lw = 0;
	for (int i=0 ; i < 65536 ; i++)
	{
		r = xc->write(&xc->lw,1);
		if (r != (int)xc->bs) cout << "error";
	}
}

int main()
{  
	xfrchk *xc = NULL;
	string dp;
	
	cout << "Device Path: ";
	cin >> dp;

	xc = new xfrchk(dp, 2048, xfrchk::cc_d);

	xc->fillbuf(0xF0F0F0F0F0F0F0F0, xc->po, xc->bs);

	//Test_constructor(xc);
	//test_read(xc);
	test_write(xc);
	//test(xc);

	delete xc;
}