/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "License").
 * You may not use this file except in compliance with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2009 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

/*
 *   diskscan:
 *   performs a verification pass over a device specified on command line;
 *   display progress on stdout, and print bad sector numbers to stderr
 */
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <stropts.h>
#include <memory.h>
#include <ctype.h>
#include <malloc.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/param.h>
#include <sys/stat.h>
#include <sys/vtoc.h>
#include <sys/dkio.h>

static void verexit();	/* signal handler and exit routine	*/
static void report();	   /* tell user how we're getting on */
static void scandisk(char *device, int devfd, int writeflag);
static void report(char *what, diskaddr_t sector);
static void verexit(int code);

#define	TRUE		1
#define	FALSE		0
#define	VER_WRITE	1
#define	VER_READ	2

static char	*progname;
static struct  dk_geom	dkg;	  /* physical device boot info */
static char	replybuf[64];	  /* used for user replies to questions */
static diskaddr_t unix_base;	  /* first sector of UNIX System partition */
static diskaddr_t unix_size;	  /* # sectors in UNIX System partition */
static long	numbadrd = 0;	  /* number of bad sectors on read */
static long	numbadwr = 0;	  /* number of bad sectors on write */
static char	eol = '\n';	  /* end-of-line char (if -n, we set to '\n') */
static int	print_warn = 1;	  /* should the warning message be printed? */
static int do_scan = VER_READ;

int
main(int argc, char *argv[]) {
	extern int	optind;
	int		devfd;	/* device file descriptor */
	struct stat 	statbuf;
	struct part_info	part_info;
	struct extpart_info	extpartinfo;
	int		c;
	int		errflag = 0;
	char		*device;
	progname = argv[0];

	/* Don't buffer stdout - we don't want to see bursts */

	setbuf(stdout, NULL);

	while ((c = getopt(argc, argv, "Wny")) != -1)
	{
		switch (c) {
		case 'W':
			do_scan = VER_WRITE;
			break;

		case 'n':
			eol = '\r';
			break;

		case 'y':
			print_warn = 0;
			break;

		default:
			++errflag;
			break;
		}
	}

	if ((argc - optind) < 1)
		errflag++;

	if (errflag) {
		(void) fprintf(stderr,
		    "\nUsage: %s [-W] [-n] [-y] <phys_device_name> \n",
			    progname);
		exit(1);
	}

	device = argv[optind];

	if (stat(device, &statbuf)) {
		(void) fprintf(stderr,
		    "%s: invalid device %s, stat failed\n", progname, device);
		perror("");
		exit(4);
	}
	if ((statbuf.st_mode & S_IFMT) != S_IFCHR) {
		(void) fprintf(stderr,
		    "%s: device %s is not character special\n",
			    progname, device);
		exit(5);
	}
	if ((devfd = open(device, O_RDWR)) == -1) {
		(void) fprintf(stderr,
		    "%s: open of %s failed\n", progname, device);
		perror("");
		exit(8);
	}

	if ((ioctl(devfd, DKIOCGGEOM, &dkg)) == -1) {
		(void) fprintf(stderr,
		    "%s: unable to get disk geometry.\n", progname);
		perror("");
		exit(9);
	}

	if ((ioctl(devfd, DKIOCEXTPARTINFO, &extpartinfo)) == 0) {
		unix_base = extpartinfo.p_start;
		unix_size = extpartinfo.p_length;
	} else {
		if ((ioctl(devfd, DKIOCPARTINFO, &part_info)) == 0) {
			unix_base = (ulong_t)part_info.p_start;
			unix_size = (uint_t)part_info.p_length;
		} else {
			(void) fprintf(stderr, "%s: unable to get partition "
			    "info.\n", progname);
			perror("");
			exit(9);
		}
	}

	scandisk(device, devfd, do_scan);
	return (0);
}

/*
 *  scandisk:
 *	  attempt to read every sector of the drive;
 *	  display bad sectors found on stderr
 */

static void
scandisk(char *device, int devfd, int writeflag)
{
	int	trksiz = 0;
	char	*verbuf;
	diskaddr_t cursec;
	int	 cylsiz =  dkg.dkg_nsect * dkg.dkg_nhead;
	int	 i;
	char	*rptr;
	diskaddr_t tmpend = 0;
	diskaddr_t tmpsec = 0;
	struct dk_minfo	mediainfo;
	uint_t  sector_size;

	if ((ioctl(devfd, DKIOCGMEDIAINFO, &mediainfo)) == 0) {
		sector_size = mediainfo.dki_lbsize;
	} else {
		sector_size = NBPSCTR;
	}
	trksiz =  sector_size * dkg.dkg_nsect;

/* #define LIBMALLOC */

#ifdef LIBMALLOC

	extern int  mallopt();

	/* This adds 5k to the binary, but it's a lot prettier */


	/* make track buffer sector aligned */
	if (mallopt(M_GRAIN, sector_size)) {
		perror("mallopt");
		exit(1);
	}
	if ((verbuf = malloc(sector_size * dkg.dkg_nsect)) ==
	    (char *)NULL) {
		perror("malloc");
		exit(1);
	}

#else

	if ((verbuf = malloc(sector_size + sector_size * dkg.dkg_nsect))
	    == (char *)NULL) {
		perror("malloc");
		exit(1);
	}
	verbuf = (char *)((((unsigned long)verbuf + sector_size)) &
	    (-sector_size));

#endif

	/* write pattern in track buffer */

	for (i = 0; i < trksiz; i++)
		verbuf[i] = (char)0xe5;

	/* Turn off retry, and set trap to turn them on again */

	(void) signal(SIGINT, verexit);
	(void) signal(SIGQUIT, verexit);

	if (writeflag == VER_READ)
		goto do_readonly;

	/*
	 *   display warning only if -n arg not passed
	 *   (otherwise the UI system will take care of it)
	 */

	if (print_warn == 1) {
		(void) printf(
		    "\nCAUTION: ABOUT TO DO DESTRUCTIVE WRITE ON %s\n", device);
		(void) printf("	 THIS WILL DESTROY ANY DATA YOU HAVE ON\n");
		(void) printf("	 THAT PARTITION OR SLICE.\n");
		(void) printf("Do you want to continue (y/n)? ");

		rptr = fgets(replybuf, 64*sizeof (char), stdin);
		if (!rptr || !((replybuf[0] == 'Y') || (replybuf[0] == 'y')))
			exit(10);
	}

	for (cursec = 0; cursec < unix_size; cursec +=  dkg.dkg_nsect) {
		if (llseek(devfd, cursec * sector_size, 0) == -1) {
			(void) fprintf(stderr,
			    "Error seeking sector %llu Cylinder %llu\n",
			    cursec, cursec / cylsiz);
			verexit(1);
		}

		/*
		 * verify sector at a time only when
		 * the whole track write fails;
		 *  (if we write a sector at a time, it takes forever)
		 */

		report("Writing", cursec);

		if (write(devfd, verbuf, trksiz) != trksiz) {
			tmpend = cursec +  dkg.dkg_nsect;
			for (tmpsec = cursec; tmpsec < tmpend; tmpsec++) {
				/*
				 *  try writing to it once; if this fails,
				 *  then announce the sector bad on stderr
				 */

				if (llseek(devfd, tmpsec * sector_size,
				    0) == -1) {
					(void) fprintf(stderr, "Error seeking "
					    "sector %llu Cylinder %llu\n",
					    tmpsec, cursec / cylsiz);
					verexit(1);
				}

				report("Writing", tmpsec);

				if (write(devfd, verbuf, sector_size)
				    != sector_size) {
					(void) fprintf(stderr,
					    "%llu\n", tmpsec + unix_base);
					numbadwr++;
				}
			}
		}
	}

	(void) putchar(eol);
	do_readonly:

	for (cursec = 0; cursec < unix_size; cursec +=  dkg.dkg_nsect) {
		if (llseek(devfd, cursec * sector_size, 0) == -1) {
			(void) fprintf(stderr,
			    "Error seeking sector %llu Cylinder %llu\n",
			    cursec, cursec / cylsiz);
			verexit(1);
		}

		/*
		 * read a sector at a time only when
		 * the whole track write fails;
		 * (if we do a sector at a time read, it takes forever)
		 */

		report("Reading", cursec);
		if (read(devfd, verbuf, trksiz) != trksiz) {
			tmpend = cursec +  dkg.dkg_nsect;
			for (tmpsec = cursec; tmpsec < tmpend; tmpsec++) {
				if (llseek(devfd, tmpsec * sector_size,
				    0) == -1) {
					(void) fprintf(stderr, "Error seeking"
					    " sector %llu Cylinder %llu\n",
					    tmpsec, cursec / cylsiz);
					verexit(1);
				}
				report("Reading", tmpsec);
				if (read(devfd, verbuf, sector_size) !=
				    sector_size) {
					(void) fprintf(stderr, "%llu\n",
					    tmpsec + unix_base);
					numbadrd++;
				}
			}
		}
	}
	(void) printf("%c%c======== Diskscan complete ========%c", eol,
	    eol, eol);

	if ((numbadrd > 0) || (numbadwr > 0)) {
		(void) printf("%cFound %ld bad sector(s) on read,"
		    " %ld bad sector(s) on write%c",
		    eol, numbadrd, numbadwr, eol);
	}
}

static void
verexit(int code)
{
	(void) printf("\n");
	exit(code);
}


/*
 *   report where we are...
 */

static void
report(char *what, diskaddr_t sector)
{
	(void) printf("%s sector %-19llu of %-19llu%c", what, sector,
	    unix_size, eol);
}