Add cgpt_wrapper binary to transparently support NAND

This CL adds a "cgpt_wrapper" binary that will forward to "cgpt" binary
as needed to transparently support NAND.

The idea is that if we are working with an MTD device, this binary will
exec out to "flashrom" to read in the GPT section, then exec out to the
actual "cgpt" binary to work on that file, and finally write it back
with "flashrom".

This CL introduces a make target "install_mtd" to install this wrapper.
The corresponding ebuild should use this make target when "mtd" USE flag
is on.

BUG=chromium:442518
BRANCH=none
CQ-DEPEND=CL:239573
TEST=unittest and some trial runs with mtdram
TEST=test on a real device with NOR flash and NAND

Change-Id: I54534afa9a970ec858f313f698c0eb210c827477
Reviewed-on: https://chromium-review.googlesource.com/239580
Tested-by: Nam Nguyen <namnguyen@chromium.org>
Reviewed-by: Daniel Ehrenberg <dehrenberg@chromium.org>
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Commit-Queue: Nam Nguyen <namnguyen@chromium.org>

4 files changed, 476 insertions, 4 deletions

diff --git a/cgpt/cgpt_find.c b/cgpt/cgpt_find.c
index c3193040..1a516d8c 100644
--- a/cgpt/cgpt_find.c
+++ b/cgpt/cgpt_find.c
@@ -8,6 +8,7 @@
 #include <unistd.h>
 
 #include "cgpt.h"
+#include "cgpt_nor.h"
 #include "cgptlib_internal.h"
 #include "vboot_host.h"
 
@@ -145,6 +146,7 @@ static int do_search(CgptFindParams *params, char *fileName) {
 }
 
 
+#define PROC_MTD "/proc/mtd"
 #define PROC_PARTITIONS "/proc/partitions"
 #define DEV_DIR "/dev"
 #define SYS_BLOCK_DIR "/sys/block"
@@ -186,23 +188,23 @@ static char *is_wholedev(const char *basename) {
   return 0;
 }
 
-
 // This scans all the physical devices it can find, looking for a match. It
 // returns true if any matches were found, false otherwise.
 static int scan_real_devs(CgptFindParams *params) {
   int found = 0;
-  char line[BUFSIZE];
   char partname[128];                   // max size for /proc/partition lines?
   FILE *fp;
   char *pathname;
 
-  fp = fopen(PROC_PARTITIONS, "r");
+  fp = fopen(PROC_PARTITIONS, "re");
   if (!fp) {
     perror("can't read " PROC_PARTITIONS);
     return found;
   }
 
-  while (fgets(line, sizeof(line), fp)) {
+  size_t line_length = 0;
+  char *line = NULL;
+  while (getline(&line, &line_length, fp) != -1) {
     int ma, mi;
     long long unsigned int sz;
 
@@ -217,6 +219,46 @@ static int scan_real_devs(CgptFindParams *params) {
   }
 
   fclose(fp);
+
+  fp = fopen(PROC_MTD, "re");
+  if (!fp) {
+    free(line);
+    return found;
+  }
+
+  while (getline(&line, &line_length, fp) != -1) {
+    uint64_t sz;
+    uint32_t erasesz;
+    char name[128];
+    // dev:  size  erasesize  name
+    if (sscanf(line, "%64[^:]: %" PRIx64 " %x \"%127[^\"]\"",
+               partname, &sz, &erasesz, name) != 4)
+      continue;
+    if (strcmp(partname, "mtd0") == 0) {
+      char temp_dir[] = "/tmp/cgpt_find.XXXXXX";
+      if (params->drive_size == 0) {
+        if (GetMtdSize("/dev/mtd0", &params->drive_size) != 0) {
+          perror("GetMtdSize");
+          goto cleanup;
+        }
+      }
+      if (ReadNorFlash(temp_dir) != 0) {
+        perror("ReadNorFlash");
+        goto cleanup;
+      }
+      char nor_file[64];
+      if (snprintf(nor_file, sizeof(nor_file), "%s/rw_gpt", temp_dir) > 0) {
+        if (do_search(params, nor_file)) {
+          found++;
+        }
+      }
+      RemoveDir(temp_dir);
+      break;
+    }
+  }
+cleanup:
+  fclose(fp);
+  free(line);
   return found;
 }
 
diff --git a/cgpt/cgpt_nor.c b/cgpt/cgpt_nor.c
new file mode 100644
index 00000000..17502a73
--- /dev/null
+++ b/cgpt/cgpt_nor.c
@@ -0,0 +1,249 @@
+/* Copyright 2015 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include <err.h>
+#include <errno.h>
+#include <fcntl.h>
+#include <ftw.h>
+#include <inttypes.h>
+#include <linux/major.h>
+#include <stdbool.h>
+#include <stdarg.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+#include "cgpt.h"
+#include "cgpt_nor.h"
+
+static const char FLASHROM_PATH[] = "/usr/sbin/flashrom";
+
+// Obtain the MTD size from its sysfs node.
+int GetMtdSize(const char *mtd_device, uint64_t *size) {
+  mtd_device = strrchr(mtd_device, '/');
+  if (mtd_device == NULL) {
+    errno = EINVAL;
+    return 1;
+  }
+  char *sysfs_name;
+  if (asprintf(&sysfs_name, "/sys/class/mtd%s/size", mtd_device) == -1) {
+    return 1;
+  }
+  FILE *fp = fopen(sysfs_name, "r");
+  free(sysfs_name);
+  if (fp == NULL) {
+    return 1;
+  }
+  int ret = (fscanf(fp, "%" PRIu64 "\n", size) != 1);
+  fclose(fp);
+  return ret;
+}
+
+int ForkExecV(const char *cwd, const char *const argv[]) {
+  pid_t pid = fork();
+  if (pid == -1) {
+    return -1;
+  }
+  int status = -1;
+  if (pid == 0) {
+    if (cwd && chdir(cwd) != 0) {
+      return -1;
+    }
+    execv(argv[0], (char *const *)argv);
+    // If this is reached, execv fails.
+    err(-1, "Cannot exec %s in %s.", argv[0], cwd);
+  } else {
+    if (waitpid(pid, &status, 0) != -1 && WIFEXITED(status))
+      return WEXITSTATUS(status);
+  }
+  return status;
+}
+
+int ForkExecL(const char *cwd, const char *cmd, ...) {
+  int argc;
+  va_list ap;
+  va_start(ap, cmd);
+  for (argc = 1; va_arg(ap, char *) != NULL; ++argc);
+  va_end(ap);
+
+  va_start(ap, cmd);
+  const char **argv = calloc(argc + 1, sizeof(char *));
+  if (argv == NULL) {
+    errno = ENOMEM;
+    return -1;
+  }
+  argv[0] = cmd;
+  int i;
+  for (i = 1; i < argc; ++i) {
+    argv[i] = va_arg(ap, char *);
+  }
+  va_end(ap);
+
+  int ret = ForkExecV(cwd, argv);
+  free(argv);
+  return ret;
+}
+
+static int read_write(int source_fd,
+                      uint64_t size,
+                      const char *src_name,
+                      int idx) {
+  int ret = 1;
+  const int bufsize = 4096;
+  char *buf = malloc(bufsize);
+  if (buf == NULL) {
+    goto clean_exit;
+  }
+
+  ret++;
+  char *dest;
+  if (asprintf(&dest, "%s_%d", src_name, idx) == -1) {
+    goto free_buf;
+  }
+
+  ret++;
+  int dest_fd = open(dest, O_WRONLY | O_CLOEXEC | O_CREAT, 0600);
+  if (dest_fd < 0) {
+    goto free_dest;
+  }
+
+  ret++;
+  uint64_t copied = 0;
+  ssize_t nr_read;
+  ssize_t nr_write;
+  while (copied < size) {
+    size_t to_read = size - copied;
+    if (to_read > bufsize) {
+      to_read = bufsize;
+    }
+    nr_read = read(source_fd, buf, to_read);
+    if (nr_read < 0) {
+      goto close_dest_fd;
+    }
+    nr_write = 0;
+    while (nr_write < nr_read) {
+      ssize_t s = write(dest_fd, buf + nr_write, nr_read - nr_write);
+      if (s < 0) {
+        goto close_dest_fd;
+      }
+      nr_write += s;
+    }
+    copied += nr_read;
+  }
+
+  ret = 0;
+
+close_dest_fd:
+  close(dest_fd);
+free_dest:
+  free(dest);
+free_buf:
+  free(buf);
+clean_exit:
+  return ret;
+}
+
+static int split_gpt(const char *dir_name, const char *file_name) {
+  int ret = 1;
+  char *source;
+  if (asprintf(&source, "%s/%s", dir_name, file_name) == -1) {
+    goto clean_exit;
+  }
+
+  ret++;
+  int fd = open(source, O_RDONLY | O_CLOEXEC);
+  if (fd < 0) {
+    goto free_source;
+  }
+
+  ret++;
+  struct stat stat;
+  if (fstat(fd, &stat) != 0 || (stat.st_size & 1) != 0) {
+    goto close_fd;
+  }
+  uint64_t half_size = stat.st_size / 2;
+
+  ret++;
+  if (read_write(fd, half_size, source, 1) != 0 ||
+      read_write(fd, half_size, source, 2) != 0) {
+    goto close_fd;
+  }
+
+  ret = 0;
+close_fd:
+  close(fd);
+free_source:
+  free(source);
+clean_exit:
+  return ret;
+}
+
+static int remove_file_or_dir(const char *fpath, const struct stat *sb,
+                              int typeflag, struct FTW *ftwbuf) {
+  return remove(fpath);
+}
+
+int RemoveDir(const char *dir) {
+  return nftw(dir, remove_file_or_dir, 20, FTW_DEPTH | FTW_PHYS);
+}
+
+// Read RW_GPT from NOR flash to "rw_gpt" in a temp dir |temp_dir_template|.
+// |temp_dir_template| is passed to mkdtemp() so it must satisfy all
+// requirements by mkdtemp.
+int ReadNorFlash(char *temp_dir_template) {
+  int ret = 0;
+
+  // Create a temp dir to work in.
+  ret++;
+  if (mkdtemp(temp_dir_template) == NULL) {
+    Error("Cannot create a temporary directory.\n");
+    return ret;
+  }
+
+  // Read RW_GPT section from NOR flash to "rw_gpt".
+  ret++;
+  if (ForkExecL(temp_dir_template, FLASHROM_PATH, "-i", "RW_GPT:rw_gpt", "-r",
+                NULL) != 0) {
+    Error("Cannot exec flashrom to read from RW_GPT section.\n");
+    RemoveDir(temp_dir_template);
+    return ret;
+  }
+
+  ret = 0;
+  return ret;
+}
+
+// Write "rw_gpt" back to NOR flash. We write the file in two parts for safety.
+int WriteNorFlash(const char *dir) {
+  int ret = 0;
+  ret++;
+  if (split_gpt(dir, "rw_gpt") != 0) {
+    Error("Cannot split rw_gpt in two.\n");
+    return ret;
+  }
+  ret++;
+  int nr_fails = 0;
+  if (ForkExecL(dir, FLASHROM_PATH, "-i", "RW_GPT_PRIMARY:rw_gpt_1",
+                "-w", NULL) != 0) {
+    Warning("Cannot write the 1st half of rw_gpt back with flashrom.\n");
+    nr_fails++;
+  }
+  if (ForkExecL(dir, FLASHROM_PATH, "-i", "RW_GPT_SECONDARY:rw_gpt_2",
+                "-w", NULL) != 0) {
+    Warning("Cannot write the 2nd half of rw_gpt back with flashrom.\n");
+    nr_fails++;
+  }
+  switch (nr_fails) {
+    case 0: ret = 0; break;
+    case 1: Warning("It might still be okay.\n"); break;
+    case 2: Error("Cannot write both parts back with flashrom.\n"); break;
+  }
+  return ret;
+}
diff --git a/cgpt/cgpt_nor.h b/cgpt/cgpt_nor.h
new file mode 100644
index 00000000..b01f88d8
--- /dev/null
+++ b/cgpt/cgpt_nor.h
@@ -0,0 +1,34 @@
+/* Copyright 2015 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ *
+ * This module provides some utility functions to use "flashrom" to read from
+ * and write to NOR flash.
+ */
+
+#ifndef VBOOT_REFERCENCE_CGPT_CGPT_NOR_H_
+#define VBOOT_REFERCENCE_CGPT_CGPT_NOR_H_
+
+// Obtain the MTD size from its sysfs node. |mtd_device| should point to
+// a dev node such as /dev/mtd0. This function returns 0 on success.
+int GetMtdSize(const char *mtd_device, uint64_t *size);
+
+// Exec |argv| in |cwd|. Return -1 on error, or exit code on success. |argv|
+// must be terminated with a NULL element as is required by execv().
+int ForkExecV(const char *cwd, const char *const argv[]);
+
+// Similar to ForkExecV but with a vararg instead of an array of pointers.
+int ForkExecL(const char *cwd, const char *cmd, ...);
+
+// Exec "rm" to remove |dir|.
+int RemoveDir(const char *dir);
+
+// Read RW_GPT from NOR flash to "rw_gpt" in a temp dir |temp_dir_template|.
+// |temp_dir_template| is passed to mkdtemp() so it must satisfy all
+// requirements by mkdtemp().
+int ReadNorFlash(char *temp_dir_template);
+
+// Write "rw_gpt" back to NOR flash. We write the file in two parts for safety.
+int WriteNorFlash(const char *dir);
+
+#endif  // VBOOT_REFERCENCE_CGPT_CGPT_NOR_H_
diff --git a/cgpt/cgpt_wrapper.c b/cgpt/cgpt_wrapper.c
new file mode 100644
index 00000000..6750d8d8
--- /dev/null
+++ b/cgpt/cgpt_wrapper.c
@@ -0,0 +1,147 @@
+/* Copyright 2015 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ *
+ * This utility wraps around "cgpt" execution to work with NAND. If the target
+ * device is an MTD device, this utility will read the GPT structures from
+ * FMAP, invokes "cgpt" on that, and writes the result back to NOR flash. */
+
+#include <errno.h>
+#include <fcntl.h>
+#include <inttypes.h>
+#include <limits.h>
+#include <linux/major.h>
+#include <stdbool.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include "cgpt.h"
+#include "cgpt_nor.h"
+
+// Check if cmdline |argv| has "-D". "-D" signifies that GPT structs are stored
+// off device, and hence we should not wrap around cgpt.
+static bool has_dash_D(int argc, const char *const argv[]) {
+  int i;
+  // We go from 2, because the second arg is a cgpt command such as "create".
+  for (i = 2; i < argc; ++i) {
+    if (strcmp("-D", argv[i]) == 0) {
+      return true;
+    }
+  }
+  return false;
+}
+
+// Check if |device_path| is an MTD device based on its major number being 90.
+static bool is_mtd(const char *device_path) {
+  struct stat stat;
+  if (lstat(device_path, &stat) != 0) {
+    return false;
+  }
+
+  if (major(stat.st_rdev) != MTD_CHAR_MAJOR) {
+    return false;
+  }
+
+  return true;
+}
+
+// Return the element in |argv| that is an MTD device.
+static const char *find_mtd_device(int argc, const char *const argv[]) {
+  int i;
+  for (i = 2; i < argc; ++i) {
+    if (is_mtd(argv[i])) {
+      return argv[i];
+    }
+  }
+  return NULL;
+}
+
+static int wrap_cgpt(int argc,
+                     const char *const argv[],
+                     const char *mtd_device) {
+  int ret = 0;
+
+  // Create a temp dir to work in.
+  ret++;
+  char temp_dir[] = "/tmp/cgpt_wrapper.XXXXXX";
+  if (ReadNorFlash(temp_dir) != 0) {
+    return ret;
+  }
+  char rw_gpt_path[PATH_MAX];
+  if (snprintf(rw_gpt_path, sizeof(rw_gpt_path), "%s/rw_gpt", temp_dir) < 0) {
+    goto cleanup;
+  }
+
+  // Obtain the MTD size.
+  ret++;
+  uint64_t drive_size = 0;
+  if (GetMtdSize(mtd_device, &drive_size) != 0) {
+    Error("Cannot get the size of %s.\n", mtd_device);
+    goto cleanup;
+  }
+
+  // Launch cgpt on "rw_gpt" with -D size.
+  ret++;
+  const char** my_argv = calloc(argc + 2 + 1, sizeof(char *));
+  if (my_argv == NULL) {
+    errno = ENOMEM;
+    goto cleanup;
+  }
+  memcpy(my_argv, argv, sizeof(char *) * argc);
+  char *real_cgpt;
+  if (asprintf(&real_cgpt, "%s.bin", argv[0]) == -1) {
+    free(my_argv);
+    goto cleanup;
+  }
+  my_argv[0] = real_cgpt;
+
+  int i;
+  for (i = 2; i < argc; ++i) {
+    if (strcmp(my_argv[i], mtd_device) == 0) {
+      my_argv[i] = rw_gpt_path;
+    }
+  }
+  my_argv[argc] = "-D";
+  char size[32];
+  snprintf(size, sizeof(size), "%" PRIu64, drive_size);
+  my_argv[argc + 1] = size;
+  i = ForkExecV(NULL, my_argv);
+  free(real_cgpt);
+  free(my_argv);
+  if (i != 0) {
+    Error("Cannot exec cgpt to modify rw_gpt.\n");
+    goto cleanup;
+  }
+
+  // Write back "rw_gpt" to NOR flash in two chunks.
+  ret++;
+  if (WriteNorFlash(temp_dir) == 0) {
+    ret = 0;
+  }
+
+cleanup:
+  RemoveDir(temp_dir);
+  return ret;
+}
+
+int main(int argc, const char *argv[]) {
+  if (argc > 2 && !has_dash_D(argc, argv)) {
+    const char *mtd_device = find_mtd_device(argc, argv);
+    if (mtd_device) {
+      return wrap_cgpt(argc, argv, mtd_device);
+    }
+  }
+
+  // Forward to cgpt as-is. Real cgpt has been renamed cgpt.bin.
+  char *real_cgpt;
+  if (asprintf(&real_cgpt, "%s.bin", argv[0]) == -1) {
+    return -1;
+  }
+  argv[0] = real_cgpt;
+  return execv(argv[0], (char * const *)argv);
+}