// code by SASANO Takayoshi, CC BY-SA

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <endian.h>
#include <fcntl.h>
#include <time.h>
#include <unistd.h>
#include <sys/ioctl.h>

#if defined(__OpenBSD__)
#include <sys/dkio.h>
#include <sys/disklabel.h>
#elif defined(__linux__)
#include <linux/fs.h>
#endif

#define	BUFFERSIZE	(1024 * 1024 * 1)
#define	SECTORSIZE	512
#define	SECTORS		(BUFFERSIZE / SECTORSIZE)
#define	ELEMENTS	(SECTORSIZE / sizeof(uint64_t))

#define	TEST_WRITE	0x1000
#define	TEST_READ	0x2000
#define	TEST_SET	0x4000
#define	TEST_ENDLESS	0x8000
#define	METHOD_MASK	0x0fff

#define	METHOD_ZERO		0x0000
#define	METHOD_ONE		0x0001
#define	METHOD_RANDOM		0x0002
#define	METHOD_64BIT		0x0004
#define	METHOD_32BIT		0x0005
#define	METHOD_16BIT		0x0006
#define	METHOD_8BIT		0x0007
#define	METHOD_NONE		0x0008
#define	METHOD_NOT		0x0400
#define	METHOD_BIGENDIAN	0x0800

static	uint8_t	buf[BUFFERSIZE];
static	uint32_t	myrandom_state;

/* pseudo random */
static	uint8_t	myrandom(void)
{
	myrandom_state *= 13;
	myrandom_state += 0x1a2b3c4d;
	return myrandom_state >> 24;
}

/* create test data from sector value */
static	uint64_t	test_data(uint64_t secno, int method)
{
	int	i;
	uint64_t	value;

	value = secno;

	switch (method & 0x07) {
	case	METHOD_8BIT:
		value &= 0xff;
		value |= (value << 8);
		/* FALLTHROUGH */
	case	METHOD_16BIT:
		value &= 0xffff;
		value |= (value << 16);
		/* FALLTHROUGH */
	case	METHOD_32BIT:
		value &= (uint64_t)0xffffffff;
		value |= (value << 32);
		/* FALLTHROUGH */
	default:
	case	METHOD_64BIT:
		break;

	case	METHOD_ZERO:
		value = 0;
		break;

	case	METHOD_ONE:
		value = ~(uint64_t)0;
		break;

	case	METHOD_RANDOM:
		for (i = 0; i < 8; i++) {
			value <<= 8;
			value |= myrandom();
		}
		break;
	}

	value = (method & METHOD_BIGENDIAN) ? be64toh(value) : le64toh(value);
	value ^= (method & METHOD_NOT) ? ~(uint64_t)0 : 0;

	return value;
}

/* fill sector with test data */
static	void	fill_sector(uint64_t *p, void **q, uint64_t secno, int method)
{
	int	i;

	for (i = 0; i < ELEMENTS; i++) {
		*p++ = test_data(secno, method);
	}

	*q = p;
}

/* fill buffer with test data */
static	void	fill_buffer(uint64_t secno, uint64_t sectors, int method)
{
	int	i;
	void	*p;

	p = buf;
	
	for (i = 0; i < sectors; i++) {
		fill_sector(p, &p, secno, method);
		secno++;
	}
}

/* check sector with test data */
static	int	check_sector(uint64_t *p, void **q, uint64_t secno, int method, int verbose)
{
	int	i, result;
	uint64_t	data;

	result = 0;

	for (i = 0; i < ELEMENTS; i++) {
		data = test_data(secno, method);

		if (data != *p) {
			if (verbose) {
				printf("offset %ld: %llx -> %llx\n",
				       i * (SECTORSIZE / ELEMENTS),
				       data, *p);
			}
			result = 1;
		}
		p++;
	}

	*q = p;

	return result;
}

/* check buffer with test data */
static	void	check_buffer(uint64_t secno, uint64_t sectors, int method)
{
	int	i;
	void	*p, *q;
	uint32_t	saved_state;

	p = buf;

	for (i = 0; i < sectors; i++) {
		saved_state = myrandom_state;
		if (check_sector(p, &q, secno, method, 0)) {
			printf("sector %lld: check failed\n", secno);
			myrandom_state = saved_state;
			check_sector(p, &q, secno, method, 1);
		}
		secno++;
		p = q;
	}
}

/* write test */
static	void	write_test(int fd, int method, uint64_t sectors)
{
	uint64_t	i, remain;

	myrandom_state = 0;

	for (i = 0; i < sectors; i += SECTORS) {
		remain = sectors - i;
		if (remain > SECTORS) {
			remain = SECTORS;
		}

		if (!(method & METHOD_NONE)) {
			fill_buffer(i, remain, method);
		}

		if (pwrite(fd, buf, remain * SECTORSIZE,
			   i * SECTORSIZE) < remain * SECTORSIZE) {
			printf("sector %lld: write error\n", i);
		}
	}

	sync();

	return;
}

/* read test */
static	void	read_test(int fd, int method, uint64_t sectors)
{
	uint64_t	i, remain;

	myrandom_state = 0;

	for (i = 0; i < sectors; i += SECTORS) {
		remain = sectors - i;
		if (remain > SECTORS) {
			remain = SECTORS;
		}

		if (pread(fd, buf, remain * SECTORSIZE,
			  i * SECTORSIZE) < remain * SECTORSIZE) {
			printf("sector %lld: read error\n", i);
		}

		if (!(method & METHOD_NONE)) {
			check_buffer(i, remain, method);
		}
	}

	return;
}

static	void	sync_sec(void)
{
	time_t	t;

	t = time(NULL);

	while (t == time(NULL));
}

/* get total sectors */
static	uint64_t	get_sectors(int fd)
{
	int	n;
	off_t	size;
	uint64_t	sectors;
#if defined(__OpenBSD__)
	struct disklabel dl;
#elif defined(__linux__)
	unsigned long	v;
#endif
	

#if defined(__OpenBSD__)
	n = ioctl(fd, DIOCGPDINFO, &dl);
#elif defined(__linux__)
	n = ioctl(fd, BLKGETSIZE, &v);
#endif

	/* get total sectors */
	if (n < 0) {
		size = lseek(fd, 0, SEEK_END);
		lseek(fd, 0, SEEK_SET);
		sectors = size / SECTORSIZE;
	} else {
#if defined(__OpenBSD__)
		sectors = DL_GETDSIZE(&dl);
		n = dl.d_secsize % SECTORSIZE;
#elif defined(__linux__)
		sectors = v;
		n = ioctl(fd, BLKSSZGET, &v);
		if (n >= 0) n = v % SECTORSIZE;
#endif
		if (n) {
			printf("device should be %d bytes/sector\n",
			       SECTORSIZE);
			sectors = 0;
			goto fin0;
		}
	}

fin0:
	return sectors;
}

static	void	show_result(char *type, uint64_t sectors, time_t seconds)
{
	uint64_t	bytes;
	double	rate;

	bytes = sectors * SECTORSIZE;

	printf("%s %lld bytes, %lld sec", type, bytes, seconds);

	if (seconds) {
		rate = (double)bytes / seconds;
		printf(" (%f bytes/sec)", rate);
	}

	printf("\n");
}

static	void	test(char *device, int method)
{
	int	fd;
	uint64_t	sectors;
	time_t	before, after;

	fd = open(device, O_RDWR);
	if (fd < 0) {
		printf("device open error\n");
		goto fin0;
	}

	sectors = get_sectors(fd);
	printf("device %s: %lld sectors method %#x\n", device, sectors, method);

	if (method & TEST_WRITE) {
		sync_sec();
		before = time(NULL);
		write_test(fd, method & METHOD_MASK, sectors);
		after = time(NULL);
		show_result("write", sectors, after - before);
	}

	if (method & TEST_READ) {
		sync_sec();
		before = time(NULL);
		read_test(fd, method & METHOD_MASK, sectors);
		after = time(NULL);
		show_result("read", sectors, after - before);
	}

fin0:
	close(fd);
	return;
}

int	main(int argc, char *argv[])
{
	int	method;

	if (argc < 3) {
		printf("usage:	%s [method] [device]\n", argv[0]);
		return 0;
	}

	memset(buf, 0, BUFFERSIZE);
	method = strtol(argv[1], NULL, 0);

	do {
		if (method & TEST_SET) {
			test(argv[2], TEST_WRITE | TEST_READ | METHOD_ONE);
			test(argv[2], TEST_WRITE | TEST_READ | METHOD_64BIT);
			test(argv[2], TEST_WRITE | TEST_READ | METHOD_RANDOM);
			test(argv[2], TEST_WRITE | TEST_READ | METHOD_ZERO);
		} else {
			test(argv[2], method);
		}
	} while (method & TEST_ENDLESS);

	return 0;
}