From d1e517eb20c2e5e14bdd8eea7168538f9f15a67c Mon Sep 17 00:00:00 2001
From: Vladimir 'phcoder' Serbinenko <phcoder@gmail.com>
Date: Sat, 25 Dec 2010 12:11:54 +0100
Subject: [PATCH] extend --disk-module to AHCI and USB

---
 grub-core/bus/usb/ohci.c |  8 ++++++--
 grub-core/disk/ahci.c    | 38 ++++++++++++++++++++++++++++++++++++++
 util/grub-install.in     | 14 +++++++++++---
 3 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/grub-core/bus/usb/ohci.c b/grub-core/bus/usb/ohci.c
index 8adaee6e0..577f35259 100644
--- a/grub-core/bus/usb/ohci.c
+++ b/grub-core/bus/usb/ohci.c
@@ -1422,18 +1422,22 @@ static struct grub_usb_controller_dev usb_controller =
   .detect_dev = grub_ohci_detect_dev
 };
 
+static void *fini_hnd;
+
 GRUB_MOD_INIT(ohci)
 {
   COMPILE_TIME_ASSERT (sizeof (struct grub_ohci_td) == 32);
   COMPILE_TIME_ASSERT (sizeof (struct grub_ohci_ed) == 16);
   grub_ohci_inithw ();
   grub_usb_controller_dev_register (&usb_controller);
-  grub_loader_register_preboot_hook (grub_ohci_fini_hw, grub_ohci_restore_hw,
-				     GRUB_LOADER_PREBOOT_HOOK_PRIO_DISK);
+  fini_hnd = grub_loader_register_preboot_hook (grub_ohci_fini_hw,
+						grub_ohci_restore_hw,
+						GRUB_LOADER_PREBOOT_HOOK_PRIO_DISK);
 }
 
 GRUB_MOD_FINI(ohci)
 {
   grub_ohci_fini_hw (0);
+  grub_loader_unregister_preboot_hook (fini_hnd);
   grub_usb_controller_dev_unregister (&usb_controller);
 }
diff --git a/grub-core/disk/ahci.c b/grub-core/disk/ahci.c
index 3160523ce..61ca7e4d6 100644
--- a/grub-core/disk/ahci.c
+++ b/grub-core/disk/ahci.c
@@ -25,6 +25,7 @@
 #include <grub/scsi.h>
 #include <grub/misc.h>
 #include <grub/list.h>
+#include <grub/loader.h>
 
 struct grub_ahci_cmd_head
 {
@@ -296,6 +297,34 @@ grub_ahci_initialize (void)
   return grub_errno;
 }
 
+static grub_err_t
+grub_ahci_fini_hw (int noreturn __attribute__ ((unused)))
+{
+  struct grub_ahci_device *dev, *next;
+
+  for (dev = grub_ahci_devices; dev; dev = next)
+    {
+      next = dev->next;
+      dev->hba->ports[dev->num].command &= ~GRUB_AHCI_HBA_PORT_CMD_FRE;
+      while ((dev->hba->ports[dev->num].command & GRUB_AHCI_HBA_PORT_CMD_FR));
+      dev->hba->ports[dev->num].command &= ~GRUB_AHCI_HBA_PORT_CMD_ST;
+      while ((dev->hba->ports[dev->num].command & GRUB_AHCI_HBA_PORT_CMD_CR));
+      grub_dma_free (dev->command_list_chunk);
+      grub_dma_free (dev->command_table_chunk);
+      grub_dma_free (dev->rfis);
+
+      grub_free (dev);
+    }
+  grub_ahci_devices = NULL;
+  return GRUB_ERR_NONE;
+}
+
+static grub_err_t
+grub_ahci_restore_hw (void)
+{
+  return grub_ahci_initialize ();
+}
+
 
 
 
@@ -533,6 +562,8 @@ static struct grub_ata_dev grub_ahci_dev =
 
 
 
+static void *fini_hnd;
+
 GRUB_MOD_INIT(ahci)
 {
   /* To prevent two drivers operating on the same disks.  */
@@ -548,9 +579,16 @@ GRUB_MOD_INIT(ahci)
 
   /* AHCI devices are handled by scsi.mod.  */
   grub_ata_dev_register (&grub_ahci_dev);
+
+  fini_hnd = grub_loader_register_preboot_hook (grub_ahci_fini_hw,
+						grub_ahci_restore_hw,
+						GRUB_LOADER_PREBOOT_HOOK_PRIO_DISK);
 }
 
 GRUB_MOD_FINI(ahci)
 {
+  grub_ahci_fini_hw (0);
+  grub_loader_unregister_preboot_hook (fini_hnd);
+
   grub_ata_dev_unregister (&grub_ahci_dev);
 }
diff --git a/util/grub-install.in b/util/grub-install.in
index b9e833360..6259b1c4a 100644
--- a/util/grub-install.in
+++ b/util/grub-install.in
@@ -77,7 +77,7 @@ if [ "${target_cpu}-${platform}" = "i386-pc" ] ; then
 elif [ "${platform}" = "ieee1275" ] || [ "${platform}" = "efi" ] ; then
     disk_module=
 else
-    disk_module=ata
+    disk_module=native
 fi
 
 # Usage: usage
@@ -114,7 +114,7 @@ Install GRUB on your drive.
 EOF
 if [ "${target_cpu}-${platform}" = "i386-pc" ] ; then
     cat <<EOF
-  --disk-module=MODULE    disk module to use
+  --disk-module=MODULE    disk module to use (biosdisk or native)
 EOF
 fi
 if [ "${target_cpu}-${platform}" = "i386-ieee1275" ] || [ "${target_cpu}-${platform}" = "powerpc-ieee1275" ] ; then
@@ -486,6 +486,14 @@ done
 # Device abstraction module, if any (lvm, raid).
 devabstraction_module="`"$grub_probe" --device-map="${device_map}" --target=abstraction --device "${grub_device}"`"
 
+if [ "x$disk_module" = xata ]; then
+    disk_module=pata
+fi
+
+if [ "x$disk_module" = xnative ]; then
+    disk_module="pata ahci ohci uhci usbms"
+fi
+
 # The order in this list is critical.  Be careful when modifying it.
 modules="$modules $disk_module"
 modules="$modules $fs_module $partmap_module $devabstraction_module"
@@ -519,7 +527,7 @@ if [ "x${devabstraction_module}" = "x" ] ; then
     # Strip partition number
     grub_partition="`echo "${grub_drive}" | sed -e 's/^[^,]*[,)]//; s/)$//'`"
     grub_drive="`echo "${grub_drive}" | sed -e s/,[a-z0-9,]*//g`"
-    if [ "$disk_module" = ata ] ; then
+    if [ "x$disk_module" != x ] && [ "x$disk_module" != xbiosdisk ]; then
         # generic method (used on coreboot and ata mod)
         uuid="`"$grub_probe" --device-map="${device_map}" --target=fs_uuid --device "${grub_device}"`"
         if [ "x${uuid}" = "x" ] ; then