file_system/mbr/
mod.rs

1/*!
2Master Boot Record (MBR) module
3
4This module provides functionality for working with MBR (Master Boot Record)
5partition tables, which are used in traditional BIOS-based systems.
6
7# Features
8
9- Parse and validate MBR structures from raw bytes or devices
10- Create and modify MBR partition tables
11- Work with individual partition entries
12- Create partition devices for accessing individual partitions
13- Utility functions for common MBR operations
14- Type-safe partition type enumeration
15
16# Examples
17
18```rust
19extern crate alloc;
20
21use file_system::*;
22
23let device = create_device!(Memory_device_type::<512>::new(512));
24
25// Create a new MBR
26let mut mbr = MBR_type::New_with_signature(0x12345678);
27// Add a FAT32 partition
28mbr.Add_partition(Partition_type_type::Fat32_lba, 2048, 204800, true).unwrap();
29
30// Write MBR to the device
31mbr.Write_to_device(&device).unwrap();
32
33// Read MBR from a device
34let mbr = MBR_type::Read_from_device(&device).unwrap();
35
36// Display MBR information
37println!("{}", mbr);
38
39// Get all valid partitions
40let partitions = mbr.get_valid_partitions();
41
42// Create a partition device
43if let Some(partition) = partitions.first() {
44    let partition_device = Create_partition_device(device.clone(), partition).unwrap();
45}
46```
47*/
48
49use alloc::vec::Vec;
50use core::fmt;
51
52use crate::{PartitionDevice, PartitionStatistics};
53
54mod utilities;
55pub use utilities::*;
56
57use crate::{Device, Error, PartitionEntry, PartitionKind, Result};
58
59#[cfg(test)]
60use crate::MemoryDevice;
61
62/// Master Boot Record structure (512 bytes)
63#[derive(Debug, Clone)]
64pub struct Mbr {
65    /// Bootstrap code (440 bytes)
66    pub bootstrap_code: [u8; 440],
67    /// Optional disk signature (4 bytes)
68    pub disk_signature: [u8; 4],
69    /// Reserved (usually 0x0000)
70    pub reserved: [u8; 2],
71    /// Partition table (4 entries × 16 bytes = 64 bytes)
72    pub partitions: [PartitionEntry; 4],
73    /// Boot signature (0x55AA)
74    pub boot_signature: [u8; 2],
75}
76
77impl Mbr {
78    /// MBR signature bytes
79    pub const SIGNATURE: [u8; 2] = [0x55, 0xAA];
80
81    /// Size of MBR in bytes
82    pub const SIZE: usize = 512;
83
84    /// Maximum number of primary partitions in MBR
85    pub const MAXIMUM_PARTITIONS_COUNT: usize = 4;
86
87    /// Create a new empty MBR
88    pub fn new() -> Self {
89        Self {
90            bootstrap_code: [0; 440],
91            disk_signature: [0; 4],
92            reserved: [0; 2],
93            partitions: [PartitionEntry::new(); 4],
94            boot_signature: Self::SIGNATURE,
95        }
96    }
97
98    /// Create a new MBR with a specific disk signature
99    pub fn new_with_signature(disk_signature: u32) -> Self {
100        let mut mbr = Self::new();
101        mbr.set_disk_signature(disk_signature);
102        mbr
103    }
104
105    /// Parse MBR from raw bytes
106    pub fn from_bytes(data: &[u8]) -> Result<Self> {
107        if data.len() < Self::SIZE {
108            return Err(Error::InvalidParameter);
109        }
110
111        // Check MBR signature
112        if data[510] != Self::SIGNATURE[0] || data[511] != Self::SIGNATURE[1] {
113            return Err(Error::Corrupted);
114        }
115
116        let mut mbr = Mbr {
117            bootstrap_code: [0; 440],
118            disk_signature: [0; 4],
119            reserved: [0; 2],
120            partitions: [PartitionEntry::new(); 4],
121            boot_signature: [0; 2],
122        };
123
124        // Copy bootstrap code
125        mbr.bootstrap_code.copy_from_slice(&data[0..440]);
126
127        // Copy disk signature
128        mbr.disk_signature.copy_from_slice(&data[440..444]);
129
130        // Copy reserved bytes
131        mbr.reserved.copy_from_slice(&data[444..446]);
132
133        // Parse partition entries
134        for (i, partition) in mbr.partitions.iter_mut().enumerate() {
135            let offset = 446 + (i * 16);
136            let partition_data = &data[offset..offset + 16];
137
138            partition.bootable = partition_data[0];
139            partition.start_head = partition_data[1];
140            partition.start_sector = partition_data[2];
141            partition.start_cylinder = partition_data[3];
142            partition.partition_type = partition_data[4];
143            partition.end_head = partition_data[5];
144            partition.end_sector = partition_data[6];
145            partition.end_cylinder = partition_data[7];
146            partition.start_lba = u32::from_le_bytes([
147                partition_data[8],
148                partition_data[9],
149                partition_data[10],
150                partition_data[11],
151            ]);
152            partition.size_sectors = u32::from_le_bytes([
153                partition_data[12],
154                partition_data[13],
155                partition_data[14],
156                partition_data[15],
157            ]);
158        }
159
160        // Copy boot signature
161        mbr.boot_signature.copy_from_slice(&data[510..512]);
162
163        Ok(mbr)
164    }
165
166    /// Read and parse MBR from a device
167    pub fn read_from_device(device: &Device) -> Result<Self> {
168        // Read the first 512 bytes (MBR sector)
169        let mut buffer = [0u8; Self::SIZE];
170
171        // Set position to the beginning of the device
172        device.set_position(&crate::Position::Start(0))?;
173
174        // Read MBR data
175        let bytes_read = device.read(&mut buffer)?;
176
177        if bytes_read.as_u64() < Self::SIZE as u64 {
178            return Err(Error::InputOutput);
179        }
180
181        Self::from_bytes(&buffer)
182    }
183
184    /// Write MBR to a device
185    pub fn write_to_device(&self, device: &Device) -> Result<()> {
186        // Set position to the beginning of the device
187        device.set_position(&crate::Position::Start(0))?;
188
189        // Convert to bytes and write
190        let buffer = self.to_bytes();
191        let bytes_written = device.write(&buffer)?;
192
193        if bytes_written.as_u64() < Self::SIZE as u64 {
194            return Err(Error::InputOutput);
195        }
196
197        device.flush()?;
198        Ok(())
199    }
200
201    /// Check if MBR has a valid signature
202    pub fn is_valid(&self) -> bool {
203        self.boot_signature == Self::SIGNATURE
204    }
205
206    /// Get all valid partitions
207    pub fn get_valid_partitions(&self) -> Vec<&PartitionEntry> {
208        self.partitions
209            .iter()
210            .filter(|partition| partition.is_valid())
211            .collect()
212    }
213
214    /// Get all valid partitions (mutable)
215    pub fn get_valid_partitions_mut(&mut self) -> Vec<&mut PartitionEntry> {
216        self.partitions
217            .iter_mut()
218            .filter(|partition| partition.is_valid())
219            .collect()
220    }
221
222    /// Get bootable partition (if any)
223    pub fn get_bootable_partition(&self) -> Option<&PartitionEntry> {
224        self.partitions
225            .iter()
226            .find(|partition| partition.is_bootable())
227    }
228
229    /// Get bootable partition (mutable, if any)
230    pub fn get_bootable_partition_mut(&mut self) -> Option<&mut PartitionEntry> {
231        self.partitions
232            .iter_mut()
233            .find(|partition| partition.is_bootable())
234    }
235
236    /// Set a partition as bootable (clears bootable flag from other partitions)
237    pub fn set_bootable_partition(&mut self, index: usize) -> Result<()> {
238        if index >= Self::MAXIMUM_PARTITIONS_COUNT {
239            return Err(Error::InvalidParameter);
240        }
241
242        // Clear bootable flag from all partitions
243        for partition in &mut self.partitions {
244            partition.set_bootable(false);
245        }
246
247        // Set the specified partition as bootable
248        self.partitions[index].set_bootable(true);
249        Ok(())
250    }
251
252    /// Check if this MBR contains a GPT protective partition
253    pub fn has_gpt_protective_partition(&self) -> bool {
254        self.partitions
255            .iter()
256            .any(|partition| partition.get_partition_type() == PartitionKind::GptProtective)
257    }
258
259    /// Get disk signature as u32
260    pub fn get_disk_signature(&self) -> u32 {
261        u32::from_le_bytes(self.disk_signature)
262    }
263
264    /// Set disk signature
265    pub fn set_disk_signature(&mut self, signature: u32) {
266        self.disk_signature = signature.to_le_bytes();
267    }
268
269    /// Get the first free partition slot
270    pub fn get_free_partition_slot(&self) -> Option<usize> {
271        self.partitions
272            .iter()
273            .position(|partition| !partition.is_valid())
274    }
275
276    /// Add a new partition
277    pub fn add_partition(
278        &mut self,
279        partition_type: crate::PartitionKind,
280        start_lba: u32,
281        size_sectors: u32,
282        bootable: bool,
283    ) -> Result<usize> {
284        let slot = self
285            .get_free_partition_slot()
286            .ok_or(Error::FileSystemFull)?;
287
288        let new_partition =
289            PartitionEntry::new_with_params(bootable, partition_type, start_lba, size_sectors);
290
291        // Check for overlaps with existing partitions
292        for existing in &self.partitions {
293            if existing.is_valid() && new_partition.overlaps_with(existing) {
294                return Err(Error::AlreadyExists);
295            }
296        }
297
298        self.partitions[slot] = new_partition;
299
300        // If this is the only bootable partition or no other bootable partition exists
301        if bootable {
302            self.set_bootable_partition(slot)?;
303        }
304
305        Ok(slot)
306    }
307
308    /// Remove a partition by index
309    pub fn remove_partition(&mut self, index: usize) -> Result<()> {
310        if index >= Self::MAXIMUM_PARTITIONS_COUNT {
311            return Err(Error::InvalidParameter);
312        }
313
314        self.partitions[index].clear();
315        Ok(())
316    }
317
318    /// Check for partition overlaps
319    pub fn has_overlapping_partitions(&self) -> bool {
320        let valid_partitions = self.get_valid_partitions();
321
322        for (i, partition1) in valid_partitions.iter().enumerate() {
323            for partition2 in valid_partitions.iter().skip(i + 1) {
324                if partition1.overlaps_with(partition2) {
325                    return true;
326                }
327            }
328        }
329
330        false
331    }
332
333    /// Get partition count
334    pub fn get_partition_count(&self) -> usize {
335        self.partitions.iter().filter(|p| p.is_valid()).count()
336    }
337
338    /// Create partition devices for all valid partitions in this MBR.
339    ///
340    /// This method iterates through all partition entries in this MBR and creates
341    /// [`PartitionDevice`] instances for each valid partition. This is useful
342    /// when you need to access all partitions on a disk programmatically.
343    ///
344    /// # Arguments
345    ///
346    /// * `base_device` - The underlying storage device containing all partitions
347    ///
348    /// # Returns
349    ///
350    /// * `Ok(Vec<Partition_device_type>)` - Vector of partition devices for all valid partitions
351    /// * `Err(Error)` - Error if any partition device creation fails
352    ///
353    /// # Examples
354    ///
355    /// ```rust
356    /// extern crate alloc;
357    /// use file_system::*;
358    ///
359    /// let device = create_device!(Memory_device_type::<512>::new(4 * 1024 * 1024));
360    /// // Create an MBR with multiple partitions
361    /// let mut mbr = MBR_type::New_with_signature(0x12345678);
362    /// mbr.Add_partition(Partition_type_type::Fat32_lba, 2048, 1024, true).unwrap();
363    /// mbr.Add_partition(Partition_type_type::Linux, 4096, 2048, false).unwrap();
364    /// mbr.Write_to_device(&device).unwrap();
365    ///
366    /// // Read it back and create all partition devices
367    /// let mbr = MBR_type::Read_from_device(&device).unwrap();
368    /// let partition_devices = mbr.Create_all_partition_devices(device).unwrap();
369    /// println!("Created {} partition devices", partition_devices.len());
370    ///
371    /// for (i, partition) in partition_devices.iter().enumerate() {
372    ///     println!("Partition {}: {} sectors", i, partition.get_sector_count());
373    /// }
374    /// ```
375    pub fn create_all_partition_devices(
376        &self,
377        base_device: Device,
378    ) -> Result<Vec<PartitionDevice>> {
379        let mut devices = Vec::new();
380
381        for partition in &self.partitions {
382            if partition.is_valid() {
383                let device = create_partition_device(base_device.clone(), partition)?;
384                devices.push(device);
385            }
386        }
387
388        Ok(devices)
389    }
390
391    /// Find partitions of a specific type within this MBR.
392    ///
393    /// This method searches through all partitions in this MBR and returns references
394    /// to those that match the specified partition type. This is useful for locating
395    /// specific types of partitions (e.g., FAT32, Linux, etc.) without creating
396    /// partition devices.
397    ///
398    /// # Arguments
399    ///
400    /// * `partition_type` - The specific partition type to find
401    ///
402    /// # Returns
403    ///
404    /// A vector of tuples containing the partition index and reference to the partition entry
405    /// for each matching partition.
406    ///
407    /// # Examples
408    ///
409    /// ```rust
410    /// extern crate alloc;
411    /// use file_system::*;
412    ///
413    /// let device = create_device!(Memory_device_type::<512>::new(4 * 1024 * 1024));
414    /// // Create an MBR with FAT32 partition
415    /// let mut mbr = MBR_type::New_with_signature(0x12345678);
416    /// mbr.Add_partition(Partition_type_type::Fat32_lba, 2048, 1024, true).unwrap();
417    /// mbr.Write_to_device(&device).unwrap();
418    ///
419    /// // Read it back and find FAT32 partitions
420    /// let mbr = MBR_type::Read_from_device(&device).unwrap();
421    /// let fat32_partitions = mbr.Find_partitions_by_type(Partition_type_type::Fat32_lba);
422    /// println!("Found {} FAT32 partitions", fat32_partitions.len());
423    /// ```
424    pub fn find_partitions_by_type(
425        &self,
426        partition_type: crate::PartitionKind,
427    ) -> Vec<(usize, &PartitionEntry)> {
428        self.partitions
429            .iter()
430            .enumerate()
431            .filter(|(_, partition)| {
432                partition.is_valid() && partition.get_partition_type() == partition_type
433            })
434            .collect()
435    }
436
437    /// Validate this MBR structure for consistency and correctness.
438    ///
439    /// This method performs comprehensive validation of this MBR structure, checking:
440    /// - MBR signature validity (0x55AA boot signature)
441    /// - Partition overlap detection
442    /// - Bootable partition count (at most one partition should be bootable)
443    ///
444    /// # Returns
445    ///
446    /// * `Ok(())` - MBR is valid and consistent
447    /// * `Err(Error::Corrupted)` - MBR is invalid or corrupted
448    ///
449    /// # Examples
450    ///
451    /// ```rust
452    /// extern crate alloc;
453    /// use file_system::*;
454    ///
455    /// let device = create_device!(Memory_device_type::<512>::new(4 * 1024 * 1024));
456    /// // First create and write a valid MBR
457    /// let mut mbr = MBR_type::New_with_signature(0x12345678);
458    /// mbr.Add_partition(Partition_type_type::Fat32_lba, 2048, 1024, true).unwrap();
459    /// mbr.Write_to_device(&device).unwrap();
460    ///
461    /// // Read it back and validate
462    /// let mbr = MBR_type::Read_from_device(&device).unwrap();
463    /// match mbr.Validate() {
464    ///     Ok(()) => println!("MBR is valid"),
465    ///     Err(Error::Corrupted) => println!("MBR is corrupted"),
466    ///     Err(e) => println!("Validation error: {}", e),
467    /// }
468    /// ```
469    pub fn validate(&self) -> Result<()> {
470        // Check MBR signature
471        if !self.is_valid() {
472            return Err(Error::Corrupted);
473        }
474
475        // Check for overlapping partitions
476        if self.has_overlapping_partitions() {
477            return Err(Error::Corrupted);
478        }
479
480        // Check that only one partition is bootable
481        let bootable_count = self.partitions.iter().filter(|p| p.is_bootable()).count();
482
483        if bootable_count > 1 {
484            return Err(Error::Corrupted);
485        }
486
487        Ok(())
488    }
489
490    /// Generate comprehensive statistics from this MBR.
491    ///
492    /// This method analyzes all partitions in this MBR and generates
493    /// detailed statistics about partition types, sizes, and other characteristics.
494    ///
495    /// # Returns
496    ///
497    /// A new `Partition_statistics_type` containing the computed statistics.
498    ///
499    /// # Examples
500    ///
501    /// ```rust
502    /// extern crate alloc;
503    /// use file_system::*;
504    ///
505    /// let device = create_device!(Memory_device_type::<512>::new(4 * 1024 * 1024));
506    /// // Create an MBR with some partitions
507    /// let mut mbr = MBR_type::New_with_signature(0x12345678);
508    /// mbr.Add_partition(Partition_type_type::Fat32_lba, 2048, 1024, true).unwrap();
509    /// mbr.Add_partition(Partition_type_type::Linux, 4096, 2048, false).unwrap();
510    /// mbr.Write_to_device(&device).unwrap();
511    ///
512    /// // Read it back and analyze
513    /// let mbr = MBR_type::Read_from_device(&device).unwrap();
514    /// let stats = mbr.Generate_statistics();
515    /// if stats.Total_partitions > 0 {
516    ///     println!("Average partition size: {} sectors",
517    ///              stats.Total_used_sectors / stats.Total_partitions as u64);
518    /// }
519    /// ```
520    pub fn generate_statistics(&self) -> PartitionStatistics {
521        let valid_partitions: Vec<_> = self.get_valid_partitions();
522
523        let total_partitions = valid_partitions.len();
524        let bootable_partitions = valid_partitions.iter().filter(|p| p.is_bootable()).count();
525
526        let fat_partitions = valid_partitions
527            .iter()
528            .filter(|p| p.get_partition_type().is_fat())
529            .count();
530
531        let linux_partitions = valid_partitions
532            .iter()
533            .filter(|p| p.get_partition_type().is_linux())
534            .count();
535
536        let hidden_partitions = valid_partitions
537            .iter()
538            .filter(|p| p.get_partition_type().is_hidden())
539            .count();
540
541        let extended_partitions = valid_partitions
542            .iter()
543            .filter(|p| p.get_partition_type().is_extended())
544            .count();
545
546        let unknown_partitions = valid_partitions
547            .iter()
548            .filter(|p| matches!(p.get_partition_type(), crate::PartitionKind::Unknown(_)))
549            .count();
550
551        let total_used_sectors = valid_partitions
552            .iter()
553            .map(|p| p.get_size_sectors() as u64)
554            .sum();
555
556        let largest_partition_sectors = valid_partitions
557            .iter()
558            .map(|p| p.get_size_sectors())
559            .max()
560            .unwrap_or(0);
561
562        let smallest_partition_sectors = valid_partitions
563            .iter()
564            .map(|p| p.get_size_sectors())
565            .min()
566            .unwrap_or(0);
567
568        PartitionStatistics {
569            total_partitions,
570            bootable_partitions,
571            fat_partitions,
572            linux_partitions,
573            hidden_partitions,
574            extended_partitions,
575            unknown_partitions,
576            total_used_sectors,
577            largest_partition_sectors,
578            smallest_partition_sectors,
579        }
580    }
581
582    /// Create a basic MBR with a single partition covering most of the disk.
583    ///
584    /// This function creates a simple MBR structure with one partition that uses
585    /// most of the available disk space. It leaves 2048 sectors at the beginning
586    /// for proper alignment, which is standard practice for modern storage devices.
587    ///
588    /// # Arguments
589    ///
590    /// * `disk_signature` - Unique 32-bit signature for the disk
591    /// * `partition_type` - Type of partition to create (e.g., FAT32, Linux, etc.)
592    /// * `total_sectors` - Total number of sectors available on the disk
593    ///
594    /// # Returns
595    ///
596    /// * `Ok(MBR_type)` - Successfully created MBR with single partition
597    /// * `Err(Error::InvalidParameter)` - Disk is too small for a partition
598    ///
599    /// # Examples
600    ///
601    /// ```rust
602    /// use file_system::*;
603    ///
604    /// // Create MBR for a 4MB device (8192 sectors)
605    /// let mbr = MBR_type::Create_basic(0x12345678, Partition_type_type::Fat32_lba, 8192).unwrap();
606    ///
607    /// // The MBR will have one FAT32 partition starting at sector 2048
608    /// let partitions = mbr.get_valid_partitions();
609    /// assert_eq!(partitions.len(), 1);
610    /// assert_eq!(partitions[0].get_start_lba(), 2048);
611    /// ```
612    pub fn create_basic(
613        disk_signature: u32,
614        partition_type: crate::PartitionKind,
615        total_sectors: u32,
616    ) -> Result<Self> {
617        let mut mbr = Self::new_with_signature(disk_signature);
618
619        // Leave some space at the beginning (typically 2048 sectors for alignment)
620        let start_lba = 2048;
621        let partition_sectors = total_sectors.saturating_sub(start_lba);
622
623        if partition_sectors == 0 {
624            return Err(Error::InvalidParameter);
625        }
626
627        mbr.add_partition(partition_type, start_lba, partition_sectors, true)?;
628
629        Ok(mbr)
630    }
631
632    /// Convert MBR back to bytes
633    pub fn to_bytes(&self) -> [u8; Self::SIZE] {
634        let mut buffer = [0u8; Self::SIZE];
635
636        // Copy bootstrap code
637        buffer[0..440].copy_from_slice(&self.bootstrap_code);
638
639        // Copy disk signature
640        buffer[440..444].copy_from_slice(&self.disk_signature);
641
642        // Copy reserved bytes
643        buffer[444..446].copy_from_slice(&self.reserved);
644
645        // Copy partition entries
646        for (i, partition) in self.partitions.iter().enumerate() {
647            let offset = 446 + (i * 16);
648            buffer[offset] = partition.bootable;
649            buffer[offset + 1] = partition.start_head;
650            buffer[offset + 2] = partition.start_sector;
651            buffer[offset + 3] = partition.start_cylinder;
652            buffer[offset + 4] = partition.partition_type;
653            buffer[offset + 5] = partition.end_head;
654            buffer[offset + 6] = partition.end_sector;
655            buffer[offset + 7] = partition.end_cylinder;
656
657            let start_lba_bytes = partition.start_lba.to_le_bytes();
658            buffer[offset + 8..offset + 12].copy_from_slice(&start_lba_bytes);
659
660            let size_bytes = partition.size_sectors.to_le_bytes();
661            buffer[offset + 12..offset + 16].copy_from_slice(&size_bytes);
662        }
663
664        // Copy boot signature
665        buffer[510..512].copy_from_slice(&self.boot_signature);
666
667        buffer
668    }
669
670    /// Find or create a partition with a specific disk signature.
671    ///
672    /// This function searches for an MBR with the specified disk signature on the device.
673    /// If found, it returns a partition device for the first valid partition. If not found,
674    /// or if no valid partitions exist, it formats the disk with a new MBR containing a
675    /// single partition that occupies the full disk space.
676    ///
677    /// # Arguments
678    ///
679    /// * `device` - The storage device to search or format
680    /// * `target_signature` - The disk signature to look for
681    /// * `partition_type` - The type of partition to create if formatting is needed
682    ///
683    /// # Returns
684    ///
685    /// * `Ok(Partition_device_type)` - Partition device for the found or created partition
686    /// * `Err(Error)` - Error if operation failed
687    ///
688    /// # Examples
689    ///
690    /// ```rust
691    /// extern crate alloc;
692    /// use file_system::*;
693    ///
694    /// let device = create_device!(Memory_device_type::<512>::new(4 * 1024 * 1024));
695    ///
696    /// // Look for partition with signature 0x12345678, create FAT32 partition if not found
697    /// let partition = MBR_type::Find_or_create_partition_with_signature(
698    ///     &device,
699    ///     0x12345678,
700    ///     Partition_type_type::Fat32_lba
701    /// ).unwrap();
702    ///
703    /// // The partition device is ready to use
704    /// assert!(partition.is_valid());
705    /// ```
706    pub fn find_or_create_partition_with_signature(
707        device: &Device,
708        target_signature: u32,
709        partition_type: crate::PartitionKind,
710    ) -> Result<PartitionDevice> {
711        // Try to read existing MBR from device
712        if let Ok(existing_mbr) = Self::read_from_device(device) {
713            // Check if the MBR has the target signature
714            if existing_mbr.get_disk_signature() == target_signature && existing_mbr.is_valid() {
715                // Get valid partitions
716                let valid_partitions = existing_mbr.get_valid_partitions();
717
718                // If we have at least one valid partition, return it
719                if !valid_partitions.is_empty() {
720                    return create_partition_device(device.clone(), valid_partitions[0]);
721                }
722            }
723        }
724
725        // Either no MBR found, wrong signature, or no valid partitions
726        // Format the disk with new MBR and create partition
727        Self::format_disk_with_signature_and_partition(device, target_signature, partition_type)
728    }
729
730    /// Format a disk with a specific signature and create a single partition.
731    ///
732    /// This function creates a new MBR with the specified disk signature and a single
733    /// partition that occupies most of the available disk space, leaving standard
734    /// alignment space at the beginning.
735    ///
736    /// # Arguments
737    ///
738    /// * `device` - The storage device to format
739    /// * `disk_signature` - The disk signature to set in the new MBR
740    /// * `partition_type` - The type of partition to create
741    ///
742    /// # Returns
743    ///
744    /// * `Ok(Partition_device_type)` - Partition device for the created partition
745    /// * `Err(Error)` - Error if operation failed
746    ///
747    /// # Examples
748    ///
749    /// ```rust
750    /// extern crate alloc;
751    /// use file_system::*;
752    ///
753    /// let device = create_device!(Memory_device_type::<512>::new(4 * 1024 * 1024));
754    ///
755    /// // Format disk and create a Linux partition with specific signature
756    /// let partition = MBR_type::Format_disk_with_signature_and_partition(
757    ///     &device,
758    ///     0xABCDEF00,
759    ///     Partition_type_type::Linux
760    /// ).unwrap();
761    ///
762    /// // Verify the partition was created correctly
763    /// assert!(partition.is_valid());
764    /// assert_eq!(partition.get_start_lba(), 2048); // Standard alignment
765    /// ```
766    pub fn format_disk_with_signature_and_partition(
767        device: &Device,
768        disk_signature: u32,
769        partition_type: crate::PartitionKind,
770    ) -> Result<PartitionDevice> {
771        // Get device size in sectors
772        let device_size = device.get_size()?;
773        let block_size = device.get_block_size()?;
774        let total_sectors = (device_size.as_u64() / block_size as u64) as u32;
775
776        // Ensure device is large enough for a meaningful partition
777        if total_sectors < 2048 {
778            return Err(Error::InvalidParameter);
779        }
780
781        // Create new MBR with the specified signature
782        let new_mbr = Self::create_basic(disk_signature, partition_type, total_sectors)?;
783
784        // Write the new MBR to device
785        new_mbr.write_to_device(device)?;
786
787        // Get the first (and only) partition
788        let valid_partitions = new_mbr.get_valid_partitions();
789        if valid_partitions.is_empty() {
790            return Err(Error::InternalError);
791        }
792
793        // Create and return partition device
794        create_partition_device(device.clone(), valid_partitions[0])
795    }
796}
797
798impl Default for Mbr {
799    fn default() -> Self {
800        Self::new()
801    }
802}
803
804impl fmt::Display for Mbr {
805    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
806        writeln!(formatter, "Master Boot Record:")?;
807        writeln!(
808            formatter,
809            "  Disk Signature: 0x{:08X}",
810            self.get_disk_signature()
811        )?;
812        writeln!(
813            formatter,
814            "  Boot Signature: 0x{:02X}{:02X}",
815            self.boot_signature[1], self.boot_signature[0]
816        )?;
817        writeln!(formatter, "  Valid: {}", self.is_valid())?;
818        writeln!(
819            formatter,
820            "  GPT Protective: {}",
821            self.has_gpt_protective_partition()
822        )?;
823        writeln!(
824            formatter,
825            "  Partition Count: {}",
826            self.get_partition_count()
827        )?;
828        writeln!(
829            formatter,
830            "  Has Overlaps: {}",
831            self.has_overlapping_partitions()
832        )?;
833        writeln!(formatter, "  Partitions:")?;
834
835        for (i, partition) in self.partitions.iter().enumerate() {
836            if partition.is_valid() {
837                writeln!(formatter, "    {}: {}", i + 1, partition)?;
838            } else {
839                writeln!(formatter, "    {}: Empty", i + 1)?;
840            }
841        }
842
843        Ok(())
844    }
845}
846
847#[cfg(test)]
848mod tests {
849    use super::*;
850    use crate::PartitionEntry;
851    use alloc::{format, vec};
852
853    fn create_sample_mbr_bytes() -> [u8; 512] {
854        let mut data = [0u8; 512];
855
856        // Set MBR signature
857        data[510] = 0x55;
858        data[511] = 0xAA;
859
860        // Set disk signature
861        data[440..444].copy_from_slice(&0x12345678u32.to_le_bytes());
862
863        // Create partition 1: FAT32 LBA, bootable, starts at LBA 2048, size 100MB
864        let partition1_offset = 446;
865        data[partition1_offset] = 0x80; // bootable
866        data[partition1_offset + 4] = 0x0C; // FAT32 LBA
867        data[partition1_offset + 8..partition1_offset + 12].copy_from_slice(&2048u32.to_le_bytes());
868        data[partition1_offset + 12..partition1_offset + 16]
869            .copy_from_slice(&204800u32.to_le_bytes());
870
871        // Create partition 2: Linux, starts after partition 1
872        let partition2_offset = 446 + 16;
873        data[partition2_offset + 4] = 0x83; // Linux
874        data[partition2_offset + 8..partition2_offset + 12]
875            .copy_from_slice(&206848u32.to_le_bytes());
876        data[partition2_offset + 12..partition2_offset + 16]
877            .copy_from_slice(&102400u32.to_le_bytes());
878
879        data
880    }
881
882    #[test]
883    fn test_mbr_new() {
884        let mbr = super::Mbr::new();
885        assert!(mbr.is_valid());
886        assert_eq!(mbr.get_disk_signature(), 0);
887        assert_eq!(mbr.get_partition_count(), 0);
888        assert!(!mbr.has_gpt_protective_partition());
889        assert!(!mbr.has_overlapping_partitions());
890    }
891
892    #[test]
893    fn test_mbr_new_with_signature() {
894        let signature = 0xDEADBEEF;
895        let mbr = Mbr::new_with_signature(signature);
896        assert!(mbr.is_valid());
897        assert_eq!(mbr.get_disk_signature(), signature);
898    }
899
900    #[test]
901    fn test_mbr_from_bytes() {
902        let data = create_sample_mbr_bytes();
903        let mbr = Mbr::from_bytes(&data).expect("Should parse MBR successfully");
904
905        assert!(mbr.is_valid());
906        assert_eq!(mbr.get_disk_signature(), 0x12345678);
907        assert_eq!(mbr.get_partition_count(), 2);
908
909        let partitions = mbr.get_valid_partitions();
910        assert_eq!(partitions.len(), 2);
911
912        // Check first partition
913        let p1 = &partitions[0];
914        assert!(p1.is_bootable());
915        assert_eq!(p1.get_partition_type(), super::PartitionKind::Fat32Lba);
916        assert_eq!(p1.get_start_lba(), 2048);
917        assert_eq!(p1.get_size_sectors(), 204800);
918
919        // Check second partition
920        let p2 = &partitions[1];
921        assert!(!p2.is_bootable());
922        assert_eq!(p2.get_partition_type(), super::PartitionKind::Linux);
923        assert_eq!(p2.get_start_lba(), 206848);
924        assert_eq!(p2.get_size_sectors(), 102400);
925    }
926
927    #[test]
928    fn test_mbr_from_bytes_invalid_signature() {
929        let mut data = create_sample_mbr_bytes();
930        data[510] = 0x00; // Invalid signature
931        data[511] = 0x00;
932
933        let result = Mbr::from_bytes(&data);
934        assert!(result.is_err());
935        assert_eq!(result.unwrap_err(), crate::Error::Corrupted);
936    }
937
938    #[test]
939    fn test_mbr_from_bytes_too_small() {
940        let data = [0u8; 256]; // Too small
941        let result = Mbr::from_bytes(&data);
942        assert!(result.is_err());
943        assert_eq!(result.unwrap_err(), crate::Error::InvalidParameter);
944    }
945
946    #[test]
947    fn test_mbr_to_bytes_round_trip() {
948        let original_data = create_sample_mbr_bytes();
949        let mbr = Mbr::from_bytes(&original_data).unwrap();
950        let serialized_data = mbr.to_bytes();
951
952        assert_eq!(original_data.len(), serialized_data.len());
953        assert_eq!(original_data, serialized_data);
954    }
955
956    #[test]
957    fn test_mbr_disk_signature() {
958        let mut mbr = super::Mbr::new();
959        assert_eq!(mbr.get_disk_signature(), 0);
960
961        mbr.set_disk_signature(0xCAFEBABE);
962        assert_eq!(mbr.get_disk_signature(), 0xCAFEBABE);
963    }
964
965    #[test]
966    fn test_mbr_get_bootable_partition() {
967        let data = create_sample_mbr_bytes();
968        let mbr = Mbr::from_bytes(&data).unwrap();
969
970        let bootable = mbr.get_bootable_partition();
971        assert!(bootable.is_some());
972        assert_eq!(
973            bootable.unwrap().get_partition_type(),
974            super::PartitionKind::Fat32Lba
975        );
976    }
977
978    #[test]
979    fn test_mbr_set_bootable_partition() {
980        let mut mbr = super::Mbr::new();
981
982        // Add two partitions
983        mbr.add_partition(super::PartitionKind::Fat32, 2048, 100000, false)
984            .unwrap();
985        mbr.add_partition(super::PartitionKind::Linux, 102048, 50000, true)
986            .unwrap();
987
988        // Second partition should be bootable
989        let bootable = mbr.get_bootable_partition().unwrap();
990        assert_eq!(bootable.get_partition_type(), super::PartitionKind::Linux);
991
992        // Set first partition as bootable
993        mbr.set_bootable_partition(0).unwrap();
994        let bootable = mbr.get_bootable_partition().unwrap();
995        assert_eq!(bootable.get_partition_type(), super::PartitionKind::Fat32);
996    }
997
998    #[test]
999    fn test_mbr_add_partition() {
1000        let mut mbr = super::Mbr::new();
1001
1002        let index = mbr
1003            .add_partition(super::PartitionKind::Fat32Lba, 2048, 204800, true)
1004            .unwrap();
1005
1006        assert_eq!(index, 0);
1007        assert_eq!(mbr.get_partition_count(), 1);
1008
1009        let partition = &mbr.partitions[index];
1010        assert!(partition.is_valid());
1011        assert!(partition.is_bootable());
1012        assert_eq!(
1013            partition.get_partition_type(),
1014            super::PartitionKind::Fat32Lba
1015        );
1016    }
1017
1018    #[test]
1019    fn test_mbr_add_overlapping_partition() {
1020        let mut mbr = super::Mbr::new();
1021
1022        // Add first partition
1023        mbr.add_partition(super::PartitionKind::Fat32, 1000, 2000, false)
1024            .unwrap();
1025
1026        // Try to add overlapping partition
1027        let result = mbr.add_partition(super::PartitionKind::Linux, 1500, 1000, false);
1028        assert!(result.is_err());
1029        assert_eq!(result.unwrap_err(), crate::Error::AlreadyExists);
1030    }
1031
1032    #[test]
1033    fn test_mbr_add_too_many_partitions() {
1034        let mut mbr = super::Mbr::new();
1035
1036        // Fill all 4 partition slots
1037        for i in 0..4 {
1038            let start = (i as u32) * 10000 + 1000;
1039            mbr.add_partition(super::PartitionKind::Linux, start, 5000, false)
1040                .unwrap();
1041        }
1042
1043        // Try to add fifth partition
1044        let result = mbr.add_partition(super::PartitionKind::Fat32, 50000, 1000, false);
1045        assert!(result.is_err());
1046        assert_eq!(result.unwrap_err(), crate::Error::FileSystemFull);
1047    }
1048
1049    #[test]
1050    fn test_mbr_remove_partition() {
1051        let mut mbr = super::Mbr::new();
1052
1053        mbr.add_partition(super::PartitionKind::Fat32, 2048, 100000, false)
1054            .unwrap();
1055        assert_eq!(mbr.get_partition_count(), 1);
1056
1057        mbr.remove_partition(0).unwrap();
1058        assert_eq!(mbr.get_partition_count(), 0);
1059    }
1060
1061    #[test]
1062    fn test_mbr_remove_invalid_partition() {
1063        let mut mbr = super::Mbr::new();
1064        let result = mbr.remove_partition(5); // Invalid index
1065        assert!(result.is_err());
1066        assert_eq!(result.unwrap_err(), crate::Error::InvalidParameter);
1067    }
1068
1069    #[test]
1070    fn test_mbr_free_partition_slot() {
1071        let mut mbr = super::Mbr::new();
1072
1073        // Should have first slot free
1074        assert_eq!(mbr.get_free_partition_slot(), Some(0));
1075
1076        // Fill first two slots
1077        mbr.add_partition(super::PartitionKind::Fat32, 1000, 1000, false)
1078            .unwrap();
1079        mbr.add_partition(super::PartitionKind::Linux, 3000, 1000, false)
1080            .unwrap();
1081
1082        // Should have third slot free
1083        assert_eq!(mbr.get_free_partition_slot(), Some(2));
1084
1085        // Fill remaining slots
1086        mbr.add_partition(super::PartitionKind::LinuxSwap, 5000, 1000, false)
1087            .unwrap();
1088        mbr.add_partition(super::PartitionKind::NtfsExfat, 7000, 1000, false)
1089            .unwrap();
1090
1091        // Should have no free slots
1092        assert_eq!(mbr.get_free_partition_slot(), None);
1093    }
1094
1095    #[test]
1096    fn test_mbr_has_gpt_protective() {
1097        let mut mbr = super::Mbr::new();
1098        assert!(!mbr.has_gpt_protective_partition());
1099
1100        mbr.add_partition(super::PartitionKind::GptProtective, 1, 0xFFFFFFFF, false)
1101            .unwrap();
1102        assert!(mbr.has_gpt_protective_partition());
1103    }
1104
1105    #[test]
1106    fn test_mbr_overlapping_partitions_detection() {
1107        let mut mbr = super::Mbr::new();
1108
1109        // Add non-overlapping partitions
1110        mbr.add_partition(super::PartitionKind::Fat32, 1000, 1000, false)
1111            .unwrap();
1112        mbr.add_partition(super::PartitionKind::Linux, 3000, 1000, false)
1113            .unwrap();
1114        assert!(!mbr.has_overlapping_partitions());
1115
1116        // Manually create overlapping partitions (bypassing validation)
1117        mbr.partitions[2] =
1118            PartitionEntry::new_with_params(false, super::PartitionKind::LinuxSwap, 1500, 1000);
1119        assert!(mbr.has_overlapping_partitions());
1120    }
1121
1122    #[test]
1123    fn test_mbr_default() {
1124        let mbr = super::Mbr::default();
1125        assert!(mbr.is_valid());
1126        assert_eq!(mbr.get_partition_count(), 0);
1127    }
1128
1129    #[test]
1130    fn test_mbr_display() {
1131        let data = create_sample_mbr_bytes();
1132        let mbr = Mbr::from_bytes(&data).unwrap();
1133
1134        let display_string = format!("{mbr}");
1135        assert!(display_string.contains("Master Boot Record"));
1136        assert!(display_string.contains("Disk Signature: 0x12345678"));
1137        assert!(display_string.contains("Valid: true"));
1138        assert!(display_string.contains("Partition Count: 2"));
1139        assert!(display_string.contains("FAT32 LBA"));
1140        assert!(display_string.contains("Linux"));
1141    }
1142
1143    #[test]
1144    fn test_mbr_constants() {
1145        assert_eq!(super::Mbr::SIZE, 512);
1146        assert_eq!(super::Mbr::MAXIMUM_PARTITIONS_COUNT, 4);
1147        assert_eq!(super::Mbr::SIGNATURE, [0x55, 0xAA]);
1148    }
1149
1150    #[test]
1151    fn test_find_or_create_partition_with_signature_existing() {
1152        let mbr_data = create_sample_mbr_bytes();
1153        // Create a larger device (4MB) and put the MBR at the beginning
1154        let mut data = vec![0u8; 4096 * 1024];
1155        data[0..512].copy_from_slice(&mbr_data);
1156
1157        let memory_device = MemoryDevice::<512>::from_vec(data);
1158        let device = crate::create_device!(memory_device);
1159
1160        // The sample MBR has signature 0x12345678
1161        let result = Mbr::find_or_create_partition_with_signature(
1162            &device,
1163            0x12345678,
1164            PartitionKind::Fat32Lba,
1165        );
1166        assert!(result.is_ok());
1167
1168        let partition_device = result.unwrap();
1169        assert!(partition_device.is_valid());
1170        assert_eq!(partition_device.get_start_lba(), 2048); // From sample data
1171    }
1172
1173    #[test]
1174    fn test_find_or_create_partition_with_signature_wrong_signature() {
1175        let mbr_data = create_sample_mbr_bytes();
1176        // Create a larger device (4MB) and put the MBR at the beginning
1177        let mut data = vec![0u8; 4096 * 1024];
1178        data[0..512].copy_from_slice(&mbr_data);
1179
1180        let memory_device = MemoryDevice::<512>::from_vec(data);
1181        let device = crate::create_device!(memory_device);
1182
1183        // Request different signature than what's in the sample MBR
1184        let result = Mbr::find_or_create_partition_with_signature(
1185            &device,
1186            0xABCDEF00, // Different from 0x12345678 in sample
1187            PartitionKind::Linux,
1188        );
1189
1190        assert!(result.is_ok());
1191
1192        let partition_device = result.unwrap();
1193        assert!(partition_device.is_valid());
1194
1195        // Verify new MBR was created with correct signature
1196        let new_mbr = Mbr::read_from_device(&device).unwrap();
1197        assert_eq!(new_mbr.get_disk_signature(), 0xABCDEF00);
1198
1199        // Check partition type
1200        let valid_partitions = new_mbr.get_valid_partitions();
1201        assert_eq!(valid_partitions.len(), 1);
1202        assert_eq!(
1203            valid_partitions[0].get_partition_type(),
1204            PartitionKind::Linux
1205        );
1206    }
1207
1208    #[test]
1209    fn test_find_or_create_partition_with_signature_no_mbr() {
1210        // Create device with no MBR
1211        let data = vec![0u8; 4096 * 1024]; // 8MB device with no MBR
1212        let memory_device = MemoryDevice::<512>::from_vec(data);
1213        let device = crate::create_device!(memory_device);
1214
1215        let result = Mbr::find_or_create_partition_with_signature(
1216            &device,
1217            0x11223344,
1218            PartitionKind::Fat32Lba,
1219        );
1220        assert!(result.is_ok());
1221
1222        let partition_device = result.unwrap();
1223        assert!(partition_device.is_valid());
1224        assert_eq!(partition_device.get_start_lba(), 2048); // Standard alignment
1225
1226        // Verify MBR was created with correct signature
1227        let mbr = Mbr::read_from_device(&device).unwrap();
1228        assert_eq!(mbr.get_disk_signature(), 0x11223344);
1229        assert!(mbr.is_valid());
1230    }
1231
1232    #[test]
1233    fn test_find_or_create_partition_with_signature_empty_mbr() {
1234        // Create device with valid MBR but no partitions
1235        let mut data = vec![0u8; 4096 * 1024];
1236        let empty_mbr = Mbr::new_with_signature(0x55667788);
1237        let mbr_bytes = empty_mbr.to_bytes();
1238        data[0..512].copy_from_slice(&mbr_bytes);
1239
1240        let memory_device = MemoryDevice::<512>::from_vec(data);
1241        let device = crate::create_device!(memory_device);
1242
1243        let result = Mbr::find_or_create_partition_with_signature(
1244            &device,
1245            0x55667788, // Same signature as existing MBR
1246            PartitionKind::NtfsExfat,
1247        );
1248        assert!(result.is_ok());
1249
1250        let partition_device = result.unwrap();
1251        assert!(partition_device.is_valid());
1252
1253        // Verify partition was created with correct type
1254        let mbr = Mbr::read_from_device(&device).unwrap();
1255        assert_eq!(mbr.get_disk_signature(), 0x55667788);
1256        let valid_partitions = mbr.get_valid_partitions();
1257        assert_eq!(valid_partitions.len(), 1);
1258        assert_eq!(
1259            valid_partitions[0].get_partition_type(),
1260            PartitionKind::NtfsExfat
1261        );
1262    }
1263
1264    #[test]
1265    fn test_format_disk_with_signature_and_partition() {
1266        let data = vec![0u8; 2048 * 1024]; // 4MB device
1267        let memory_device = MemoryDevice::<512>::from_vec(data);
1268        let device = crate::create_device!(memory_device);
1269
1270        let result = Mbr::format_disk_with_signature_and_partition(
1271            &device,
1272            0x99887766,
1273            PartitionKind::LinuxSwap,
1274        );
1275        assert!(result.is_ok());
1276
1277        let partition_device = result.unwrap();
1278        assert!(partition_device.is_valid());
1279        assert_eq!(partition_device.get_start_lba(), 2048);
1280
1281        // Verify MBR
1282        let mbr = Mbr::read_from_device(&device).unwrap();
1283        assert_eq!(mbr.get_disk_signature(), 0x99887766);
1284        assert!(mbr.is_valid());
1285
1286        let valid_partitions = mbr.get_valid_partitions();
1287        assert_eq!(valid_partitions.len(), 1);
1288        assert_eq!(
1289            valid_partitions[0].get_partition_type(),
1290            PartitionKind::LinuxSwap
1291        );
1292        assert!(valid_partitions[0].is_bootable()); // Should be bootable by default
1293    }
1294
1295    #[test]
1296    fn test_format_disk_with_signature_and_partition_device_too_small() {
1297        // Create very small device (less than 2048 sectors)
1298        let data = vec![0u8; 1024]; // 2 sectors only
1299        let memory_device = MemoryDevice::<512>::from_vec(data);
1300        let device = crate::create_device!(memory_device);
1301
1302        let result = Mbr::format_disk_with_signature_and_partition(
1303            &device,
1304            0x12345678,
1305            PartitionKind::Fat32Lba,
1306        );
1307        assert!(result.is_err());
1308        assert_eq!(result.unwrap_err(), Error::InvalidParameter);
1309    }
1310}
file_system/mbr/mod.rs

file_system/mbr/
mod.rs
