1use crate::disk;
7use crate::DiskDevice;
8use std::{fmt, io};
9
10use simple_bytes::{Bytes, BytesArray, BytesRead, BytesWrite};
11
12#[non_exhaustive]
13#[derive(Debug)]
14pub enum MBRError {
16 Io(io::Error),
18 InvalidMBRLength,
20 InvalidMBRSignature,
22 InvalidPartitionLength,
24 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
53pub 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 pub fn new() -> Self {
88 Self::default()
89 }
90
91 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 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 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 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 pub fn bootcode(&self) -> &[u8; 440] {
175 &self.bootcode
176 }
177
178 pub fn set_bootcode(&mut self, bootcode: [u8; 440]) -> &Self {
183 self.bootcode = bootcode;
184 self
185 }
186
187 pub fn disk_signature(&self) -> &[u8; 4] {
189 &self.disk_signature
190 }
191
192 pub fn set_disk_signature(&mut self, sig: [u8; 4]) -> &Self {
197 self.disk_signature = sig;
198 self
199 }
200
201 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 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 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 pub fn update_conservative<D: DiskDevice>(&self, device: &mut D) -> Result<usize, MBRError> {
249 let cur = device.stream_position()?;
250 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#[derive(Copy, Clone, Debug, Eq, PartialEq)]
267pub struct PartRecord {
268 pub boot_indicator: u8,
270 pub start_head: u8,
272 pub start_sector: u8,
274 pub start_track: u8,
276 pub os_type: u8,
278 pub end_head: u8,
280 pub end_sector: u8,
282 pub end_track: u8,
284 pub lb_start: u32,
286 pub lb_size: u32,
288}
289
290impl PartRecord {
291 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 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 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 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
371pub 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
383pub 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
395pub 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
407pub 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}