Skip to main content

disk_types_aosc/
table.rs

1use crate::{device::BlockDeviceExt, partition::PartitionType};
2
3/// Specifies whether the partition table on the disk is **MSDOS** or **GPT**.
4#[derive(Debug, PartialEq, Clone, Copy, Hash)]
5pub enum PartitionTable {
6    Msdos,
7    Gpt,
8}
9
10/// A possible error when validating the partition table.
11#[derive(Debug, thiserror::Error, PartialEq)]
12pub enum PartitionTableError {
13    #[error("primary partitions exceeded on partition table")]
14    PrimaryPartitionsExceeded,
15    #[error("partition table not found")]
16    NotFound,
17}
18
19/// Methods for block devices that may have a partition table.
20pub trait PartitionTableExt: BlockDeviceExt {
21    /// Fetch the partition table info on this device, if it exists.
22    fn get_partition_table(&self) -> Option<PartitionTable>;
23
24    /// Obtain the number of primary and logical partitions, in that order.
25    fn get_partition_type_count(&self) -> (usize, usize, bool);
26
27    /// Checks if the additional partition type can be added to the partition table.
28    fn supports_additional_partition_type(
29        &self,
30        new_type: PartitionType,
31    ) -> Result<(), PartitionTableError> {
32        match self.get_partition_table() {
33            Some(PartitionTable::Gpt) => (),
34            Some(PartitionTable::Msdos) => {
35                let (primary, logical, extended) = self.get_partition_type_count();
36                if new_type == PartitionType::Primary {
37                    if primary >= 4 || (primary >= 3 && (extended || logical != 0)) {
38                        return Err(PartitionTableError::PrimaryPartitionsExceeded);
39                    }
40                } else if primary >= 4 {
41                    return Err(PartitionTableError::PrimaryPartitionsExceeded);
42                }
43            }
44            None => return Err(PartitionTableError::NotFound),
45        }
46
47        Ok(())
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    use std::path::Path;
55
56    pub struct FictionalBlock {
57        partitions: Vec<PartitionType>,
58    }
59
60    impl BlockDeviceExt for FictionalBlock {
61        fn get_device_name(&self) -> &str {
62            "fictional"
63        }
64
65        fn get_device_path(&self) -> &Path {
66            Path::new("/dev/fictional")
67        }
68
69        fn get_mount_point(&self) -> Option<&Path> {
70            None
71        }
72    }
73
74    impl PartitionTableExt for FictionalBlock {
75        fn get_partition_table(&self) -> Option<PartitionTable> {
76            Some(PartitionTable::Msdos)
77        }
78
79        fn get_partition_type_count(&self) -> (usize, usize, bool) {
80            self.partitions
81                .iter()
82                .fold((0, 0, false), |sum, &part| match part {
83                    PartitionType::Logical => (sum.0, sum.1 + 1, sum.2),
84                    PartitionType::Primary => (sum.0 + 1, sum.1, sum.2),
85                    PartitionType::Extended => (sum.0, sum.1, true),
86                })
87        }
88    }
89
90    #[test]
91    fn partition_table_msdos_checks() {
92        let maxed_block = FictionalBlock {
93            partitions: vec![
94                PartitionType::Primary,
95                PartitionType::Primary,
96                PartitionType::Primary,
97                PartitionType::Primary,
98            ],
99        };
100
101        assert_eq!(
102            maxed_block.supports_additional_partition_type(PartitionType::Primary),
103            Err(PartitionTableError::PrimaryPartitionsExceeded)
104        );
105
106        assert_eq!(
107            maxed_block.supports_additional_partition_type(PartitionType::Logical),
108            Err(PartitionTableError::PrimaryPartitionsExceeded)
109        );
110
111        let max_extended = FictionalBlock {
112            partitions: vec![
113                PartitionType::Primary,
114                PartitionType::Primary,
115                PartitionType::Primary,
116                PartitionType::Extended,
117                PartitionType::Logical,
118                PartitionType::Logical,
119            ],
120        };
121
122        assert_eq!(
123            max_extended.supports_additional_partition_type(PartitionType::Primary),
124            Err(PartitionTableError::PrimaryPartitionsExceeded)
125        );
126
127        assert_eq!(
128            max_extended.supports_additional_partition_type(PartitionType::Logical),
129            Ok(())
130        );
131
132        let free = FictionalBlock {
133            partitions: vec![
134                PartitionType::Primary,
135                PartitionType::Primary,
136                PartitionType::Extended,
137                PartitionType::Logical,
138                PartitionType::Logical,
139            ],
140        };
141
142        assert_eq!(
143            free.supports_additional_partition_type(PartitionType::Primary),
144            Ok(())
145        );
146
147        assert_eq!(
148            free.supports_additional_partition_type(PartitionType::Logical),
149            Ok(())
150        );
151    }
152}