agent:storage: Refactor storage handler registration

- Added `driver_types` method to `StorageHandler` trait to return driver
  types managed by each handler.
- Implemented driver_types method for all storage handlers.
- Updated `STORAGE_HANDLERS` initialization to use `driver_types` for
  handler registration.

Fixes: #10242

Signed-off-by: ChengyuZhu6 <chengyu.zhu@intel.com>

src/agent/src/storage/bind_watcher_handler.rs

@@ -4,6 +4,8 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+use crate::device::DRIVER_WATCHABLE_BIND_TYPE;
+use crate::storage::{new_device, StorageContext, StorageHandler};
 use anyhow::Result;
 use kata_types::mount::StorageDevice;
 use protocols::agent::Storage;
@@ -11,13 +13,16 @@ use std::iter;
 use std::sync::Arc;
 use tracing::instrument;
 
-use crate::storage::{new_device, StorageContext, StorageHandler};
-
 #[derive(Debug)]
 pub struct BindWatcherHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for BindWatcherHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_WATCHABLE_BIND_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,

src/agent/src/storage/block_handler.rs

@@ -10,6 +10,10 @@ use std::path::Path;
 use std::str::FromStr;
 use std::sync::Arc;
 
+use crate::device::{
+    DRIVER_BLK_CCW_TYPE, DRIVER_BLK_MMIO_TYPE, DRIVER_BLK_PCI_TYPE, DRIVER_NVDIMM_TYPE,
+    DRIVER_SCSI_TYPE,
+};
 use anyhow::{anyhow, Context, Result};
 use kata_types::mount::StorageDevice;
 use protocols::agent::Storage;
@@ -29,6 +33,11 @@ pub struct VirtioBlkMmioHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for VirtioBlkMmioHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_BLK_MMIO_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,
@@ -50,6 +59,11 @@ pub struct VirtioBlkPciHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for VirtioBlkPciHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_BLK_PCI_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,
@@ -81,6 +95,11 @@ pub struct VirtioBlkCcwHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for VirtioBlkCcwHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_BLK_CCW_TYPE]
+    }
+
     #[cfg(target_arch = "s390x")]
     #[instrument]
     async fn create_device(
@@ -111,6 +130,11 @@ pub struct ScsiHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for ScsiHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_SCSI_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,
@@ -131,6 +155,11 @@ pub struct PmemHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for PmemHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_NVDIMM_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,

src/agent/src/storage/ephemeral_handler.rs

@@ -35,6 +35,11 @@ pub struct EphemeralHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for EphemeralHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_EPHEMERAL_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,

src/agent/src/storage/fs_handler.rs

@@ -8,18 +8,23 @@ use std::fs;
 use std::path::Path;
 use std::sync::Arc;
 
+use crate::device::{DRIVER_9P_TYPE, DRIVER_OVERLAYFS_TYPE, DRIVER_VIRTIOFS_TYPE};
+use crate::storage::{common_storage_handler, new_device, StorageContext, StorageHandler};
 use anyhow::{anyhow, Context, Result};
 use kata_types::mount::StorageDevice;
 use protocols::agent::Storage;
 use tracing::instrument;
 
-use crate::storage::{common_storage_handler, new_device, StorageContext, StorageHandler};
-
 #[derive(Debug)]
 pub struct OverlayfsHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for OverlayfsHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_OVERLAYFS_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,
@@ -61,6 +66,11 @@ pub struct Virtio9pHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for Virtio9pHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_9P_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,
@@ -77,6 +87,11 @@ pub struct VirtioFsHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for VirtioFsHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_VIRTIOFS_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,

src/agent/src/storage/image_pull_handler.rs

@@ -32,6 +32,11 @@ impl ImagePullHandler {
 
 #[async_trait::async_trait]
 impl StorageHandler for ImagePullHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[KATA_VIRTUAL_VOLUME_IMAGE_GUEST_PULL]
+    }
+
     #[instrument]
     async fn create_device(
         &self,

src/agent/src/storage/local_handler.rs

@@ -8,6 +8,7 @@ use std::fs;
 use std::os::unix::fs::PermissionsExt;
 use std::sync::Arc;
 
+use crate::device::DRIVER_LOCAL_TYPE;
 use anyhow::{Context, Result};
 use kata_types::mount::{StorageDevice, KATA_MOUNT_OPTION_FS_GID};
 use nix::unistd::Gid;
@@ -21,6 +22,11 @@ pub struct LocalHandler {}
 
 #[async_trait::async_trait]
 impl StorageHandler for LocalHandler {
+    #[instrument]
+    fn driver_types(&self) -> &[&str] {
+        &[DRIVER_LOCAL_TYPE]
+    }
+
     #[instrument]
     async fn create_device(
         &self,

src/agent/src/storage/mod.rs

@@ -12,8 +12,6 @@ use std::sync::Arc;
 
 use anyhow::{anyhow, Context, Result};
 use kata_sys_util::mount::{create_mount_destination, parse_mount_options};
-#[cfg(feature = "guest-pull")]
-use kata_types::mount::KATA_VIRTUAL_VOLUME_IMAGE_GUEST_PULL;
 use kata_types::mount::{StorageDevice, StorageHandlerManager, KATA_SHAREDFS_GUEST_PREMOUNT_TAG};
 use nix::unistd::{Gid, Uid};
 use protocols::agent::Storage;
@@ -29,11 +27,6 @@ use self::fs_handler::{OverlayfsHandler, Virtio9pHandler, VirtioFsHandler};
 #[cfg(feature = "guest-pull")]
 use self::image_pull_handler::ImagePullHandler;
 use self::local_handler::LocalHandler;
-use crate::device::{
-    DRIVER_9P_TYPE, DRIVER_BLK_MMIO_TYPE, DRIVER_BLK_PCI_TYPE, DRIVER_EPHEMERAL_TYPE,
-    DRIVER_LOCAL_TYPE, DRIVER_NVDIMM_TYPE, DRIVER_OVERLAYFS_TYPE, DRIVER_SCSI_TYPE,
-    DRIVER_VIRTIOFS_TYPE, DRIVER_WATCHABLE_BIND_TYPE,
-};
 use crate::mount::{baremount, is_mounted, remove_mounts};
 use crate::sandbox::Sandbox;
 
@@ -133,26 +126,36 @@ pub trait StorageHandler: Send + Sync {
         storage: Storage,
         ctx: &mut StorageContext,
     ) -> Result<Arc<dyn StorageDevice>>;
+
+    /// Return the driver types that the handler manages.
+    fn driver_types(&self) -> &[&str];
 }
 
 #[rustfmt::skip]
 lazy_static! {
     pub static ref STORAGE_HANDLERS: StorageHandlerManager<Arc<dyn StorageHandler>> = {
         let mut manager: StorageHandlerManager<Arc<dyn StorageHandler>> = StorageHandlerManager::new();
-        manager.add_handler(&[DRIVER_9P_TYPE], Arc::new(Virtio9pHandler{})).unwrap();
-        #[cfg(target_arch = "s390x")]
-        manager.add_handler(&[crate::device::DRIVER_BLK_CCW_TYPE], Arc::new(self::block_handler::VirtioBlkCcwHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_BLK_MMIO_TYPE], Arc::new(VirtioBlkMmioHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_BLK_PCI_TYPE], Arc::new(VirtioBlkPciHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_EPHEMERAL_TYPE], Arc::new(EphemeralHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_LOCAL_TYPE], Arc::new(LocalHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_NVDIMM_TYPE], Arc::new(PmemHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_OVERLAYFS_TYPE], Arc::new(OverlayfsHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_SCSI_TYPE], Arc::new(ScsiHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_VIRTIOFS_TYPE], Arc::new(VirtioFsHandler{})).unwrap();
-        manager.add_handler(&[DRIVER_WATCHABLE_BIND_TYPE], Arc::new(BindWatcherHandler{})).unwrap();
-        #[cfg(feature = "guest-pull")]
-        manager.add_handler(&[KATA_VIRTUAL_VOLUME_IMAGE_GUEST_PULL], Arc::new(ImagePullHandler{})).unwrap();
+        let handlers: Vec<Arc<dyn StorageHandler>> = vec![
+            Arc::new(Virtio9pHandler {}),
+            Arc::new(VirtioBlkMmioHandler {}),
+            Arc::new(VirtioBlkPciHandler {}),
+            Arc::new(EphemeralHandler {}),
+            Arc::new(LocalHandler {}),
+            Arc::new(PmemHandler {}),
+            Arc::new(OverlayfsHandler {}),
+            Arc::new(ScsiHandler {}),
+            Arc::new(VirtioFsHandler {}),
+            Arc::new(BindWatcherHandler {}),
+            #[cfg(target_arch = "s390x")]
+            Arc::new(self::block_handler::VirtioBlkCcwHandler {}),
+            #[cfg(feature = "guest-pull")]
+            Arc::new(ImagePullHandler {}),
+        ];
+
+        for handler in handlers {
+            manager.add_handler(handler.driver_types(), handler.clone()).unwrap();
+        }
+
         manager
     };
 }