mbrman___bitvec_1_0/
lib.rs

1//! A library that allows managing MBR partition tables.
2//!
3//! ## Features
4//!
5//!  *  Create primary partitions and logical volumes
6//!  *  Delete primary partitions and logical volumes
7//!  *  Automatically generate logical volume's EBR (or can be provided manually)
8//!  *  If the disk geometry is set, the partition CHS addresses will be calculated
9//!     automatically when writing to disk
10//!
11//! ## Examples
12//!
13//! ### Read all the partitions of a disk
14//!
15//! ```rust
16//! let mut f = std::fs::File::open("tests/fixtures/disk1.img")
17//!     .expect("could not open disk");
18//! let mbr = mbrman::MBR::read_from(&mut f, 512)
19//!     .expect("could not find MBR");
20//!
21//! println!("Disk signature: {:?}", mbr.header.disk_signature);
22//!
23//! for (i, p) in mbr.iter() {
24//!     // NOTE: The first four partitions are always provided by iter()
25//!     if p.is_used() {
26//!         println!("Partition #{}: type = {:?}, size = {} bytes, starting lba = {}",
27//!             i,
28//!             p.sys,
29//!             p.sectors * mbr.sector_size,
30//!             p.starting_lba);
31//!     }
32//! }
33//! ```
34//!
35//! ### Create and delete primary partitions
36//!
37//! ```rust
38//! let mut f = std::fs::File::open("tests/fixtures/disk1.img")
39//!     .expect("could not open disk");
40//! let mut mbr = mbrman::MBR::read_from(&mut f, 512)
41//!     .expect("could not find MBR");
42//!
43//! let free_partition_number = mbr.iter().find(|(i, p)| p.is_unused()).map(|(i, _)| i)
44//!     .expect("no more places available");
45//! let sectors = mbr.get_maximum_partition_size()
46//!     .expect("no more space available");
47//! let starting_lba = mbr.find_optimal_place(sectors)
48//!     .expect("could not find a place to put the partition");
49//!
50//! mbr[free_partition_number] = mbrman::MBRPartitionEntry {
51//!     boot: mbrman::BOOT_INACTIVE,        // boot flag
52//!     first_chs: mbrman::CHS::empty(),    // first CHS address (only useful for old computers)
53//!     sys: 0x83,                          // Linux filesystem
54//!     last_chs: mbrman::CHS::empty(),     // last CHS address (only useful for old computers)
55//!     starting_lba,                       // the sector where the partition starts
56//!     sectors,                            // the number of sectors in that partition
57//! };
58//!
59//! mbr[free_partition_number] = mbrman::MBRPartitionEntry::empty();
60//!
61//! // NOTE: no modification is committed to the disk until we call mbr.write_into()
62//! ```
63//!
64//! ### Create a new partition table from an empty disk
65//!
66//! ```rust
67//! let ss = 512; // sector size
68//! let data = vec![0; 100 * ss as usize];
69//! let mut cur = std::io::Cursor::new(data);
70//!
71//! let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
72//!     .expect("could not create partition table");
73//!
74//! // NOTE: commit the change to the in-memory buffer
75//! mbr.write_into(&mut cur);
76//! ```
77//!
78//! ### Add a new logical volume to the disk
79//!
80//! ```rust
81//! let ss = 512; // sector size
82//! let data = vec![0; 100 * ss as usize];
83//! let mut cur = std::io::Cursor::new(data);
84//!
85//! let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
86//!     .expect("could not create partition table");
87//!
88//! mbr[1] = mbrman::MBRPartitionEntry {
89//!     boot: mbrman::BOOT_INACTIVE,        // boot flag
90//!     first_chs: mbrman::CHS::empty(),    // first CHS address (only useful for old computers)
91//!     sys: 0x0f,                          // extended partition with LBA
92//!     last_chs: mbrman::CHS::empty(),     // last CHS address (only useful for old computers)
93//!     starting_lba: 1,                    // the sector where the partition starts
94//!     sectors: mbr.disk_size - 1,         // the number of sectors in that partition
95//! };
96//!
97//! // this helper function will do all the hard work for you
98//! // here it creates a logical volume with Linux filesystem that occupies the entire disk
99//! // NOTE: you will lose 1 sector because it is used by the EBR
100//! mbr.push(0x83, 1, mbr.disk_size - 1);
101//!
102//! // NOTE: commit the change to the in-memory buffer
103//! mbr.write_into(&mut cur);
104//! ```
105//!
106//! ### Add a new logical volume manually to the disk
107//!
108//! This is useful only if you need to specify exactly where goes the EBR and the partition itself.
109//!
110//! ```rust
111//! let ss = 512; // sector size
112//! let data = vec![0; 100 * ss as usize];
113//! let mut cur = std::io::Cursor::new(data);
114//!
115//! let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
116//!     .expect("could not create partition table");
117//!
118//! mbr[1] = mbrman::MBRPartitionEntry {
119//!     boot: mbrman::BOOT_INACTIVE,        // boot flag
120//!     first_chs: mbrman::CHS::empty(),    // first CHS address (only useful for old computers)
121//!     sys: 0x0f,                          // extended partition with LBA
122//!     last_chs: mbrman::CHS::empty(),     // last CHS address (only useful for old computers)
123//!     starting_lba: 1,                    // the sector where the partition starts
124//!     sectors: mbr.disk_size - 1,         // the number of sectors in that partition
125//! };
126//!
127//! // NOTE: mbrman won't check the consistency of the partition you have created manually
128//! mbr.logical_partitions.push(
129//!     mbrman::LogicalPartition {
130//!         // this is the actual partition entry for the logical volume
131//!         partition: mbrman::MBRPartitionEntry {
132//!             boot: mbrman::BOOT_INACTIVE,
133//!             first_chs: mbrman::CHS::empty(),
134//!             sys: 0x83,
135//!             last_chs: mbrman::CHS::empty(),
136//!             starting_lba: 2,                    // the sector index 1 is used by the EBR
137//!             sectors: mbr.disk_size - 2,
138//!         },
139//!         // this is the absolute LBA address of the EBR
140//!         absolute_ebr_lba: 1,
141//!         // the number of sectors in the first EBR is never known
142//!         ebr_sectors: None,
143//!         // empty boot sector in the EBR
144//!         bootstrap_code: [0; 446],
145//!         // this is the absolute CHS address of the EBR (only used by old computers)
146//!         ebr_first_chs: mbrman::CHS::empty(),                // only for old computers
147//!         // this is the absolute CHS address of the last EBR (only used by old computers)
148//!         // NOTE: this is not know the first EBR
149//!         ebr_last_chs: None,
150//!     }
151//! );
152//!
153//! // NOTE: commit the change to the in-memory buffer
154//! mbr.write_into(&mut cur);
155//! ```
156
157#![deny(missing_docs)]
158
159use bincode::{deserialize_from, serialize_into};
160use bitvec::prelude::*;
161use serde::de::{SeqAccess, Visitor};
162use serde::ser::SerializeTuple;
163use serde::{Deserialize, Deserializer, Serialize, Serializer};
164use serde_big_array::BigArray;
165use std::convert::TryFrom;
166use std::io::{Read, Seek, SeekFrom, Write};
167use std::iter::{once, repeat};
168use std::ops::{Index, IndexMut};
169use thiserror::Error;
170
171const DEFAULT_ALIGN: u32 = 2048;
172const MAX_ALIGN: u32 = 16384;
173const FIRST_USABLE_LBA: u32 = 1;
174const BOOT_SIGNATURE: [u8; 2] = [0x55, 0xaa];
175
176/// Boot flag for a bootable partition
177pub const BOOT_ACTIVE: u8 = 0x80;
178/// Boot flag for a non-bootable partition
179pub const BOOT_INACTIVE: u8 = 0x00;
180
181/// The result of reading, writing or managing a MBR.
182pub type Result<T> = std::result::Result<T, Error>;
183
184/// An error
185#[derive(Debug, Error)]
186#[non_exhaustive]
187pub enum Error {
188    /// The CHS address requested cannot be represented in CHS
189    ///
190    /// # Remark
191    ///
192    /// There is a hard limit around 8GB for CHS addressing.
193    #[error("exceeded the maximum limit of CHS")]
194    LBAExceedsMaximumCHS,
195    /// The CHS address requested exceeds the number of cylinders in the disk
196    #[error("exceeded the maximum number of cylinders on disk")]
197    LBAExceedsMaximumCylinders,
198    /// Deserialization errors.
199    #[error("deserialization failed")]
200    Deserialize(#[from] bincode::Error),
201    /// I/O errors.
202    #[error("generic I/O error")]
203    Io(#[from] std::io::Error),
204    /// Inconsistent extended boot record
205    #[error("inconsistent extended boot record")]
206    InconsistentEBR,
207    /// No extended partition
208    #[error("no extended partition")]
209    NoExtendedPartition,
210    /// The EBR starts before the extended partition
211    #[error("EBR starts before the extended partition")]
212    EBRStartsBeforeExtendedPartition,
213    /// The EBR starts too close to the extended partition
214    #[error("EBR starts too close to the end of the extended partition")]
215    EBRStartsTooCloseToTheEndOfExtendedPartition,
216    /// The EBR ends after the extended partition
217    #[error("EBR ends after the extended partition")]
218    EBREndsAfterExtendedPartition,
219    /// Not enough sectors to create a logical partition
220    #[error("not enough sectors to create a logical partition")]
221    NotEnoughSectorsToCreateLogicalPartition,
222    /// An operation that required to find a partition, was unable to find that partition.
223    #[error("partition not found")]
224    PartitionNotFound,
225    /// An error that occurs when there is not enough space left on the table to continue.
226    #[error("no space left")]
227    NoSpaceLeft,
228    /// MBR doesn't have the expected signature value
229    #[error("invalid MBR signature")]
230    InvalidSignature,
231    /// Partition has invalid boot flag
232    #[error("partition has invalid boot flag")]
233    InvalidBootFlag,
234}
235
236/// A type representing a MBR partition table including its partition, the sector size of the disk
237/// and the alignment of the partitions to the sectors.
238///
239/// # Examples:
240/// Read an existing MBR on a reader and list its partitions:
241/// ```
242/// let mut f = std::fs::File::open("tests/fixtures/disk1.img")
243///     .expect("could not open disk");
244/// let mbr = mbrman::MBR::read_from(&mut f, 512)
245///     .expect("could not find MBR");
246///
247/// println!("Disk signature: {:?}", mbr.header.disk_signature);
248///
249/// for (i, p) in mbr.iter() {
250///     if p.is_used() {
251///         println!("Partition #{}: type = {:?}, size = {} bytes, starting lba = {}",
252///             i,
253///             p.sys,
254///             p.sectors * mbr.sector_size,
255///             p.starting_lba);
256///     }
257/// }
258/// ```
259#[derive(Debug, Clone, PartialEq, Eq)]
260pub struct MBR {
261    /// Sector size of the disk.
262    ///
263    /// You should not change this, otherwise the starting locations of your partitions will be
264    /// different in bytes.
265    pub sector_size: u32,
266    /// MBR partition header (disk GUID, first/last usable LBA, etc...)
267    pub header: MBRHeader,
268    /// A vector with all the logical partitions. You can push new ones (even empty ones)
269    pub logical_partitions: Vec<LogicalPartition>,
270    /// Partitions alignment (in sectors)
271    ///
272    /// This field change the behavior of the methods `get_maximum_partition_size()`,
273    /// `find_free_sectors()`, `find_first_place()`, `find_last_place()` and `find_optimal_place()`
274    /// so they return only values aligned to the alignment.
275    ///
276    /// # Panics
277    /// The value must be greater than 0, otherwise you will encounter divisions by zero.
278    pub align: u32,
279    /// Disk geometry: number of cylinders
280    pub cylinders: u16,
281    /// Disk geometry: number of heads
282    pub heads: u8,
283    /// Disk geometry: number of sectors
284    pub sectors: u8,
285    /// Disk size in sectors
286    pub disk_size: u32,
287}
288
289impl MBR {
290    /// Get an iterator over the partition entries and their index. The
291    /// index always starts at 1.
292    pub fn iter(&self) -> impl Iterator<Item = (usize, &MBRPartitionEntry)> {
293        self.header.iter().chain(
294            self.logical_partitions
295                .iter()
296                .map(|x| &x.partition)
297                .enumerate()
298                .map(|(i, x)| (i + 5, x)),
299        )
300    }
301
302    /// Get a mutable iterator over the partition entries and their
303    /// index. The index always starts at 1.
304    pub fn iter_mut(&mut self) -> impl Iterator<Item = (usize, &mut MBRPartitionEntry)> {
305        let mut partitions: Vec<_> = vec![
306            &mut self.header.partition_1,
307            &mut self.header.partition_2,
308            &mut self.header.partition_3,
309            &mut self.header.partition_4,
310        ];
311        partitions.extend(self.logical_partitions.iter_mut().map(|x| &mut x.partition));
312        partitions.into_iter().enumerate().map(|(i, x)| (i + 1, x))
313    }
314
315    /// Get `Some(&MBRPartitionEntry)` if it exists, None otherwise.
316    ///
317    /// # Remarks
318    ///
319    ///  -  The partitions start at index 1
320    ///  -  The first 4 partitions always exist
321    pub fn get(&self, i: usize) -> Option<&MBRPartitionEntry> {
322        match i {
323            0 => None,
324            1 => Some(&self.header.partition_1),
325            2 => Some(&self.header.partition_2),
326            3 => Some(&self.header.partition_3),
327            4 => Some(&self.header.partition_4),
328            i => self.logical_partitions.get(i - 5).map(|x| &x.partition),
329        }
330    }
331
332    /// Get `Some(&mut MBRPartitionEntry)` if it exists, None otherwise.
333    ///
334    /// # Remarks
335    ///
336    ///  -  The partitions start at index 1
337    ///  -  The first 4 partitions always exist
338    pub fn get_mut(&mut self, i: usize) -> Option<&mut MBRPartitionEntry> {
339        match i {
340            0 => None,
341            1 => Some(&mut self.header.partition_1),
342            2 => Some(&mut self.header.partition_2),
343            3 => Some(&mut self.header.partition_3),
344            4 => Some(&mut self.header.partition_4),
345            i => self
346                .logical_partitions
347                .get_mut(i - 5)
348                .map(|x| &mut x.partition),
349        }
350    }
351
352    /// The total number of partitions on the disk: primary partitions and logical partitions.
353    ///
354    /// # Remark
355    ///
356    /// The primary partitions are always counted even if they are empty.
357    pub fn len(&self) -> usize {
358        4 + self.logical_partitions.len()
359    }
360
361    /// Always false: primary partitions are always counted even if they are empty.
362    pub fn is_empty(&self) -> bool {
363        false
364    }
365
366    /// Make a new MBR
367    ///
368    /// # Examples
369    /// Basic usage:
370    /// ```
371    /// let mut f = std::fs::File::open("tests/fixtures/disk1.img")
372    ///     .expect("could not open disk");
373    /// let mbr = mbrman::MBR::new_from(&mut f, 512, [0x01, 0x02, 0x03, 0x04])
374    ///     .expect("could not make a partition table");
375    /// ```
376    pub fn new_from<S>(seeker: &mut S, sector_size: u32, disk_signature: [u8; 4]) -> Result<MBR>
377    where
378        S: Seek,
379    {
380        let disk_size = u32::try_from(seeker.seek(SeekFrom::End(0))? / u64::from(sector_size))
381            .unwrap_or(u32::max_value());
382        let header = MBRHeader::new(disk_signature);
383
384        Ok(MBR {
385            sector_size,
386            header,
387            logical_partitions: Vec::new(),
388            align: DEFAULT_ALIGN,
389            cylinders: 0,
390            heads: 0,
391            sectors: 0,
392            disk_size,
393        })
394    }
395
396    /// Read the MBR on a reader.
397    ///
398    /// # Examples
399    /// Basic usage:
400    /// ```
401    /// let mut f = std::fs::File::open("tests/fixtures/disk1.img")
402    ///     .expect("could not open disk");
403    /// let mbr = mbrman::MBR::read_from(&mut f, 512)
404    ///     .expect("could not read the partition table");
405    /// ```
406    pub fn read_from<R: ?Sized>(mut reader: &mut R, sector_size: u32) -> Result<MBR>
407    where
408        R: Read + Seek,
409    {
410        let disk_size = u32::try_from(reader.seek(SeekFrom::End(0))? / u64::from(sector_size))
411            .unwrap_or(u32::max_value());
412        let header = MBRHeader::read_from(&mut reader)?;
413
414        let mut logical_partitions = Vec::new();
415        if let Some(extended) = header.get_extended_partition() {
416            // NOTE: The number of sectors is an index field; thus, the zero
417            //       value is invalid, reserved and must not be used in normal
418            //       partition entries. The entry is used by operating systems
419            //       in certain circumstances; in such cases the CHS addresses
420            //       are ignored.
421            let mut relative_ebr_lba = 0;
422            let mut ebr_sectors = None;
423            let mut ebr_first_chs = extended.first_chs;
424            let mut ebr_last_chs = None;
425            loop {
426                reader.seek(SeekFrom::Start(u64::from(
427                    (extended.starting_lba + relative_ebr_lba) * sector_size,
428                )))?;
429                let (partition, next, bootstrap_code) = match EBRHeader::read_from(&mut reader) {
430                    Ok(ebr) => ebr.unwrap(),
431                    Err(err) => {
432                        if relative_ebr_lba == 0 {
433                            // NOTE: if the extended partition is empty, it is not required that an
434                            //       EBR exists
435                            break;
436                        } else {
437                            return Err(err);
438                        }
439                    }
440                };
441                let absolute_ebr_lba = extended.starting_lba + relative_ebr_lba;
442                logical_partitions.push(LogicalPartition {
443                    partition: MBRPartitionEntry {
444                        starting_lba: partition.starting_lba + absolute_ebr_lba,
445                        ..partition
446                    },
447                    absolute_ebr_lba,
448                    ebr_sectors,
449                    ebr_first_chs,
450                    ebr_last_chs,
451                    bootstrap_code,
452                });
453
454                if next.starting_lba > 0 && relative_ebr_lba >= next.starting_lba {
455                    return Err(Error::InconsistentEBR);
456                }
457
458                relative_ebr_lba = next.starting_lba;
459                ebr_sectors = Some(next.sectors);
460                ebr_first_chs = next.first_chs;
461                ebr_last_chs = Some(next.last_chs);
462
463                if relative_ebr_lba == 0 {
464                    break;
465                }
466            }
467        }
468
469        let align = MBR::find_alignment(&header, &logical_partitions);
470
471        Ok(MBR {
472            sector_size,
473            header,
474            logical_partitions,
475            align,
476            cylinders: 0,
477            heads: 0,
478            sectors: 0,
479            disk_size,
480        })
481    }
482
483    fn find_alignment(header: &MBRHeader, logical_partitions: &[LogicalPartition]) -> u32 {
484        let lbas = header
485            .iter()
486            .map(|(_, x)| x)
487            .chain(logical_partitions.iter().map(|x| &x.partition))
488            .filter(|x| x.is_used())
489            .map(|x| x.starting_lba)
490            .collect::<Vec<_>>();
491
492        if lbas.is_empty() {
493            return DEFAULT_ALIGN;
494        }
495
496        if lbas.len() == 1 && lbas[0] == FIRST_USABLE_LBA {
497            return FIRST_USABLE_LBA;
498        }
499
500        (1..=MAX_ALIGN.min(*lbas.iter().max().unwrap_or(&1)))
501            .filter(|div| lbas.iter().all(|x| x % div == 0))
502            .max()
503            .unwrap()
504    }
505
506    /// Return `true` if the MBR has a valid geometry. The geometry can be set by setting
507    /// the fiels `cylinders`, `heads` and `sectors`.
508    ///
509    /// Remarks
510    ///
511    /// The cylinders, heads and sectors must have a value greater than zero.
512    ///
513    /// The cylinders cannot exceed 1023.
514    ///
515    /// The sectors cannot exceed 63.
516    pub fn check_geometry(&self) -> bool {
517        self.cylinders > 0
518            && self.cylinders <= 1023
519            && self.heads > 0
520            && self.sectors > 0
521            && self.sectors <= 63
522    }
523
524    /// Write the MBR to a writer. This function will seek automatically in the writer.
525    /// This function will update the CHS address of the partitions automatically if a valid
526    /// geometry has been set. See `check_geometry`.
527    ///
528    /// # Examples
529    /// Basic usage:
530    /// ```
531    /// let ss = 512;
532    /// let data = vec![0; 100 * ss as usize];
533    /// let mut cur = std::io::Cursor::new(data);
534    /// let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
535    ///     .expect("could not make a partition table");
536    ///
537    /// // actually write:
538    /// mbr.write_into(&mut cur)
539    ///     .expect("could not write MBR to disk")
540    /// ```
541    pub fn write_into<W: ?Sized>(&mut self, mut writer: &mut W) -> Result<()>
542    where
543        W: Write + Seek,
544    {
545        self.header.write_into(&mut writer)?;
546
547        if let Some(extended) = self.header.get_extended_partition() {
548            if let Some(first) = self.logical_partitions.get_mut(0) {
549                first.absolute_ebr_lba = extended.starting_lba;
550                first.ebr_sectors = None;
551                first.ebr_last_chs = None;
552            }
553
554            if self.check_geometry() {
555                for l in self.logical_partitions.iter_mut() {
556                    l.update_chs(self.cylinders, self.heads, self.sectors)?;
557                }
558            }
559
560            // newtype for bootstrap code so we can serialize it via BigArray
561            // pending https://github.com/serde-rs/serde/issues/1937
562            #[derive(Serialize)]
563            struct BootstrapCode446(#[serde(with = "BigArray")] [u8; 446]);
564
565            let next_logical_partitions = self
566                .logical_partitions
567                .iter()
568                .skip(1)
569                .map(Some)
570                .chain(once(None));
571            for (l, next) in self.logical_partitions.iter().zip(next_logical_partitions) {
572                let partition = MBRPartitionEntry {
573                    starting_lba: l.partition.starting_lba.saturating_sub(l.absolute_ebr_lba),
574                    ..l.partition
575                };
576                partition.check()?;
577                writer.seek(SeekFrom::Start(u64::from(
578                    l.absolute_ebr_lba * self.sector_size,
579                )))?;
580                serialize_into(&mut writer, &BootstrapCode446(l.bootstrap_code))?;
581                serialize_into(&mut writer, &partition)?;
582                if let Some(next) = next {
583                    serialize_into(
584                        &mut writer,
585                        &MBRPartitionEntry {
586                            boot: BOOT_INACTIVE,
587                            first_chs: next.ebr_first_chs,
588                            sys: extended.sys,
589                            last_chs: next.ebr_last_chs.unwrap(),
590                            starting_lba: next
591                                .absolute_ebr_lba
592                                .saturating_sub(extended.starting_lba),
593                            sectors: next.ebr_sectors.unwrap(),
594                        },
595                    )?;
596                } else {
597                    serialize_into(&mut writer, &MBRPartitionEntry::empty())?;
598                }
599                writer.write_all(&[0; 16 * 2])?;
600                serialize_into(&mut writer, &BOOT_SIGNATURE)?;
601            }
602        }
603
604        Ok(())
605    }
606
607    /// Get a cylinder size in sectors. This function is useful if you want to
608    /// align your partitions to the cylinder.
609    pub fn get_cylinder_size(&self) -> u32 {
610        u32::from(self.heads) * u32::from(self.sectors)
611    }
612
613    /// Finds the primary partition (ignoring extended partitions) or logical
614    /// partition where the given sector resides.
615    pub fn find_at_sector(&self, sector: u32) -> Option<usize> {
616        let between = |sector, start, len| sector >= start && sector < start + len;
617
618        let primary = self
619            .header
620            .iter()
621            .find(|(_, x)| x.is_used() && between(sector, x.starting_lba, x.sectors));
622
623        match primary {
624            Some((_, x)) if x.is_extended() => self
625                .logical_partitions
626                .iter()
627                .enumerate()
628                .find(|(_, x)| {
629                    x.partition.is_used()
630                        && between(sector, x.partition.starting_lba, x.partition.sectors)
631                })
632                .map(|(i, _)| 5 + i),
633            Some((i, _)) => Some(i),
634            None => None,
635        }
636    }
637
638    /// Remove a partition entry that resides at a given sector. If the partition is the extended
639    /// partition, it will delete also all the logical partitions.
640    ///
641    /// # Errors
642    /// It is an error to provide a sector which does not belong to a partition.
643    pub fn remove_at_sector(&mut self, sector: u32) -> Result<()> {
644        let i = self
645            .find_at_sector(sector)
646            .ok_or(Error::PartitionNotFound)?;
647
648        if i >= 5 {
649            self.remove(i);
650        } else {
651            if self[i].is_extended() {
652                self.logical_partitions.clear();
653            }
654            self[i] = MBRPartitionEntry::empty();
655        }
656
657        Ok(())
658    }
659
660    /// Find free spots in the partition table.
661    /// This function will return a vector of tuple with on the left: the starting LBA of the free
662    /// spot; and on the right: the size (in sectors) of the free spot.
663    /// This function will automatically align with the alignment defined in the `MBR`.
664    ///
665    /// # Examples
666    /// Basic usage:
667    /// ```
668    /// let ss = 512;
669    /// let data = vec![0; 100 * ss as usize];
670    /// let mut cur = std::io::Cursor::new(data);
671    /// let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
672    ///     .expect("could not create partition table");
673    ///
674    /// mbr[1] = mbrman::MBRPartitionEntry {
675    ///     boot: mbrman::BOOT_INACTIVE,
676    ///     first_chs: mbrman::CHS::empty(),
677    ///     sys: 0x83,
678    ///     last_chs: mbrman::CHS::empty(),
679    ///     starting_lba: 6,
680    ///     sectors: mbr.disk_size - 11,
681    /// };
682    ///
683    /// // NOTE: align to the sectors, so we can use every last one of them
684    /// // NOTE: this is only for the demonstration purpose, this is not recommended
685    /// mbr.align = 1;
686    ///
687    /// assert_eq!(
688    ///     mbr.find_free_sectors(),
689    ///     vec![(1, 5), (mbr.disk_size - 5, 5)]
690    /// );
691    /// ```
692    pub fn find_free_sectors(&self) -> Vec<(u32, u32)> {
693        assert!(self.align > 0, "align must be greater than 0");
694
695        let collect_free_sectors = |positions: Vec<u32>| {
696            positions
697                .chunks(2)
698                .map(|x| (x[0] + 1, x[1] - x[0] - 1))
699                .filter(|(_, l)| *l > 0)
700                .map(|(i, l)| (i, l, ((i - 1) / self.align + 1) * self.align - i))
701                .map(|(i, l, s)| (i + s, l.saturating_sub(s)))
702                .filter(|(_, l)| *l > 0)
703                .collect::<Vec<_>>()
704        };
705
706        let mut positions = vec![0];
707        for (_, partition) in self.header.iter().filter(|(_, x)| x.is_used()) {
708            positions.push(partition.starting_lba);
709            positions.push(partition.starting_lba + partition.sectors - 1);
710        }
711        positions.push(self.disk_size);
712        positions.sort_unstable();
713
714        let mut res = collect_free_sectors(positions);
715
716        if let Some(extended) = self.header.get_extended_partition() {
717            let mut positions = vec![extended.starting_lba];
718            for l in self
719                .logical_partitions
720                .iter()
721                .filter(|x| x.partition.is_used())
722            {
723                let starting_lba = l.absolute_ebr_lba + l.partition.starting_lba;
724                positions.push(starting_lba);
725                positions.push(starting_lba + l.partition.sectors - 1);
726            }
727            positions.push(extended.starting_lba + extended.sectors);
728            positions.sort_unstable();
729            res.extend(collect_free_sectors(positions));
730        }
731
732        res
733    }
734
735    /// Find the first place (most on the left) where you could start a new partition of the size
736    /// given in parameter.
737    /// This function will automatically align with the alignment defined in the `MBR`.
738    ///
739    /// # Examples:
740    /// Basic usage:
741    /// ```
742    /// let ss = 512;
743    /// let data = vec![0; 100 * ss as usize];
744    /// let mut cur = std::io::Cursor::new(data);
745    /// let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
746    ///     .expect("could not create partition table");
747    ///
748    /// mbr[1] = mbrman::MBRPartitionEntry {
749    ///     boot: mbrman::BOOT_INACTIVE,
750    ///     first_chs: mbrman::CHS::empty(),
751    ///     sys: 0x83,
752    ///     last_chs: mbrman::CHS::empty(),
753    ///     starting_lba: 6,
754    ///     sectors: mbr.disk_size - 6,
755    /// };
756    ///
757    /// // NOTE: align to the sectors, so we can use every last one of them
758    /// // NOTE: this is only for the demonstration purpose, this is not recommended
759    /// mbr.align = 1;
760    ///
761    /// assert_eq!(mbr.find_first_place(5), Some(1));
762    /// ```
763    pub fn find_first_place(&self, size: u32) -> Option<u32> {
764        self.find_free_sectors()
765            .iter()
766            .find(|(_, l)| *l >= size)
767            .map(|(i, _)| *i)
768    }
769
770    /// Find the last place (most on the right) where you could start a new partition of the size
771    /// given in parameter.
772    /// This function will automatically align with the alignment defined in the `MBR`.
773    ///
774    /// # Examples:
775    /// Basic usage:
776    /// ```
777    /// let ss = 512;
778    /// let data = vec![0; 100 * ss as usize];
779    /// let mut cur = std::io::Cursor::new(data);
780    /// let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
781    ///     .expect("could not create partition table");
782    ///
783    /// mbr[1] = mbrman::MBRPartitionEntry {
784    ///     boot: mbrman::BOOT_INACTIVE,
785    ///     first_chs: mbrman::CHS::empty(),
786    ///     sys: 0x83,
787    ///     last_chs: mbrman::CHS::empty(),
788    ///     starting_lba: 6,
789    ///     sectors: 5,
790    /// };
791    ///
792    /// // NOTE: align to the sectors, so we can use every last one of them
793    /// // NOTE: this is only for the demonstration purpose, this is not recommended
794    /// mbr.align = 1;
795    ///
796    /// assert_eq!(mbr.find_last_place(5), Some(mbr.disk_size - 5));
797    /// ```
798    pub fn find_last_place(&self, size: u32) -> Option<u32> {
799        self.find_free_sectors()
800            .iter()
801            .filter(|(_, l)| *l >= size)
802            .last()
803            .map(|(i, l)| (i + l - size) / self.align * self.align)
804    }
805
806    /// Find the most optimal place (in the smallest free space) where you could start a new
807    /// partition of the size given in parameter.
808    /// This function will automatically align with the alignment defined in the `MBR`.
809    ///
810    /// # Examples:
811    /// Basic usage:
812    /// ```
813    /// let ss = 512;
814    /// let data = vec![0; 100 * ss as usize];
815    /// let mut cur = std::io::Cursor::new(data);
816    /// let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
817    ///     .expect("could not create partition table");
818    ///
819    /// mbr[1] = mbrman::MBRPartitionEntry {
820    ///     boot: mbrman::BOOT_INACTIVE,
821    ///     first_chs: mbrman::CHS::empty(),
822    ///     sys: 0x83,
823    ///     last_chs: mbrman::CHS::empty(),
824    ///     starting_lba: 11,
825    ///     sectors: mbr.disk_size - 11 - 5,
826    /// };
827    ///
828    /// // NOTE: align to the sectors, so we can use every last one of them
829    /// // NOTE: this is only for the demonstration purpose, this is not recommended
830    /// mbr.align = 1;
831    ///
832    /// // NOTE: the space as the end is more optimal because it will allow you to still be able to
833    /// //       insert a bigger partition later
834    /// assert_eq!(mbr.find_optimal_place(5), Some(mbr.disk_size - 5));
835    /// ```
836    pub fn find_optimal_place(&self, size: u32) -> Option<u32> {
837        let mut slots = self
838            .find_free_sectors()
839            .into_iter()
840            .filter(|(_, l)| *l >= size)
841            .collect::<Vec<_>>();
842        slots.sort_by(|(_, l1), (_, l2)| l1.cmp(l2));
843        slots.first().map(|&(i, _)| i)
844    }
845
846    /// Get the maximum size (in sectors) of a partition you could create in the MBR.
847    /// This function will automatically align with the alignment defined in the `MBR`.
848    ///
849    /// # Examples:
850    /// Basic usage:
851    /// ```
852    /// let ss = 512;
853    /// let data = vec![0; 100 * ss as usize];
854    /// let mut cur = std::io::Cursor::new(data);
855    /// let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
856    ///     .expect("could not create partition table");
857    ///
858    /// // NOTE: align to the sectors, so we can use every last one of them
859    /// // NOTE: this is only for the demonstration purpose, this is not recommended
860    /// mbr.align = 1;
861    ///
862    /// assert_eq!(
863    ///     mbr.get_maximum_partition_size().unwrap_or(0),
864    ///     mbr.disk_size - 1
865    /// );
866    /// ```
867    pub fn get_maximum_partition_size(&self) -> Result<u32> {
868        self.find_free_sectors()
869            .into_iter()
870            .map(|(_, l)| l / self.align * self.align)
871            .max()
872            .ok_or(Error::NoSpaceLeft)
873    }
874
875    /// Push a new logical partition to the end of the extended partition list. This function will
876    /// take care of creating the EBR for you. The EBR will be located at `starting_lba` (provided
877    /// in input) and the logical partition itself will be located a block further to stay
878    /// aligned. The size of the logical partition will be one block smaller than the `sectors`
879    /// provided in input.
880    pub fn push(
881        &mut self,
882        sys: u8,
883        mut starting_lba: u32,
884        mut sectors: u32,
885    ) -> Result<&mut LogicalPartition> {
886        let extended = self
887            .header
888            .get_extended_partition()
889            .ok_or(Error::NoExtendedPartition)?;
890
891        starting_lba = ((starting_lba - 1) / self.align + 1) * self.align;
892        sectors = ((sectors - 1) / self.align + 1) * self.align;
893        if sectors < 2 * self.align {
894            return Err(Error::NotEnoughSectorsToCreateLogicalPartition);
895        }
896
897        let mut l = LogicalPartition {
898            partition: MBRPartitionEntry {
899                boot: BOOT_INACTIVE,
900                first_chs: CHS::empty(),
901                sys,
902                last_chs: CHS::empty(),
903                starting_lba: starting_lba + self.align,
904                sectors: sectors - self.align,
905            },
906            absolute_ebr_lba: starting_lba,
907            ebr_sectors: if self.logical_partitions.is_empty() {
908                None
909            } else {
910                Some(sectors)
911            },
912            ebr_first_chs: CHS::empty(),
913            ebr_last_chs: if self.logical_partitions.is_empty() {
914                None
915            } else {
916                Some(CHS::empty())
917            },
918            bootstrap_code: [0; 446],
919        };
920
921        if l.absolute_ebr_lba < extended.starting_lba {
922            return Err(Error::EBRStartsBeforeExtendedPartition);
923        }
924
925        if l.absolute_ebr_lba > extended.starting_lba + extended.sectors - 2 * self.align {
926            return Err(Error::EBRStartsTooCloseToTheEndOfExtendedPartition);
927        }
928
929        if let Some(ebr_sectors) = l.ebr_sectors {
930            let ending_ebr_lba = l.absolute_ebr_lba + ebr_sectors - 1;
931
932            if ending_ebr_lba > extended.starting_lba + extended.sectors - 1 {
933                return Err(Error::EBREndsAfterExtendedPartition);
934            }
935        }
936
937        if self.check_geometry() {
938            l.update_chs(self.cylinders, self.heads, self.sectors)?;
939        }
940        self.logical_partitions.push(l);
941
942        Ok(self.logical_partitions.last_mut().unwrap())
943    }
944
945    /// Remove a logical partition. This will remove a logical partition in the array.
946    ///
947    /// # Remark
948    ///
949    /// This operation will decrease by one the index of every logical partition after the one that
950    /// has been removed.
951    ///
952    /// # Panics
953    ///
954    /// Panics if `index` is out of bounds.
955    pub fn remove(&mut self, index: usize) -> LogicalPartition {
956        assert!(index >= 5, "logical partitions start at 5");
957        self.logical_partitions.remove(index - 5)
958    }
959}
960
961impl Index<usize> for MBR {
962    type Output = MBRPartitionEntry;
963
964    fn index(&self, i: usize) -> &Self::Output {
965        assert!(i != 0, "invalid partition index: 0");
966        self.get(i).expect("invalid partition")
967    }
968}
969
970impl IndexMut<usize> for MBR {
971    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
972        assert!(i != 0, "invalid partition index: 0");
973        self.get_mut(i).expect("invalid partition")
974    }
975}
976
977/// An MBR partition table header
978#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
979pub struct MBRHeader {
980    /// Bootstrap code area
981    #[serde(with = "BigArray")]
982    pub bootstrap_code: [u8; 440],
983    /// 32-bit disk signature
984    pub disk_signature: [u8; 4],
985    /// `[0x5a, 0x5a]` if protected, `[0x00, 0x00]` if not
986    pub copy_protected: [u8; 2],
987    /// Partition 1
988    pub partition_1: MBRPartitionEntry,
989    /// Partition 2
990    pub partition_2: MBRPartitionEntry,
991    /// Partition 3
992    pub partition_3: MBRPartitionEntry,
993    /// Partition 4
994    pub partition_4: MBRPartitionEntry,
995    /// Boot signature
996    pub boot_signature: [u8; 2],
997}
998
999impl MBRHeader {
1000    /// Check if the partition table is copy-protected
1001    pub fn is_copy_protected(&self) -> Option<bool> {
1002        match self.copy_protected {
1003            [0x00, 0x00] => Some(false),
1004            [0x5a, 0x5a] => Some(true),
1005            _ => None,
1006        }
1007    }
1008
1009    /// Get `Some(&MBRPartitionEntry)` if it exists, None otherwise.
1010    ///
1011    /// # Remarks
1012    ///
1013    ///  -  The partitions start at index 1
1014    ///  -  The first 4 partitions always exist
1015    ///  -  This function does not return logical partitions
1016    pub fn get(&self, i: usize) -> Option<&MBRPartitionEntry> {
1017        match i {
1018            1 => Some(&self.partition_1),
1019            2 => Some(&self.partition_2),
1020            3 => Some(&self.partition_3),
1021            4 => Some(&self.partition_4),
1022            _ => None,
1023        }
1024    }
1025
1026    /// Get `Some(&mut MBRPartitionEntry)` if it exists, None otherwise.
1027    ///
1028    /// # Remarks
1029    ///
1030    ///  -  The partitions start at index 1
1031    ///  -  The first 4 partitions always exist
1032    ///  -  This function does not return logical partitions
1033    pub fn get_mut(&mut self, i: usize) -> Option<&mut MBRPartitionEntry> {
1034        match i {
1035            1 => Some(&mut self.partition_1),
1036            2 => Some(&mut self.partition_2),
1037            3 => Some(&mut self.partition_3),
1038            4 => Some(&mut self.partition_4),
1039            _ => None,
1040        }
1041    }
1042
1043    /// Get an iterator over the primary partition entries and their index. The
1044    /// index always starts at 1.
1045    pub fn iter(&self) -> impl Iterator<Item = (usize, &MBRPartitionEntry)> {
1046        vec![
1047            &self.partition_1,
1048            &self.partition_2,
1049            &self.partition_3,
1050            &self.partition_4,
1051        ]
1052        .into_iter()
1053        .enumerate()
1054        .map(|(i, x)| (i + 1, x))
1055    }
1056
1057    /// Get a mutable iterator over the primary partition entries and their
1058    /// index. The index always starts at 1.
1059    pub fn iter_mut(&mut self) -> impl Iterator<Item = (usize, &mut MBRPartitionEntry)> {
1060        vec![
1061            &mut self.partition_1,
1062            &mut self.partition_2,
1063            &mut self.partition_3,
1064            &mut self.partition_4,
1065        ]
1066        .into_iter()
1067        .enumerate()
1068        .map(|(i, x)| (i + 1, x))
1069    }
1070
1071    fn get_extended_partition(&self) -> Option<&MBRPartitionEntry> {
1072        self.iter().find(|(_, x)| x.is_extended()).map(|(_, x)| x)
1073    }
1074
1075    /// Attempt to read a MBR header from a reader.  This operation will seek at the
1076    /// correct location before trying to write to disk.
1077    pub fn read_from<R: ?Sized>(mut reader: &mut R) -> Result<MBRHeader>
1078    where
1079        R: Read + Seek,
1080    {
1081        reader.seek(SeekFrom::Start(0))?;
1082        let header: Self = deserialize_from(&mut reader)?;
1083        header.check()?;
1084        Ok(header)
1085    }
1086
1087    /// Make a new MBR header
1088    pub fn new(disk_signature: [u8; 4]) -> MBRHeader {
1089        MBRHeader {
1090            bootstrap_code: [0; 440],
1091            disk_signature,
1092            copy_protected: [0x00, 0x00],
1093            partition_1: MBRPartitionEntry::empty(),
1094            partition_2: MBRPartitionEntry::empty(),
1095            partition_3: MBRPartitionEntry::empty(),
1096            partition_4: MBRPartitionEntry::empty(),
1097            boot_signature: BOOT_SIGNATURE,
1098        }
1099    }
1100
1101    /// Write the MBR header into a writer. This operation will seek at the
1102    /// correct location before trying to write to disk.
1103    pub fn write_into<W: ?Sized>(&self, mut writer: &mut W) -> Result<()>
1104    where
1105        W: Write + Seek,
1106    {
1107        self.check()?;
1108        writer.seek(SeekFrom::Start(0))?;
1109        serialize_into(&mut writer, &self)?;
1110
1111        Ok(())
1112    }
1113
1114    /// Validate invariants
1115    fn check(&self) -> Result<()> {
1116        if self.boot_signature != BOOT_SIGNATURE {
1117            return Err(Error::InvalidSignature);
1118        }
1119        self.iter().try_for_each(|(_, partition)| partition.check())
1120    }
1121}
1122
1123impl Index<usize> for MBRHeader {
1124    type Output = MBRPartitionEntry;
1125
1126    fn index(&self, i: usize) -> &Self::Output {
1127        assert!(i != 0, "invalid partition index: 0");
1128        self.get(i).expect("invalid partition")
1129    }
1130}
1131
1132impl IndexMut<usize> for MBRHeader {
1133    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
1134        assert!(i != 0, "invalid partition index: 0");
1135        self.get_mut(i).expect("invalid partition")
1136    }
1137}
1138
1139#[derive(Debug, Clone, Deserialize)]
1140struct EBRHeader {
1141    #[serde(with = "BigArray")]
1142    bootstrap_code: [u8; 446],
1143    partition_1: MBRPartitionEntry,
1144    partition_2: MBRPartitionEntry,
1145    _unused_partition_3: [u8; 16],
1146    _unused_partition_4: [u8; 16],
1147    boot_signature: [u8; 2],
1148}
1149
1150impl EBRHeader {
1151    fn read_from<R: ?Sized>(mut reader: &mut R) -> Result<EBRHeader>
1152    where
1153        R: Read,
1154    {
1155        let header: Self = deserialize_from(&mut reader)?;
1156        header.check()?;
1157        Ok(header)
1158    }
1159
1160    /// Validate invariants
1161    fn check(&self) -> Result<()> {
1162        if self.boot_signature != BOOT_SIGNATURE {
1163            return Err(Error::InvalidSignature);
1164        }
1165        self.partition_1.check()?;
1166        self.partition_2.check()
1167    }
1168
1169    fn unwrap(self) -> (MBRPartitionEntry, MBRPartitionEntry, [u8; 446]) {
1170        (self.partition_1, self.partition_2, self.bootstrap_code)
1171    }
1172}
1173
1174/// An MBR partition entry
1175#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
1176pub struct MBRPartitionEntry {
1177    /// Boot flag
1178    pub boot: u8,
1179    /// CHS address of the first sector in the partition
1180    pub first_chs: CHS,
1181    /// Partition type (file system ID)
1182    pub sys: u8,
1183    /// CHS address of the last sector in the partition
1184    pub last_chs: CHS,
1185    /// Starting LBA of the partition
1186    pub starting_lba: u32,
1187    /// Number of sectors allocated to the partition
1188    pub sectors: u32,
1189}
1190
1191impl MBRPartitionEntry {
1192    /// Creates an empty partition entry
1193    ///
1194    /// # Examples
1195    /// Basic usage:
1196    /// ```
1197    /// let ss = 512;
1198    /// let data = vec![0; 100 * ss as usize];
1199    /// let mut cur = std::io::Cursor::new(data);
1200    /// let mut mbr = mbrman::MBR::new_from(&mut cur, ss as u32, [0xff; 4])
1201    ///     .expect("could not create partition table");
1202    ///
1203    /// mbr[1] = mbrman::MBRPartitionEntry::empty();
1204    ///
1205    /// // NOTE: an empty partition entry is considered as not allocated
1206    /// assert!(mbr[1].is_unused());
1207    /// ```
1208    pub fn empty() -> MBRPartitionEntry {
1209        MBRPartitionEntry {
1210            boot: BOOT_INACTIVE,
1211            first_chs: CHS::empty(),
1212            sys: 0,
1213            last_chs: CHS::empty(),
1214            starting_lba: 0,
1215            sectors: 0,
1216        }
1217    }
1218
1219    /// Returns `true` if the partition entry is used (type (sys) != 0)
1220    pub fn is_used(&self) -> bool {
1221        self.sys > 0
1222    }
1223
1224    /// Returns `true` if the partition entry is not used (type (sys) == 0)
1225    pub fn is_unused(&self) -> bool {
1226        !self.is_used()
1227    }
1228
1229    /// Returns `true` if the partition is an extended type partition
1230    pub fn is_extended(&self) -> bool {
1231        self.sys == 0x05
1232            || self.sys == 0x0f
1233            || self.sys == 0x85
1234            || self.sys == 0xc5
1235            || self.sys == 0xd5
1236    }
1237
1238    /// Returns `true` if the partition is marked active (bootable)
1239    pub fn is_active(&self) -> bool {
1240        self.boot == BOOT_ACTIVE
1241    }
1242
1243    /// Validate invariants
1244    fn check(&self) -> Result<()> {
1245        if self.boot != BOOT_ACTIVE && self.boot != BOOT_INACTIVE {
1246            return Err(Error::InvalidBootFlag);
1247        }
1248        Ok(())
1249    }
1250}
1251
1252/// An abstraction struct for a logical partition
1253#[derive(Debug, Clone, PartialEq, Eq)]
1254pub struct LogicalPartition {
1255    /// MBR partition entry of the logical partition
1256    pub partition: MBRPartitionEntry,
1257    /// Absolute LBA of the EBR partition table
1258    pub absolute_ebr_lba: u32,
1259    /// Number of sectors in the EBR
1260    ///
1261    /// # Remark
1262    ///
1263    /// This information is known for all the EBR except the first logical partition
1264    pub ebr_sectors: Option<u32>,
1265    /// CHS address of the EBR header
1266    ///
1267    /// # Remark
1268    ///
1269    /// This information is copied from the extended partition for the first logical partition
1270    pub ebr_first_chs: CHS,
1271    /// CHS address of the last sector of the extended partition in the EBR header
1272    ///
1273    /// # Remark
1274    ///
1275    /// This information is known for all the EBR except the first logical partition
1276    pub ebr_last_chs: Option<CHS>,
1277    /// Bootstrap code area
1278    pub bootstrap_code: [u8; 446],
1279}
1280
1281impl LogicalPartition {
1282    /// Update the fields `partition.first_chs`, `partition.last_chs`.
1283    /// `ebr_first_chs` and `ebr_last_chs` using the disk geometry provided in parameter.
1284    pub fn update_chs(&mut self, cylinders: u16, heads: u8, sectors: u8) -> Result<()> {
1285        self.partition.first_chs =
1286            CHS::from_lba_exact(self.partition.starting_lba, cylinders, heads, sectors)?;
1287        self.partition.last_chs = CHS::from_lba_exact(
1288            self.partition.starting_lba + self.partition.sectors - 1,
1289            cylinders,
1290            heads,
1291            sectors,
1292        )?;
1293        self.ebr_first_chs = CHS::from_lba_exact(self.absolute_ebr_lba, cylinders, heads, sectors)?;
1294        self.ebr_last_chs = if let Some(ebr_sectors) = self.ebr_sectors {
1295            Some(CHS::from_lba_exact(
1296                self.absolute_ebr_lba + ebr_sectors - 1,
1297                cylinders,
1298                heads,
1299                sectors,
1300            )?)
1301        } else {
1302            None
1303        };
1304
1305        Ok(())
1306    }
1307}
1308
1309/// A CHS address (cylinder/head/sector)
1310#[derive(Debug, Copy, Clone, PartialEq, Eq)]
1311pub struct CHS {
1312    /// Cylinder
1313    pub cylinder: u16,
1314    /// Head
1315    pub head: u8,
1316    /// Sector
1317    pub sector: u8,
1318}
1319
1320impl CHS {
1321    /// Creates a new CHS address based on input parameters.
1322    ///
1323    /// # Remark
1324    ///
1325    /// The values entered in input are not checked.
1326    pub fn new(cylinder: u16, head: u8, sector: u8) -> CHS {
1327        CHS {
1328            cylinder,
1329            head,
1330            sector,
1331        }
1332    }
1333
1334    /// Creates an empty CHS addressing (0/0/0).
1335    ///
1336    /// # Remark
1337    ///
1338    /// This is what you need on recent hardware because CHS is never used.
1339    pub fn empty() -> CHS {
1340        CHS {
1341            cylinder: 0,
1342            head: 0,
1343            sector: 0,
1344        }
1345    }
1346
1347    /// Creates a CHS address calculated from the number of cylinders, heads
1348    /// and sectors per track of the hard disk.
1349    ///
1350    /// # Remarks
1351    ///
1352    ///  *  You probably don't need to do this at all! This is only useful if you
1353    ///     intend to use partitions that use the CHS addressing. Check the column
1354    ///     "Access" of [this table on Wikipedia](https://en.wikipedia.org/wiki/Partition_type).
1355    ///  *  On old systems, partitions must be aligned on cylinders.
1356    pub fn from_lba_exact(lba: u32, cylinders: u16, heads: u8, sectors: u8) -> Result<CHS> {
1357        // NOTE: code inspired from libfdisk (long2chs)
1358        let cylinders = u32::from(cylinders);
1359        let heads = u32::from(heads);
1360        let sectors = u32::from(sectors);
1361        let cylinder_size = heads * sectors;
1362
1363        let cylinder = lba / cylinder_size;
1364        let rem = lba % cylinder_size;
1365        let head = rem / sectors;
1366        let sector = rem % sectors + 1;
1367
1368        if cylinder > 1023 {
1369            return Err(Error::LBAExceedsMaximumCHS);
1370        }
1371
1372        if cylinder > cylinders {
1373            return Err(Error::LBAExceedsMaximumCylinders);
1374        }
1375
1376        Ok(CHS {
1377            cylinder: u16::try_from(cylinder).unwrap(),
1378            head: u8::try_from(head).unwrap(),
1379            sector: u8::try_from(sector).unwrap(),
1380        })
1381    }
1382
1383    /// Creates a CHS address, aligned to the nearest cylinder. The cylinder
1384    /// chosen will always be the exact cylinder (if the LBA is exactly at the
1385    /// beginning of a cylinder); or the next cylinder. But it will never
1386    /// choose the previous cylinder.
1387    pub fn from_lba_aligned(lba: u32, cylinders: u16, heads: u8, sectors: u8) -> Result<CHS> {
1388        let cylinders = u32::from(cylinders);
1389        let heads = u32::from(heads);
1390        let sectors = u32::from(sectors);
1391        let cylinder_size = heads * sectors;
1392
1393        let cylinder = ((lba - 1) / cylinder_size) + 1;
1394
1395        if cylinder > 1023 {
1396            return Err(Error::LBAExceedsMaximumCHS);
1397        }
1398
1399        if cylinder > cylinders {
1400            return Err(Error::LBAExceedsMaximumCylinders);
1401        }
1402
1403        // NOTE: In CHS addressing the sector numbers always start at 1, there is no sector 0
1404        //       https://en.wikipedia.org/wiki/Cylinder-head-sector
1405        Ok(CHS {
1406            cylinder: u16::try_from(cylinder).unwrap(),
1407            head: 0,
1408            sector: 1,
1409        })
1410    }
1411
1412    /// Convert a CHS address to LBA
1413    pub fn to_lba(self, heads: u8, sectors: u8) -> u32 {
1414        let heads = u32::from(heads);
1415        let sectors = u32::from(sectors);
1416        let c = u32::from(self.cylinder);
1417        let h = u32::from(self.head);
1418        let s = u32::from(self.sector);
1419
1420        // NOTE: In CHS addressing the sector numbers always start at 1, there is no sector 0
1421        //       https://en.wikipedia.org/wiki/Cylinder-head-sector
1422        c * (heads * sectors) + h * sectors + s - 1
1423    }
1424
1425    /// Check if the CHS address is empty
1426    ///
1427    /// # Remark
1428    ///
1429    /// This function does not check if the CHS address is withing range of
1430    /// possible values for a provided hard disk.
1431    pub fn is_empty(self) -> bool {
1432        self.cylinder == 0 && self.head == 0 && self.sector == 0
1433    }
1434
1435    /// Check if the CHS address is valid and within range of the possible
1436    /// values for the hard disk geometry provided in argument.
1437    pub fn is_valid(self, cylinders: u16, heads: u8, sectors: u8) -> bool {
1438        // NOTE: In CHS addressing the sector numbers always start at 1, there is no sector 0
1439        //       https://en.wikipedia.org/wiki/Cylinder-head-sector
1440        self.sector > 0 && self.sector <= sectors && self.head < heads && self.cylinder < cylinders
1441    }
1442}
1443
1444struct CHSVisitor;
1445
1446impl<'de> Visitor<'de> for CHSVisitor {
1447    type Value = CHS;
1448
1449    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
1450        formatter.write_str("CHS addressing")
1451    }
1452
1453    fn visit_seq<A>(self, mut seq: A) -> std::result::Result<CHS, A::Error>
1454    where
1455        A: SeqAccess<'de>,
1456    {
1457        let head = BitVec::<u8, Msb0>::from_vec(vec![seq.next_element::<u8>()?.unwrap()]);
1458        let mut bv = BitVec::<u8, Msb0>::from_vec(vec![seq.next_element::<u8>()?.unwrap()]);
1459        let mut cylinder = BitVec::<u16, Msb0>::with_capacity(10);
1460        cylinder.extend(repeat(false).take(6));
1461        cylinder.extend(bv.drain(..2));
1462        cylinder.extend(BitVec::<u8, Msb0>::from_vec(vec![seq
1463            .next_element::<u8>()?
1464            .unwrap()]));
1465        let mut sector = BitVec::<u8, Msb0>::with_capacity(8);
1466        sector.push(false);
1467        sector.push(false);
1468        sector.extend(bv.drain(..));
1469
1470        Ok(CHS {
1471            cylinder: cylinder.as_raw_slice()[0],
1472            head: head.as_raw_slice()[0],
1473            sector: sector.as_raw_slice()[0],
1474        })
1475    }
1476}
1477
1478impl<'de> Deserialize<'de> for CHS {
1479    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
1480    where
1481        D: Deserializer<'de>,
1482    {
1483        deserializer.deserialize_tuple(3, CHSVisitor)
1484    }
1485}
1486
1487impl Serialize for CHS {
1488    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
1489    where
1490        S: Serializer,
1491    {
1492        let mut bv = BitVec::<u8, Msb0>::from_vec(vec![self.head]);
1493        let mut sector = BitVec::<u8, Msb0>::from_vec(vec![self.sector]);
1494        let mut cylinder = BitVec::<u16, Msb0>::from_vec(vec![self.cylinder]);
1495        bv.extend(cylinder.drain(..8).skip(6));
1496        bv.extend(sector.drain(2..));
1497        bv.extend(cylinder.drain(..));
1498
1499        let mut seq = serializer.serialize_tuple(3)?;
1500        for x in bv.as_raw_slice() {
1501            seq.serialize_element(&x)?;
1502        }
1503        seq.end()
1504    }
1505}
1506
1507#[cfg(test)]
1508#[allow(clippy::cognitive_complexity)]
1509mod tests {
1510    use super::*;
1511    use std::fs::File;
1512    use std::io::Cursor;
1513
1514    const DISK1: &str = "tests/fixtures/disk1.img";
1515    const DISK2: &str = "tests/fixtures/disk2.img";
1516
1517    #[test]
1518    fn deserialize_maximum_chs_value() {
1519        let chs: CHS = bincode::deserialize(&[0xff, 0xff, 0xff]).unwrap();
1520        assert_eq!(
1521            chs,
1522            CHS {
1523                cylinder: 1023,
1524                head: 255,
1525                sector: 63,
1526            }
1527        );
1528    }
1529
1530    #[test]
1531    fn serialize_maximum_chs_value() {
1532        let chs = CHS {
1533            cylinder: 1023,
1534            head: 255,
1535            sector: 63,
1536        };
1537        let out = bincode::serialize(&chs).unwrap();
1538        assert_eq!(out, &[0xff, 0xff, 0xff]);
1539    }
1540
1541    #[test]
1542    fn serialize_and_deserialize_some_chs_value() {
1543        let chs: CHS = bincode::deserialize(&[0xaa, 0xaa, 0xaa]).unwrap();
1544        assert_eq!(
1545            chs,
1546            CHS {
1547                cylinder: 682,
1548                head: 170,
1549                sector: 42,
1550            }
1551        );
1552        let out = bincode::serialize(&chs).unwrap();
1553        assert_eq!(out, &[0xaa, 0xaa, 0xaa]);
1554    }
1555
1556    #[test]
1557    fn align_chs_to_cylinder() {
1558        fn lba2c(lba: u32) -> u16 {
1559            let chs = CHS::from_lba_aligned(lba, 100, 2, 2).unwrap();
1560
1561            assert_eq!(chs.head, 0);
1562            assert_eq!(chs.sector, 1);
1563
1564            chs.cylinder
1565        }
1566
1567        assert_eq!(lba2c(12), 3);
1568        assert_eq!(lba2c(10), 3);
1569    }
1570
1571    #[test]
1572    fn convert_chs_to_lba_and_back() {
1573        // NOTE: 2484/16/63 is taken from a real life example of hard disk of 1280MB
1574        // LBA address 666666 is around 341MB for a sector size of 512 bytes
1575        let chs = CHS::from_lba_exact(666_666, 2484, 16, 63).unwrap();
1576        assert_eq!(chs.to_lba(16, 63), 666_666);
1577
1578        let chs = CHS::from_lba_aligned(666_666, 2484, 16, 63).unwrap();
1579        assert_eq!(chs.to_lba(16, 63), 667_296);
1580
1581        let chs = CHS::from_lba_exact(667_296, 2484, 16, 63).unwrap();
1582        assert_eq!(chs.head, 0);
1583        assert_eq!(chs.sector, 1);
1584    }
1585
1586    #[test]
1587    #[allow(clippy::bool_assert_comparison)]
1588    fn read_disk1() {
1589        let mut mbr = MBR::read_from(&mut File::open(DISK1).unwrap(), 512).unwrap();
1590        assert!(mbr.header.partition_1.is_used());
1591        assert!(mbr.header.partition_2.is_used());
1592        assert!(mbr.header.partition_3.is_unused());
1593        assert!(mbr.header.partition_4.is_unused());
1594        assert!(mbr.header.partition_1.is_active());
1595        assert!(!mbr.header.partition_2.is_active());
1596        assert!(!mbr.header.partition_3.is_active());
1597        assert!(!mbr.header.partition_4.is_active());
1598        assert_eq!(mbr.len(), 4);
1599        assert_eq!(mbr.header.iter().count(), 4);
1600        assert_eq!(mbr.header.iter_mut().count(), 4);
1601        assert_eq!(mbr.iter_mut().count(), 4);
1602        assert_eq!(mbr.header.partition_1.boot, BOOT_ACTIVE);
1603        assert_eq!(mbr.header.partition_1.sys, 0x06);
1604        assert_eq!(mbr.header.partition_1.starting_lba, 1);
1605        assert_eq!(mbr.header.partition_1.sectors, 1);
1606        assert_eq!(mbr.header.partition_2.boot, BOOT_INACTIVE);
1607        assert_eq!(mbr.header.partition_2.sys, 0x0b);
1608        assert_eq!(mbr.header.partition_2.starting_lba, 3);
1609        assert_eq!(mbr.header.partition_2.sectors, 1);
1610    }
1611
1612    #[test]
1613    fn read_disk2() {
1614        let mut mbr = MBR::read_from(&mut File::open(DISK2).unwrap(), 512).unwrap();
1615        assert!(mbr.header.partition_1.is_used());
1616        assert!(mbr.header.partition_2.is_used());
1617        assert!(mbr.header.partition_3.is_unused());
1618        assert!(mbr.header.partition_4.is_unused());
1619        assert!(!mbr.header.partition_1.is_active());
1620        assert!(!mbr.header.partition_2.is_active());
1621        assert!(!mbr.header.partition_3.is_active());
1622        assert!(!mbr.header.partition_4.is_active());
1623        assert_eq!(mbr.header.partition_2.sys, 0x05);
1624        assert_eq!(mbr.header.partition_2.starting_lba, 5);
1625        assert_eq!(mbr.header.partition_2.first_chs, CHS::new(0, 1, 3));
1626        assert_eq!(mbr.header.partition_2.last_chs, CHS::new(3, 0, 2));
1627        assert_eq!(mbr.header.partition_2.sectors, 15);
1628        assert_eq!(mbr.len(), 9);
1629        assert_eq!(mbr.iter().count(), 9);
1630        assert_eq!(mbr.iter_mut().count(), 9);
1631        assert_eq!(mbr.iter().filter(|(_, x)| x.is_used()).count(), 7);
1632        assert!(mbr
1633            .iter()
1634            .filter(|(_, x)| x.is_used() && !x.is_extended())
1635            .all(|(_, x)| x.sys == 0x83));
1636        assert!(mbr.get(9).is_some());
1637        assert!(mbr.get(10).is_none());
1638        assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 5);
1639        assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
1640        assert_eq!(mbr.logical_partitions[1].absolute_ebr_lba, 7);
1641        assert_eq!(mbr.logical_partitions[1].ebr_sectors, Some(3));
1642        assert_eq!(mbr.logical_partitions[2].absolute_ebr_lba, 10);
1643        assert_eq!(mbr.logical_partitions[2].ebr_sectors, Some(4));
1644        // NOTE: this is actually for testing the CHS conversion to and from LBA
1645        assert_eq!(
1646            mbr.logical_partitions[2].partition.first_chs,
1647            CHS::new(1, 1, 3)
1648        );
1649        assert_eq!(
1650            mbr.logical_partitions[2].partition.first_chs.to_lba(2, 3),
1651            mbr.logical_partitions[2].partition.starting_lba
1652        );
1653        assert_eq!(
1654            CHS::from_lba_exact(
1655                mbr.logical_partitions[2].partition.starting_lba,
1656                u16::max_value(),
1657                2,
1658                3
1659            )
1660            .unwrap(),
1661            mbr.logical_partitions[2].partition.first_chs
1662        );
1663        assert_eq!(mbr.logical_partitions[3].absolute_ebr_lba, 14);
1664        assert_eq!(mbr.logical_partitions[3].ebr_sectors, Some(2));
1665        assert_eq!(mbr.logical_partitions[4].absolute_ebr_lba, 16);
1666        assert_eq!(mbr.logical_partitions[4].ebr_sectors, Some(2));
1667        assert!(mbr.logical_partitions.get(5).is_none());
1668    }
1669
1670    #[test]
1671    fn read_empty_extended_partition() {
1672        let ss = 512_u32;
1673        let data = vec![0; 10 * ss as usize];
1674        let mut cur = Cursor::new(data);
1675        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1676        mbr.header.partition_1.sys = 0x0f;
1677        mbr.header.partition_1.starting_lba = 1;
1678        mbr.header.partition_1.sectors = 10;
1679
1680        mbr.write_into(&mut cur).unwrap();
1681
1682        let mbr = MBR::read_from(&mut cur, ss).unwrap();
1683        assert_eq!(mbr.header.partition_1.sys, 0x0f);
1684        assert_eq!(mbr.header.partition_1.starting_lba, 1);
1685        assert_eq!(mbr.header.partition_1.sectors, 10);
1686        assert!(mbr.logical_partitions.is_empty());
1687    }
1688
1689    #[test]
1690    #[allow(clippy::bool_assert_comparison)]
1691    fn new_mbr_then_write_then_read_twice() {
1692        let ss = 512_u32;
1693        let data = vec![0; 12 * ss as usize];
1694        let mut cur = Cursor::new(data);
1695        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1696        mbr.header.partition_1.sys = 0x83;
1697        mbr.header.partition_1.starting_lba = 1;
1698        mbr.header.partition_1.sectors = 4;
1699        mbr.header.partition_3.sys = 0x0f;
1700        mbr.header.partition_3.starting_lba = 5;
1701        mbr.header.partition_3.sectors = 6;
1702        let mut empty = MBRPartitionEntry::empty();
1703        empty.starting_lba = 6;
1704        empty.sectors = 1;
1705        mbr.logical_partitions.push(LogicalPartition {
1706            partition: empty,
1707            absolute_ebr_lba: 5,
1708            ebr_sectors: None,
1709            ebr_first_chs: CHS::empty(),
1710            ebr_last_chs: None,
1711            bootstrap_code: [0; 446],
1712        });
1713        mbr.logical_partitions.push(LogicalPartition {
1714            partition: MBRPartitionEntry {
1715                boot: BOOT_ACTIVE,
1716                first_chs: CHS::empty(),
1717                sys: 0x83,
1718                last_chs: CHS::empty(),
1719                starting_lba: 9,
1720                sectors: 1,
1721            },
1722            absolute_ebr_lba: 7,
1723            ebr_sectors: Some(3),
1724            ebr_first_chs: CHS::empty(),
1725            ebr_last_chs: Some(CHS::empty()),
1726            bootstrap_code: [0; 446],
1727        });
1728
1729        // NOTE: don't put this in a loop otherwise it's hard to guess if the error come from
1730        //       reading, writing or reading after writing again
1731        mbr.write_into(&mut cur).unwrap();
1732
1733        let mut mbr = MBR::read_from(&mut cur, ss).unwrap();
1734        assert_eq!(mbr.header.partition_1.boot, BOOT_INACTIVE);
1735        assert_eq!(mbr.header.partition_1.sys, 0x83);
1736        assert_eq!(mbr.header.partition_1.starting_lba, 1);
1737        assert_eq!(mbr.header.partition_1.sectors, 4);
1738        assert_eq!(mbr.logical_partitions.len(), 2);
1739        assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 5);
1740        assert_eq!(mbr.logical_partitions[0].partition.boot, BOOT_INACTIVE);
1741        assert_eq!(mbr.logical_partitions[0].partition.starting_lba, 6);
1742        assert_eq!(mbr.logical_partitions[0].partition.sectors, 1);
1743        assert_eq!(mbr.logical_partitions[0].partition.sys, 0);
1744        assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
1745        assert_eq!(mbr.logical_partitions[1].absolute_ebr_lba, 7);
1746        assert_eq!(mbr.logical_partitions[1].partition.boot, BOOT_ACTIVE);
1747        assert_eq!(mbr.logical_partitions[1].partition.starting_lba, 9);
1748        assert_eq!(mbr.logical_partitions[1].partition.sectors, 1);
1749        assert_eq!(mbr.logical_partitions[1].partition.sys, 0x83);
1750        assert_eq!(mbr.logical_partitions[1].ebr_sectors, Some(3));
1751
1752        mbr.write_into(&mut cur).unwrap();
1753
1754        let mbr = MBR::read_from(&mut cur, ss).unwrap();
1755        assert_eq!(mbr.header.partition_1.boot, BOOT_INACTIVE);
1756        assert_eq!(mbr.header.partition_1.sys, 0x83);
1757        assert_eq!(mbr.header.partition_1.starting_lba, 1);
1758        assert_eq!(mbr.header.partition_1.sectors, 4);
1759        assert_eq!(mbr.logical_partitions.len(), 2);
1760        assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 5);
1761        assert_eq!(mbr.logical_partitions[0].partition.boot, BOOT_INACTIVE);
1762        assert_eq!(mbr.logical_partitions[0].partition.starting_lba, 6);
1763        assert_eq!(mbr.logical_partitions[0].partition.sectors, 1);
1764        assert_eq!(mbr.logical_partitions[0].partition.sys, 0);
1765        assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
1766        assert_eq!(mbr.logical_partitions[1].absolute_ebr_lba, 7);
1767        assert_eq!(mbr.logical_partitions[1].partition.boot, BOOT_ACTIVE);
1768        assert_eq!(mbr.logical_partitions[1].partition.starting_lba, 9);
1769        assert_eq!(mbr.logical_partitions[1].partition.sectors, 1);
1770        assert_eq!(mbr.logical_partitions[1].partition.sys, 0x83);
1771        assert_eq!(mbr.logical_partitions[1].ebr_sectors, Some(3));
1772    }
1773
1774    #[test]
1775    fn find_at_sector() {
1776        let mbr = MBR::read_from(&mut File::open(DISK2).unwrap(), 512).unwrap();
1777        assert_eq!(mbr.find_at_sector(2), Some(1));
1778        assert_eq!(mbr.find_at_sector(4), None);
1779        assert_eq!(mbr.find_at_sector(8), Some(6));
1780        assert_eq!(mbr.find_at_sector(7), None);
1781    }
1782
1783    #[test]
1784    fn find_free_sectors() {
1785        let ss = 512_u32;
1786        let data = vec![0; 10 * ss as usize];
1787        let mut cur = Cursor::new(data);
1788        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1789        mbr.align = 1;
1790        mbr.header.partition_1.sys = 0x83;
1791        mbr.header.partition_1.starting_lba = 1;
1792        mbr.header.partition_1.sectors = 1;
1793        mbr.header.partition_3.sys = 0x0f;
1794        mbr.header.partition_3.starting_lba = 5;
1795        mbr.header.partition_3.sectors = 5;
1796        mbr.logical_partitions.push(LogicalPartition {
1797            bootstrap_code: [0; 446],
1798            partition: MBRPartitionEntry {
1799                boot: BOOT_INACTIVE,
1800                first_chs: CHS::empty(),
1801                sys: 0x83,
1802                last_chs: CHS::empty(),
1803                starting_lba: 2,
1804                sectors: 1,
1805            },
1806            absolute_ebr_lba: 5,
1807            ebr_sectors: None,
1808            ebr_first_chs: CHS::empty(),
1809            ebr_last_chs: None,
1810        });
1811
1812        assert_eq!(mbr.find_free_sectors(), vec![(2, 3), (6, 1), (8, 2)]);
1813    }
1814
1815    #[test]
1816    fn push_logical_partition_aligned_to_cylinder() {
1817        let ss = 512_u32;
1818        let data = vec![0; 54 * ss as usize];
1819        let mut cur = Cursor::new(data);
1820        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1821        mbr.cylinders = 6;
1822        mbr.heads = 3;
1823        mbr.sectors = 3;
1824        let align = 9;
1825        mbr.align = mbr.get_cylinder_size();
1826        assert_eq!(mbr.align, align);
1827        mbr.header.partition_1.sys = 0x05;
1828        mbr.header.partition_1.first_chs = CHS::new(1, 0, 1);
1829        mbr.header.partition_1.last_chs = CHS::new(5, 0, 1);
1830        mbr.header.partition_1.starting_lba = align;
1831        mbr.header.partition_1.sectors = mbr.disk_size;
1832        let p = mbr.push(0x00, 2, 2 * align).unwrap();
1833
1834        assert_eq!(p.absolute_ebr_lba, align);
1835        assert_eq!(p.partition.starting_lba, 2 * align);
1836        assert_eq!(p.ebr_sectors, None);
1837        assert_eq!(p.partition.sectors, align);
1838        assert_eq!(p.ebr_first_chs, CHS::new(1, 0, 1));
1839        assert_eq!(p.ebr_last_chs, None);
1840        assert_eq!(p.partition.first_chs, CHS::new(2, 0, 1));
1841        assert_eq!(p.partition.last_chs, CHS::new(2, 2, 3));
1842
1843        let p = mbr
1844            .push(
1845                0x83,
1846                CHS::new(3, 0, 1).to_lba(mbr.heads, mbr.sectors),
1847                align * 3,
1848            )
1849            .unwrap();
1850
1851        assert_eq!(p.absolute_ebr_lba, align * 3);
1852        assert_eq!(p.partition.starting_lba, 4 * align);
1853        assert_eq!(p.ebr_sectors, Some(align * 3));
1854        assert_eq!(p.partition.sectors, align * 2);
1855        assert_eq!(p.ebr_first_chs, CHS::new(3, 0, 1));
1856        assert_eq!(p.ebr_last_chs, Some(CHS::new(5, 2, 3)));
1857        assert_eq!(p.partition.first_chs, CHS::new(4, 0, 1));
1858        assert_eq!(p.partition.last_chs, CHS::new(5, 2, 3));
1859
1860        mbr.write_into(&mut cur).unwrap();
1861        let mut same_mbr = MBR::read_from(&mut cur, ss).unwrap();
1862        same_mbr.cylinders = mbr.cylinders;
1863        same_mbr.heads = mbr.heads;
1864        same_mbr.sectors = mbr.sectors;
1865
1866        assert_eq!(mbr, same_mbr);
1867    }
1868
1869    #[test]
1870    fn push_logical_partition_check_within_range() {
1871        let ss = 512_u32;
1872        let data = vec![0; 10 * ss as usize];
1873        let mut cur = Cursor::new(data);
1874        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1875        mbr.align = 1;
1876        mbr.header.partition_1.sys = 0x0f;
1877        mbr.header.partition_1.starting_lba = 4;
1878        mbr.header.partition_1.sectors = 2;
1879        assert!(mbr.push(0x83, 3, 2).is_err());
1880        assert!(mbr.push(0x83, 6, 2).is_err());
1881        mbr.push(0x83, 4, 2).unwrap();
1882        assert!(mbr.push(0x83, 4, 3).is_err());
1883    }
1884
1885    #[test]
1886    fn remove_logical_partition() {
1887        let ss = 512_u32;
1888        let data = vec![0; 10 * ss as usize];
1889        let mut cur = Cursor::new(data);
1890        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1891        mbr.align = 1;
1892        mbr.header.partition_1.sys = 0x0f;
1893        mbr.header.partition_1.starting_lba = 1;
1894        mbr.header.partition_1.sectors = mbr.disk_size - 1;
1895        mbr.push(0x00, 1, 2).unwrap();
1896        mbr.push(0x83, 4, 3).unwrap();
1897
1898        mbr.logical_partitions.remove(0);
1899        mbr.write_into(&mut cur).unwrap();
1900
1901        let same_mbr = MBR::read_from(&mut cur, ss).unwrap();
1902
1903        assert_eq!(mbr, same_mbr);
1904        assert_eq!(mbr.logical_partitions[0].partition.starting_lba, 5);
1905        assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 1);
1906        assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
1907    }
1908
1909    fn read_corrupt_mbr(path: &str, bad_offset: usize, bad_data: u8) -> Error {
1910        let mut data = Vec::new();
1911        File::open(path).unwrap().read_to_end(&mut data).unwrap();
1912        data[bad_offset] = bad_data;
1913        let mut cur = Cursor::new(&data);
1914        MBR::read_from(&mut cur, 512).unwrap_err()
1915    }
1916
1917    #[test]
1918    fn read_invalid_boot_signature() {
1919        // MBR
1920        assert!(matches!(
1921            read_corrupt_mbr(DISK2, 0xffe, 0xaa),
1922            Error::InvalidSignature
1923        ));
1924        // EBR
1925        assert!(matches!(
1926            read_corrupt_mbr(DISK2, 0x21fe, 0xaa),
1927            Error::InvalidSignature
1928        ));
1929    }
1930
1931    #[test]
1932    fn write_invalid_boot_signature() {
1933        let ss = 512_u32;
1934        let data = vec![0; 10 * ss as usize];
1935        let mut cur = Cursor::new(data);
1936
1937        // invalid in MBRHeader struct
1938        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1939        mbr.header.partition_1.sys = 0x0a;
1940        mbr.header.partition_1.starting_lba = 1;
1941        mbr.header.partition_1.sectors = 10;
1942        mbr.header.boot_signature = [0, 0];
1943        assert!(matches!(
1944            mbr.write_into(&mut cur).unwrap_err(),
1945            Error::InvalidSignature
1946        ));
1947    }
1948
1949    #[test]
1950    fn read_invalid_boot_flag() {
1951        // primary partition
1952        assert!(matches!(
1953            read_corrupt_mbr(DISK2, 0x1be, BOOT_ACTIVE | 0x01),
1954            Error::InvalidBootFlag
1955        ));
1956        // EBR first partition
1957        assert!(matches!(
1958            read_corrupt_mbr(DISK2, 0xfbe, BOOT_ACTIVE | 0x01),
1959            Error::InvalidBootFlag
1960        ));
1961        // EBR second partition
1962        assert!(matches!(
1963            read_corrupt_mbr(DISK2, 0xfce, BOOT_ACTIVE | 0x01),
1964            Error::InvalidBootFlag
1965        ));
1966    }
1967
1968    #[test]
1969    fn write_invalid_boot_flag() {
1970        let ss = 512_u32;
1971        let data = vec![0; 10 * ss as usize];
1972        let mut cur = Cursor::new(data);
1973
1974        // primary partition
1975        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1976        mbr.header.partition_2.boot = BOOT_ACTIVE | 0x01;
1977        mbr.header.partition_2.sys = 0x0a;
1978        mbr.header.partition_2.starting_lba = 1;
1979        mbr.header.partition_2.sectors = 10;
1980        assert!(matches!(
1981            mbr.write_into(&mut cur).unwrap_err(),
1982            Error::InvalidBootFlag
1983        ));
1984
1985        // logical partition
1986        let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
1987        mbr.align = 1;
1988        mbr.header.partition_1.sys = 0x0f;
1989        mbr.header.partition_1.starting_lba = 1;
1990        mbr.header.partition_1.sectors = 10;
1991        mbr.push(0x0f, 1, 9).unwrap().partition.boot = BOOT_ACTIVE | 0x01;
1992        assert!(matches!(
1993            mbr.write_into(&mut cur).unwrap_err(),
1994            Error::InvalidBootFlag
1995        ));
1996    }
1997}
1998
1999#[cfg(doctest)]
2000mod test_readme {
2001    // for Rust < 1.54
2002    // https://blog.guillaume-gomez.fr/articles/2021-08-03+Improvements+for+%23%5Bdoc%5D+attributes+in+Rust
2003    macro_rules! check_doc {
2004        ($x:expr) => {
2005            #[doc = $x]
2006            extern "C" {}
2007        };
2008    }
2009    check_doc!(include_str!("../README.md"));
2010}