/*
 * 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 (c) 1984, 1986, 1987, 1988, 1989 AT&T */
/*        All Rights Reserved */


/*      Copyright (c) 1984 AT&T */
/*        All Rights Reserved   */


/*
 * Copyright 2008 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 * Copyright 2021 Jason King
 */

/*
 * Print a disk partition map (volume table of contents, or VTOC).
 */

#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <limits.h>
#include <err.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/dkio.h>
#include <sys/vtoc.h>
#include <sys/mnttab.h>
#include <sys/vfstab.h>
#include <sys/mkdev.h>

#include <sys/efi_partition.h>
/*
 * Assumes V_NUMPAR must be a power of 2.
 *
 * for V_NUMPAR = 8, we have
 *      parttn(x)=(x & 0x07)    noparttn(x)=(x & 0x3fff8)
 *
 * for V_NUMPAR = 16, we have
 *      parttn(x)=(x & 0x0f)    noparttn(x)=(x & 0x3fff0)
 */
#define parttn(x)       (x % V_NUMPAR)
#define noparttn(x)     (x & (MAXMIN & ~(V_NUMPAR-1)))

/*
 * Disk freespace structure.
 */
typedef struct {
        u_longlong_t    fr_start;       /* Start of free space */
        u_longlong_t    fr_size;        /* Length of free space */
} freemap_t;

static  freemap_t       *findfree(struct dk_geom *, struct extvtoc *);
static  int     partcmp(const void *, const void *);
static  int     partcmp64(const void *, const void *);
static  int     prtvtoc(char *);
static  void    putfree(struct extvtoc *, freemap_t *);
static  void    putfree64(struct dk_gpt *, freemap_t *);
static  void    puttable(struct dk_geom *, struct extvtoc *, freemap_t *,
                        char *, char **);
static  void    puttable64(struct dk_gpt *, freemap_t *,
                        char *, char **);
static  int     readgeom(int, char *, struct dk_geom *);
static  int     readvtoc(int, char *, struct extvtoc *);
static  int     readefi(int, char *, struct dk_gpt **);
static  void    usage(void);
static  char    *safe_strdup(const char *, const char *);
static  void    *safe_calloc(const char *, size_t, size_t);

#define SAFE_STRDUP(a)          safe_strdup(__func__, (a))
#define SAFE_CALLOC(a, b)       safe_calloc(__func__, (a), (b))

/*
 * External variables.
 */
extern char     *getfullrawname();
/*
 * Static variables.
 */
static short    fflag;                  /* Print freespace shell assignments */
static short    hflag;                  /* Omit headers */
static short    sflag;                  /* Omit all but the column header */
static char     *fstab = VFSTAB;        /* Fstab pathname */
static char     *mnttab = MNTTAB;       /* mnttab pathname */

int
main(int argc, char *argv[])
{
        int status = EXIT_SUCCESS;
        int c;

        while ((c = getopt(argc, argv, "fhst:m:")) != -1) {
                switch (c) {
                case 'f':
                        ++fflag;
                        break;
                case 'h':
                        ++hflag;
                        break;
                case 's':
                        ++sflag;
                        break;
                case 't':
                        fstab = optarg;
                        break;
                case 'm':
                        mnttab = optarg;
                        break;
                default:
                        usage();
                }
        }

        if (optind >= argc)
                usage();

        for (int i = optind; i < argc; i++) {
                if (prtvtoc(argv[i]) != 0) {
                        status = EXIT_FAILURE;
                }
        }

        return (status);
}

static freemap_t        *freemap;
/*
 * findfree(): Find free space on a disk.
 */
static freemap_t *
findfree(struct dk_geom *geom, struct extvtoc *vtoc)
{
        struct extpartition *part;
        struct extpartition **list;
        freemap_t *freeidx;
        diskaddr_t fullsize;
        ulong_t cylsize;
        struct extpartition *sorted[V_NUMPAR + 1];

        if (vtoc->v_nparts > V_NUMPAR) {
                errx(EXIT_FAILURE, "putfree(): Too many partitions on disk!");
        }

        freemap = SAFE_CALLOC(sizeof (freemap_t), V_NUMPAR + 1);
        cylsize = (geom->dkg_nsect) * (geom->dkg_nhead);
        fullsize = (diskaddr_t)(geom->dkg_ncyl) * cylsize;
        list = sorted;
        for (part = vtoc->v_part; part < vtoc->v_part + vtoc->v_nparts;
            ++part) {
                if (part->p_size && part->p_tag != V_BACKUP)
                        *list++ = part;
        }
        *list = 0;
        qsort(sorted, list - sorted, sizeof (*sorted), partcmp);
        freeidx = freemap;
        freeidx->fr_start = 0;
        for (list = sorted; (part = *list) != NULL; ++list) {
                if (part->p_start <= freeidx->fr_start) {
                        freeidx->fr_start += part->p_size;
                } else {
                        freeidx->fr_size = part->p_start - freeidx->fr_start;
                        (++freeidx)->fr_start = part->p_start + part->p_size;
                }
        }
        if (freeidx->fr_start < fullsize) {
                freeidx->fr_size = fullsize - freeidx->fr_start;
                ++freeidx;
        }
        freeidx->fr_start = freeidx->fr_size = 0;
        return (freemap);
}

/*
 * findfree64(): Find free space on a disk.
 */
static freemap_t *
findfree64(struct dk_gpt *efi)
{
        struct dk_part *part;
        struct dk_part **list;
        freemap_t *freeidx;
        diskaddr_t fullsize;
        struct dk_part **sorted;

        freemap = SAFE_CALLOC(sizeof (freemap_t), efi->efi_nparts + 1);
        sorted = SAFE_CALLOC(sizeof (struct dk_part), efi->efi_nparts + 1);
        fullsize = efi->efi_last_u_lba;
        list = sorted;
        for (part = efi->efi_parts; part < efi->efi_parts + efi->efi_nparts;
            ++part) {
                if (part->p_size && part->p_tag != V_BACKUP)
                        *list++ = part;
        }
        *list = 0;
        qsort(sorted, list - sorted, sizeof (*sorted), partcmp64);
        freeidx = freemap;
        freeidx->fr_start = efi->efi_first_u_lba;
        for (list = sorted; (part = *list) != NULL; ++list) {
                if (part->p_start == freeidx->fr_start) {
                        freeidx->fr_start += part->p_size;
                } else {
                        freeidx->fr_size = part->p_start - freeidx->fr_start;
                        (++freeidx)->fr_start = part->p_start + part->p_size;
                }
        }
        if (freeidx->fr_start < fullsize) {
                freeidx->fr_size = fullsize - freeidx->fr_start;
                ++freeidx;
        }
        freeidx->fr_start = freeidx->fr_size = 0;
        return (freemap);
}

/*
 * getmntpt()
 *
 * Get the filesystem mountpoint of each partition on the disk
 * from the fstab or mnttab. Returns a pointer to an array of pointers to
 * directory names (indexed by partition number).
 */
static char **
getmntpt(major_t slot, minor_t nopartminor)
{
        FILE *file;
        char devbuf[PATH_MAX], *item;
        static char *list[V_NUMPAR];
        struct stat sb;
        struct mnttab mtab;
        struct vfstab vtab;

        for (unsigned idx = 0; idx < V_NUMPAR; ++idx)
                list[idx] = NULL;

        /* read mnttab for partition mountpoints */
        if ((file = fopen(mnttab, "r")) == NULL) {
                warn("failed to open %s", mnttab);
        } else {
                while (getmntent(file, &mtab) == 0) {
                        item = mtab.mnt_special;
                        if (item == NULL || mtab.mnt_mountp == NULL)
                                continue;

                        /*
                         * Is it from /dev?
                         */
                        if (strncmp(item, "/dev/", strlen("/dev/")) != 0)
                                continue;

                        /*
                         * Is it a character device?
                         */
                        (void) snprintf(devbuf, sizeof (devbuf), "/dev/r%s",
                            item + strlen("/dev/"));

                        if (stat(devbuf, &sb) != 0 ||
                            (sb.st_mode & S_IFMT) != S_IFCHR) {
                                continue;
                        }

                        /*
                         * device must match input slot and nopartminor
                         */
                        if (major(sb.st_rdev) != slot ||
                            noparttn(minor(sb.st_rdev)) != nopartminor) {
                                continue;
                        }

                        list[parttn(minor(sb.st_rdev))] =
                            SAFE_STRDUP(mtab.mnt_mountp);
                }
                (void) fclose(file);
        }

        if ((file = fopen(fstab, "r")) == NULL) {
                warn("failed to open %s", fstab);
                return (list);
        }

        /*
         * Look for the disk in the vfstab so that we can report its mount
         * point even if it isn't currently mounted.
         */
        while (getvfsent(file, &vtab) == 0) {
                item = vtab.vfs_special;
                if (item == NULL || vtab.vfs_mountp == NULL)
                        continue;

                if (strncmp(item, "/dev/", strlen("/dev/")) != 0)
                        continue;

                /*
                 * Is it a character device?
                 */
                (void) snprintf(devbuf, sizeof (devbuf), "/dev/r%s",
                    item + strlen("/dev/"));

                if (stat(devbuf, &sb) != 0 ||
                    (sb.st_mode & S_IFMT) != S_IFCHR) {
                        continue;
                }

                /*
                 * device must match input slot and nopartminor
                 */
                if (major(sb.st_rdev) != slot ||
                    noparttn(minor(sb.st_rdev)) != nopartminor) {
                        continue;
                }

                /*
                 * use mnttab entry if both tables have entries
                 */
                if (list[parttn(minor(sb.st_rdev))] != NULL)
                        continue;

                list[parttn(minor(sb.st_rdev))] = SAFE_STRDUP(vtab.vfs_mountp);
        }
        (void) fclose(file);

        return (list);
}

/*
 * partcmp(): Qsort() key comparison of partitions by starting sector numbers.
 */
static int
partcmp(const void *one, const void *two)
{
        struct partition *p1 = *(struct partition **)one;
        struct partition *p2 = *(struct partition **)two;

        if (p1->p_start > p2->p_start) {
                return (1);
        } else if (p1->p_start < p2->p_start) {
                return (-1);
        } else {
                return (0);
        }
}

static int
partcmp64(const void *one, const void *two)
{
        dk_part_t *p1 = *(dk_part_t **)one;
        dk_part_t *p2 = *(dk_part_t **)two;

        if (p1->p_start > p2->p_start) {
                return (1);
        } else if (p1->p_start < p2->p_start) {
                return (-1);
        } else {
                return (0);
        }
}

/*
 * prtvtoc(): Read and print a VTOC.
 */
static int
prtvtoc(char *devname)
{
        int fd;
        int idx = 0;
        freemap_t *freemap;
        struct stat sb;
        struct extvtoc vtoc;
        int geo;
        struct dk_geom geom;
        char *name;
        int newvtoc = 0;
        struct dk_gpt *efi;

        name = getfullrawname(devname);
        if (name == NULL) {
                warnx("%s: internal administrative call (getfullrawname) "
                    "failed", devname);
                return (-1);
        }
        if (strcmp(name, "") == 0)
                name = devname;
        if ((fd = open(name, O_NONBLOCK|O_RDONLY)) < 0) {
                warn("%s: failed to open device", name);
                return (-1);
        }
        if (fstat(fd, &sb) < 0) {
                warn("%s: failed to stat device", name);
                return (-1);
        }
        if ((sb.st_mode & S_IFMT) != S_IFCHR) {
                warnx("%s: Not a raw device", name);
                return (-1);
        }

        geo = (readgeom(fd, name, &geom) == 0);
        if (geo) {
                if ((idx = readvtoc(fd, name, &vtoc)) == VT_ENOTSUP) {
                        idx = (readefi(fd, name, &efi) == 0);
                        newvtoc = 1;
                } else {
                        idx = (idx == 0);
                }
        }
        (void) close(fd);
        if ((!geo) || (!idx))
                return (-1);
        if (!newvtoc)
                freemap = findfree(&geom, &vtoc);
        else
                freemap = findfree64(efi);
        if (fflag) {
                if (!newvtoc)
                        putfree(&vtoc, freemap);
                else
                        putfree64(efi, freemap);
        } else {
                if (!newvtoc) {
                        puttable(&geom, &vtoc, freemap, devname,
                            getmntpt(major(sb.st_rdev),
                            noparttn(minor(sb.st_rdev))));
                } else {
                        puttable64(efi, freemap, devname,
                            getmntpt(major(sb.st_rdev),
                            noparttn(minor(sb.st_rdev))));
                }
        }
        if (newvtoc)
                efi_free(efi);
        return (0);
}

/*
 * putfree():
 *
 * Print shell assignments for disk free space. FREE_START and FREE_SIZE
 * represent the starting block and number of blocks of the first chunk
 * of free space. FREE_PART lists the unassigned partitions.
 */
static void
putfree(struct extvtoc *vtoc, freemap_t *freemap)
{
        freemap_t *freeidx;
        ushort_t idx;
        int free_count = 0;

        for (freeidx = freemap; freeidx->fr_size; ++freeidx)
                free_count++;

        (void) printf("FREE_START=%llu FREE_SIZE=%llu FREE_COUNT=%d FREE_PART=",
            freemap->fr_start, freemap->fr_size, free_count);

        for (idx = 0; idx < vtoc->v_nparts; ++idx) {
                if (vtoc->v_part[idx].p_size == 0 && idx != 2)
                        (void) printf("%x", idx);
        }
        (void) printf("\n");
}

static void
putfree64(struct dk_gpt *efi, freemap_t *freemap)
{
        freemap_t *freeidx;
        ushort_t idx;
        int free_count = 0;

        for (freeidx = freemap; freeidx->fr_size; ++freeidx)
                free_count++;

        (void) printf("FREE_START=%llu FREE_SIZE=%llu FREE_COUNT=%d FREE_PART=",
            freemap->fr_start, freemap->fr_size, free_count);

        for (idx = 0; idx < efi->efi_nparts; ++idx) {
                if (efi->efi_parts[idx].p_size == 0 && idx != 2)
                        (void) printf("%x", idx);
        }
        (void) printf("\n");
}

static void
print_table_header()
{
        (void) printf("*                            First       Sector"
            "      Last\n");
        (void) printf("* Partition  Tag  Flags      Sector       Count"
            "      Sector  Mount Directory\n");
}

static void
print_table_row(uint_t partition, uint_t tag, uint_t flag,
    u_longlong_t first_sector, u_longlong_t sector_count,
    u_longlong_t last_sector, const char *mount_dir)
{
        (void) printf("  %6u   %4u    %02x  %11llu %11llu %11llu",
            partition, tag, flag, first_sector, sector_count, last_sector);
        if (mount_dir != NULL) {
                (void) printf("   %s", mount_dir);
        }
        (void) printf("\n");
}

static void
print_freemap(freemap_t *freemap)
{
        if (freemap->fr_size == 0) {
                /*
                 * The freemap is completely empty, so do not print the header.
                 */
                return;
        }

        (void) printf("* Unallocated space:\n"
            "*         First       Sector      Last\n"
            "*         Sector       Count      Sector\n");

        do {
                (void) printf("*   %11llu %11llu %11llu\n",
                    freemap->fr_start, freemap->fr_size,
                    freemap->fr_size + freemap->fr_start - 1);
        } while ((++freemap)->fr_size != 0);

        (void) printf("*\n");
}

/*
 * puttable(): Print a human-readable VTOC.
 */
static void
puttable(struct dk_geom *geom, struct extvtoc *vtoc, freemap_t *freemap,
    char *name, char **mtab)
{
        ushort_t idx;
        ulong_t cylsize;

        cylsize = (geom->dkg_nsect) * (geom->dkg_nhead);
        if (!hflag && !sflag) {
                u_longlong_t asectors = (u_longlong_t)cylsize * geom->dkg_ncyl;
                u_longlong_t sectors = (u_longlong_t)cylsize * geom->dkg_pcyl;

                (void) printf("* %s", name);
                if (vtoc->v_volume[0] != '\0')
                        (void) printf(" (volume \"%.8s\")", vtoc->v_volume);

                (void) printf(" partition map\n");
                (void) printf("*\n* Dimensions:\n");
                (void) printf("* %11u bytes/sector\n", vtoc->v_sectorsz);
                (void) printf("* %11u sectors/track\n", geom->dkg_nsect);
                (void) printf("* %11u tracks/cylinder\n", geom->dkg_nhead);
                (void) printf("* %11lu sectors/cylinder\n", cylsize);
                (void) printf("* %11u cylinders\n", geom->dkg_pcyl);
                (void) printf("* %11u accessible cylinders\n", geom->dkg_ncyl);
                (void) printf("* %11llu sectors\n", sectors);
                (void) printf("* %11llu accessible sectors\n", asectors);
                (void) printf("*\n* Flags:\n");
                (void) printf("*   1: unmountable\n");
                (void) printf("*  10: read-only\n*\n");

                print_freemap(freemap);
        }

        if (!hflag) {
                print_table_header();
        }

        for (idx = 0; idx < vtoc->v_nparts; ++idx) {
                const char *mount_dir = NULL;
                struct extpartition *p = &vtoc->v_part[idx];

                if (p->p_size == 0)
                        continue;

                if (mtab != NULL) {
                        mount_dir = mtab[idx];
                }

                print_table_row(idx, p->p_tag, p->p_flag, p->p_start,
                    p->p_size, p->p_start + p->p_size - 1, mount_dir);
        }
}

/*
 * puttable(): Print a human-readable VTOC.
 */
static void
puttable64(struct dk_gpt *efi, freemap_t *freemap, char *name, char **mtab)
{
        if (!hflag && !sflag) {
                (void) printf("* %s", name);
                for (uint_t idx = 0; idx < efi->efi_nparts; idx++) {
                        if (efi->efi_parts[idx].p_tag == V_RESERVED &&
                            efi->efi_parts[idx].p_name[0] != '\0') {
                                (void) printf(" (volume \"%.8s\")",
                                    efi->efi_parts[idx].p_name);
                        }
                }
                (void) printf(" partition map\n");
                (void) printf("*\n* Dimensions:\n");
                (void) printf("* %11u bytes/sector\n", efi->efi_lbasize);
                (void) printf("* %11llu sectors\n", efi->efi_last_lba + 1);
                (void) printf("* %11llu accessible sectors\n",
                    efi->efi_last_u_lba - efi->efi_first_u_lba + 1);
                (void) printf("*\n* Flags:\n");
                (void) printf("*   1: unmountable\n");
                (void) printf("*  10: read-only\n*\n");

                print_freemap(freemap);
        }

        if (!hflag) {
                print_table_header();
        }

        for (uint_t idx = 0; idx < efi->efi_nparts; ++idx) {
                const char *mount_dir = NULL;
                dk_part_t *p = &efi->efi_parts[idx];

                if (p->p_size == 0)
                        continue;

                if (idx < 7 && mtab != NULL) {
                        mount_dir = mtab[idx];
                }

                print_table_row(idx, p->p_tag, p->p_flag, p->p_start,
                    p->p_size, p->p_start + p->p_size - 1, mount_dir);
        }
}

/*
 * readgeom(): Read the disk geometry.
 */
static int
readgeom(int fd, char *name, struct dk_geom *geom)
{
        if (ioctl(fd, DKIOCGGEOM, geom) < 0) {
                if (errno != ENOTSUP) {
                        warnx("%s: Unable to read Disk geometry errno = 0x%x",
                            name, errno);
                        return (-1);
                }

                (void) memset(geom, 0, sizeof (struct dk_geom));
        }

        return (0);
}

/*
 * readvtoc(): Read a partition map.
 */
static int
readvtoc(int fd, char *name, struct extvtoc *vtoc)
{
        int retval;

        if ((retval = read_extvtoc(fd, vtoc)) >= 0)
                return (0);

        switch (retval) {
        case VT_EIO:
                warnx("%s: Unable to read VTOC", name);
                return (-1);
        case VT_EINVAL:
                warnx("%s: Invalid VTOC", name);
                return (-1);
        case VT_ERROR:
                warnx("%s: Unknown problem reading VTOC", name);
                return (-1);
        }

        return (retval);
}

/*
 * readefi(): Read a partition map.
 */
static int
readefi(int fd, char *name, struct dk_gpt **efi)
{
        int     retval;

        if ((retval = efi_alloc_and_read(fd, efi)) >= 0)
                return (0);

        switch (retval) {
        case VT_EIO:
                warnx("%s: Unable to read VTOC", name);
                return (-1);
        case VT_EINVAL:
                warnx("%s: Invalid VTOC", name);
                return (-1);
        case VT_ERROR:
                warnx("%s: Unknown problem reading VTOC", name);
                return (-1);
        }

        return (retval);
}

static void
memory_err(size_t l, int e, const char *fname)
{
        const char *reason;

        switch (e) {
        case EAGAIN:
                reason = "not enough memory was available, please try again";
                break;
        case ENOMEM:
                reason = "allocation size was too large";
                break;
        default:
                reason = strerror(e);
                break;
        }

        errx(EXIT_FAILURE, "%s: failed to allocate %llu bytes of memory: %s",
            fname, (u_longlong_t)l, reason);
}

static void *
safe_calloc(const char *fname, size_t nelem, size_t elsize)
{
        void *r;

        if ((r = calloc(nelem, elsize)) == NULL) {
                memory_err(nelem * elsize, errno, fname);
        }

        return (r);
}

static char *
safe_strdup(const char *fname, const char *str)
{
        size_t l = strlen(str);
        char *r;

        if ((r = strndup(str, l)) == NULL) {
                memory_err(l + 1, errno, fname);
        }

        return (r);
}

/*
 * usage(): Print a helpful message and exit.
 */
static void
usage()
{
        (void) fprintf(stderr, "Usage:\t%s [ -fhs ] [ -t fstab ] [ -m mnttab ] "
            "rawdisk ...\n", getprogname());
        exit(1);
}