From 81980388d478bc6d323ab3068aa657433cdeebd8 Mon Sep 17 00:00:00 2001
From: Jiang Liu <gerry@linux.alibaba.com>
Date: Sat, 12 Aug 2023 11:02:50 +0800
Subject: [PATCH] agent/image: export the image service singleton instance

Export the image service singleton instance.

Signed-off-by: Jiang Liu <gerry@linux.alibaba.com>
---
 src/agent/src/image_rpc.rs | 10 ++++++++--
 src/agent/src/rpc.rs       |  9 ++++++---
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/agent/src/image_rpc.rs b/src/agent/src/image_rpc.rs
index 4f02355ac..4ce4a14fa 100644
--- a/src/agent/src/image_rpc.rs
+++ b/src/agent/src/image_rpc.rs
@@ -36,14 +36,20 @@ const KATA_CC_IMAGE_WORK_DIR: &str = "/run/image/";
 const KATA_CC_PAUSE_BUNDLE: &str = "/pause_bundle";
 const CONFIG_JSON: &str = "config.json";
 
+#[rustfmt::skip]
+lazy_static! {
+    pub static ref IMAGE_SERVICE: Mutex<Option<ImageService>> = Mutex::new(None);
+}
+
 // Convenience function to obtain the scope logger.
 fn sl() -> slog::Logger {
     slog_scope::logger().new(o!("subsystem" => "cgroups"))
 }
 
+#[derive(Clone)]
 pub struct ImageService {
     sandbox: Arc<Mutex<Sandbox>>,
-    attestation_agent_started: AtomicBool,
+    attestation_agent_started: Arc<AtomicBool>,
     image_client: Arc<Mutex<ImageClient>>,
     container_count: Arc<AtomicU16>,
 }
@@ -67,7 +73,7 @@ impl ImageService {
 
         Self {
             sandbox,
-            attestation_agent_started: AtomicBool::new(false),
+            attestation_agent_started: Arc::new(AtomicBool::new(false)),
             image_client: Arc::new(Mutex::new(image_client)),
             container_count: Arc::new(AtomicU16::new(0)),
         }
diff --git a/src/agent/src/rpc.rs b/src/agent/src/rpc.rs
index 2548c72e8..3d79434cd 100644
--- a/src/agent/src/rpc.rs
+++ b/src/agent/src/rpc.rs
@@ -207,7 +207,8 @@ impl AgentService {
             "receive createcontainer, storages: {:?}", &req.storages
         );
 
-        // Merge the image bundle OCI spec into the container creation request OCI spec.
+        // In case of pulling image inside guest, we need to merge the image bundle OCI spec
+        // into the container creation request OCI spec.
         self.merge_bundle_oci(&mut oci).await?;
 
         // Some devices need some extra processing (the ones invoked with
@@ -1773,9 +1774,11 @@ pub async fn start(
     let health_service = Box::new(HealthService {}) as Box<dyn health_ttrpc::Health + Send + Sync>;
     let hservice = health_ttrpc::create_health(Arc::new(health_service));
 
+    let image_service = image_rpc::ImageService::new(s);
+    *image_rpc::IMAGE_SERVICE.lock().await = Some(image_service.clone());
     let image_service =
-        Box::new(image_rpc::ImageService::new(s)) as Box<dyn image_ttrpc::Image + Send + Sync>;
-    let iservice = image_ttrpc::create_image(Arc::new(image_service));
+        Arc::new(Box::new(image_service) as Box<dyn image_ttrpc::Image + Send + Sync>);
+    let iservice = image_ttrpc::create_image(image_service);
 
     let server = TtrpcServer::new()
         .bind(server_address)?