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