gpt/
mbr.rs

1//! MBR-related types and helper functions.
2//!
3//! This module provides access to low-level primitives
4//! to work with Master Boot Record (MBR), also known as LBA0.
5
6use crate::disk;
7use crate::DiskDevice;
8use std::{fmt, io};
9
10use simple_bytes::{Bytes, BytesArray, BytesRead, BytesWrite};
11
12#[non_exhaustive]
13#[derive(Debug)]
14/// Errors returned when interacting with a Gpt Disk.
15pub enum MBRError {
16    /// Generic IO Error
17    Io(io::Error),
18    /// The provided buffer does not match the expected mbr length
19    InvalidMBRLength,
20    /// invalid MBR signature
21    InvalidMBRSignature,
22    /// Invalid Partition Length != 16
23    InvalidPartitionLength,
24    /// Somthing Overflowed or Underflowed
25    /// This will never occur when dealing with sane values
26    Overflow(&'static str),
27}
28
29impl From<io::Error> for MBRError {
30    fn from(e: io::Error) -> Self {
31        Self::Io(e)
32    }
33}
34
35impl std::error::Error for MBRError {}
36
37impl fmt::Display for MBRError {
38    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
39        use MBRError::*;
40        let desc = match self {
41            Io(e) => return write!(fmt, "MBR IO Error: {e}"),
42            InvalidMBRLength => "The provided buffer does not match the expected mbr length",
43            InvalidMBRSignature => "Invalid MBR signature",
44            InvalidPartitionLength => "Invalid Partition length expected 16",
45            Overflow(m) => return write!(fmt, "MBR error Overflow: {m}"),
46        };
47        write!(fmt, "{desc}")
48    }
49}
50
51const MBR_SIGNATURE: [u8; 2] = [0x55, 0xAA];
52
53/// Protective MBR, as defined by GPT.
54pub struct ProtectiveMBR {
55    bootcode: [u8; 440],
56    disk_signature: [u8; 4],
57    unknown: u16,
58    partitions: [PartRecord; 4],
59    signature: [u8; 2],
60}
61
62impl fmt::Debug for ProtectiveMBR {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        write!(f, "Protective MBR, partitions: {:#?}", self.partitions)
65    }
66}
67
68impl Default for ProtectiveMBR {
69    fn default() -> Self {
70        Self {
71            bootcode: [0x00; 440],
72            disk_signature: [0x00; 4],
73            unknown: 0,
74            partitions: [
75                PartRecord::new_protective(None),
76                PartRecord::zero(),
77                PartRecord::zero(),
78                PartRecord::zero(),
79            ],
80            signature: MBR_SIGNATURE,
81        }
82    }
83}
84
85impl ProtectiveMBR {
86    /// Create a default protective-MBR object.
87    pub fn new() -> Self {
88        Self::default()
89    }
90
91    /// Create a protective-MBR object with a specific protective partition size (in LB).
92    /// The protective partition size should be the size of the disk - 1 (because the protective
93    /// partition always begins at LBA 1 (the second sector)).
94    pub fn with_lb_size(lb_size: u32) -> Self {
95        Self {
96            bootcode: [0x00; 440],
97            disk_signature: [0x00; 4],
98            unknown: 0,
99            partitions: [
100                PartRecord::new_protective(Some(lb_size)),
101                PartRecord::zero(),
102                PartRecord::zero(),
103                PartRecord::zero(),
104            ],
105            signature: MBR_SIGNATURE,
106        }
107    }
108
109    /// Parse input bytes into a protective-MBR object.
110    pub fn from_bytes(buf: &[u8], sector_size: disk::LogicalBlockSize) -> Result<Self, MBRError> {
111        let mut pmbr = Self::new();
112        let totlen: u64 = sector_size.into();
113
114        if buf.len() != (totlen as usize) {
115            return Err(MBRError::InvalidMBRLength);
116        }
117
118        let mut bytes = Bytes::from(buf);
119
120        pmbr.bootcode.copy_from_slice(bytes.read(440));
121        pmbr.disk_signature.copy_from_slice(bytes.read(4));
122        pmbr.unknown = bytes.read_le_u16();
123
124        for p in pmbr.partitions.iter_mut() {
125            *p = PartRecord::from_bytes(bytes.read(16))?;
126        }
127
128        assert_eq!(simple_bytes::BytesSeek::position(&bytes), 510);
129
130        pmbr.signature.copy_from_slice(bytes.read(2));
131        if pmbr.signature == MBR_SIGNATURE {
132            Ok(pmbr)
133        } else {
134            Err(MBRError::InvalidMBRSignature)
135        }
136    }
137
138    /// Read the LBA0 of a disk device and parse it into a protective-MBR object.
139    pub fn from_disk<D: DiskDevice>(
140        device: &mut D,
141        sector_size: disk::LogicalBlockSize,
142    ) -> Result<Self, MBRError> {
143        let totlen: u64 = sector_size.into();
144        let mut buf = vec![0_u8; totlen as usize];
145        let cur = device.stream_position()?;
146
147        device.seek(io::SeekFrom::Start(0))?;
148        device.read_exact(&mut buf)?;
149        let pmbr = Self::from_bytes(&buf, sector_size);
150        device.seek(io::SeekFrom::Start(cur))?;
151        pmbr
152    }
153
154    /// Return the memory representation of this MBR as a byte vector.
155    ///
156    /// This will always be 512
157    pub fn to_bytes(&self) -> [u8; 512] {
158        let mut bytes = BytesArray::from([0u8; 512]);
159
160        bytes.write(self.bootcode);
161        bytes.write(self.disk_signature);
162        bytes.write_le_u16(self.unknown);
163
164        for p in &self.partitions {
165            bytes.write(p.to_bytes());
166        }
167
168        bytes.write(self.signature);
169
170        bytes.into_array()
171    }
172
173    /// Return the 440 bytes of BIOS bootcode.
174    pub fn bootcode(&self) -> &[u8; 440] {
175        &self.bootcode
176    }
177
178    /// Set the 440 bytes of BIOS bootcode.
179    ///
180    /// This only changes the in-memory state, without overwriting
181    /// any on-disk data.
182    pub fn set_bootcode(&mut self, bootcode: [u8; 440]) -> &Self {
183        self.bootcode = bootcode;
184        self
185    }
186
187    /// Return the 4 bytes of MBR disk signature.
188    pub fn disk_signature(&self) -> &[u8; 4] {
189        &self.disk_signature
190    }
191
192    /// Set the 4 bytes of MBR disk signature.
193    ///
194    /// This only changes the in-memory state, without overwriting
195    /// any on-disk data.
196    pub fn set_disk_signature(&mut self, sig: [u8; 4]) -> &Self {
197        self.disk_signature = sig;
198        self
199    }
200
201    /// Returns the given partition (0..=3) or None if the partition index is invalid.
202    pub fn partition(&self, partition_index: usize) -> Option<PartRecord> {
203        if partition_index >= self.partitions.len() {
204            None
205        } else {
206            Some(self.partitions[partition_index])
207        }
208    }
209
210    /// Set the data for the given partition.
211    /// Returns the previous partition record or None if the partition index is invalid.
212    ///
213    /// This only changes the in-memory state, without overwriting
214    /// any on-disk data.
215    pub fn set_partition(
216        &mut self,
217        partition_index: usize,
218        partition: PartRecord,
219    ) -> Option<PartRecord> {
220        if partition_index >= self.partitions.len() {
221            None
222        } else {
223            Some(std::mem::replace(
224                &mut self.partitions[partition_index],
225                partition,
226            ))
227        }
228    }
229
230    /// Write a protective MBR to LBA0, overwriting any existing data.
231    ///
232    /// This will not write the entire lba0 if the sector size is 4096
233    pub fn overwrite_lba0<D: DiskDevice>(&self, device: &mut D) -> Result<usize, MBRError> {
234        let cur = device.stream_position()?;
235        let _ = device.seek(io::SeekFrom::Start(0))?;
236        let data = self.to_bytes();
237        device.write_all(&data)?;
238        device.flush()?;
239
240        device.seek(io::SeekFrom::Start(cur))?;
241        Ok(data.len())
242    }
243
244    /// Update LBA0, preserving most bytes of any existing MBR.
245    ///
246    /// This overwrites the four MBR partition records and the
247    /// well-known signature, leaving all other MBR bits as-is.
248    pub fn update_conservative<D: DiskDevice>(&self, device: &mut D) -> Result<usize, MBRError> {
249        let cur = device.stream_position()?;
250        // Seek to first partition record.
251        // (GPT spec 2.7 - sec. 5.2.3 - table 15)
252        let _ = device.seek(io::SeekFrom::Start(446))?;
253        for p in &self.partitions {
254            device.write_all(&p.to_bytes())?;
255        }
256        device.write_all(&self.signature)?;
257        device.flush()?;
258
259        device.seek(io::SeekFrom::Start(cur))?;
260        let bytes_updated: usize = (16 * 4) + 2;
261        Ok(bytes_updated)
262    }
263}
264
265/// A partition record, MBR-style.
266#[derive(Copy, Clone, Debug, Eq, PartialEq)]
267pub struct PartRecord {
268    /// Bit 7 set if partition is active (bootable)
269    pub boot_indicator: u8,
270    /// CHS address of partition start: 8-bit value of head in CHS address
271    pub start_head: u8,
272    /// CHS address of partition start: Upper 2 bits are 8th-9th bits of cylinder, lower 6 bits are sector
273    pub start_sector: u8,
274    /// CHS address of partition start: Lower 8 bits of cylinder
275    pub start_track: u8,
276    /// Partition type. See <https://www.win.tue.nl/~aeb/partitions/partition_types-1.html>
277    pub os_type: u8,
278    /// CHS address of partition end: 8-bit value of head in CHS address
279    pub end_head: u8,
280    /// CHS address of partition end: Upper 2 bits are 8th-9th bits of cylinder, lower 6 bits are sector
281    pub end_sector: u8,
282    /// CHS address of partition end: Lower 8 bits of cylinder
283    pub end_track: u8,
284    /// LBA of start of partition
285    pub lb_start: u32,
286    /// Number of sectors in partition
287    pub lb_size: u32,
288}
289
290impl PartRecord {
291    /// Create a protective Partition Record object with a specific disk size (in LB).
292    pub fn new_protective(lb_size: Option<u32>) -> Self {
293        let size = lb_size.unwrap_or(0xFF_FF_FF_FF);
294        Self {
295            boot_indicator: 0x00,
296            start_head: 0x00,
297            start_sector: 0x02,
298            start_track: 0x00,
299            os_type: 0xEE,
300            end_head: 0xFF,
301            end_sector: 0xFF,
302            end_track: 0xFF,
303            lb_start: 1,
304            lb_size: size,
305        }
306    }
307
308    /// Create an all-zero Partition Record.
309    pub fn zero() -> Self {
310        Self {
311            boot_indicator: 0x00,
312            start_head: 0x00,
313            start_sector: 0x00,
314            start_track: 0x00,
315            os_type: 0x00,
316            end_head: 0x00,
317            end_sector: 0x00,
318            end_track: 0x00,
319            lb_start: 0,
320            lb_size: 0,
321        }
322    }
323
324    /// Parse input bytes into a Partition Record.
325    pub fn from_bytes(buf: &[u8]) -> Result<Self, MBRError> {
326        if buf.len() != 16 {
327            return Err(MBRError::InvalidPartitionLength);
328        }
329
330        let mut bytes = Bytes::from(buf);
331
332        let pr = Self {
333            boot_indicator: bytes.read_u8(),
334            start_head: bytes.read_u8(),
335            start_sector: bytes.read_u8(),
336            start_track: bytes.read_u8(),
337            os_type: bytes.read_u8(),
338            end_head: bytes.read_u8(),
339            end_sector: bytes.read_u8(),
340            end_track: bytes.read_u8(),
341            lb_start: bytes.read_le_u32(),
342            lb_size: bytes.read_le_u32(),
343        };
344
345        Ok(pr)
346    }
347
348    /// Return the memory representation of this Partition Record as a byte vector.
349    pub fn to_bytes(&self) -> [u8; 16] {
350        let mut bytes = BytesArray::from([0u8; 16]);
351
352        bytes.write_u8(self.boot_indicator);
353
354        bytes.write_u8(self.start_head);
355        bytes.write_u8(self.start_sector);
356        bytes.write_u8(self.start_track);
357
358        bytes.write_u8(self.os_type);
359
360        bytes.write_u8(self.end_head);
361        bytes.write_u8(self.end_sector);
362        bytes.write_u8(self.end_track);
363
364        bytes.write_le_u32(self.lb_start);
365        bytes.write_le_u32(self.lb_size);
366
367        bytes.into_array()
368    }
369}
370
371/// Return the 440 bytes of BIOS bootcode.
372pub fn read_bootcode<D: DiskDevice>(device: &mut D) -> io::Result<[u8; 440]> {
373    let bootcode_offset = 0;
374    let cur = device.stream_position()?;
375    let _ = device.seek(io::SeekFrom::Start(bootcode_offset))?;
376    let mut bootcode = [0x00; 440];
377    device.read_exact(&mut bootcode)?;
378
379    device.seek(io::SeekFrom::Start(cur))?;
380    Ok(bootcode)
381}
382
383/// Write the 440 bytes of BIOS bootcode.
384pub fn write_bootcode<D: DiskDevice>(device: &mut D, bootcode: &[u8; 440]) -> io::Result<()> {
385    let bootcode_offset = 0;
386    let cur = device.stream_position()?;
387    let _ = device.seek(io::SeekFrom::Start(bootcode_offset))?;
388    device.write_all(bootcode)?;
389    device.flush()?;
390
391    device.seek(io::SeekFrom::Start(cur))?;
392    Ok(())
393}
394
395/// Read the 4 bytes of MBR disk signature.
396pub fn read_disk_signature<D: DiskDevice>(device: &mut D) -> io::Result<[u8; 4]> {
397    let dsig_offset = 440;
398    let cur = device.stream_position()?;
399    let _ = device.seek(io::SeekFrom::Start(dsig_offset))?;
400    let mut dsig = [0x00; 4];
401    device.read_exact(&mut dsig)?;
402
403    device.seek(io::SeekFrom::Start(cur))?;
404    Ok(dsig)
405}
406
407/// Write the 4 bytes of MBR disk signature.
408pub fn write_disk_signature<D: DiskDevice>(device: &mut D, sig: &[u8; 4]) -> io::Result<()> {
409    let dsig_offset = 440;
410    let cur = device.stream_position()?;
411    let _ = device.seek(io::SeekFrom::Start(dsig_offset))?;
412    device.write_all(sig)?;
413    device.flush()?;
414
415    device.seek(io::SeekFrom::Start(cur))?;
416    Ok(())
417}