1use crate::{device::BlockDeviceExt, partition::PartitionType};
2
3#[derive(Debug, PartialEq, Clone, Copy, Hash)]
5pub enum PartitionTable {
6 Msdos,
7 Gpt,
8}
9
10#[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
19pub trait PartitionTableExt: BlockDeviceExt {
21 fn get_partition_table(&self) -> Option<PartitionTable>;
23
24 fn get_partition_type_count(&self) -> (usize, usize, bool);
26
27 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}