Add --selftest command to citadel_updater

Bug: b/111597461
Test: Ran on red board and C1DVT. Works.

Change-Id: Iae66f83e8601e00e4d16d5a03e2bbe55e51861f8
Signed-off-by: Bill Richardson <wfrichar@google.com>
diff --git a/citadel/updater/updater.cpp b/citadel/updater/updater.cpp
index 618a9d1..9316499 100644
--- a/citadel/updater/updater.cpp
+++ b/citadel/updater/updater.cpp
@@ -76,6 +76,8 @@
   int change_pw;
   uint32_t erase_code;
   int ap_uart;
+  int selftest;
+  char **selftest_args;
   /* generic connection options */
   const char *device;
 } options;
@@ -94,6 +96,7 @@
   OPT_CHANGE_PW,
   OPT_ERASE,
   OPT_AP_UART,
+  OPT_SELFTEST,
 };
 
 const char *short_opts = ":hvlV:fF:";
@@ -120,6 +123,7 @@
   {"erase",         1, NULL, OPT_ERASE},
   {"ap_uart",       0, NULL, OPT_AP_UART},
   {"ap-uart",       0, NULL, OPT_AP_UART},
+  {"selftest",      0, NULL, OPT_SELFTEST},
 #ifndef ANDROID
   {"device",        1, NULL, OPT_DEVICE},
 #endif
@@ -175,6 +179,10 @@
     "\n\n"
     "  --erase=CODE         Erase all user secrets and reboot.\n"
     "                       This skips all other actions.\n"
+    "\n\n"
+    "  --selftest [ARGS]    Run one or more selftests. With no ARGS, it asks\n"
+    "                       Citdel for the available tests. This should be\n"
+    "                       run by itself for best results.\n"
 #ifndef ANDROID
     "\n"
     "Options:\n"
@@ -732,7 +740,6 @@
   return rv;
 }
 
-
 static uint32_t do_erase(AppClient &app)
 {
   std::vector<uint8_t> data(sizeof(uint32_t));
@@ -746,6 +753,72 @@
   return rv;
 }
 
+#define MAX_SELFTEST_REPLY_LEN 4096
+static uint32_t do_selftest(AppClient &app, int argc, char *argv[])
+{
+  int i = options.selftest;
+  char *e = 0;
+  uint32_t rv = APP_ERROR_BOGUS_ARGS;;
+
+  /* No args at all should just request the selftest version */
+  if (i >= argc) {
+    std::vector<uint8_t> data;
+    data.reserve(MAX_SELFTEST_REPLY_LEN);
+
+    rv = app.Call(NUGGET_PARAM_SELFTEST, data, &data);
+
+    if (is_app_success(rv)) {
+      /* Make SURE it's null-terminated */
+      size_t len = data.size();
+      if (len) {
+        data[len - 1] = '\0';
+        printf("%s\n", data.data());
+      }
+    }
+
+    return rv;
+  }
+
+  /* The only test right now is the trng test */
+  if (*argv[i] == 't' || *argv[i] == 'T') {
+    i++;
+
+    /* It takes uint32_t args */
+    std::vector<uint32_t> args;
+    args.push_back(SELFTEST_TRNGSTATS); /* First arg is command */
+
+    for (; i < argc; i++) {
+      uint32_t tmp = (uint32_t)strtoul(argv[i], &e, 0);
+      if (e && *e) {
+        Error("Invalid arg: \"%s\"\n", argv[i]);
+        return -1;
+      }
+      args.push_back(tmp);
+    }
+
+    size_t arg_bytes = args.size() * sizeof(uint32_t);
+    std::vector<uint8_t> data(arg_bytes); // sending this much
+    memcpy(data.data(), args.data(), arg_bytes);
+    data.reserve(MAX_SELFTEST_REPLY_LEN);
+
+    rv = app.Call(NUGGET_PARAM_SELFTEST, data, &data);
+
+    if (is_app_success(rv)) {
+      /* Make SURE it's null-terminated */
+      size_t len = data.size();
+      if (len) {
+        data[len - 1] = '\0';
+        printf("%s\n", data.data());
+      }
+    }
+
+  } else {
+    Error("Unrecognized test name");
+  }
+
+  return rv;
+}
+
 // This is currently device-specific, but could be abstracted further
 #ifdef ANDROID
 static uint32_t do_force_reset(CitadeldProxyClient &client)
@@ -764,7 +837,8 @@
 #endif
 
 int execute_commands(const std::vector<uint8_t> &image,
-                     const char *old_passwd, const char *passwd)
+                     const char *old_passwd, const char *passwd,
+                     int argc, char *argv[])
 {
 #ifdef ANDROID
   CitadeldProxyClient client;
@@ -855,6 +929,11 @@
     return 7;
   }
 
+  if (options.selftest &&
+      do_selftest(app, argc, argv) != APP_SUCCESS) {
+    return 1;
+  }
+
   if (options.force_reset &&
       do_force_reset(client) != APP_SUCCESS) {
     return 1;
@@ -963,7 +1042,6 @@
       options.erase_code = (uint32_t)strtoul(optarg, &e, 0);
       if (!*optarg || (e && *e)) {
         Error("Invalid argument: \"%s\"\n", optarg);
-        errorcnt++;
       }
       got_action = 1;
       break;
@@ -971,6 +1049,11 @@
       options.ap_uart = 1;
       got_action = 1;
       break;
+    case OPT_SELFTEST:
+      options.selftest = optind;
+      options.selftest_args = argv;
+      got_action = 1;
+      break;
 
       /* generic options below */
     case OPT_DEVICE:
@@ -1044,7 +1127,7 @@
   }
 
   /* Okay, let's do it! */
-  (void) execute_commands(image, old_passwd, passwd);
+  (void) execute_commands(image, old_passwd, passwd, argc, argv);
   /* This is the last action, so fall through either way */
 
 out: