/*
 * Copyright (c) 2007 Bruce M. Simpson.
 * 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.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``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 THE AUTHOR OR CONTRIBUTORS 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/types.h>
#include <sys/ioctl.h>
#include <sys/pciio.h>
#include <sys/mman.h>
#include <sys/memrange.h>
#include <sys/stat.h>
#include <machine/endian.h>

#include <stddef.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <libgen.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>

#define _PATH_DEVPCI    "/dev/pci"
#define _PATH_DEVMEM    "/dev/mem"

#define PCI_CFG_CMD             0x04            /* command register */
#define PCI_CFG_ROM_BAR         0x30            /* rom base register */

#define PCI_ROM_ADDR_MASK       0xFFFFFC00      /* the 21 MSBs form the BAR */
#define PCI_ROM_RESERVED_MASK   0x03FE          /* mask for reserved bits */
#define PCI_ROM_ACTIVATE        0x01            /* mask for activation bit */

#define PCI_CMD_MEM_SPACE       0x02            /* memory space bit */
#define PCI_HDRTYPE_MFD         0x80            /* MFD bit in HDRTYPE reg. */

#define MAX_PCI_DEVS            64              /* # of devices in system */

typedef enum {
        PRINT = 0,
        SAVE = 1
} action_t;

/*
 * This is set to a safe physical base address in PCI range for my Vaio.
 * YOUR MACHINE *WILL* VARY, I SUGGEST YOU LOOK UP YOUR MACHINE'S MEMORY
 * MAP IN DETAIL IF YOU PLAN ON SAVING ROMS.
 *
 * This is the hole between the APIC and the BIOS (FED00000-FEDFFFFF);
 * should be a safe range on the i815 Solano chipset.
 */
#define PCI_DEFAULT_ROM_ADDR    0xFED00000

static char *progname = NULL;
static uintptr_t base_addr = PCI_DEFAULT_ROM_ADDR;

static void     usage(void);
static void     banner(void);
static void     pci_enum_devs(int pci_fd, action_t action);
static uint32_t pci_testrombar(int pci_fd, struct pci_conf *dev);
static int      pci_enable_bars(int pci_fd, struct pci_conf *dev,
    uint16_t *oldcmd);
static int      pci_disable_bars(int pci_fd, struct pci_conf *dev,
    uint16_t *oldcmd);
static int      pci_save_rom(char *filename, int romsize);

int
main(int argc, char *argv[])
{
        int              pci_fd;
        int              err;
        int              ch;
        action_t         action;
        char            *base_addr_string;
        char            *ep;

        err = -1;
        pci_fd = -1;
        action = PRINT;
        base_addr_string = NULL;
        ep = NULL;
        progname = basename(argv[0]);

        while ((ch = getopt(argc, argv, "sb:h")) != -1)
                switch (ch) {
                case 's':
                        action = SAVE;
                        break;
                case 'b':
                        base_addr_string = optarg;
                        break;
                case 'h':
                default:
                     usage();
        }
        argc -= optind;
        argv += optind;

        if (base_addr_string != NULL) {
                uintmax_t base_addr_max;

                base_addr_max = strtoumax(base_addr_string, &ep, 16);
                if (*ep != '\0') {
                        fprintf(stderr, "Invalid base address.\r\n");
                        usage();
                }
                /* XXX: TODO: deal with 64-bit PCI. */
                base_addr = (uintptr_t)base_addr_max;
                base_addr &= ~PCI_ROM_RESERVED_MASK;
        }

        if (argc > 0)
                usage();

        if ((pci_fd = open(_PATH_DEVPCI, O_RDWR)) == -1) {
                perror("open");
                goto cleanup;
        }

        banner();
        pci_enum_devs(pci_fd, action);

        err = 0;
cleanup:
        if (pci_fd != -1)
                close(pci_fd);

        exit ((err == 0) ? EXIT_SUCCESS : EXIT_FAILURE);
}

static void
usage(void)
{

        fprintf(stderr, "usage: %s [-s] [-b <base-address>]\r\n", progname);
        exit(EXIT_FAILURE);
}

static void
banner(void)
{

        fprintf(stderr,
                "WARNING: You are advised to run this program in single\r\n"
                "user mode, with few or no processes running.\r\n\r\n");
}

/*
 * Enumerate PCI device list to a limit of MAX_PCI_DEVS devices.
 */
static void
pci_enum_devs(int pci_fd, action_t action)
{
        struct pci_conf          devs[MAX_PCI_DEVS];
        char                     filename[16];
        struct pci_conf_io       pc;
        struct pci_conf         *p;
        int                      result;
        int                      romsize;
        uint16_t                 oldcmd;

        result = -1;
        romsize = 0;

        bzero(&pc, sizeof(pc));
        pc.match_buf_len = sizeof(devs);
        pc.matches = devs;

        if (ioctl(pci_fd, PCIOCGETCONF, &pc) == -1) {
                perror("ioctl PCIOCGETCONF");
                return;
        }

        if (pc.status == PCI_GETCONF_ERROR) {
                fprintf(stderr,
                    "Error fetching PCI device list from kernel.\r\n");
                return;
        }

        if (pc.status == PCI_GETCONF_MORE_DEVS) {
                fprintf(stderr,
"More than %d devices exist. Only the first %d will be inspected.\r\n",
                    MAX_PCI_DEVS, MAX_PCI_DEVS);
        }

        for (p = devs ; p < &devs[pc.num_matches]; p++) {

                /* No PCI bridges; only PCI devices. */
                if (p->pc_hdr != 0x00)
                        continue;

                romsize = pci_testrombar(pci_fd, p);

                switch (action) {
                case PRINT:
                        printf(
"Domain %04Xh Bus %02Xh Device %02Xh Function %02Xh: ",
                                p->pc_sel.pc_domain, p->pc_sel.pc_bus,
                                p->pc_sel.pc_dev, p->pc_sel.pc_func);
                        printf((romsize ? "%dKB ROM aperture detected."
                                        : "No ROM present."), romsize/1024);
                        printf("\r\n");
                        break;
                case SAVE:
                        if (romsize == 0)
                                continue;       /* XXX */

                        snprintf(filename, sizeof(filename), "%08X.rom",
                            ((p->pc_device << 16) | p->pc_vendor));

                        fprintf(stderr, "Saving %dKB ROM image to %s...\r\n",
                            romsize, filename);

                        if (pci_enable_bars(pci_fd, p, &oldcmd) == 0)
                                result = pci_save_rom(filename, romsize);

                        pci_disable_bars(pci_fd, p, &oldcmd);

                        if (result == 0)  {
                                fprintf(stderr, "Done.\r\n");
                        } else  {
                                fprintf(stderr,
"An error occurred whilst saving the ROM.\r\n");
                        }
                        break;
                } /* switch */
        } /* for */
}

/*
 * Return: size of ROM aperture off dev, 0 if no ROM exists.
 */
static uint32_t
pci_testrombar(int pci_fd, struct pci_conf *dev)
{
        struct pci_io    io;
        uint32_t         romsize;

        romsize = 0;

        /*
         * Only attempt to discover ROMs on Header Type 0x00 devices.
         */
        if (dev->pc_hdr != 0x00)
                return romsize;

        /*
         * Activate ROM BAR
         */
        io.pi_sel = dev->pc_sel;
        io.pi_reg = PCI_CFG_ROM_BAR;
        io.pi_width = 4;
        io.pi_data = 0xFFFFFFFF;
        if (ioctl(pci_fd, PCIOCWRITE, &io) == -1)
                return romsize;

        /*
         * Read back ROM BAR and compare with mask
         */
        if (ioctl(pci_fd, PCIOCREAD, &io) == -1)
                return 0;

        /*
         * Calculate ROM aperture if one was set.
         */
        if (io.pi_data & PCI_ROM_ADDR_MASK)
                romsize = -(io.pi_data & PCI_ROM_ADDR_MASK);

        /*
         * Disable the ROM BAR when done.
         */
        io.pi_data = 0;
        if (ioctl(pci_fd, PCIOCWRITE, &io) == -1)
                return 0;

        return romsize;
}

static int
pci_save_rom(char *filename, int romsize)
{
        int      fd, mem_fd, err;
        void    *map_addr;

        fd = err = mem_fd = -1;
        map_addr = MAP_FAILED;

        if ((mem_fd = open(_PATH_DEVMEM, O_RDONLY)) == -1) {
                perror("open");
                return -1;
        }

        map_addr = mmap(NULL, romsize, PROT_READ, MAP_SHARED|MAP_NOCORE,
            mem_fd, base_addr);

        /* Dump ROM aperture to a file. */
        if ((fd = open(filename, O_CREAT|O_RDWR|O_TRUNC|O_NOFOLLOW,
            S_IRUSR|S_IWUSR)) == -1) {
                perror("open");
                goto cleanup;
        }

        if (write(fd, map_addr, romsize) != romsize)
                perror("write");

        err = 0;
cleanup:
        if (fd != -1)
                close(fd);

        if (map_addr != MAP_FAILED)
                munmap((void *)base_addr, romsize);

        if (mem_fd != -1)
                close(mem_fd);

        return err;
}

static int
pci_enable_bars(int pci_fd, struct pci_conf *dev, uint16_t *oldcmd)
{
        struct pci_io io;

        /* Don't grok bridges. */
        if (dev->pc_hdr != 0x00)
                return -1;

        /* Save command register. */
        io.pi_sel = dev->pc_sel;
        io.pi_reg = PCI_CFG_CMD;
        io.pi_width = 2;
        if (ioctl(pci_fd, PCIOCREAD, &io) == -1)
                return -1;
        *oldcmd = (uint16_t)io.pi_data;

        io.pi_data |= PCI_CMD_MEM_SPACE;
        if (ioctl(pci_fd, PCIOCWRITE, &io) == -1)
                return -1;

        /*
         * Activate ROM BAR and map at the specified base address.
         */
        io.pi_sel = dev->pc_sel;
        io.pi_reg = PCI_CFG_ROM_BAR;
        io.pi_width = 4;
        io.pi_data = (base_addr | PCI_ROM_ACTIVATE);
        if (ioctl(pci_fd, PCIOCWRITE, &io) == -1)
                return -1;

        return 0;
}

static int
pci_disable_bars(int pci_fd, struct pci_conf *dev, uint16_t *oldcmd)
{
        struct pci_io    io;

        /*
         * Clear ROM BAR to deactivate the mapping.
         */
        io.pi_sel = dev->pc_sel;
        io.pi_reg = PCI_CFG_ROM_BAR;
        io.pi_width = 4;
        io.pi_data = 0;
        if (ioctl(pci_fd, PCIOCWRITE, &io) == -1)
                return 0;

        /*
         * Restore state of the command register.
         */
        io.pi_sel = dev->pc_sel;
        io.pi_reg = PCI_CFG_CMD;
        io.pi_width = 2;
        io.pi_data = *oldcmd;
        if (ioctl(pci_fd, PCIOCWRITE, &io) == -1) {
                perror("ioctl");
                return 0;
        }

        return 0;
}