Implement dm-sworndisk with core crate

core/src/lib.rs

@@ -37,7 +37,7 @@ mod util;
 extern crate alloc;
 
 #[cfg(feature = "linux")]
-pub use self::os::{Arc, Mutex};
+pub use self::os::{Arc, Mutex, Vec};
 
 #[cfg(feature = "occlum")]
 #[macro_use]

core/src/os/linux/mod.rs

@@ -125,6 +125,18 @@ impl<T: ?Sized> AsMut<T> for Box<T> {
     }
 }
 
+impl<I: Iterator + ?Sized> Iterator for Box<I> {
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<I::Item> {
+        self.inner.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.inner.size_hint()
+    }
+}
+
 impl<T: fmt::Debug + ?Sized> fmt::Debug for Box<T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Debug::fmt(&**self, f)

linux/dm-sworndisk/src/lib.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0
 // Copyright (c) 2023 Ant Group CO., Ltd.
 
-//! Dummy dm-sworndisk module.
+//! A device mapper target based on `SwornDisk`.
 
 #![feature(allocator_api)]
 
@@ -16,7 +16,9 @@ use core::{
 };
 use crossbeam_queue::SegQueue;
 use kernel::{c_str, new_condvar, prelude::*, sync::CondVar};
-use sworndisk::{Arc, BlockId, BlockSet, BufMut, BufRef, Errno, Error, Mutex};
+use sworndisk::{
+    AeadKey, Arc, BlockId, BlockSet, Buf, BufMut, BufRef, Errno, Error, Mutex, SwornDisk, Vec,
+};
 
 use crate::bio::{Bio, BioOp, BioVec};
 use crate::block_device::{HostBlockDevice, BLOCK_SECTORS, BLOCK_SIZE};
@@ -46,6 +48,8 @@ impl kernel::Module for TargetManager {
         // test_thread();
         // test_aead();
         // test_skcipher();
+        // test_rawdisk(c_str!("/dev/sdb"), 1 * 1024 * 1024 * 1024);
+        // test_sworndisk(c_str!("/dev/sdb"), 10 * 1024 * 1024 * 1024);
 
         let dm_sworndisk = Box::pin_init(TargetType::register::<DmSwornDisk>(
             c_str!("sworndisk"),
@@ -64,17 +68,16 @@ impl Drop for TargetManager {
 }
 
 /// A request queue, dispatching bios from device mapper to `RawDisk`.
-struct ReqQueue {
+struct ReqQueue<D: BlockSet> {
     bios: Mutex<SegQueue<Bio>>,
-    // TODO: replace `RawDisk` with `SwornDisk`.
-    disk: RawDisk,
+    disk: SwornDisk<D>,
     should_stop: AtomicBool,
     new_bio_condvar: Pin<Box<CondVar>>,
 }
 
-impl ReqQueue {
+impl<D: BlockSet + 'static> ReqQueue<D> {
     /// Constructs a `ReqQueue`.
-    pub fn new(disk: RawDisk) -> Self {
+    pub fn new(disk: SwornDisk<D>) -> Self {
         Self {
             bios: Mutex::new(SegQueue::new()),
             disk,
@@ -125,29 +128,41 @@ impl ReqQueue {
     fn process(&self, bio: Bio) {
         if bio.start_sector() % BLOCK_SECTORS != 0 || bio.len() % BLOCK_SIZE != 0 {
             pr_warn!(
-                "bio not aligned to BLOCK_SIZE, start_sector: {}, len: {}",
+                "bio not aligned to BLOCK_SIZE, {:?}, start_sector: {}, len: {}",
+                bio.op(),
                 bio.start_sector(),
                 bio.len(),
             );
+
+            match bio.op() {
+                BioOp::Read => self.read_unaligned(&bio),
+                BioOp::Write => self.write_unaligned(&bio),
+                _ => unreachable!(),
+            }
             bio.end();
             return;
         }
 
+        // TODO: we may need a block range iterator abstraction to deal with
+        // blocks, including the unaligned ones.
+
         match bio.op() {
             BioOp::Read => {
-                let mut pos = bio.start_sector() / BLOCK_SECTORS;
-                for mut bio_vec in bio.iter() {
-                    let buf = BufMut::try_from(bio_vec.deref_mut()).unwrap();
-                    let nblocks = buf.nblocks();
-                    if let Err(err) = self.disk.read(pos, buf) {
-                        pr_info!(
-                            "read sworndisk failed, block_id: {}, nblocks: {}, err: {:?}",
-                            pos,
-                            nblocks,
-                            err,
-                        );
-                    }
-                    pos += nblocks;
+                let mut bio_vecs = Vec::new();
+                bio.iter().for_each(|bio_vec| bio_vecs.push(bio_vec));
+                let mut bufs = Vec::new();
+                bio_vecs.iter_mut().for_each(|bio_vec| {
+                    bufs.push(BufMut::try_from(bio_vec.deref_mut()).unwrap());
+                });
+
+                let start_block = bio.start_sector() / BLOCK_SECTORS;
+                if let Err(err) = self.disk.readv(start_block, &mut bufs) {
+                    pr_info!(
+                        "read sworndisk failed, block_id: {}, nblocks: {}, err: {:?}",
+                        start_block,
+                        bio.len() / BLOCK_SIZE,
+                        err,
+                    );
                 }
             }
             BioOp::Write => {
@@ -171,23 +186,91 @@ impl ReqQueue {
         bio.end();
     }
 
+    fn read_unaligned(&self, bio: &Bio) {
+        let start_bytes = bio.start_sector() * (BLOCK_SIZE / BLOCK_SECTORS);
+        let end_bytes = start_bytes + bio.len();
+
+        let start_block = start_bytes / BLOCK_SIZE;
+        let end_block = (end_bytes + BLOCK_SIZE - 1) / BLOCK_SIZE;
+        let nblocks = end_block - start_block;
+        let mut buf = Buf::alloc(nblocks).expect("alloc read buffer failed");
+
+        if let Err(err) = self.disk.read(start_block, buf.as_mut()) {
+            pr_info!(
+                "read sworndisk failed, block_id: {}, nblocks: {}, err: {:?}",
+                start_block,
+                nblocks,
+                err,
+            );
+        }
+
+        let mut offset = start_bytes % BLOCK_SIZE;
+        for mut bio_vec in bio.iter() {
+            let vec_len = bio_vec.len();
+            bio_vec.copy_from_slice(&buf.as_slice()[offset..offset + vec_len]);
+            offset += vec_len;
+        }
+    }
+
+    fn write_unaligned(&self, bio: &Bio) {
+        let start_bytes = bio.start_sector() * (BLOCK_SIZE / BLOCK_SECTORS);
+        let end_bytes = start_bytes + bio.len();
+
+        let start_block = start_bytes / BLOCK_SIZE;
+        let end_block = (end_bytes + BLOCK_SIZE - 1) / BLOCK_SIZE;
+        let nblocks = end_block - start_block;
+        let mut buf = Buf::alloc(nblocks).expect("alloc read buffer failed");
+
+        // TODO: we could only read the first and last blocks.
+        if let Err(err) = self.disk.read(start_block, buf.as_mut()) {
+            pr_info!(
+                "read sworndisk failed, block_id: {}, nblocks: {}, err: {:?}",
+                start_block,
+                nblocks,
+                err,
+            );
+        }
+
+        let mut offset = start_bytes % BLOCK_SIZE;
+        for mut bio_vec in bio.iter() {
+            let vec_len = bio_vec.len();
+            buf.as_mut_slice()[offset..offset + vec_len].copy_from_slice(bio_vec.deref());
+            offset += vec_len;
+        }
+
+        if let Err(err) = self.disk.write(start_block, buf.as_ref()) {
+            pr_info!(
+                "write sworndisk failed, block_id: {}, nblocks: {}, err: {:?}",
+                start_block,
+                nblocks,
+                err,
+            );
+        }
+    }
+
     /// Processes all the pending `Bio`s in the queue.
     fn clear(&self) {
         while let Some(bio) = self.bios.lock().pop() {
             self.process(bio);
         }
     }
+
+    /// Sync the disk before handle next bio.
+    pub fn sync(&self) -> Result<(), Error> {
+        let _ = self.bios.lock();
+        self.disk.sync()
+    }
 }
 
 /// A struct represent a `dm_target` type, which should impl `DmTargetOps`.
 struct DmSwornDisk {
-    queue: Arc<ReqQueue>,
+    queue: Arc<ReqQueue<RawDisk>>,
     worker: Mutex<Option<JoinHandle<()>>>,
 }
 
 impl DmSwornDisk {
     /// Returns an in-place initializer.
-    fn new(queue: Arc<ReqQueue>, worker: JoinHandle<()>) -> impl Init<Self> {
+    fn new(queue: Arc<ReqQueue<RawDisk>>, worker: JoinHandle<()>) -> impl Init<Self> {
         init!(Self {
             queue,
             worker: Mutex::new(Some(worker))
@@ -298,38 +381,72 @@ impl DmTargetOps for DmSwornDisk {
     type Private = DmSwornDisk;
 
     fn ctr(target: &mut Target<Self>, args: Args) -> Result<Box<Self>> {
-        // TODO: accept more arguments, e.g., root key.
-        if args.len() != 1 {
+        if args.len() != 3 {
             target.set_error(c_str!("Invalid argument count"));
             return Err(EINVAL);
         }
 
-        let Ok(raw_disk) = RawDisk::open(&args[0]) else {
+        let should_format = match args[0].to_str() {
+            Ok("create") => true,
+            Ok("open") => false,
+            _ => {
+                target.set_error(c_str!("Unsupported operation"));
+                return Err(EINVAL);
+            }
+        };
+
+        let Ok(raw_disk) = RawDisk::open(&args[1]) else {
             target.set_error(c_str!("Device lookup failed"));
             return Err(ENODEV);
         };
 
-        // TODO: use raw_disk to construct a sworndisk instance.
-        let queue = Arc::new(ReqQueue::new(raw_disk));
+        let root_key = {
+            let mut key = AeadKey::default();
+            // SAFETY: `key` and `args` contain valid pointers.
+            let err =
+                unsafe { bindings::hex2bin(key.as_mut_ptr(), args[2].as_char_ptr(), key.len()) };
+            if err != 0 {
+                target.set_error(c_str!("Invalid root key"));
+                return Err(EINVAL);
+            };
+            key
+        };
+
+        let sworndisk = match should_format {
+            true => SwornDisk::create(raw_disk, root_key, None).map_err(|_| {
+                target.set_error(c_str!("Create sworndisk failed"));
+                ENODEV
+            })?,
+            // TODO: open with a `SyncIdStore`.
+            false => SwornDisk::open(raw_disk, root_key, None).map_err(|_| {
+                target.set_error(c_str!("Open sworndisk failed"));
+                ENODEV
+            })?,
+        };
+        target.set_region(0..sworndisk.total_blocks());
+
+        let queue = Arc::new(ReqQueue::new(sworndisk));
         let worker = ReqQueue::spawn_req_worker(queue.clone());
 
         Box::init(DmSwornDisk::new(queue, worker))
     }
 
     fn dtr(target: &mut Target<Self>) {
         let Some(dm_sworndisk) = target.private() else {
-            pr_warn!("Error, found no dm_sworndisk\n");
+            pr_err!("Error, found no dm_sworndisk\n");
             return;
         };
 
         dm_sworndisk.queue.set_stopped();
         let worker = dm_sworndisk.worker.lock().take().unwrap();
         worker.join().unwrap();
+
+        dm_sworndisk.queue.sync().unwrap();
     }
 
     fn map(target: &Target<Self>, bio: Bio) -> MapState {
         let Some(dm_sworndisk) = target.private() else {
-            pr_warn!("Error, found no dm_sworndisk\n");
+            pr_err!("Error, found no dm_sworndisk\n");
             return MapState::Kill;
         };
 
@@ -559,3 +676,77 @@ fn test_skcipher() {
     }
     assert_eq!(decrypted, plaintext);
 }
+
+fn test_rawdisk(dev_path: &'static CStr, test_size: usize) {
+    spawn(move || {
+        pr_info!("----test_rawdisk begin");
+        let rawdisk = RawDisk::open(dev_path).unwrap();
+        let nblocks = test_size / BLOCK_SIZE;
+        if rawdisk.region.end < nblocks {
+            pr_err!("raw disk is too small for test_size: {test_size}");
+            return;
+        }
+
+        let mut buf = Buf::alloc(1).unwrap();
+        pr_info!("----write begin");
+        for i in 0..nblocks {
+            buf.as_mut_slice().fill(i as u8);
+            match rawdisk.write(i, buf.as_ref()) {
+                Err(err) => pr_info!("write sworndisk failed, block_id: {}, err: {:?}", i, err),
+                Ok(_) => continue,
+            }
+        }
+        pr_info!("----write end");
+
+        rawdisk.flush().unwrap();
+
+        pr_info!("----read begin");
+        for i in 0..nblocks {
+            match rawdisk.read(i, buf.as_mut()) {
+                Err(err) => pr_info!("read sworndisk failed, block_id: {}, err: {:?}", i, err),
+                Ok(_) => continue,
+            }
+            assert_eq!(buf.as_slice()[0], i as u8);
+        }
+        pr_info!("----read end");
+        pr_info!("----test_rawdisk end");
+    });
+}
+
+fn test_sworndisk(dev_path: &'static CStr, test_size: usize) {
+    spawn(move || {
+        pr_info!("----test_sworndisk begin");
+        let root_key = AeadKey::default();
+        let raw_disk = RawDisk::open(dev_path).unwrap();
+        let sworndisk = SwornDisk::create(raw_disk, root_key, None).unwrap();
+        let nblocks = test_size / BLOCK_SIZE;
+        if sworndisk.total_blocks() < nblocks {
+            pr_err!("sworndisk is too small for test_size: {test_size}");
+            return;
+        }
+
+        let mut buf = Buf::alloc(1).unwrap();
+        pr_info!("----write begin");
+        for i in 0..nblocks {
+            buf.as_mut_slice().fill(i as u8);
+            match sworndisk.write(i, buf.as_ref()) {
+                Err(err) => pr_info!("write sworndisk failed, block_id: {}, err: {:?}", i, err),
+                Ok(_) => continue,
+            }
+        }
+        pr_info!("----write end");
+
+        sworndisk.sync().unwrap();
+
+        pr_info!("----read begin");
+        for i in 0..nblocks {
+            match sworndisk.read(i, buf.as_mut()) {
+                Err(err) => pr_info!("read sworndisk failed, block_id: {}, err: {:?}", i, err),
+                Ok(_) => continue,
+            }
+            assert_eq!(buf.as_slice()[0], i as u8);
+        }
+        pr_info!("----read end");
+        pr_info!("----test_sworndisk end");
+    });
+}