From 881c84d01113be1406f25a0cf00b494ae100ee1c Mon Sep 17 00:00:00 2001
From: reagentoo <reagentoo@gmail.com>
Date: Fri, 30 Oct 2020 12:21:24 +0300
Subject: [PATCH] search: Support searching for partition UUID with --part-uuid

These changes are made in the same way as in the probe command.
---
 docs/grub.texi                       | 11 ++---
 grub-core/Makefile.core.def          |  5 +++
 grub-core/commands/search.c          | 61 +++++++++++++++++++++++++++-
 grub-core/commands/search_partuuid.c |  5 +++
 grub-core/commands/search_wrap.c     |  8 +++-
 include/grub/search.h                |  2 +
 6 files changed, 84 insertions(+), 8 deletions(-)
 create mode 100644 grub-core/commands/search_partuuid.c

diff --git a/docs/grub.texi b/docs/grub.texi
index 37f7ce7da..fbeabf835 100644
--- a/docs/grub.texi
+++ b/docs/grub.texi
@@ -4938,8 +4938,8 @@ unbootable. @xref{Using digital signatures}, for more information.
  [@option{--file}|@option{--label}|@option{--fs-uuid}] @
  [@option{--set} [var]] [@option{--no-floppy}] name
 Search devices by file (@option{-f}, @option{--file}), filesystem label
-(@option{-l}, @option{--label}), or filesystem UUID (@option{-u},
-@option{--fs-uuid}).
+(@option{-l}, @option{--label}), filesystem UUID (@option{-u},
+@option{--fs-uuid}), or partition UUID (@option{-p}, @option{--part-uuid})
 
 If the @option{--set} option is used, the first device found is set as the
 value of environment variable @var{var}.  The default variable is
@@ -4948,9 +4948,10 @@ value of environment variable @var{var}.  The default variable is
 The @option{--no-floppy} option prevents searching floppy devices, which can
 be slow.
 
-The @samp{search.file}, @samp{search.fs_label}, and @samp{search.fs_uuid}
-commands are aliases for @samp{search --file}, @samp{search --label}, and
-@samp{search --fs-uuid} respectively.
+The @samp{search.file}, @samp{search.fs_label}, @samp{search.fs_uuid},
+and @samp{search.part_uuid}
+commands are aliases for @samp{search --file}, @samp{search --label},
+@samp{search --fs-uuid}, and @samp{search --part-uuid} respectively.
 @end deffn
 
 
diff --git a/grub-core/Makefile.core.def b/grub-core/Makefile.core.def
index b5f47fc41..a40a1afbc 100644
--- a/grub-core/Makefile.core.def
+++ b/grub-core/Makefile.core.def
@@ -1080,6 +1080,11 @@ module = {
   common = commands/search_label.c;
 };
 
+module = {
+  name = search_part_uuid;
+  common = commands/search_partuuid.c;
+};
+
 module = {
   name = setpci;
   common = commands/setpci.c;
diff --git a/grub-core/commands/search.c b/grub-core/commands/search.c
index ed090b3af..3dafa273c 100644
--- a/grub-core/commands/search.c
+++ b/grub-core/commands/search.c
@@ -30,6 +30,8 @@
 #include <grub/i18n.h>
 #include <grub/disk.h>
 #include <grub/partition.h>
+#include <grub/gpt_partition.h>
+#include <grub/i386/pc/boot.h>
 
 GRUB_MOD_LICENSE ("GPLv3+");
 
@@ -66,13 +68,13 @@ iterate_device (const char *name, void *data)
       name[0] == 'f' && name[1] == 'd' && name[2] >= '0' && name[2] <= '9')
     return 1;
 
-#ifdef DO_SEARCH_FS_UUID
+#if defined (DO_SEARCH_FS_UUID) || defined (DO_SEARCH_PART_UUID)
 #define compare_fn grub_strcasecmp
 #else
 #define compare_fn grub_strcmp
 #endif
 
-#ifdef DO_SEARCH_FILE
+#if defined (DO_SEARCH_FILE)
     {
       char *buf;
       grub_file_t file;
@@ -90,6 +92,57 @@ iterate_device (const char *name, void *data)
 	}
       grub_free (buf);
     }
+#elif defined (DO_SEARCH_PART_UUID)
+    {
+      grub_device_t dev;
+
+      /* AAAABBBB-CCCC-DDDD-EEEE-FFFFFFFFFFFF + null terminator */
+      char val[37] = "none";
+
+      dev = grub_device_open (name);
+      if (dev && dev->disk && dev->disk->partition)
+	{
+	  struct grub_partition *p = dev->disk->partition;
+	  grub_disk_t disk = grub_disk_open(dev->disk->name);
+
+	  if (!disk)
+	    return 1;
+	  if (grub_strcmp(dev->disk->partition->partmap->name, "gpt") == 0)
+	    {
+	      struct grub_gpt_partentry entry;
+	      grub_gpt_part_guid_t *guid;
+
+	      if (grub_disk_read(disk, p->offset, p->index, sizeof(entry), &entry))
+		return 1;
+	      guid = &entry.guid;
+	      grub_snprintf (val, sizeof(val),
+			     "%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x",
+			     grub_le_to_cpu32 (guid->data1),
+			     grub_le_to_cpu16 (guid->data2),
+			     grub_le_to_cpu16 (guid->data3),
+			     guid->data4[0], guid->data4[1], guid->data4[2],
+			     guid->data4[3], guid->data4[4], guid->data4[5],
+			     guid->data4[6], guid->data4[7]);
+
+	      if (compare_fn (val, ctx->key) == 0)
+		found = 1;
+	    }
+	  else if (grub_strcmp(dev->disk->partition->partmap->name, "msdos") == 0)
+	    {
+	      grub_uint32_t nt_disk_sig;
+
+	      if (grub_disk_read(disk, 0, GRUB_BOOT_MACHINE_WINDOWS_NT_MAGIC,
+				 sizeof(nt_disk_sig), &nt_disk_sig) == 0)
+		grub_snprintf (val, sizeof(val), "%08x-%02x",
+			       grub_le_to_cpu32(nt_disk_sig), 1 + p->number);
+
+	      if (compare_fn (val, ctx->key) == 0)
+		found = 1;
+	    }
+	  grub_disk_close(disk);
+	  grub_device_close(dev);
+	}
+    }
 #else
     {
       /* SEARCH_FS_UUID or SEARCH_LABEL */
@@ -315,6 +368,8 @@ static grub_command_t cmd;
 GRUB_MOD_INIT(search_fs_file)
 #elif defined (DO_SEARCH_FS_UUID)
 GRUB_MOD_INIT(search_fs_uuid)
+#elif defined (DO_SEARCH_PART_UUID)
+GRUB_MOD_INIT(search_part_uuid)
 #else
 GRUB_MOD_INIT(search_label)
 #endif
@@ -329,6 +384,8 @@ GRUB_MOD_INIT(search_label)
 GRUB_MOD_FINI(search_fs_file)
 #elif defined (DO_SEARCH_FS_UUID)
 GRUB_MOD_FINI(search_fs_uuid)
+#elif defined (DO_SEARCH_PART_UUID)
+GRUB_MOD_FINI(search_part_uuid)
 #else
 GRUB_MOD_FINI(search_label)
 #endif
diff --git a/grub-core/commands/search_partuuid.c b/grub-core/commands/search_partuuid.c
new file mode 100644
index 000000000..2bbfac3a7
--- /dev/null
+++ b/grub-core/commands/search_partuuid.c
@@ -0,0 +1,5 @@
+#define DO_SEARCH_PART_UUID 1
+#define FUNC_NAME grub_search_part_uuid
+#define COMMAND_NAME "search.part_uuid"
+#define HELP_MESSAGE N_("Search devices by PARTUUID. If VARIABLE is specified, the first device found is set to a variable.")
+#include "search.c"
diff --git a/grub-core/commands/search_wrap.c b/grub-core/commands/search_wrap.c
index 47fc8eb99..ece6dd0f6 100644
--- a/grub-core/commands/search_wrap.c
+++ b/grub-core/commands/search_wrap.c
@@ -36,6 +36,8 @@ static const struct grub_arg_option options[] =
      0, 0},
     {"fs-uuid",		'u', 0, N_("Search devices by a filesystem UUID."),
      0, 0},
+    {"part-uuid",	'p', 0, N_("Search devices by a partition UUID."),
+     0, 0},
     {"set",		's', GRUB_ARG_OPTION_OPTIONAL,
      N_("Set a variable to the first device found."), N_("VARNAME"),
      ARG_TYPE_STRING},
@@ -71,6 +73,7 @@ enum options
     SEARCH_FILE,
     SEARCH_LABEL,
     SEARCH_FS_UUID,
+    SEARCH_PART_UUID,
     SEARCH_SET,
     SEARCH_NO_FLOPPY,
     SEARCH_HINT,
@@ -183,6 +186,9 @@ grub_cmd_search (grub_extcmd_context_t ctxt, int argc, char **args)
   if (state[SEARCH_LABEL].set)
     grub_search_label (id, var, state[SEARCH_NO_FLOPPY].set, 
 		       hints, nhints);
+  else if (state[SEARCH_PART_UUID].set)
+    grub_search_part_uuid (id, var, state[SEARCH_NO_FLOPPY].set,
+			 hints, nhints);
   else if (state[SEARCH_FS_UUID].set)
     grub_search_fs_uuid (id, var, state[SEARCH_NO_FLOPPY].set,
 			 hints, nhints);
@@ -204,7 +210,7 @@ GRUB_MOD_INIT(search)
   cmd =
     grub_register_extcmd ("search", grub_cmd_search,
 			  GRUB_COMMAND_FLAG_EXTRACTOR | GRUB_COMMAND_ACCEPT_DASH,
-			  N_("[-f|-l|-u|-s|-n] [--hint HINT [--hint HINT] ...]"
+			  N_("[-f|-l|-u|-p|-s|-n] [--hint HINT [--hint HINT] ...]"
 			     " NAME"),
 			  N_("Search devices by file, filesystem label"
 			     " or filesystem UUID."
diff --git a/include/grub/search.h b/include/grub/search.h
index d80347df3..c0a9d5485 100644
--- a/include/grub/search.h
+++ b/include/grub/search.h
@@ -25,5 +25,7 @@ void grub_search_fs_uuid (const char *key, const char *var, int no_floppy,
 			  char **hints, unsigned nhints);
 void grub_search_label (const char *key, const char *var, int no_floppy,
 			char **hints, unsigned nhints);
+void grub_search_part_uuid (const char *key, const char *var, int no_floppy,
+			  char **hints, unsigned nhints);
 
 #endif
-- 
2.29.1

