disk_hook/
lib.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Disk decorators with operation hooks for storage tests.
5//!
6//! [`HookDisk::new`] provides sparse RAM storage using [`disklayer_ram`], and
7//! [`HookDisk::wrap`] decorates an existing [`Disk`] while preserving its
8//! capabilities.
9//!
10//! Use [`HookDisk::with_sync_hook`] to observe operation parameters or inject
11//! errors, and [`HookDisk::with_async_hook`] to asynchronously gate I/O. Hooks
12//! run before the operation; returning an error prevents it from reaching
13//! storage. Hooks do not receive data buffers.
14//!
15//! ```
16//! use disk_hook::{Disk, HookDisk};
17//!
18//! let disk: Disk = HookDisk::new(4096, 512).into_disk();
19//! assert_eq!(disk.sector_count(), 8);
20//! ```
21//!
22//! Use this crate as a dev-dependency, not as a production storage backend.
23
24#![forbid(unsafe_code)]
25
26pub use disk_backend::Disk;
27pub use disk_backend::DiskError;
28use disk_backend::DiskIo;
29use disk_backend::UnmapBehavior;
30use disklayer_ram::ram_disk_with_sector_size;
31use inspect::Inspect;
32use scsi_buffers::RequestBuffers;
33use std::future::Future;
34use std::pin::Pin;
35
36/// An operation about to be issued to a test disk.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum DiskOperation {
39    /// Read `len` bytes starting at `sector`.
40    Read {
41        /// Starting logical sector.
42        sector: u64,
43        /// Transfer length in bytes.
44        len: usize,
45    },
46    /// Write `len` bytes starting at `sector`, with the requested FUA flag.
47    Write {
48        /// Starting logical sector.
49        sector: u64,
50        /// Transfer length in bytes.
51        len: usize,
52        /// Whether force unit access was requested.
53        fua: bool,
54    },
55    /// Flush the disk's cache.
56    SyncCache,
57    /// Unmap a range of sectors.
58    Unmap {
59        /// Starting logical sector.
60        sector: u64,
61        /// Number of logical sectors.
62        count: u64,
63        /// Whether only whole allocation blocks should be unmapped.
64        block_level_only: bool,
65    },
66    /// Eject the medium.
67    Eject,
68}
69
70type HookFuture = Pin<Box<dyn Future<Output = Result<(), DiskError>> + Send>>;
71type Hook = dyn Fn(DiskOperation) -> HookFuture + Send + Sync;
72
73/// A RAM disk or disk decorator with an optional operation hook.
74///
75/// Configure metadata and hooks before calling [`into_disk`](Self::into_disk).
76/// The resulting disk handle is cheap to clone and shares storage, capacity,
77/// and hook state.
78#[derive(Inspect)]
79pub struct HookDisk {
80    sector_size: u32,
81    physical_sector_size: u32,
82    disk_id: Option<[u8; 16]>,
83    read_only: bool,
84    fua_respected: bool,
85    unmap_behavior: UnmapBehavior,
86    optimal_unmap_sectors: u32,
87    disk: Disk,
88    #[inspect(skip)]
89    hook: Option<Box<Hook>>,
90}
91
92impl HookDisk {
93    /// Creates a sparse RAM disk of `size` bytes using [`disklayer_ram`].
94    ///
95    /// Unwritten sectors read as zeroes. Storage is allocated only on writes.
96    ///
97    /// Panics if the size is zero, the sector size is not a power of two
98    /// >= 512, or the size is not a multiple of the sector size.
99    pub fn new(size: u64, sector_size: u32) -> Self {
100        Self::wrap(
101            ram_disk_with_sector_size(size, false, sector_size)
102                .expect("invalid hook disk specification"),
103        )
104    }
105
106    /// Wraps a disk, preserving its metadata, reservations, and resize support.
107    pub fn wrap(disk: Disk) -> Self {
108        Self {
109            sector_size: disk.sector_size(),
110            physical_sector_size: disk.physical_sector_size(),
111            disk_id: disk.disk_id(),
112            read_only: disk.is_read_only(),
113            fua_respected: disk.is_fua_respected(),
114            unmap_behavior: disk.unmap_behavior(),
115            optimal_unmap_sectors: disk.optimal_unmap_sectors(),
116            disk,
117            hook: None,
118        }
119    }
120
121    /// Converts this builder into a disk handle, preserving hooks and metadata.
122    ///
123    /// The returned handle can be cheaply cloned to share storage and hook state.
124    ///
125    /// Panics if the configured physical sector size is not a power of two or
126    /// is smaller than the logical sector size.
127    pub fn into_disk(self) -> Disk {
128        Disk::new(self).expect("invalid hook disk metadata")
129    }
130
131    /// Sets the physical sector size reported by the disk.
132    pub fn with_physical_sector_size(mut self, sector_size: u32) -> Self {
133        self.physical_sector_size = sector_size;
134        self
135    }
136
137    /// Sets the identifier reported by the disk.
138    pub fn with_disk_id(mut self, disk_id: [u8; 16]) -> Self {
139        self.disk_id = Some(disk_id);
140        self
141    }
142
143    /// Sets whether the disk rejects writes and unmaps that reach the backend.
144    ///
145    /// Setting this to false does not make a read-only backing disk writable.
146    pub fn with_read_only(mut self, read_only: bool) -> Self {
147        self.read_only = read_only;
148        self
149    }
150
151    /// Sets whether the disk advertises FUA support.
152    pub fn with_fua_respected(mut self, respected: bool) -> Self {
153        self.fua_respected = respected;
154        self
155    }
156
157    /// Sets the advertised unmap behavior.
158    ///
159    /// This does not change the backing disk's behavior. Advertise `Zeroes`
160    /// only if it actually zeroes unmapped sectors. [`Disk`] does not forward
161    /// ignored unmaps, so they do not invoke the hook.
162    pub fn with_unmap_behavior(mut self, behavior: UnmapBehavior) -> Self {
163        self.unmap_behavior = behavior;
164        self
165    }
166
167    /// Sets the advertised optimal unmap granularity in sectors.
168    pub fn with_optimal_unmap_sectors(mut self, sectors: u32) -> Self {
169        self.optimal_unmap_sectors = sectors;
170        self
171    }
172
173    /// Installs a synchronous hook, replacing any existing hook.
174    ///
175    /// The hook runs before each I/O operation and may fail it. Concurrent
176    /// operations may invoke the hook concurrently.
177    pub fn with_sync_hook<F>(mut self, hook: F) -> Self
178    where
179        F: Fn(DiskOperation) -> Result<(), DiskError> + Send + Sync + 'static,
180    {
181        self.hook = Some(Box::new(move |op| Box::pin(std::future::ready(hook(op)))));
182        self
183    }
184
185    /// Installs an async hook, replacing any existing hook.
186    ///
187    /// The hook runs before each I/O operation and may fail or suspend it.
188    /// Concurrent operations may invoke the hook concurrently.
189    pub fn with_async_hook<F, Fut>(mut self, hook: F) -> Self
190    where
191        F: Fn(DiskOperation) -> Fut + Send + Sync + 'static,
192        Fut: Future<Output = Result<(), DiskError>> + Send + 'static,
193    {
194        self.hook = Some(Box::new(move |op| Box::pin(hook(op))));
195        self
196    }
197
198    async fn before(&self, operation: DiskOperation) -> Result<(), DiskError> {
199        if let Some(hook) = &self.hook {
200            hook(operation).await?;
201        }
202        Ok(())
203    }
204}
205
206impl DiskIo for HookDisk {
207    fn disk_type(&self) -> &str {
208        "hook"
209    }
210
211    fn sector_count(&self) -> u64 {
212        self.disk.sector_count()
213    }
214
215    fn sector_size(&self) -> u32 {
216        self.sector_size
217    }
218
219    fn physical_sector_size(&self) -> u32 {
220        self.physical_sector_size
221    }
222
223    fn disk_id(&self) -> Option<[u8; 16]> {
224        self.disk_id
225    }
226
227    fn is_read_only(&self) -> bool {
228        self.read_only
229    }
230
231    fn is_fua_respected(&self) -> bool {
232        self.fua_respected
233    }
234
235    fn unmap_behavior(&self) -> UnmapBehavior {
236        self.unmap_behavior
237    }
238
239    fn optimal_unmap_sectors(&self) -> u32 {
240        self.optimal_unmap_sectors
241    }
242
243    fn pr(&self) -> Option<&dyn disk_backend::pr::PersistentReservation> {
244        self.disk.pr()
245    }
246
247    async fn read_vectored(
248        &self,
249        buffers: &RequestBuffers<'_>,
250        sector: u64,
251    ) -> Result<(), DiskError> {
252        self.before(DiskOperation::Read {
253            sector,
254            len: buffers.len(),
255        })
256        .await?;
257        self.disk.read_vectored(buffers, sector).await
258    }
259
260    async fn write_vectored(
261        &self,
262        buffers: &RequestBuffers<'_>,
263        sector: u64,
264        fua: bool,
265    ) -> Result<(), DiskError> {
266        self.before(DiskOperation::Write {
267            sector,
268            len: buffers.len(),
269            fua,
270        })
271        .await?;
272        if self.read_only {
273            return Err(DiskError::ReadOnly);
274        }
275        self.disk.write_vectored(buffers, sector, fua).await
276    }
277
278    async fn sync_cache(&self) -> Result<(), DiskError> {
279        self.before(DiskOperation::SyncCache).await?;
280        self.disk.sync_cache().await
281    }
282
283    async fn unmap(
284        &self,
285        sector: u64,
286        count: u64,
287        block_level_only: bool,
288    ) -> Result<(), DiskError> {
289        self.before(DiskOperation::Unmap {
290            sector,
291            count,
292            block_level_only,
293        })
294        .await?;
295        if self.read_only {
296            return Err(DiskError::ReadOnly);
297        }
298        self.disk.unmap(sector, count, block_level_only).await
299    }
300
301    async fn eject(&self) -> Result<(), DiskError> {
302        self.before(DiskOperation::Eject).await?;
303        self.disk.eject().await
304    }
305
306    async fn wait_resize(&self, sector_count: u64) -> u64 {
307        self.disk.wait_resize(sector_count).await
308    }
309}
310
311#[cfg(test)]
312mod tests;
disk_hook/lib.rs

disk_hook/
lib.rs
