imago/macros/
on_disk_struct.rs

1//! Macro to allow binary structures be read/stored from/to disk.
2
3use std::io;
4use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
5
6/// A structure to be loaded from or stored on disk (e.g. qcow2 image header), in binary format.
7pub(crate) trait OnDiskStruct: Sized {
8    /// The on-disk size of this structure, in bytes
9    const ON_DISK_SIZE: usize;
10
11    /// Convenience alias for [`Self::ON_DISK_SIZE`] for objects
12    fn on_disk_size(&self) -> usize {
13        Self::ON_DISK_SIZE
14    }
15
16    /// Load this structure from the given slice.
17    ///
18    /// Use little endian, unless the structure has a specific inherent endianness.
19    #[allow(dead_code)]
20    fn load_from_le(bytes: &[u8]) -> io::Result<Self> {
21        Self::load_from(bytes)
22    }
23
24    /// Store this structure in the given slice.
25    ///
26    /// Use little endian, unless the structure has a specific inherent endianness.
27    #[allow(dead_code)]
28    fn store_to_le(&self, bytes: &mut [u8]) -> io::Result<()> {
29        self.store_to(bytes)
30    }
31
32    /// Load this structure from the given slice.
33    ///
34    /// Use big endian, unless the structure has a specific inherent endianness.
35    fn load_from_be(bytes: &[u8]) -> io::Result<Self> {
36        Self::load_from(bytes)
37    }
38
39    /// Store this structure in the given slice.
40    ///
41    /// Use big endian, unless the structure has a specific inherent endianness.
42    fn store_to_be(&self, bytes: &mut [u8]) -> io::Result<()> {
43        self.store_to(bytes)
44    }
45
46    /// Load this structure from the given slice, in default inherent endianness.
47    fn load_from(bytes: &[u8]) -> io::Result<Self>;
48
49    /// Store this structure in the given slice, in default inherent endianness.
50    fn store_to(&self, bytes: &mut [u8]) -> io::Result<()>;
51}
52
53/// Implement [`OnDiskStruct`] for a type wrapping a primitive type.
54///
55/// First argument: The wrapping type, e.g. `AtomicU16`; second argument: The inner primitive type
56/// (e.g. [`std::sync::atomic::AtomicPrimitive::Storage`], which just is not yet stable); third
57/// argument: The wrapping function (for loading); fourth argument: The unwrapping function (for
58/// storing).
59macro_rules! impl_on_disk_struct_wrapped_primitive {
60    ($type:ty, $raw_type:ty, $wrap:expr, $unwrap:expr) => {
61        impl $crate::macros::on_disk_struct::OnDiskStruct for $type {
62            const ON_DISK_SIZE: usize = std::mem::size_of::<$raw_type>();
63
64            fn load_from_le(bytes: &[u8]) -> std::io::Result<Self> {
65                Ok($wrap(<$raw_type>::load_from_le(bytes)?))
66            }
67
68            fn store_to_le(&self, bytes: &mut [u8]) -> std::io::Result<()> {
69                let raw = $unwrap(self);
70                raw.store_to_le(bytes)
71            }
72
73            fn load_from_be(bytes: &[u8]) -> std::io::Result<Self> {
74                Ok($wrap(<$raw_type>::load_from_be(bytes)?))
75            }
76
77            fn store_to_be(&self, bytes: &mut [u8]) -> std::io::Result<()> {
78                let raw = $unwrap(self);
79                raw.store_to_be(bytes)
80            }
81
82            fn load_from(bytes: &[u8]) -> std::io::Result<Self> {
83                Ok($wrap(<$raw_type>::load_from(bytes)?))
84            }
85
86            fn store_to(&self, bytes: &mut [u8]) -> std::io::Result<()> {
87                let raw = $unwrap(self);
88                raw.store_to(bytes)
89            }
90        }
91    };
92}
93
94/// Implement [`OnDiskStruct`] for a primitive integer type.
95macro_rules! impl_on_disk_struct_primitive {
96    ($type:ty) => {
97        impl $crate::macros::on_disk_struct::OnDiskStruct for $type {
98            const ON_DISK_SIZE: usize = std::mem::size_of::<$type>();
99
100            fn load_from_le(bytes: &[u8]) -> std::io::Result<Self> {
101                Ok(<$type>::from_le_bytes(bytes.try_into().map_err(|_| {
102                    $crate::misc_helpers::invalid_data(format!(
103                        "Cannot load type {} (length {}) from buffer of length {}",
104                        std::any::type_name::<$type>(),
105                        Self::ON_DISK_SIZE,
106                        bytes.len(),
107                    ))
108                })?))
109            }
110
111            fn store_to_le(&self, bytes: &mut [u8]) -> std::io::Result<()> {
112                if bytes.len() != Self::ON_DISK_SIZE {
113                    return Err($crate::misc_helpers::invalid_data(format!(
114                        "Cannot write type {} (length {}) into buffer of length {}",
115                        std::any::type_name::<$type>(),
116                        Self::ON_DISK_SIZE,
117                        bytes.len(),
118                    )));
119                }
120
121                bytes.copy_from_slice(&self.to_le_bytes());
122                Ok(())
123            }
124
125            fn load_from_be(bytes: &[u8]) -> std::io::Result<Self> {
126                Ok(<$type>::from_be_bytes(bytes.try_into().map_err(|_| {
127                    $crate::misc_helpers::invalid_data(format!(
128                        "Cannot load type {} (length {}) from buffer of length {}",
129                        std::any::type_name::<$type>(),
130                        Self::ON_DISK_SIZE,
131                        bytes.len(),
132                    ))
133                })?))
134            }
135
136            fn store_to_be(&self, bytes: &mut [u8]) -> std::io::Result<()> {
137                if bytes.len() != Self::ON_DISK_SIZE {
138                    return Err($crate::misc_helpers::invalid_data(format!(
139                        "Cannot write type {} (length {}) into buffer of length {}",
140                        std::any::type_name::<$type>(),
141                        Self::ON_DISK_SIZE,
142                        bytes.len(),
143                    )));
144                }
145
146                bytes.copy_from_slice(&self.to_be_bytes());
147                Ok(())
148            }
149
150            fn load_from(bytes: &[u8]) -> std::io::Result<Self> {
151                Ok(<$type>::from_ne_bytes(bytes.try_into().map_err(|_| {
152                    $crate::misc_helpers::invalid_data(format!(
153                        "Cannot load type {} (length {}) from buffer of length {}",
154                        std::any::type_name::<$type>(),
155                        Self::ON_DISK_SIZE,
156                        bytes.len(),
157                    ))
158                })?))
159            }
160
161            fn store_to(&self, bytes: &mut [u8]) -> std::io::Result<()> {
162                if bytes.len() != Self::ON_DISK_SIZE {
163                    return Err($crate::misc_helpers::invalid_data(format!(
164                        "Cannot write type {} (length {}) into buffer of length {}",
165                        std::any::type_name::<$type>(),
166                        Self::ON_DISK_SIZE,
167                        bytes.len(),
168                    )));
169                }
170
171                bytes.copy_from_slice(&self.to_ne_bytes());
172                Ok(())
173            }
174        }
175    };
176}
177
178impl_on_disk_struct_primitive!(u16);
179impl_on_disk_struct_primitive!(u32);
180impl_on_disk_struct_primitive!(u64);
181impl_on_disk_struct_wrapped_primitive!(AtomicU32, u32, Into::into, |x: &AtomicU32| x
182    .load(Ordering::Relaxed));
183impl_on_disk_struct_wrapped_primitive!(AtomicU64, u64, Into::into, |x: &AtomicU64| x
184    .load(Ordering::Relaxed));
185
186/// Implement [`OnDiskStruct`] for the contained `struct` definition.
187///
188/// The struct name must be followed by a specification of endianness and whether to allow gaps in
189/// the field offsets, i.e.: `struct <Struct>/<LE, BE, NE>, <no_gaps, allow_gaps>`.
190///
191/// All field types must be annotated with their byte offset in the structure, e.g. `foo: u32[42]`.
192macro_rules! on_disk_struct {
193    (
194        $(#[$attr:meta])*
195        struct $struct_name:ident/$endianness:ident, $packed:ident {
196            $(
197                $(#[$id_attr:meta])*
198                $identifier:ident: $type:ty[$offset:literal],
199            )+
200        }
201    ) => {
202        $(#[$attr])*
203        struct $struct_name {
204            $(
205                $(#[$id_attr])*
206                $identifier: $type,
207            )+
208        }
209
210        // Verify strict field ordering
211        const _: () = const {
212            let mut next = 0;
213            $(
214                $crate::macros::on_disk_struct::on_disk_struct_helper!(@check_layout $packed, $offset, next);
215                next = $offset + <$type>::ON_DISK_SIZE;
216            )+
217            let _ = next;
218        };
219
220        impl $crate::macros::on_disk_struct::OnDiskStruct for $struct_name {
221            const ON_DISK_SIZE: usize = $crate::macros::last_element!($($offset + <$type>::ON_DISK_SIZE),+);
222
223            fn load_from(bytes: &[u8]) -> std::io::Result<Self> {
224                if bytes.len() < Self::ON_DISK_SIZE {
225                    return Err($crate::misc_helpers::invalid_data(format!(
226                        "Cannot read struct {} (length {}) from buffer of length {}",
227                        std::any::type_name::<$struct_name>(),
228                        Self::ON_DISK_SIZE,
229                        bytes.len(),
230                    )));
231                }
232
233                Ok($struct_name {
234                    $(
235                        $identifier: $crate::macros::on_disk_struct::on_disk_struct_helper!(
236                            @load $endianness,
237                            $type,
238                            &bytes[$offset..($offset + <$type>::ON_DISK_SIZE)]
239                        )?,
240                    )+
241                })
242            }
243
244            fn store_to(&self, bytes: &mut [u8]) -> std::io::Result<()> {
245                if bytes.len() < Self::ON_DISK_SIZE {
246                    return Err($crate::misc_helpers::invalid_data(format!(
247                        "Cannot write struct {} (length {}) into buffer of length {}",
248                        std::any::type_name::<$struct_name>(),
249                        Self::ON_DISK_SIZE,
250                        bytes.len(),
251                    )));
252                }
253
254                $(
255                    $crate::macros::on_disk_struct::on_disk_struct_helper!(
256                        @store $endianness,
257                        self.$identifier,
258                        &mut bytes[$offset..($offset + <$type>::ON_DISK_SIZE)]
259                    )?;
260                )+
261
262                Ok(())
263            }
264        }
265    }
266}
267
268pub(crate) use on_disk_struct;
269
270/// Helper macro for [`on_disk_struct`].
271///
272/// Various functionalities, selected by the first parameter.
273macro_rules! on_disk_struct_helper {
274    // Check a field offset against the actual offset within the struct, either allowing gaps or not
275    (@check_layout no_gaps, $field_offset:literal, $actual_offset:expr) => {
276        assert!($field_offset == $actual_offset);
277    };
278    (@check_layout allow_gaps, $field_offset:literal, $actual_offset:expr) => {
279        assert!($field_offset >= $actual_offset);
280    };
281
282    // Load with an endianness specified
283    (@load NE, $type:ty, $slice:expr) => {
284        <$type>::load_from($slice)
285    };
286    (@load LE, $type:ty, $slice:expr) => {
287        <$type>::load_from_le($slice)
288    };
289    (@load BE, $type:ty, $slice:expr) => {
290        <$type>::load_from_be($slice)
291    };
292
293    // Store with an endianness specified
294    (@store NE, $value:expr, $slice:expr) => {
295        $value.store_to($slice)
296    };
297    (@store LE, $value:expr, $slice:expr) => {
298        $value.store_to_le($slice)
299    };
300    (@store BE, $value:expr, $slice:expr) => {
301        $value.store_to_be($slice)
302    };
303}
304
305pub(crate) use on_disk_struct_helper;
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use std::sync::atomic::Ordering;
311
312    // -- Primitive round-trips --
313
314    #[test]
315    fn u16_be_round_trip() {
316        let val: u16 = 0x0102;
317        let mut buf = [0u8; 2];
318        val.store_to_be(&mut buf).unwrap();
319        assert_eq!(buf, [0x01, 0x02]);
320        assert_eq!(u16::load_from_be(&buf).unwrap(), val);
321    }
322
323    #[test]
324    fn u16_le_round_trip() {
325        let val: u16 = 0x0102;
326        let mut buf = [0u8; 2];
327        val.store_to_le(&mut buf).unwrap();
328        assert_eq!(buf, [0x02, 0x01]);
329        assert_eq!(u16::load_from_le(&buf).unwrap(), val);
330    }
331
332    #[test]
333    fn u32_be_round_trip() {
334        let val: u32 = 0x01020304;
335        let mut buf = [0u8; 4];
336        val.store_to_be(&mut buf).unwrap();
337        assert_eq!(buf, [0x01, 0x02, 0x03, 0x04]);
338        assert_eq!(u32::load_from_be(&buf).unwrap(), val);
339    }
340
341    #[test]
342    fn u32_le_round_trip() {
343        let val: u32 = 0x01020304;
344        let mut buf = [0u8; 4];
345        val.store_to_le(&mut buf).unwrap();
346        assert_eq!(buf, [0x04, 0x03, 0x02, 0x01]);
347        assert_eq!(u32::load_from_le(&buf).unwrap(), val);
348    }
349
350    #[test]
351    fn u64_be_round_trip() {
352        let val: u64 = 0x0102030405060708;
353        let mut buf = [0u8; 8];
354        val.store_to_be(&mut buf).unwrap();
355        assert_eq!(buf, [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]);
356        assert_eq!(u64::load_from_be(&buf).unwrap(), val);
357    }
358
359    #[test]
360    fn u64_le_round_trip() {
361        let val: u64 = 0x0102030405060708;
362        let mut buf = [0u8; 8];
363        val.store_to_le(&mut buf).unwrap();
364        assert_eq!(buf, [0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01]);
365        assert_eq!(u64::load_from_le(&buf).unwrap(), val);
366    }
367
368    // -- Primitive error cases --
369
370    #[test]
371    fn primitive_store_wrong_size() {
372        let val: u32 = 42;
373        assert!(val.store_to_be(&mut [0u8; 3]).is_err());
374        assert!(val.store_to_be(&mut [0u8; 5]).is_err());
375        assert!(val.store_to_le(&mut [0u8; 0]).is_err());
376    }
377
378    #[test]
379    fn primitive_load_wrong_size() {
380        assert!(u32::load_from_be(&[0u8; 3]).is_err());
381        assert!(u32::load_from_be(&[0u8; 5]).is_err());
382        assert!(u64::load_from_le(&[0u8; 7]).is_err());
383    }
384
385    // -- Atomic round-trips --
386
387    #[test]
388    fn atomic_u32_be_round_trip() {
389        let val = AtomicU32::new(0xdeadbeef);
390        let mut buf = [0u8; 4];
391        val.store_to_be(&mut buf).unwrap();
392        assert_eq!(buf, [0xde, 0xad, 0xbe, 0xef]);
393        let loaded = AtomicU32::load_from_be(&buf).unwrap();
394        assert_eq!(loaded.load(Ordering::Relaxed), 0xdeadbeef);
395    }
396
397    #[test]
398    fn atomic_u64_le_round_trip() {
399        let val = AtomicU64::new(0x0102030405060708);
400        let mut buf = [0u8; 8];
401        val.store_to_le(&mut buf).unwrap();
402        assert_eq!(buf, [0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01]);
403        let loaded = AtomicU64::load_from_le(&buf).unwrap();
404        assert_eq!(loaded.load(Ordering::Relaxed), 0x0102030405060708);
405    }
406
407    // -- on_disk_struct: big-endian, no gaps --
408
409    on_disk_struct! {
410        struct TestBE/BE, no_gaps {
411            a: u32[0],
412            b: u64[4],
413            c: u16[12],
414        }
415    }
416
417    #[test]
418    fn be_struct_on_disk_size() {
419        assert_eq!(TestBE::ON_DISK_SIZE, 14);
420    }
421
422    #[test]
423    fn be_struct_round_trip() {
424        let s = TestBE {
425            a: 0x01020304,
426            b: 0x0506070809101112,
427            c: 0x1314,
428        };
429        let mut buf = [0u8; 14];
430        s.store_to(&mut buf).unwrap();
431
432        // Verify big-endian byte layout
433        assert_eq!(&buf[0..4], &[0x01, 0x02, 0x03, 0x04]);
434        assert_eq!(
435            &buf[4..12],
436            &[0x05, 0x06, 0x07, 0x08, 0x09, 0x10, 0x11, 0x12]
437        );
438        assert_eq!(&buf[12..14], &[0x13, 0x14]);
439
440        let loaded = TestBE::load_from(&buf).unwrap();
441        assert_eq!(loaded.a, s.a);
442        assert_eq!(loaded.b, s.b);
443        assert_eq!(loaded.c, s.c);
444    }
445
446    #[test]
447    fn be_struct_load_from_larger_buffer() {
448        let mut buf = [0u8; 20];
449        buf[0..4].copy_from_slice(&0x11223344u32.to_be_bytes());
450        let loaded = TestBE::load_from(&buf).unwrap();
451        assert_eq!(loaded.a, 0x11223344);
452    }
453
454    #[test]
455    fn be_struct_load_too_short() {
456        assert!(TestBE::load_from(&[0u8; 13]).is_err());
457        assert!(TestBE::load_from(&[0u8; 0]).is_err());
458    }
459
460    #[test]
461    fn be_struct_store_too_short() {
462        let s = TestBE { a: 0, b: 0, c: 0 };
463        assert!(s.store_to(&mut [0u8; 13]).is_err());
464    }
465
466    // -- on_disk_struct: little-endian, no gaps --
467
468    on_disk_struct! {
469        struct TestLE/LE, no_gaps {
470            x: u32[0],
471            y: u32[4],
472        }
473    }
474
475    #[test]
476    fn le_struct_round_trip() {
477        let s = TestLE {
478            x: 0x01020304,
479            y: 0x05060708,
480        };
481        let mut buf = [0u8; 8];
482        s.store_to(&mut buf).unwrap();
483
484        // Verify little-endian byte layout
485        assert_eq!(&buf[0..4], &[0x04, 0x03, 0x02, 0x01]);
486        assert_eq!(&buf[4..8], &[0x08, 0x07, 0x06, 0x05]);
487
488        let loaded = TestLE::load_from(&buf).unwrap();
489        assert_eq!(loaded.x, s.x);
490        assert_eq!(loaded.y, s.y);
491    }
492
493    // -- on_disk_struct: with atomics --
494
495    on_disk_struct! {
496        struct TestAtomics/BE, no_gaps {
497            plain: u32[0],
498            atomic32: AtomicU32[4],
499            atomic64: AtomicU64[8],
500        }
501    }
502
503    #[test]
504    fn atomic_struct_on_disk_size() {
505        assert_eq!(TestAtomics::ON_DISK_SIZE, 16);
506    }
507
508    #[test]
509    fn atomic_struct_round_trip() {
510        let s = TestAtomics {
511            plain: 0xaaaaaaaa,
512            atomic32: AtomicU32::new(0xbbbbbbbb),
513            atomic64: AtomicU64::new(0xccccccccdddddddd),
514        };
515        let mut buf = [0u8; 16];
516        s.store_to(&mut buf).unwrap();
517
518        assert_eq!(&buf[0..4], &[0xaa, 0xaa, 0xaa, 0xaa]);
519        assert_eq!(&buf[4..8], &[0xbb, 0xbb, 0xbb, 0xbb]);
520        assert_eq!(
521            &buf[8..16],
522            &[0xcc, 0xcc, 0xcc, 0xcc, 0xdd, 0xdd, 0xdd, 0xdd]
523        );
524
525        let loaded = TestAtomics::load_from(&buf).unwrap();
526        assert_eq!(loaded.plain, 0xaaaaaaaa);
527        assert_eq!(loaded.atomic32.load(Ordering::Relaxed), 0xbbbbbbbb);
528        assert_eq!(loaded.atomic64.load(Ordering::Relaxed), 0xccccccccdddddddd);
529    }
530
531    // -- on_disk_struct: allow_gaps --
532
533    on_disk_struct! {
534        struct TestGaps/BE, allow_gaps {
535            first: u32[0],
536            // 4-byte gap at offset 4
537            second: u16[8],
538        }
539    }
540
541    #[test]
542    fn gaps_struct_on_disk_size() {
543        assert_eq!(TestGaps::ON_DISK_SIZE, 10);
544    }
545
546    #[test]
547    fn gaps_struct_round_trip() {
548        let s = TestGaps {
549            first: 0x11223344,
550            second: 0x5566,
551        };
552        let mut buf = [0xffu8; 10];
553        s.store_to(&mut buf).unwrap();
554
555        assert_eq!(&buf[0..4], &[0x11, 0x22, 0x33, 0x44]);
556        // Gap bytes at [4..8] are untouched
557        assert_eq!(&buf[4..8], &[0xff, 0xff, 0xff, 0xff]);
558        assert_eq!(&buf[8..10], &[0x55, 0x66]);
559
560        // Round-trip ignores gap bytes
561        let loaded = TestGaps::load_from(&buf).unwrap();
562        assert_eq!(loaded.first, 0x11223344);
563        assert_eq!(loaded.second, 0x5566);
564    }
565
566    // -- on_disk_struct: single field --
567
568    on_disk_struct! {
569        struct TestSingle/BE, no_gaps {
570            only: u64[0],
571        }
572    }
573
574    #[test]
575    fn single_field_struct() {
576        assert_eq!(TestSingle::ON_DISK_SIZE, 8);
577        let s = TestSingle {
578            only: 0x0102030405060708,
579        };
580        let mut buf = [0u8; 8];
581        s.store_to(&mut buf).unwrap();
582        assert_eq!(buf, [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]);
583        let loaded = TestSingle::load_from(&buf).unwrap();
584        assert_eq!(loaded.only, s.only);
585    }
586
587    // -- on_disk_struct: inherent endianness overrides le/be calls --
588
589    #[test]
590    fn be_struct_inherent_endianness() {
591        let s = TestBE {
592            a: 0x01020304,
593            b: 0,
594            c: 0,
595        };
596        let mut buf_via_store = [0u8; 14];
597        let mut buf_via_be = [0u8; 14];
598        let mut buf_via_le = [0u8; 14];
599
600        s.store_to(&mut buf_via_store).unwrap();
601        s.store_to_be(&mut buf_via_be).unwrap();
602        s.store_to_le(&mut buf_via_le).unwrap();
603
604        // All three should produce identical output (inherent BE endianness)
605        assert_eq!(buf_via_store, buf_via_be);
606        assert_eq!(buf_via_store, buf_via_le);
607    }
608
609    // -- on_disk_size convenience method --
610
611    #[test]
612    fn on_disk_size_method() {
613        let s = TestBE { a: 0, b: 0, c: 0 };
614        assert_eq!(s.on_disk_size(), TestBE::ON_DISK_SIZE);
615        assert_eq!(s.on_disk_size(), 14);
616    }
617
618    // -- Known byte pattern: decode hand-crafted bytes --
619
620    #[test]
621    fn decode_known_be_bytes() {
622        // Hand-crafted big-endian buffer
623        let buf: [u8; 14] = [
624            0x00, 0x00, 0x00, 0x01, // a = 1
625            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, // b = 255
626            0x01, 0x00, // c = 256
627        ];
628        let s = TestBE::load_from(&buf).unwrap();
629        assert_eq!(s.a, 1);
630        assert_eq!(s.b, 255);
631        assert_eq!(s.c, 256);
632    }
633}
imago/macros/on_disk_struct.rs

imago/macros/
on_disk_struct.rs
