/*      $NetBSD: dkctl.c,v 1.27 2024/09/14 08:30:44 mlelstv Exp $       */

/*
 * Copyright 2001 Wasabi Systems, Inc.
 * All rights reserved.
 *
 * Written by Jason R. Thorpe for Wasabi Systems, Inc.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. All advertising materials mentioning features or use of this software
 *    must display the following acknowledgement:
 *      This product includes software developed for the NetBSD Project by
 *      Wasabi Systems, Inc.
 * 4. The name of Wasabi Systems, Inc. may not be used to endorse
 *    or promote products derived from this software without specific prior
 *    written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY WASABI SYSTEMS, INC. ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL WASABI SYSTEMS, INC
 * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

/*
 * dkctl(8) -- a program to manipulate disks.
 */
#include <sys/cdefs.h>

#ifndef lint
__RCSID("$NetBSD: dkctl.c,v 1.27 2024/09/14 08:30:44 mlelstv Exp $");
#endif

#include <sys/param.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/dkio.h>
#include <sys/disk.h>
#include <sys/queue.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>     
#include <unistd.h>
#include <util.h>

#define YES     1
#define NO      0

/* I don't think nl_langinfo is suitable in this case */
#define YES_STR "yes"
#define NO_STR  "no"
#define YESNO_ARG       YES_STR " | " NO_STR

#ifndef PRIdaddr
#define PRIdaddr PRId64
#endif

struct command {
        const char *cmd_name;
        const char *arg_names;
        void (*cmd_func)(int, char *[]);
        int open_flags;
};

static struct command *lookup(const char *);
__dead static void      usage(void);
static void     run(int, char *[]);
static void     showall(void);

static int      fd;                             /* file descriptor for device */
static const    char *dvname;                   /* device name */
static char     dvname_store[MAXPATHLEN];       /* for opendisk(3) */

static int dkw_sort(const void *, const void *);
static int yesno(const char *);

static void     disk_getcache(int, char *[]);
static void     disk_setcache(int, char *[]);
static void     disk_synccache(int, char *[]);
static void     disk_keeplabel(int, char *[]);
static void     disk_badsectors(int, char *[]);

static void     disk_addwedge(int, char *[]);
static void     disk_delwedge(int, char *[]);
static void     disk_getwedgeinfo(int, char *[]);
static void     disk_getgeometry(int, char *[]);
static void     disk_listwedges(int, char *[]);
static void     disk_makewedges(int, char *[]);
static void     disk_strategy(int, char *[]);

static struct command commands[] = {
        { "addwedge",
          "name startblk blkcnt ptype",
          disk_addwedge,
          O_RDWR },

        { "badsector",
          "flush | list | retry",
           disk_badsectors,
           O_RDWR },

        { "delwedge",
          "dk",
          disk_delwedge,
          O_RDWR },

        { "getcache",
          "",
          disk_getcache,
          O_RDONLY },

        { "getwedgeinfo",
          "",
          disk_getwedgeinfo,
          O_RDONLY },

        { "getgeometry",
          "",
          disk_getgeometry,
          O_RDONLY },

        { "keeplabel",
          YESNO_ARG,
          disk_keeplabel,
          O_RDWR },

        { "listwedges",
          "",
          disk_listwedges,
          O_RDONLY },

        { "makewedges",
          "",
          disk_makewedges,
          O_RDWR },

        { "setcache",
          "none | r | w | rw [save]",
          disk_setcache,
          O_RDWR },

        { "strategy",
          "[name]",
          disk_strategy,
          O_RDWR },

        { "synccache",
          "[force]",
          disk_synccache,
          O_RDWR },

        { NULL,
          NULL,
          NULL,
          0 },
};

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

        /* Must have at least: device command */
        if (argc < 2)
                usage();

        dvname = argv[1];
        if (argc == 2)
                showall();
        else
                run(argc - 2, argv + 2);

        return EXIT_SUCCESS;
}

static void
run(int argc, char *argv[])
{
        struct command *command;

        command = lookup(argv[0]);

        /* Open the device. */
        fd = opendisk(dvname, command->open_flags, dvname_store,
            sizeof(dvname_store), 0);
        if (fd == -1)
                err(EXIT_FAILURE, "%s", dvname);
        dvname = dvname_store;

        (*command->cmd_func)(argc, argv);

        /* Close the device. */
        (void)close(fd);
}

static struct command *
lookup(const char *name)
{
        int i;

        /* Look up the command. */
        for (i = 0; commands[i].cmd_name != NULL; i++)
                if (strcmp(name, commands[i].cmd_name) == 0)
                        break;
        if (commands[i].cmd_name == NULL)
                errx(EXIT_FAILURE, "unknown command: %s", name);

        return &commands[i];
}

static void
usage(void)
{
        int i;

        fprintf(stderr,
            "Usage: %s device\n"
            "       %s device command [arg [...]]\n",
            getprogname(), getprogname());

        fprintf(stderr, "   Available commands:\n");
        for (i = 0; commands[i].cmd_name != NULL; i++)
                fprintf(stderr, "\t%s %s\n", commands[i].cmd_name,
                    commands[i].arg_names);

        exit(EXIT_FAILURE);
}

static void
showall(void)
{
        static const char *cmds[] = { "strategy", "getcache", "listwedges" };
        size_t i;
        char *args[2];

        args[1] = NULL;
        for (i = 0; i < __arraycount(cmds); i++) {
                printf("%s:\n", cmds[i]);
                args[0] = __UNCONST(cmds[i]);
                run(1, args);
                putchar('\n');
        }
}

static void
disk_strategy(int argc, char *argv[])
{
        struct disk_strategy odks;
        struct disk_strategy dks;

        memset(&dks, 0, sizeof(dks));
        if (ioctl(fd, DIOCGSTRATEGY, &odks) == -1) {
                err(EXIT_FAILURE, "%s: DIOCGSTRATEGY", dvname);
        }

        memset(&dks, 0, sizeof(dks));
        switch (argc) {
        case 1:
                /* show the buffer queue strategy used */
                printf("%s: %s\n", dvname, odks.dks_name);
                return;
        case 2:
                /* set the buffer queue strategy */
                strlcpy(dks.dks_name, argv[1], sizeof(dks.dks_name));
                if (ioctl(fd, DIOCSSTRATEGY, &dks) == -1) {
                        err(EXIT_FAILURE, "%s: DIOCSSTRATEGY", dvname);
                }
                printf("%s: %s -> %s\n", dvname, odks.dks_name, argv[1]);
                break;
        default:
                usage();
                /* NOTREACHED */
        }
}

static void
disk_getcache(int argc, char *argv[])
{
        int bits;

        if (ioctl(fd, DIOCGCACHE, &bits) == -1)
                err(EXIT_FAILURE, "%s: getcache", dvname);

        if ((bits & (DKCACHE_READ|DKCACHE_WRITE)) == 0)
                printf("%s: No caches enabled\n", dvname);
        else {
                if (bits & DKCACHE_READ)
                        printf("%s: read cache enabled\n", dvname);
                if (bits & DKCACHE_WRITE)
                        printf("%s: write-back cache enabled\n", dvname);
        }

        printf("%s: read cache enable is %schangeable\n", dvname,
            (bits & DKCACHE_RCHANGE) ? "" : "not ");
        printf("%s: write cache enable is %schangeable\n", dvname,
            (bits & DKCACHE_WCHANGE) ? "" : "not ");

        printf("%s: cache parameters are %ssavable\n", dvname,
            (bits & DKCACHE_SAVE) ? "" : "not ");

#ifdef DKCACHE_FUA
        printf("%s: cache Force Unit Access (FUA) %ssupported\n", dvname,
            (bits & DKCACHE_FUA) ? "" : "not ");
#endif /* DKCACHE_FUA */

#ifdef DKCACHE_DPO
        printf("%s: cache Disable Page Out (DPO) %ssupported\n", dvname,
            (bits & DKCACHE_DPO) ? "" : "not ");
#endif /* DKCACHE_DPO */
}

static void
disk_setcache(int argc, char *argv[])
{
        int bits;

        if (argc > 3 || argc == 1)
                usage();

        if (strcmp(argv[1], "none") == 0)
                bits = 0;
        else if (strcmp(argv[1], "r") == 0)
                bits = DKCACHE_READ;
        else if (strcmp(argv[1], "w") == 0)
                bits = DKCACHE_WRITE;
        else if (strcmp(argv[1], "rw") == 0)
                bits = DKCACHE_READ|DKCACHE_WRITE;
        else
                usage();

        if (argc == 3) {
                if (strcmp(argv[2], "save") == 0)
                        bits |= DKCACHE_SAVE;
                else
                        usage();
        }

        if (ioctl(fd, DIOCSCACHE, &bits) == -1)
                err(EXIT_FAILURE, "%s: %s", dvname, argv[0]);
}

static void
disk_synccache(int argc, char *argv[])
{
        int force;

        switch (argc) {
        case 1:
                force = 0;
                break;

        case 2:
                if (strcmp(argv[1], "force") == 0)
                        force = 1;
                else
                        usage();
                break;

        default:
                usage();
        }

        if (ioctl(fd, DIOCCACHESYNC, &force) == -1)
                err(EXIT_FAILURE, "%s: %s", dvname, argv[0]);
}

static void
disk_keeplabel(int argc, char *argv[])
{
        int keep;
        int yn;

        if (argc != 2)
                usage();

        yn = yesno(argv[1]);
        if (yn < 0)
                usage();

        keep = yn == YES;

        if (ioctl(fd, DIOCKLABEL, &keep) == -1)
                err(EXIT_FAILURE, "%s: %s", dvname, argv[0]);
}


static void
disk_badsectors(int argc, char *argv[])
{
        struct disk_badsectors *dbs, *dbs2, buffer[200];
        SLIST_HEAD(, disk_badsectors) dbstop;
        struct disk_badsecinfo dbsi;
        daddr_t blk, totbad, bad;
        u_int32_t count;
        struct stat sb;
        u_char *block;
        time_t tm;

        if (argc != 2)
                usage();

        if (strcmp(argv[1], "list") == 0) {
                /*
                 * Copy the list of kernel bad sectors out in chunks that fit
                 * into buffer[].  Updating dbsi_skip means we don't sit here
                 * forever only getting the first chunk that fit in buffer[].
                 */
                dbsi.dbsi_buffer = (caddr_t)buffer;
                dbsi.dbsi_bufsize = sizeof(buffer);
                dbsi.dbsi_skip = 0;
                dbsi.dbsi_copied = 0;
                dbsi.dbsi_left = 0;

                do {
                        if (ioctl(fd, DIOCBSLIST, (caddr_t)&dbsi) == -1)
                                err(EXIT_FAILURE, "%s: badsectors list", dvname);

                        dbs = (struct disk_badsectors *)dbsi.dbsi_buffer;
                        for (count = dbsi.dbsi_copied; count > 0; count--) {
                                tm = dbs->dbs_failedat.tv_sec;
                                printf("%s: blocks %" PRIdaddr " - %" PRIdaddr " failed at %s",
                                        dvname, dbs->dbs_min, dbs->dbs_max,
                                        ctime(&tm));
                                dbs++;
                        }
                        dbsi.dbsi_skip += dbsi.dbsi_copied;
                } while (dbsi.dbsi_left != 0);

        } else if (strcmp(argv[1], "flush") == 0) {
                if (ioctl(fd, DIOCBSFLUSH) == -1)
                        err(EXIT_FAILURE, "%s: badsectors flush", dvname);

        } else if (strcmp(argv[1], "retry") == 0) {
                /*
                 * Enforce use of raw device here because the block device
                 * causes access to blocks to be clustered in a larger group,
                 * making it impossible to determine which individual sectors
                 * are the cause of a problem.
                 */ 
                if (fstat(fd, &sb) == -1)
                        err(EXIT_FAILURE, "fstat");

                if (!S_ISCHR(sb.st_mode)) {
                        fprintf(stderr, "'badsector retry' must be used %s\n",
                                "with character device");
                        exit(1);
                }

                SLIST_INIT(&dbstop);

                /*
                 * Build up a copy of the in-kernel list in a number of stages.
                 * That the list we build up here is in the reverse order to
                 * the kernel's is of no concern.
                 */
                dbsi.dbsi_buffer = (caddr_t)buffer;
                dbsi.dbsi_bufsize = sizeof(buffer);
                dbsi.dbsi_skip = 0;
                dbsi.dbsi_copied = 0;
                dbsi.dbsi_left = 0;

                do {
                        if (ioctl(fd, DIOCBSLIST, (caddr_t)&dbsi) == -1)
                                err(EXIT_FAILURE, "%s: badsectors list", dvname);

                        dbs = (struct disk_badsectors *)dbsi.dbsi_buffer;
                        for (count = dbsi.dbsi_copied; count > 0; count--) {
                                dbs2 = malloc(sizeof *dbs2);
                                if (dbs2 == NULL)
                                        err(EXIT_FAILURE, NULL);
                                *dbs2 = *dbs;
                                SLIST_INSERT_HEAD(&dbstop, dbs2, dbs_next);
                                dbs++;
                        }
                        dbsi.dbsi_skip += dbsi.dbsi_copied;
                } while (dbsi.dbsi_left != 0);

                /*
                 * Just calculate and print out something that will hopefully
                 * provide some useful information about what's going to take
                 * place next (if anything.)
                 */
                bad = 0;
                totbad = 0;
                if ((block = calloc(1, DEV_BSIZE)) == NULL)
                        err(EXIT_FAILURE, NULL);
                SLIST_FOREACH(dbs, &dbstop, dbs_next) {
                        bad++;
                        totbad += dbs->dbs_max - dbs->dbs_min + 1;
                }

                printf("%s: bad sector clusters %"PRIdaddr
                    " total sectors %"PRIdaddr"\n", dvname, bad, totbad);

                /*
                 * Clear out the kernel's list of bad sectors, ready for us
                 * to test all those it thought were bad.
                 */
                if (ioctl(fd, DIOCBSFLUSH) == -1)
                        err(EXIT_FAILURE, "%s: badsectors flush", dvname);

                printf("%s: bad sectors flushed\n", dvname);

                /*
                 * For each entry we obtained from the kernel, retry each
                 * individual sector recorded as bad by seeking to it and
                 * attempting to read it in.  Print out a line item for each
                 * bad block we verify.
                 *
                 * PRIdaddr is used here because the type of dbs_max is daddr_t
                 * and that may be either a 32bit or 64bit number(!)
                 */
                SLIST_FOREACH(dbs, &dbstop, dbs_next) {
                        printf("%s: Retrying %"PRIdaddr" - %"
                            PRIdaddr"\n", dvname, dbs->dbs_min, dbs->dbs_max);

                        for (blk = dbs->dbs_min; blk <= dbs->dbs_max; blk++) {
                                if (lseek(fd, (off_t)blk * DEV_BSIZE,
                                    SEEK_SET) == -1) {
                                        warn("%s: lseek block %" PRIdaddr "",
                                            dvname, blk);
                                        continue;
                                }
                                printf("%s: block %"PRIdaddr" - ", dvname, blk);
                                if (read(fd, block, DEV_BSIZE) != DEV_BSIZE)
                                        printf("failed\n");
                                else
                                        printf("ok\n");
                                fflush(stdout);
                        }
                }
        }
}

static void
disk_addwedge(int argc, char *argv[])
{
        struct dkwedge_info dkw;
        char *cp;
        daddr_t start;
        uint64_t size;

        if (argc != 5)
                usage();

        /* XXX Unicode: dkw_wname is supposed to be utf-8 */
        if (strlcpy((char *)dkw.dkw_wname, argv[1], sizeof(dkw.dkw_wname)) >=
            sizeof(dkw.dkw_wname))
                errx(EXIT_FAILURE, "Wedge name too long; max %zd characters",
                    sizeof(dkw.dkw_wname) - 1);

        if (strlcpy(dkw.dkw_ptype, argv[4], sizeof(dkw.dkw_ptype)) >=
            sizeof(dkw.dkw_ptype))
                errx(EXIT_FAILURE, "Wedge partition type too long; max %zd characters",
                    sizeof(dkw.dkw_ptype) - 1);

        errno = 0;
        start = strtoll(argv[2], &cp, 0);
        if (*cp != '\0')
                errx(EXIT_FAILURE, "Invalid start block: %s", argv[2]);
        if (errno == ERANGE && (start == LLONG_MAX ||
                                start == LLONG_MIN))
                errx(EXIT_FAILURE, "Start block out of range.");
        if (start < 0)
                errx(EXIT_FAILURE, "Start block must be >= 0.");

        errno = 0;
        size = strtoull(argv[3], &cp, 0);
        if (*cp != '\0')
                errx(EXIT_FAILURE, "Invalid block count: %s", argv[3]);
        if (errno == ERANGE && (size == ULLONG_MAX))
                errx(EXIT_FAILURE, "Block count out of range.");

        dkw.dkw_offset = start;
        dkw.dkw_size = size;

        if (ioctl(fd, DIOCAWEDGE, &dkw) == -1)
                err(EXIT_FAILURE, "%s: %s", dvname, argv[0]);
        else
                printf("%s created successfully.\n", dkw.dkw_devname);

}

static void
disk_delwedge(int argc, char *argv[])
{
        struct dkwedge_info dkw;

        if (argc != 2)
                usage();

        if (strlcpy(dkw.dkw_devname, argv[1], sizeof(dkw.dkw_devname)) >=
            sizeof(dkw.dkw_devname))
                errx(EXIT_FAILURE, "Wedge dk name too long; max %zd characters",
                    sizeof(dkw.dkw_devname) - 1);

        if (ioctl(fd, DIOCDWEDGE, &dkw) == -1)
                err(EXIT_FAILURE, "%s: %s", dvname, argv[0]);
}

static void
disk_getwedgeinfo(int argc, char *argv[])
{
        struct dkwedge_info dkw;

        if (argc != 1)
                usage();

        if (ioctl(fd, DIOCGWEDGEINFO, &dkw) == -1)
                err(EXIT_FAILURE, "%s: getwedgeinfo", dvname);

        printf("%s at %s: %s\n", dkw.dkw_devname, dkw.dkw_parent,
            dkw.dkw_wname);     /* XXX Unicode */
        printf("%s: %"PRIu64" blocks at %"PRId64", type: %s\n",
            dkw.dkw_devname, dkw.dkw_size, dkw.dkw_offset, dkw.dkw_ptype);
}

static void
disk_getgeometry(int argc, char *argv[])
{
        off_t bytes;
        u_int secsize;

        if (argc != 1)
                usage();

        if (ioctl(fd, DIOCGMEDIASIZE, &bytes) == -1)
                err(EXIT_FAILURE, "%s: getmediasize", dvname);

        secsize = DEV_BSIZE;
        if (ioctl(fd, DIOCGSECTORSIZE, &secsize) == -1)
                warn("%s: getsectorsize", dvname);

        printf("%s: %"PRIu64" bytes in %"PRIu64" blocks of %u bytes\n",
            dvname, bytes, bytes/secsize, secsize);
}

static void
disk_listwedges(int argc, char *argv[])
{
        struct dkwedge_info *dkw;
        struct dkwedge_list dkwl;
        size_t bufsize;
        u_int i;
        int c;
        bool error, quiet;

        optreset = 1;
        optind = 1;
        quiet = error = false;
        while ((c = getopt(argc, argv, "qe")) != -1)
                switch (c) {
                case 'e':
                        error = true;
                        break;
                case 'q':
                        quiet = true;
                        break;
                default:
                        usage();
                }

        argc -= optind;
        argv += optind;

        if (argc != 0)
                usage();

        dkw = NULL;
        dkwl.dkwl_buf = dkw;
        dkwl.dkwl_bufsize = 0;

        for (;;) {
                if (ioctl(fd, DIOCLWEDGES, &dkwl) == -1)
                        err(EXIT_FAILURE, "%s: listwedges", dvname);
                if (dkwl.dkwl_nwedges == dkwl.dkwl_ncopied)
                        break;
                bufsize = dkwl.dkwl_nwedges * sizeof(*dkw);
                if (dkwl.dkwl_bufsize < bufsize) {
                        dkw = realloc(dkwl.dkwl_buf, bufsize);
                        if (dkw == NULL)
                                errx(EXIT_FAILURE, "%s: listwedges: unable to "
                                    "allocate wedge info buffer", dvname);
                        dkwl.dkwl_buf = dkw;
                        dkwl.dkwl_bufsize = bufsize;
                }
        }

        if (dkwl.dkwl_nwedges == 0) {
                if (!quiet)
                        printf("%s: no wedges configured\n", dvname);
                if (error)
                        exit(EXIT_FAILURE);
                return;
        }

        if (quiet)
                return;

        qsort(dkw, dkwl.dkwl_nwedges, sizeof(*dkw), dkw_sort);

        printf("%s: %u wedge%s:\n", dvname, dkwl.dkwl_nwedges,
            dkwl.dkwl_nwedges == 1 ? "" : "s");
        for (i = 0; i < dkwl.dkwl_nwedges; i++) {
                printf("%s: %s, %"PRIu64" blocks at %"PRId64", type: %s\n",
                    dkw[i].dkw_devname,
                    dkw[i].dkw_wname,   /* XXX Unicode */
                    dkw[i].dkw_size, dkw[i].dkw_offset, dkw[i].dkw_ptype);
        }
}

static void
disk_makewedges(int argc, char *argv[])
{
        int bits;

        if (argc != 1)
                usage();

        if (ioctl(fd, DIOCMWEDGES, &bits) == -1)
                err(EXIT_FAILURE, "%s: %s", dvname, argv[0]);
        else
                printf("successfully scanned %s.\n", dvname);
}

static int
dkw_sort(const void *a, const void *b)
{
        const struct dkwedge_info *dkwa = a, *dkwb = b;
        const daddr_t oa = dkwa->dkw_offset, ob = dkwb->dkw_offset;

        return (oa < ob) ? -1 : (oa > ob) ? 1 : 0;
}

/*
 * return YES, NO or -1.
 */
static int
yesno(const char *p)
{

        if (!strcmp(p, YES_STR))
                return YES;
        if (!strcmp(p, NO_STR))
                return NO;
        return -1;
}