set-verity-state: Add -h -v -R (automatic reboot) options

-h: print help
-v: noisy output
-R: automatic reboot if needed

Also remove the "argc != 0" check, as this shouldn't happen most of the
time anyway. The only possible way (I can think of) for this to happen
is to explicitly call execve() with an empty argv, which is against the
calling convention, thus a footgun, to begin with.

Bug: 241688845
Test: adb shell [disable-verity|enable-verity] [-R] [-v] [-h]
Test: adb shell set-verity-state [-R] [-v] [-h] [0|1]
Change-Id: I27fc18b0958650c5926322a3f4bd1df758fd96c8
diff --git a/set-verity-state/Android.bp b/set-verity-state/Android.bp
index 34bded8..f40118b 100644
--- a/set-verity-state/Android.bp
+++ b/set-verity-state/Android.bp
@@ -18,6 +18,9 @@
     static_libs: [
         "libavb_user",
     ],
+    header_libs: [
+        "libcutils_headers",
+    ],
 
     cflags: ["-Werror"],
     cppflags: [
diff --git a/set-verity-state/set-verity-state.cpp b/set-verity-state/set-verity-state.cpp
index 52757b6..de9a452 100644
--- a/set-verity-state/set-verity-state.cpp
+++ b/set-verity-state/set-verity-state.cpp
@@ -14,12 +14,14 @@
  * limitations under the License.
  */
 
+#include <getopt.h>
 #include <stdio.h>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/properties.h>
 #include <binder/ProcessState.h>
+#include <cutils/android_reboot.h>
 #include <fs_mgr_overlayfs.h>
 #include <libavb_user/libavb_user.h>
 
@@ -27,6 +29,18 @@
 
 namespace {
 
+void print_usage() {
+  printf(
+      "Usage:\n"
+      "\tdisable-verity\n"
+      "\tenable-verity\n"
+      "\tset-verity-state [0|1]\n"
+      "Options:\n"
+      "\t-h --help\tthis help\n"
+      "\t-R --reboot\tautomatic reboot if needed for new settings to take effect\n"
+      "\t-v --verbose\tbe noisy\n");
+}
+
 #ifdef ALLOW_DISABLE_VERITY
 const bool kAllowDisableVerity = true;
 #else
@@ -56,6 +70,15 @@
   return !android::base::GetProperty("ro.boot.vbmeta.digest", "").empty();
 }
 
+[[noreturn]] void reboot(const std::string& name) {
+  LOG(INFO) << "Rebooting device for new settings to take effect";
+  ::sync();
+  android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot," + name);
+  ::sleep(60);
+  LOG(ERROR) << "Failed to reboot";
+  ::exit(1);
+}
+
 bool overlayfs_setup(bool enable) {
   auto change = false;
   errno = 0;
@@ -110,35 +133,66 @@
   return {.success = true, .want_reboot = true};
 }
 
-void MyLogger(android::base::LogId id, android::base::LogSeverity severity, const char* tag,
-              const char* file, unsigned int line, const char* message) {
-  // Hide log starting with '[fs_mgr]' unless it's an error.
-  if (severity == android::base::ERROR || message[0] != '[') {
-    fprintf(stderr, "%s\n", message);
+class MyLogger {
+ public:
+  explicit MyLogger(bool verbose) : verbose_(verbose) {}
+
+  void operator()(android::base::LogId id, android::base::LogSeverity severity, const char* tag,
+                  const char* file, unsigned int line, const char* message) {
+    // Hide log starting with '[fs_mgr]' unless it's an error.
+    if (verbose_ || severity >= android::base::ERROR || message[0] != '[') {
+      fprintf(stderr, "%s\n", message);
+    }
+    logd_(id, severity, tag, file, line, message);
   }
-  static auto logd = android::base::LogdLogger();
-  logd(id, severity, tag, file, line, message);
-}
+
+ private:
+  android::base::LogdLogger logd_;
+  bool verbose_;
+};
 
 }  // namespace
 
 int main(int argc, char* argv[]) {
-  android::base::InitLogging(argv, MyLogger);
+  bool auto_reboot = false;
+  bool verbose = false;
 
-  if (argc == 0) {
-    LOG(FATAL) << "set-verity-state called with empty argv";
+  struct option longopts[] = {
+      {"help", no_argument, nullptr, 'h'},
+      {"reboot", no_argument, nullptr, 'R'},
+      {"verbose", no_argument, nullptr, 'v'},
+      {0, 0, nullptr, 0},
+  };
+  for (int opt; (opt = ::getopt_long(argc, argv, "hRv", longopts, nullptr)) != -1;) {
+    switch (opt) {
+      case 'h':
+        print_usage();
+        return 0;
+      case 'R':
+        auto_reboot = true;
+        break;
+      case 'v':
+        verbose = true;
+        break;
+      default:
+        print_usage();
+        return 1;
+    }
   }
 
+  android::base::InitLogging(argv, MyLogger(verbose));
+
   bool enable_verity = false;
-  std::string procname = android::base::Basename(argv[0]);
-  if (procname == "enable-verity") {
+  const std::string progname = getprogname();
+  if (progname == "enable-verity") {
     enable_verity = true;
-  } else if (procname == "disable-verity") {
+  } else if (progname == "disable-verity") {
     enable_verity = false;
-  } else if (argc == 2 && (argv[1] == "1"s || argv[1] == "0"s)) {
-    enable_verity = (argv[1] == "1"s);
+  } else if (optind < argc && (argv[optind] == "1"s || argv[optind] == "0"s)) {
+    // progname "set-verity-state"
+    enable_verity = (argv[optind] == "1"s);
   } else {
-    printf("usage: %s [1|0]\n", argv[0]);
+    print_usage();
     return 1;
   }
 
@@ -179,6 +233,9 @@
   }
 
   if (want_reboot) {
+    if (auto_reboot) {
+      reboot(progname);
+    }
     printf("Reboot the device for new settings to take effect\n");
   }