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