cgpt: Properly show ChromeOS-scheme MTD partitions

When working on NAND, we do not actually work with one device name. We
work on a temporary file instead. Moreover, depending on the type of the
partition, we need to show different devices.

BUG=None
BRANCH=None
TEST=All commands must be run on storm_nand
TEST=/usr/bin/cgpt.bin find -t kernel should print out /dev/mtd2
TEST=/usr/bin/cgpt.bin find -t rootfs should print out /dev/ubiblock5_0
TEST=/usr/bin/cgpt.bin find -t data should print out /dev/ubi1_0

Change-Id: Ia36777ffa6a9cfc7c8ec4b128e49ece140428238
Reviewed-on: https://chromium-review.googlesource.com/242291
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Tested-by: Nam Nguyen <namnguyen@chromium.org>
Commit-Queue: Nam Nguyen <namnguyen@google.com>

2 files changed, 32 insertions, 4 deletions

diff --git a/cgpt/cgpt_find.c b/cgpt/cgpt_find.c
index 1a516d8c..c099a9c4 100644
--- a/cgpt/cgpt_find.c
+++ b/cgpt/cgpt_find.c
@@ -72,18 +72,37 @@ static int match_content(CgptFindParams *params, struct drive *drive,
 
 // This needs to handle /dev/mmcblk0 -> /dev/mmcblk0p3, /dev/sda -> /dev/sda3
 static void showmatch(CgptFindParams *params, char *filename,
-                           int partnum, GptEntry *entry) {
+                      int partnum, GptEntry *entry) {
   char * format = "%s%d\n";
   if (strncmp("/dev/mmcblk", filename, 11) == 0)
     format = "%sp%d\n";
-  if (params->numeric)
+
+  if (params->numeric) {
     printf("%d\n", partnum);
-  else
-    printf(format, filename, partnum);
+  } else {
+    if (params->show_fn) {
+      params->show_fn(params, filename, partnum, entry);
+    } else {
+      printf(format, filename, partnum);
+    }
+  }
   if (params->verbose > 0)
     EntryDetails(entry, partnum - 1, params->numeric);
 }
 
+// This handles the MTD devices. ChromeOS uses /dev/mtdX for kernel partitions,
+// /dev/ubiblockX_0 for root partitions, and /dev/ubiX for stateful partition.
+static void chromeos_mtd_show(CgptFindParams *params, char *filename,
+                              int partnum, GptEntry *entry) {
+  if (GuidEqual(&guid_chromeos_kernel, &entry->type)) {
+    printf("/dev/mtd%d\n", partnum);
+  } else if (GuidEqual(&guid_chromeos_rootfs, &entry->type)) {
+    printf("/dev/ubiblock%d_0\n", partnum);
+  } else {
+    printf("/dev/ubi%d_0\n", partnum);
+  }
+}
+
 // This returns true if a GPT partition matches the search criteria. If a match
 // isn't found (or if the file doesn't contain a GPT), it returns false. The
 // filename and partition number that matched is left in a global, since we
@@ -248,9 +267,11 @@ static int scan_real_devs(CgptFindParams *params) {
       }
       char nor_file[64];
       if (snprintf(nor_file, sizeof(nor_file), "%s/rw_gpt", temp_dir) > 0) {
+        params->show_fn = chromeos_mtd_show;
         if (do_search(params, nor_file)) {
           found++;
         }
+        params->show_fn = NULL;
       }
       RemoveDir(temp_dir);
       break;
diff --git a/host/include/cgpt_params.h b/host/include/cgpt_params.h
index e7ac7327..97d89e6a 100644
--- a/host/include/cgpt_params.h
+++ b/host/include/cgpt_params.h
@@ -78,6 +78,9 @@ typedef struct CgptPrioritizeParams {
   int orig_priority;
 } CgptPrioritizeParams;
 
+struct CgptFindParams;
+typedef void (*CgptFindShowFn)(struct CgptFindParams *params, char *filename,
+                               int partnum, GptEntry *entry);
 typedef struct CgptFindParams {
   char *drive_name;
   uint64_t drive_size;
@@ -96,6 +99,10 @@ typedef struct CgptFindParams {
   char *label;
   int hits;
   int match_partnum;           /* 1-based; 0 means no match */
+  /* when working with MTD, we actually work on a temp file, but we still need
+   * to print the device name. so this parameter is here to properly show the
+   * correct device name in that special case. */
+  CgptFindShowFn show_fn;
 } CgptFindParams;
 
 typedef struct CgptLegacyParams {