imago/
raw.rs

1//! Access generic files as images.
2//!
3//! Allows accessing generic storage objects (`Storage`) as images (i.e. `FormatAccess`).
4
5use crate::format::builder::{
6    FormatCreateBuilder, FormatCreateBuilderBase, FormatDriverBuilder, FormatDriverBuilderBase,
7};
8use crate::format::drivers::FormatDriverInstance;
9use crate::format::gate::ImplicitOpenGate;
10use crate::format::{Format, PreallocateMode};
11use crate::{
12    storage, DenyImplicitOpenGate, ShallowMapping, Storage, StorageExt, StorageOpenOptions,
13};
14use maybe_async::maybe_async;
15use std::fmt::{self, Display, Formatter};
16use std::io;
17use std::path::{Path, PathBuf};
18use std::sync::atomic::{AtomicU64, Ordering};
19
20/// Wraps a storage object without any translation.
21#[derive(Debug)]
22pub struct Raw<S: Storage + 'static> {
23    /// Wrapped storage object.
24    inner: S,
25
26    /// Whether this image may be modified.
27    writable: bool,
28
29    /// Disk size, which is the file size when this object was created.
30    size: AtomicU64,
31}
32
33#[maybe_async]
34impl<S: Storage + 'static> Raw<S> {
35    /// Create a new [`FormatDriverBuilder`] instance for the given image.
36    pub fn builder(image: S) -> RawOpenBuilder<S> {
37        RawOpenBuilder::new(image)
38    }
39
40    /// Create a new [`FormatDriverBuilder`] instance for an image under the given path.
41    pub fn builder_path<P: AsRef<Path>>(image_path: P) -> RawOpenBuilder<S> {
42        RawOpenBuilder::new_path(image_path)
43    }
44
45    /// Create a new [`FormatCreateBuilder`] instance for the given file.
46    pub fn create_builder(image: S) -> RawCreateBuilder<S> {
47        RawCreateBuilder::new(image)
48    }
49
50    /// Wrap `inner`, allowing it to be used as a disk image in raw format.
51    pub async fn open_image(inner: S, writable: bool) -> io::Result<Self> {
52        let size = inner.size()?;
53        Ok(Raw {
54            inner,
55            writable,
56            size: size.into(),
57        })
58    }
59
60    /// Open the given path as a storage object, and wrap it in `Raw`.
61    pub async fn open_path<P: AsRef<Path>>(path: P, writable: bool) -> io::Result<Self> {
62        let storage_opts = StorageOpenOptions::new().write(writable).filename(path);
63        let inner = S::open(storage_opts).await?;
64        Self::open_image(inner, writable).await
65    }
66
67    /// Wrap `inner`, allowing it to be used as a disk image in raw format.
68    #[cfg(feature = "sync-wrappers")]
69    pub fn open_image_sync(inner: S, writable: bool) -> io::Result<Self> {
70        let size = inner.size()?;
71        Ok(Raw {
72            inner,
73            writable,
74            size: size.into(),
75        })
76    }
77
78    #[cfg(feature = "sync-wrappers")]
79    /// Synchronous wrapper around [`Raw::open_path()`].
80    pub fn open_path_sync<P: AsRef<Path>>(path: P, writable: bool) -> io::Result<Self> {
81        tokio::runtime::Builder::new_current_thread()
82            .build()?
83            .block_on(Self::open_path(path, writable))
84    }
85}
86
87#[maybe_async(?Send)]
88impl<S: Storage + 'static> FormatDriverInstance for Raw<S> {
89    type Storage = S;
90
91    fn format(&self) -> Format {
92        Format::Raw
93    }
94
95    async unsafe fn probe(_storage: &S) -> io::Result<bool>
96    where
97        Self: Sized,
98    {
99        Ok(true)
100    }
101
102    fn size(&self) -> u64 {
103        self.size.load(Ordering::Relaxed)
104    }
105
106    fn zero_granularity(&self) -> Option<u64> {
107        None
108    }
109
110    fn collect_storage_dependencies(&self) -> Vec<&S> {
111        vec![&self.inner]
112    }
113
114    fn writable(&self) -> bool {
115        self.writable
116    }
117
118    #[allow(clippy::needless_lifetimes)] // Elidable in sync, but async needs a named lifetime for the boxed future bound
119    async fn get_mapping<'a>(
120        &'a self,
121        offset: u64,
122        max_length: u64,
123    ) -> io::Result<(ShallowMapping<'a, S>, u64)> {
124        let remaining = match self.size().checked_sub(offset) {
125            None | Some(0) => return Ok((ShallowMapping::Eof {}, 0)),
126            Some(remaining) => remaining,
127        };
128
129        Ok((
130            ShallowMapping::Raw {
131                storage: &self.inner,
132                offset,
133                writable: true,
134            },
135            std::cmp::min(max_length, remaining),
136        ))
137    }
138
139    #[allow(clippy::needless_lifetimes)] // Elidable in sync, but async needs a named lifetime for the boxed future bound
140    async fn ensure_data_mapping<'a>(
141        &'a self,
142        offset: u64,
143        length: u64,
144        _overwrite: bool,
145    ) -> io::Result<(&'a S, u64, u64)> {
146        let Some(remaining) = self.size().checked_sub(offset) else {
147            return Err(io::Error::other("Cannot allocate past the end of file"));
148        };
149        if length > remaining {
150            return Err(io::Error::other("Cannot allocate past the end of file"));
151        }
152
153        Ok((&self.inner, offset, length))
154    }
155
156    async fn ensure_zero_mapping(&self, offset: u64, length: u64) -> io::Result<(u64, u64)> {
157        let zero_align = self.inner.zero_align();
158        assert!(zero_align.is_power_of_two());
159
160        let zero_align_mask = zero_align as u64 - 1;
161
162        let aligned_end = (offset + length) & !zero_align_mask;
163        let aligned_offset = (offset + zero_align_mask) & !zero_align_mask;
164        let aligned_length = aligned_end.saturating_sub(aligned_offset);
165        if aligned_length == 0 {
166            return Ok((aligned_offset, 0));
167        }
168
169        // FIXME: Introduce request flags, and request no fallback
170        self.inner
171            .write_zeroes(aligned_offset, aligned_length)
172            .await?;
173        Ok((aligned_offset, aligned_length))
174    }
175
176    async unsafe fn discard_to_zero_unsafe(
177        &self,
178        offset: u64,
179        length: u64,
180    ) -> io::Result<(u64, u64)> {
181        self.ensure_zero_mapping(offset, length).await
182    }
183
184    async unsafe fn discard_to_any_unsafe(
185        &self,
186        offset: u64,
187        length: u64,
188    ) -> io::Result<(u64, u64)> {
189        let discard_align = self.inner.discard_align();
190        assert!(discard_align.is_power_of_two());
191
192        let discard_align_mask = discard_align as u64 - 1;
193
194        let aligned_end = (offset + length) & !discard_align_mask;
195        let aligned_offset = (offset + discard_align_mask) & !discard_align_mask;
196        let aligned_length = aligned_end.saturating_sub(aligned_offset);
197        if aligned_length == 0 {
198            return Ok((aligned_offset, 0));
199        }
200
201        self.inner.discard(aligned_offset, aligned_length).await?;
202        Ok((aligned_offset, aligned_length))
203    }
204
205    async unsafe fn discard_to_backing_unsafe(
206        &self,
207        offset: u64,
208        length: u64,
209    ) -> io::Result<(u64, u64)> {
210        unsafe { self.discard_to_zero_unsafe(offset, length).await }
211    }
212
213    async fn flush(&self) -> io::Result<()> {
214        // No internal buffers to flush
215        self.inner.flush().await
216    }
217
218    async fn sync(&self) -> io::Result<()> {
219        self.inner.sync().await
220    }
221
222    async unsafe fn invalidate_cache(&self) -> io::Result<()> {
223        // No internal buffers to drop
224        // Safe: Caller says we should do this
225        unsafe { self.inner.invalidate_cache() }.await
226    }
227
228    async fn resize_grow(
229        &self,
230        new_size: u64,
231        format_prealloc_mode: PreallocateMode,
232    ) -> io::Result<()> {
233        #[allow(deprecated)]
234        if self
235            .size
236            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
237                (new_size > old).then_some(new_size)
238            })
239            .is_err()
240        {
241            return Ok(()); // only grow, else do nothing
242        }
243
244        let storage_prealloc_mode = match format_prealloc_mode {
245            PreallocateMode::None => storage::PreallocateMode::None,
246            PreallocateMode::Zero | PreallocateMode::FormatAllocate => {
247                storage::PreallocateMode::Zero
248            }
249            PreallocateMode::FullAllocate => storage::PreallocateMode::Allocate,
250            PreallocateMode::WriteData => storage::PreallocateMode::WriteData,
251        };
252        self.inner.resize(new_size, storage_prealloc_mode).await
253    }
254
255    async fn resize_shrink(&mut self, new_size: u64) -> io::Result<()> {
256        #[allow(deprecated)]
257        if self
258            .size
259            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
260                (new_size < old).then_some(new_size)
261            })
262            .is_err()
263        {
264            return Ok(()); // only shrink, else do nothing
265        }
266
267        self.inner
268            .resize(new_size, storage::PreallocateMode::None)
269            .await
270    }
271}
272
273impl<S: Storage + 'static> Display for Raw<S> {
274    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
275        write!(f, "raw[{}]", self.inner)
276    }
277}
278
279/// Options builder for opening a raw image.
280pub struct RawOpenBuilder<S: Storage + 'static>(FormatDriverBuilderBase<S>);
281
282#[maybe_async(AFIT)]
283impl<S: Storage + 'static> FormatDriverBuilder<S> for RawOpenBuilder<S> {
284    type Format = Raw<S>;
285    const FORMAT: Format = Format::Raw;
286
287    fn new(image: S) -> Self {
288        RawOpenBuilder(FormatDriverBuilderBase::new(image))
289    }
290
291    fn new_path<P: AsRef<Path>>(path: P) -> Self {
292        RawOpenBuilder(FormatDriverBuilderBase::new_path(path))
293    }
294
295    fn write(mut self, writable: bool) -> Self {
296        self.0.set_write(writable);
297        self
298    }
299
300    fn storage_open_options(mut self, options: StorageOpenOptions) -> Self {
301        self.0.set_storage_open_options(options);
302        self
303    }
304
305    async fn open<G: ImplicitOpenGate<S>>(self, mut gate: G) -> io::Result<Self::Format> {
306        let writable = self.0.get_writable();
307        let file = self.0.open_image(&mut gate).await?;
308        Raw::open_image(file, writable).await
309    }
310
311    fn get_image_path(&self) -> Option<PathBuf> {
312        self.0.get_image_path()
313    }
314
315    fn get_writable(&self) -> bool {
316        self.0.get_writable()
317    }
318
319    fn get_storage_open_options(&self) -> Option<&StorageOpenOptions> {
320        self.0.get_storage_opts()
321    }
322}
323
324/// Creation builder for a new raw image.
325pub struct RawCreateBuilder<S: Storage + 'static>(FormatCreateBuilderBase<S>);
326
327#[maybe_async(AFIT)]
328impl<S: Storage + 'static> FormatCreateBuilder<S> for RawCreateBuilder<S> {
329    const FORMAT: Format = Format::Raw;
330    type DriverBuilder = RawOpenBuilder<S>;
331
332    fn new(image: S) -> Self {
333        RawCreateBuilder(FormatCreateBuilderBase::new(image))
334    }
335
336    fn size(mut self, size: u64) -> Self {
337        self.0.set_size(size);
338        self
339    }
340
341    fn preallocate(mut self, prealloc_mode: PreallocateMode) -> Self {
342        self.0.set_preallocate(prealloc_mode);
343        self
344    }
345
346    fn get_size(&self) -> u64 {
347        self.0.get_size()
348    }
349
350    fn get_preallocate(&self) -> PreallocateMode {
351        self.0.get_preallocate()
352    }
353
354    async fn create(self) -> io::Result<()> {
355        self.create_open(DenyImplicitOpenGate::default(), |image| {
356            Ok(Raw::builder(image))
357        })
358        .await?;
359        Ok(())
360    }
361
362    async fn create_open<
363        G: ImplicitOpenGate<S>,
364        F: FnOnce(S) -> io::Result<Self::DriverBuilder>,
365    >(
366        self,
367        open_gate: G,
368        open_builder_fn: F,
369    ) -> io::Result<Raw<S>> {
370        let size = self.0.get_size();
371        let prealloc = self.0.get_preallocate();
372        let image = self.0.get_image();
373
374        // Clear of data (and allow for full preallocation, if requested)
375        if image.size()? > 0 {
376            image.resize(size, storage::PreallocateMode::None).await?;
377        }
378
379        let img = open_builder_fn(image)?.write(true).open(open_gate).await?;
380        if size > 0 {
381            img.resize_grow(size, prealloc).await?;
382        }
383
384        Ok(img)
385    }
386}
imago/raw.rs

imago/
raw.rs
