/*      $NetBSD: mbrlabel.c,v 1.29 2018/03/30 13:14:25 mlelstv Exp $    */

/*
 * Copyright (C) 1998 Wolfgang Solfrank.
 * Copyright (C) 1998 TooLs GmbH.
 * All rights reserved.
 *
 * 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 by TooLs GmbH.
 * 4. The name of TooLs GmbH may not be used to endorse or promote products
 *    derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY TOOLS GMBH ``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 TOOLS GMBH 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.
 */

#include <sys/cdefs.h>
#ifndef lint
__RCSID("$NetBSD: mbrlabel.c,v 1.29 2018/03/30 13:14:25 mlelstv Exp $");
#endif /* not lint */

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

#include <sys/param.h>
#define FSTYPENAMES
#include <sys/disklabel.h>
#include <sys/bootblock.h>
#include <sys/ioctl.h>

#include "dkcksum.h"
#include "extern.h"

__dead static void      usage(void);
static void     getlabel(int);
static void     setlabel(int, int);
static int      getparts(int, u_int32_t, u_int32_t, int);
static u_int16_t        getshort(void *);
static u_int32_t        getlong(void *);

struct disklabel label;

static void
getlabel(int sd)
{

        if (ioctl(sd, DIOCGDINFO, &label) < 0) {
                perror("get label");
                exit(1);
        }
        /*
         * Some ports seem to not set the number of partitions
         * correctly, albeit they seem to set the raw partition ok!
         */
        if (label.d_npartitions <= getrawpartition())
                label.d_npartitions = getrawpartition() + 1;
}

static void
setlabel(int sd, int doraw)
{
        int one = 1;

        label.d_checksum = 0;
        label.d_checksum = dkcksum(&label);
        if (ioctl(sd, doraw ? DIOCWDINFO : DIOCSDINFO, &label) < 0) {
                perror("set label");
                exit(1);
        }
        if (!doraw)
                /* If we haven't written to the disk, don't discard on close */
                ioctl(sd, DIOCKLABEL, &one);

}

static u_int16_t
getshort(void *p)
{
        unsigned char *cp = p;

        return (cp[0] | (cp[1] << 8));
}

static u_int32_t
getlong(void *p)
{
        unsigned char *cp = p;

        return (cp[0] | (cp[1] << 8) | (cp[2] << 16) | (cp[3] << 24));
}

static int
getparts(int sd, u_int32_t off, u_int32_t extoff, int verbose)
{
        unsigned char           *buf;
        struct mbr_partition    parts[MBR_PART_COUNT];
        struct partition        npe;
        off_t                   loff;
        int                     i, j, unused, changed;
        unsigned                bsize = label.d_secsize;

        if (bsize < DEV_BSIZE) {
                fprintf(stderr,"Invalid sector size %u\n", bsize);
                exit(1);
        }

        buf = malloc(bsize);
        if (buf == NULL) {
                perror("malloc I/O buffer");
                exit(1);
        }

        changed = 0;
        loff = (off_t)off * bsize;

        if (lseek(sd, loff, SEEK_SET) != loff) {
                perror("seek label");
                exit(1);
        }
        if (read(sd, buf, bsize) != (ssize_t)bsize) {
                if (off != MBR_BBSECTOR)
                        perror("read label (sector is possibly out of "
                            "range)");
                else
                        perror("read label");
                exit(1);
        }
        if (getshort(buf + MBR_MAGIC_OFFSET) != MBR_MAGIC)
                return (changed);
        memcpy(parts, buf + MBR_PART_OFFSET, sizeof parts);

                                /* scan partition table */
        for (i = 0; i < MBR_PART_COUNT; i++) {
                if (parts[i].mbrp_type == 0 ||
                                /* extended partitions are handled below */
                    MBR_IS_EXTENDED(parts[i].mbrp_type))
                        continue;

                memset((void *)&npe, 0, sizeof(npe));
                npe.p_size = getlong(&parts[i].mbrp_size);
                npe.p_offset = getlong(&parts[i].mbrp_start) + off;
                npe.p_fstype = xlat_mbr_fstype(parts[i].mbrp_type);

                                /* find existing entry, or first free slot */
                unused = -1;    /* flag as no free slot */
                if (verbose)
                        printf(
                            "Found %s partition; size %u (%u MB), offset %u\n",
                            fstypenames[npe.p_fstype],
                            npe.p_size, npe.p_size / 2048, npe.p_offset);

                for (j = 0; j < label.d_npartitions; j++) {
                        struct partition *lpe;

                        if (j == RAW_PART)
                                continue;
                        lpe = &label.d_partitions[j];

                        if (lpe->p_size == npe.p_size &&
                            lpe->p_offset == npe.p_offset) {
                                if (verbose)
                                        printf(
                                            "  skipping existing %s partition at slot %c.\n",
                                            fstypenames[lpe->p_fstype],
                                            j + 'a');
                                unused = -2;    /* flag as existing */
                                break;
                        }

                        if (unused == -1 && lpe->p_size == 0 &&
                            lpe->p_fstype == FS_UNUSED)
                                unused = j;
                }
                if (unused == -1) {
                        for (j = 0; j < label.d_npartitions; j++) {
                                struct partition *lpe;

                                if (j == RAW_PART)
                                        continue;
                                lpe = &label.d_partitions[j];

                                if ((npe.p_offset >= lpe->p_offset &&
                                    npe.p_offset < lpe->p_offset + lpe->p_size) ||
                                   (npe.p_offset + npe.p_size - 1 >= lpe->p_offset &&
                                    npe.p_offset + npe.p_size - 1 < lpe->p_offset + lpe->p_size)) {
                                        printf(
                                            "  skipping overlapping %s partition at slot %c.\n",
                                            fstypenames[lpe->p_fstype],
                                            j + 'a');
                                        unused = -2;    /* flag as existing */
                                        break;
                                }
                        }
                }
                if (unused == -2)
                        continue;       /* entry exists, skip... */
                if (unused == -1) {
                        if (label.d_npartitions < MAXPARTITIONS) {
                                unused = label.d_npartitions;
                                label.d_npartitions++;
                        } else {
                                printf(
                                "  WARNING: no slots free for %s partition.\n",
                                    fstypenames[npe.p_fstype]);
                                continue;
                        }
                }

                if (verbose)
                        printf("  adding %s partition to slot %c.\n",
                            fstypenames[npe.p_fstype], unused + 'a');

                /*
                 * XXX guess some filesystem parameters, these should be
                 * scanned from the superblocks
                 */
                switch (npe.p_fstype) {
                case FS_BSDFFS:
                case FS_APPLEUFS:
                        npe.p_fsize = 1024;
                        npe.p_frag = 8;
                        npe.p_cpg = 16;
                        break;
                case FS_BSDLFS:
                        npe.p_fsize = 1024;
                        npe.p_frag = 8;
                        npe.p_sgs = 7;
                        break;
                }

                changed++;
                label.d_partitions[unused] = npe;
        }

                                /* recursively scan extended partitions */
        for (i = 0; i < MBR_PART_COUNT; i++) {
                u_int32_t poff;

                if (MBR_IS_EXTENDED(parts[i].mbrp_type)) {
                        poff = getlong(&parts[i].mbrp_start) + extoff;
                        changed += getparts(sd, poff,
                            extoff ? extoff : poff, verbose);
                }
        }

        free(buf);
        return (changed);
}

static void
usage(void)
{
        fprintf(stderr, "usage: %s [-fqrw] [-s sector] device\n",
            getprogname());
        exit(1);
}


int
main(int argc, char **argv)
{
        int     sd, ch, changed;
        char    *ep, name[MAXPATHLEN];
        int     force;                  /* force label update */
        int     raw;                    /* update on-disk label as well */
        int     verbose;                /* verbose output */
        int     write_it;               /* update in-core label if changed */
        uint32_t sector;                /* sector that contains the MBR */
        ulong   ulsector;

        force = 0;
        raw = 0;
        verbose = 1;
        write_it = 0;
        sector = MBR_BBSECTOR;
        while ((ch = getopt(argc, argv, "fqrs:w")) != -1) {
                switch (ch) {
                case 'f':
                        force = 1;
                        break;
                case 'q':
                        verbose = 0;
                        break;
                case 'r':
                        raw = 1;
                        break;
                case 's':
                        errno = 0;
                        ulsector = strtoul(optarg, &ep, 10);
                        if (optarg[0] == '\0' || *ep != '\0')
                                errx(EXIT_FAILURE,
                                    "sector number (%s) incorrectly specified",
                                    optarg);
                        if ((errno == ERANGE && ulsector == ULONG_MAX) ||
                            ulsector > UINT32_MAX)
                                errx(EXIT_FAILURE,
                                    "sector number (%s) out of range",
                                    optarg);
                        sector = (uint32_t)ulsector;
                        break;
                case 'w':
                        write_it = 1;
                        break;
                default:
                        usage();
                }
        }
        argc -= optind;
        argv += optind;
        if (argc != 1)
                usage();

        if ((sd = opendisk(argv[0], write_it ? O_RDWR : O_RDONLY, name,
            (size_t)MAXPATHLEN, 0)) < 0) {
                perror(argv[0]);
                exit(1);
        }
        getlabel(sd);
        changed = getparts(sd, sector, 0, verbose);

        if (verbose) {
                putchar('\n');
                showpartitions(stdout, &label, 0);
                putchar('\n');
        }
        if (write_it) {
                if (! changed && ! force)
                        printf("No change; not updating disk label.\n");
                else {
                        if (verbose)
                                printf("Updating in-core %sdisk label.\n",
                                    raw ? "and on-disk " : "");
                        setlabel(sd, raw);
                }
        } else {
                printf("Not updating disk label.\n");
        }
        close(sd);
        return (0);
}