Merge "Create and Start secure_env snapshot control server" into main am: 02d6815f17

Original change: https://android-review.googlesource.com/c/device/google/cuttlefish/+/2746805

Change-Id: I39584712e7933c1d35749393b19fab747460676c
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

6 files changed, 82 insertions, 27 deletions

diff --git a/host/commands/secure_env/Android.bp b/host/commands/secure_env/Android.bp
index 7ac681265..616a3d214 100644
--- a/host/commands/secure_env/Android.bp
+++ b/host/commands/secure_env/Android.bp
@@ -112,6 +112,7 @@ cc_library {
         "device_tpm.cpp",
         "oemlock/oemlock.cpp",
         "oemlock/oemlock_responder.cpp",
+        "snapshot_control.cpp",
         "storage/insecure_json_storage.cpp",
     ],
     target: {
@@ -156,7 +157,6 @@ cc_binary_host {
         not_windows: {
             srcs: [
                 "secure_env_not_windows_main.cpp",
-                "snapshot_control.cpp",
             ],
             static_libs: [
                 "libsecure_env_not_windows",
diff --git a/host/commands/secure_env/confui_sign_server.cpp b/host/commands/secure_env/confui_sign_server.cpp
index f82ee8503..4b48e0404 100644
--- a/host/commands/secure_env/confui_sign_server.cpp
+++ b/host/commands/secure_env/confui_sign_server.cpp
@@ -15,6 +15,9 @@
 
 #include "confui_sign_server.h"
 
+#include <mutex>
+#include <shared_mutex>
+
 #include <android-base/logging.h>
 
 #include "host/commands/secure_env/primary_key_builder.h"
@@ -31,9 +34,12 @@ constexpr const char kConfirmationTokenMessageTag[] = "confirmation token";
 }  // namespace
 
 namespace cuttlefish {
-ConfUiSignServer::ConfUiSignServer(TpmResourceManager& tpm_resource_manager,
-                                   SharedFD server_fd)
-    : tpm_resource_manager_(tpm_resource_manager), server_fd_(server_fd) {
+ConfUiSignServer::ConfUiSignServer(
+    TpmResourceManager& tpm_resource_manager,
+    std::shared_ptr<SnapshotController> snapshot_ctrl, SharedFD server_fd)
+    : tpm_resource_manager_(tpm_resource_manager),
+      snapshot_ctrl_{std::move(snapshot_ctrl)},
+      server_fd_(server_fd) {
   auto config = cuttlefish::CuttlefishConfig::Get();
   CHECK(config) << "Config must not be null";
   auto instance = config->ForDefaultInstance();
@@ -47,6 +53,10 @@ ConfUiSignServer::ConfUiSignServer(TpmResourceManager& tpm_resource_manager,
                                                SOCK_STREAM, 0600);
     }
     auto accepted_socket_fd = SharedFD::Accept(*server_fd_);
+    std::shared_lock<std::shared_mutex> reader_lock;
+    if (snapshot_ctrl_->Enabled()) {
+      reader_lock = std::move(snapshot_ctrl_->WaitInitializedOrResumed());
+    }
     if (!accepted_socket_fd->IsOpen()) {
       LOG(ERROR) << "Confirmation UI host signing client socket is broken.";
       continue;
diff --git a/host/commands/secure_env/confui_sign_server.h b/host/commands/secure_env/confui_sign_server.h
index 531a3ed85..efbb302d8 100644
--- a/host/commands/secure_env/confui_sign_server.h
+++ b/host/commands/secure_env/confui_sign_server.h
@@ -15,21 +15,25 @@
 
 #pragma once
 
+#include <memory>
 #include <string>
 
 #include "common/libs/fs/shared_fd.h"
 #include "common/libs/security/confui_sign.h"
+#include "host/commands/secure_env/snapshot_control.h"
 #include "host/commands/secure_env/tpm_resource_manager.h"
 
 namespace cuttlefish {
 class ConfUiSignServer {
  public:
   ConfUiSignServer(TpmResourceManager& tpm_resource_manager,
+                   std::shared_ptr<SnapshotController> snapshot_ctrl,
                    SharedFD server_fd);
   [[noreturn]] void MainLoop();
 
  private:
   TpmResourceManager& tpm_resource_manager_;
+  std::shared_ptr<SnapshotController> snapshot_ctrl_;
   std::string server_socket_path_;
   SharedFD server_fd_;
 };
diff --git a/host/commands/secure_env/secure_env_not_windows_main.cpp b/host/commands/secure_env/secure_env_not_windows_main.cpp
index defda777f..be4f6b241 100644
--- a/host/commands/secure_env/secure_env_not_windows_main.cpp
+++ b/host/commands/secure_env/secure_env_not_windows_main.cpp
@@ -41,6 +41,7 @@
 #include "host/commands/secure_env/oemlock/oemlock_responder.h"
 #include "host/commands/secure_env/proxy_keymaster_context.h"
 #include "host/commands/secure_env/rust/kmr_ta.h"
+#include "host/commands/secure_env/snapshot_control.h"
 #include "host/commands/secure_env/soft_gatekeeper.h"
 #include "host/commands/secure_env/storage/insecure_json_storage.h"
 #include "host/commands/secure_env/storage/storage.h"
@@ -292,51 +293,89 @@ Result<void> SecureEnvMain(int argc, char** argv) {
       keymaster::MessageVersion(keymaster::KmVersion::KEYMINT_3,
                                 0 /* km_date */)));
 
+  SharedFD channel_to_run_cvd = DupFdFlag(FLAGS_snapshot_control_fd);
+  std::shared_ptr<SnapshotController> snapshot_controller = std::move(CF_EXPECT(
+      SnapshotController::CreateSnapshotController(channel_to_run_cvd)));
+  CF_EXPECT(snapshot_controller != nullptr);
+  threads.emplace_back([&snapshot_controller]() {
+    // will send suspend/resume commands to all the other worker threads
+    snapshot_controller->ControllerLoop();
+  });
+
   auto keymaster_in = DupFdFlag(FLAGS_keymaster_fd_in);
   auto keymaster_out = DupFdFlag(FLAGS_keymaster_fd_out);
   keymaster::AndroidKeymaster* borrowed_km = keymaster.get();
-  threads.emplace_back([keymaster_in, keymaster_out, borrowed_km]() {
+  threads.emplace_back([keymaster_in, keymaster_out, borrowed_km,
+                        &snapshot_controller]() {
     while (true) {
       SharedFdKeymasterChannel keymaster_channel(keymaster_in, keymaster_out);
 
       KeymasterResponder keymaster_responder(keymaster_channel, *borrowed_km);
 
-      while (keymaster_responder.ProcessMessage()) {
-      }
+      bool loop_continue = true;
+      do {
+        std::shared_lock<std::shared_mutex> reader_lock;
+        if (snapshot_controller->Enabled()) {
+          reader_lock =
+              std::move(snapshot_controller->WaitInitializedOrResumed());
+        }
+        loop_continue = keymaster_responder.ProcessMessage();
+        // release the reader lock.
+      } while (loop_continue);
     }
   });
 
   auto gatekeeper_in = DupFdFlag(FLAGS_gatekeeper_fd_in);
   auto gatekeeper_out = DupFdFlag(FLAGS_gatekeeper_fd_out);
-  threads.emplace_back([gatekeeper_in, gatekeeper_out, &gatekeeper]() {
+  threads.emplace_back([gatekeeper_in, gatekeeper_out, &gatekeeper,
+                        &snapshot_controller]() {
     while (true) {
       SharedFdGatekeeperChannel gatekeeper_channel(gatekeeper_in,
                                                    gatekeeper_out);
 
       GatekeeperResponder gatekeeper_responder(gatekeeper_channel, *gatekeeper);
 
-      while (gatekeeper_responder.ProcessMessage()) {
-      }
+      bool loop_continue = true;
+      do {
+        std::shared_lock<std::shared_mutex> reader_lock;
+        if (snapshot_controller->Enabled()) {
+          reader_lock =
+              std::move(snapshot_controller->WaitInitializedOrResumed());
+        }
+        loop_continue = gatekeeper_responder.ProcessMessage();
+        // release the reader lock.
+      } while (loop_continue);
     }
   });
 
   auto oemlock_in = DupFdFlag(FLAGS_oemlock_fd_in);
   auto oemlock_out = DupFdFlag(FLAGS_oemlock_fd_out);
-  threads.emplace_back([oemlock_in, oemlock_out, &oemlock]() {
-    while (true) {
-      transport::SharedFdChannel channel(oemlock_in, oemlock_out);
-      oemlock::OemLockResponder responder(channel, *oemlock);
-      while (responder.ProcessMessage().ok()) {
-      }
-    }
-  });
+  threads.emplace_back(
+      [oemlock_in, oemlock_out, &oemlock, &snapshot_controller]() {
+        while (true) {
+          transport::SharedFdChannel channel(oemlock_in, oemlock_out);
+          oemlock::OemLockResponder responder(channel, *oemlock);
+          bool loop_continue = true;
+          do {
+            std::shared_lock<std::shared_mutex> reader_lock;
+            if (snapshot_controller->Enabled()) {
+              reader_lock =
+                  std::move(snapshot_controller->WaitInitializedOrResumed());
+            }
+            loop_continue = responder.ProcessMessage().ok();
+            // release the reader lock.
+          } while (loop_continue);
+        }
+      });
 
   auto confui_server_fd = DupFdFlag(FLAGS_confui_server_fd);
-  threads.emplace_back([confui_server_fd, resource_manager]() {
-    ConfUiSignServer confui_sign_server(*resource_manager, confui_server_fd);
-    // no return, infinite loop
-    confui_sign_server.MainLoop();
-  });
+  threads.emplace_back(
+      [confui_server_fd, resource_manager, &snapshot_controller]() {
+        ConfUiSignServer confui_sign_server(
+            *resource_manager, snapshot_controller, confui_server_fd);
+        // no return, infinite loop
+        confui_sign_server.MainLoop();
+      });
 
   auto kernel_events_fd = DupFdFlag(FLAGS_kernel_events_fd);
   threads.emplace_back(StartKernelEventMonitor(kernel_events_fd));
diff --git a/host/commands/secure_env/snapshot_control.cpp b/host/commands/secure_env/snapshot_control.cpp
index 833fbdb4c..80b368f54 100644
--- a/host/commands/secure_env/snapshot_control.cpp
+++ b/host/commands/secure_env/snapshot_control.cpp
@@ -17,7 +17,6 @@
 
 #include <unistd.h>
 
-#include <mutex>
 #include <thread>
 
 #include "host/libs/config/cuttlefish_config.h"
@@ -72,12 +71,14 @@ Result<void> SnapshotController::ControllerLoop() {
   return {};
 }
 
-void SnapshotController::WaitInitializedOrResumed() {
+std::shared_lock<std::shared_mutex>
+SnapshotController::WaitInitializedOrResumed() {
   std::shared_lock reader_lock(reader_writer_mutex_);
   std::atomic<bool>* suspended_atomic_ptr = &suspended_;
   suspended_cv_.wait(reader_lock, [suspended_atomic_ptr]() {
     return !(suspended_atomic_ptr->load());
   });
+  return std::move(reader_lock);
 }
 
 }  // namespace cuttlefish
diff --git a/host/commands/secure_env/snapshot_control.h b/host/commands/secure_env/snapshot_control.h
index 1b32fe45f..94d7466d8 100644
--- a/host/commands/secure_env/snapshot_control.h
+++ b/host/commands/secure_env/snapshot_control.h
@@ -18,6 +18,7 @@
 #include <atomic>
 #include <condition_variable>
 #include <memory>
+#include <mutex>
 #include <shared_mutex>
 
 #include "common/libs/fs/shared_fd.h"
@@ -35,9 +36,9 @@ class SnapshotController {
   SnapshotController& operator=(SnapshotController&&) = delete;
 
   /*
-   * waits until the "suspended_" is false.
+   * waits until the "suspended_" is false and returns reader lock
    */
-  void WaitInitializedOrResumed();
+  std::shared_lock<std::shared_mutex> WaitInitializedOrResumed();
 
   bool Enabled() const { return is_crosvm_; }