/* drivemap.c - command to manage the BIOS drive mappings.  */
/*
 *  GRUB  --  GRand Unified Bootloader
 *  Copyright (C) 2008  Free Software Foundation, Inc.
 *
 *  GRUB is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  GRUB is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with GRUB.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <grub/normal.h>
#include <grub/dl.h>
#include <grub/mm.h>
#include <grub/misc.h>
#include <grub/disk.h>
#include <grub/loader.h>
#include <grub/machine/loader.h>
#include <grub/machine/biosdisk.h>

static const struct grub_arg_option options[] =
  {
    {"list", 'l', 0, "show the current mappings", 0, 0},
    {"reset", 'r', 0, "reset all mappings to the default values", 0, 0},
    {0, 0, 0, 0, 0, 0}
  };

static grub_preboot_hookid insthandler_hook = 0;

typedef struct drivemap_node
{
  grub_uint8_t newdrive;
  grub_uint8_t redirto;
  struct drivemap_node *next;
} drivemap_node_t;

static drivemap_node_t *drivemap = 0;
static grub_err_t drivemap_install_int13_handler(void);

static grub_err_t
drivemap_set (grub_uint8_t newdrive, grub_uint8_t redirto)
  /* Puts the specified mapping into the table, replacing an existing mapping
   * for newdrive or adding a new one if required. */
{
  drivemap_node_t *mapping = 0, *search = drivemap;
  while (search)
    {
      if (search->newdrive == newdrive)
        {
          mapping = search;
          break;
        }
      search = search->next;
    }

  if (mapping)  /* There was a mapping already in place, modify it */
    mapping->redirto = redirto;
  else  /* Create a new mapping and add it to the head of the list */
    {
      mapping = grub_malloc (sizeof (drivemap_node_t));
      if (!mapping)
        return grub_error (GRUB_ERR_OUT_OF_MEMORY, "cannot allocate map entry");
      mapping->newdrive = newdrive;
      mapping->redirto = redirto;
      mapping->next = drivemap;
      drivemap = mapping;
    }
  return GRUB_ERR_NONE;
}

static void
drivemap_remove (grub_uint8_t newdrive)
  /* Removes the mapping for newdrive from the table. If there is no mapping,
   * then this function behaves like a no-op on the map. */
{
  drivemap_node_t *mapping = 0, *search = drivemap, *previous = 0;
  while (search)
    {
      if (search->newdrive == newdrive)
        {
          mapping = search;
          break;
        }
      previous = search;
      search = search->next;
    }
  if (mapping) /* Found */
    {
      if (previous)
        previous->next = mapping->next;
      else drivemap = mapping->next; /* Entry was head of list */
      grub_free (mapping);
    }
}

static grub_err_t parse_biosdisk (const char *name, grub_uint8_t *disknum)
{
  if (!name) return GRUB_ERR_BAD_ARGUMENT;
  if (name[0] == '(')
    name++; /* Skip the first ( in (hd0) - disk_open wants just the name! */
  grub_disk_t disk = grub_disk_open (name);
  if (!disk)
    return GRUB_ERR_UNKNOWN_DEVICE;
  else
    {
      enum grub_disk_dev_id id = disk->dev->id;
      if (disknum)
        *disknum = disk->id;   /* Only valid, of course if it's a biosdisk */
      grub_disk_close (disk);
      return (GRUB_DISK_DEVICE_BIOSDISK_ID != id) ?
          GRUB_ERR_BAD_DEVICE : GRUB_ERR_NONE;
    }
}

static grub_err_t revparse_biosdisk(const grub_uint8_t dnum, const char **output)
{
  grub_err_t retval = GRUB_ERR_UNKNOWN_DEVICE;
  auto int find (const char *name);
  int find (const char *name)
  {
    grub_disk_t disk = grub_disk_open (name);
    if (!disk)
      return 0;
    else
      {
        int found = 0;
        if (dnum == disk->id && GRUB_DISK_DEVICE_BIOSDISK_ID == disk->dev->id)
          {
            found = 1;
            *output = name;
            retval = GRUB_ERR_NONE;
          }
        grub_disk_close (disk);
        return found;
      }
  }

  grub_disk_dev_iterate (&find);
  return retval;
}

static grub_err_t
grub_cmd_drivemap (struct grub_arg_list *state, int argc, char **args)
{
  if (state[0].set) /* Show: list mappings */
    {
      if (!drivemap)
        grub_printf ("No drives have been remapped");
      else
        {
          grub_printf ("Showing only remapped drives. Drives that have had "
                       "their slot assigned to another one and have not been "
                       "themselves remapped will become inaccessible through "
                       "the BIOS routines to the booted OS.\n\n");
          grub_printf ("Mapped\tGRUB\n");
          drivemap_node_t *curnode = drivemap;
          while (curnode)
            {
              const char *dname = 0;
              grub_err_t err = revparse_biosdisk (curnode->redirto, &dname);
              if (err != GRUB_ERR_NONE)
                return grub_error (err, "invalid mapping: non-existent disk or"
                                        "not managed by the BIOS");
              grub_printf("0x%02x\t%4s\n", curnode->newdrive, dname);
              curnode = curnode->next;
            }
        }
    }
  else if (state[1].set) /* Reset: just delete all mappings */
    {
      if (drivemap)
        {
          drivemap_node_t *curnode = drivemap, *prevnode = 0;
          while (curnode)
            {
              prevnode = curnode;
              curnode = curnode->next;
              grub_free (prevnode);
            }
          drivemap = 0;
        }
    }
  else
    {
      if (argc != 2)
        return grub_error (GRUB_ERR_BAD_ARGUMENT, "two arguments required");
      grub_uint8_t mapfrom = 0;
      grub_err_t err = parse_biosdisk (args[0], &mapfrom);
      if (err != GRUB_ERR_NONE)
        return grub_error (err, "invalid disk or not managed by the BIOS");

      grub_errno = GRUB_ERR_NONE;
      unsigned long mapto = grub_strtoul (args[1], 0, 0);
      if (grub_errno != GRUB_ERR_NONE)
        return grub_error (grub_errno,
                          "BIOS disk number must be between 0 and 255");
      else if (mapto == mapfrom)  /* Reset to default */      
        {
          grub_printf ("Removing the mapping for %s (%02x)", args[0], mapfrom);
          drivemap_remove (mapfrom);
        }
      else  /* Map */
        {
          grub_printf ("Mapping %s (%02x) to %02x\n", args[0], mapfrom, (grub_uint8_t)mapto);
          return drivemap_set ((grub_uint8_t)mapto, mapfrom);
        }
    }

  return GRUB_ERR_NONE;
}

typedef struct __attribute__ ((packed)) int13map_node
{
  grub_uint8_t disknum;
  grub_uint8_t mapto;
} int13map_node_t;

#define INT13H_OFFSET(x) ( ((grub_uint8_t*)(x)) - ((grub_uint8_t*)&grub_drivemap_int13_handler_base) )
#define INT13H_REBASE(x) ( (void*) (((grub_uint8_t*)handler_base) + (x)) )
#define INT13H_TONEWADDR(x) INT13H_REBASE( INT13H_OFFSET( x ) )
/* This code rests on the assumption that GRUB does not activate any kind of
 * memory mapping, since it accesses realmode structures by their absolute
 * addresses, like the IVT at 0 or the BDA at 0x400 */
static grub_err_t drivemap_install_int13_handler(void)
{
  grub_size_t entries = 0;
  drivemap_node_t *curentry = drivemap;
  while (curentry)  /* Count entries to prepare a contiguous map block */
    {
      entries++;
      curentry = curentry->next;
    }
  if (0 == entries)
    return GRUB_ERR_NONE;  /* No need to install the int13h handler */
  else
    {
      
      grub_uint32_t *ivtslot = (grub_uint32_t*)0x0000004c;
      
      /* Save the pointer to the old int13h handler */    
      grub_drivemap_int13_oldhandler = *ivtslot;
      grub_printf ("Old int13 handler at %08x\n", grub_drivemap_int13_oldhandler);
      /* Reserve a section of conventional memory as "BIOS memory" for handler:
       * BDA offset 0x13 contains the top of such memory */
      grub_uint16_t *bpaMemInKb = (grub_uint16_t*)0x00000413;
      grub_printf ("Top of conventional memory: %u\n", *bpaMemInKb);
      grub_size_t totalSize = grub_drivemap_int13_size
                            + (entries + 1) * sizeof(int13map_node_t);
      grub_printf ("Payload is %u bytes long, reserving %u Kb\n", totalSize, (totalSize >> 10) + (totalSize % 1024) == 0 ? 0 : 1);
      *bpaMemInKb -= (totalSize >> 10) +
                     ((totalSize % 1024) == 0) ? 0 : 1;
      /* Install the int13h handler (copy to the reserved area) */
      grub_uint8_t *handler_base = (grub_uint8_t*)(*bpaMemInKb << 10);
      grub_printf ("Copying int13 handler to: %p\n", handler_base);
      grub_memcpy (handler_base, &grub_drivemap_int13_handler_base, grub_drivemap_int13_size);
      /* Copy the mappings to the reserved area */
      curentry = drivemap;
      grub_size_t i;
      int13map_node_t *handler_map = (int13map_node_t*) INT13H_TONEWADDR(&grub_drivemap_int13_mapstart);
      grub_printf ("Target map at %p, copying mappings...\n", handler_map);
      for (i = 0; i < entries && curentry; i++, curentry = curentry->next)
        {
          handler_map[i].disknum = curentry->newdrive;
          handler_map[i].mapto = curentry->redirto;
          grub_printf ("\t#%d: 0x%02x <- 0x%02x\n", i, handler_map[i].disknum, handler_map[i].mapto);
        }
      /* Signal end-of-map */
      handler_map[i].disknum = 0;
      handler_map[i].mapto = 0;
      grub_printf ("\t#%d: 0x%02x <- 0x%02x (end)\n", i, handler_map[i].disknum, handler_map[i].mapto);
      /* Install the int13h handler (set the IVT entry for it) */
      grub_uint32_t ivtentry = ((grub_uint32_t)handler_base) << 12; /* Segment address */
      ivtentry |= (grub_uint16_t) INT13H_OFFSET(grub_drivemap_int13_handler);
      grub_printf ("New int13 handler IVT pointer: %08x\n", ivtentry);
      *ivtslot = ivtentry;
      
      return GRUB_ERR_NONE;
    }
}
#undef INT13H_TONEWADDR
#undef INT13H_REBASE
#undef INT13H_OFFSET

GRUB_MOD_INIT(drivemap)
{
  (void)mod;			/* To stop warning. */
  grub_register_command ("drivemap", grub_cmd_drivemap, GRUB_COMMAND_FLAG_BOTH,
			 "drivemap -s | -r | (hdX) newdrivenum", "Manage the BIOS drive mappings", options);
  insthandler_hook = grub_loader_register_preboot (&drivemap_install_int13_handler, 1);
}

GRUB_MOD_FINI(drivemap)
{
  grub_loader_unregister_preboot (insthandler_hook);
  insthandler_hook = 0;
  grub_unregister_command ("drivemap");
}