// 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	WRITE_TEST	0x1000
#define	READ_TEST	0x2000
#define	METHOD_MASK	0x0fff

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

static	uint8_t	buf[BUFFERSIZE];


/* create test data from sector value */
static	uint64_t	test_data(uint64_t secno, int method)
{
	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;
	}

	value = (method & METHOD_BIGENDIAN) ? be64toh(value) : le64toh(value);

	return value;
}

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

	for (i = 0; i < ELEMENTS; i++) {
		*p++ = data;
	}

	*q = p;
}

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

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

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

	result = 0;

	for (i = 0; i < ELEMENTS; i++) {
		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;
	uint64_t	data;

	p = buf;

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

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

	fd = open(device, O_WRONLY);
	if (fd < 0) goto fin0;

	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();
	close(fd);
fin0:
	return;
}

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

	fd = open(device, O_RDONLY);
	if (fd < 0) goto fin0;

	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);
		}
	}

	close(fd);
fin0:
	return;
}

static	void	sync_sec(void)
{
	time_t	t;

	t = time(NULL);

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

/* get total sectors */
static	uint64_t	get_sectors(char *device)
{
	int	fd, n;
	off_t	size;
	uint64_t	sectors;
#if defined(__OpenBSD__)
	struct disklabel dl;
#elif defined(__linux__)
	unsigned long	v;
#endif
	
	/* open device */
	fd = open(device, O_RDONLY);
	if (fd < 0) {
		sectors = 0;
		goto fin0;
	}

#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 fin1;
		}
	}

fin1:
	close(fd);
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)
{
	uint64_t	sectors;
	time_t	before, after;

	sectors = get_sectors(device);
	if (sectors == 0) {
		printf("device open error\n");
		goto fin0;
	}

	printf("device %s: %lld sectors\n", device, sectors);

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

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

fin0:
	return;
}

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

	memset(buf, 0, BUFFERSIZE);
	test(argv[2], strtol(argv[1], NULL, 0));

	return 0;
}