1#![cfg_attr(not(feature = "std"), no_std)]
2
3extern crate alloc;
4
5#[cfg(test)]
6extern crate std;
7
8use alloc::{boxed::Box, format, string::String, sync::Arc};
9use async_trait::async_trait;
10use bytemuck::try_from_bytes;
11use core::fmt;
12use gibblox_core::{
13 AlignedByteReader, BlockReader, BlockReaderConfigIdentity, GibbloxError, GibbloxErrorKind,
14 GibbloxResult, ReadContext, WindowBlockReader,
15};
16use hadris_part::mbr::MasterBootRecord;
17use tracing::{info, trace};
18
19const MBR_SECTOR_SIZE: usize = 512;
20#[cfg(test)]
21const MBR_SIGNATURE: [u8; 2] = [0x55, 0xaa];
22#[cfg(test)]
23const MBR_SIGNATURE_OFFSET: usize = 510;
24const MBR_DISK_SIGNATURE_OFFSET: usize = 0x1b8;
25#[cfg(test)]
26const MBR_PARTITION_TABLE_OFFSET: usize = 0x1be;
27#[cfg(test)]
28const MBR_PARTITION_ENTRY_LEN: usize = 16;
29const MBR_PRIMARY_PARTITION_COUNT: usize = 4;
30
31const MBR_EXTENDED_TYPE_CHS: u8 = 0x05;
32const MBR_EXTENDED_TYPE_LBA: u8 = 0x0f;
33const MBR_EXTENDED_TYPE_LINUX: u8 = 0x85;
34
35#[derive(Clone, Debug, PartialEq, Eq)]
36pub enum MbrPartitionSelector {
37 PartUuid(String),
38 Index(u32),
39}
40
41impl MbrPartitionSelector {
42 pub fn part_uuid(value: impl Into<String>) -> Self {
43 Self::PartUuid(value.into())
44 }
45
46 pub const fn index(value: u32) -> Self {
47 Self::Index(value)
48 }
49}
50
51impl fmt::Display for MbrPartitionSelector {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 match self {
54 Self::PartUuid(value) => write!(f, "partuuid={value}"),
55 Self::Index(value) => write!(f, "index={value}"),
56 }
57 }
58}
59
60#[derive(Clone, Debug, PartialEq, Eq)]
61pub struct MbrBlockReaderConfig {
62 pub selector: MbrPartitionSelector,
63 pub block_size: u32,
64 pub source_identity: Option<String>,
65}
66
67impl MbrBlockReaderConfig {
68 pub fn new(selector: MbrPartitionSelector, block_size: u32) -> Self {
69 Self {
70 selector,
71 block_size,
72 source_identity: None,
73 }
74 }
75
76 pub fn with_source_identity(mut self, source_identity: impl Into<String>) -> Self {
77 self.source_identity = Some(source_identity.into());
78 self
79 }
80
81 fn validate(&self) -> GibbloxResult<()> {
82 if self.block_size == 0 || !self.block_size.is_power_of_two() {
83 return Err(GibbloxError::with_message(
84 GibbloxErrorKind::InvalidInput,
85 "block size must be non-zero power of two",
86 ));
87 }
88 Ok(())
89 }
90}
91
92impl BlockReaderConfigIdentity for MbrBlockReaderConfig {
93 fn write_identity(&self, out: &mut dyn fmt::Write) -> fmt::Result {
94 out.write_str("mbr-partition:(")?;
95 out.write_str(
96 self.source_identity
97 .as_deref()
98 .unwrap_or("<unknown-source>"),
99 )?;
100 write!(
101 out,
102 "):selector={}:block_size={}",
103 self.selector, self.block_size
104 )
105 }
106}
107
108pub struct MbrBlockReader {
109 partition_size_bytes: u64,
110 partition_index: u32,
111 partition_partuuid: String,
112 partition_reader: WindowBlockReader<Arc<dyn BlockReader>>,
113 config: MbrBlockReaderConfig,
114}
115
116impl MbrBlockReader {
117 pub async fn new<S: BlockReader + 'static>(
118 source: S,
119 selector: MbrPartitionSelector,
120 block_size: u32,
121 ) -> GibbloxResult<Self> {
122 Self::open_with_config(source, MbrBlockReaderConfig::new(selector, block_size)).await
123 }
124
125 pub async fn open_with_config<S: BlockReader + 'static>(
126 source: S,
127 config: MbrBlockReaderConfig,
128 ) -> GibbloxResult<Self> {
129 config.validate()?;
130 info!(selector = %config.selector, block_size = config.block_size, "constructing MBR-backed reader");
131
132 let source_block_size = source.block_size();
133 if source_block_size == 0 || !source_block_size.is_power_of_two() {
134 return Err(GibbloxError::with_message(
135 GibbloxErrorKind::InvalidInput,
136 "source block size must be non-zero power of two",
137 ));
138 }
139
140 let source_total_blocks = source.total_blocks().await?;
141 let source_size_bytes = source_total_blocks
142 .checked_mul(source_block_size as u64)
143 .ok_or_else(|| {
144 GibbloxError::with_message(GibbloxErrorKind::OutOfRange, "image size overflow")
145 })?;
146 if source_size_bytes < MBR_SECTOR_SIZE as u64 {
147 return Err(GibbloxError::with_message(
148 GibbloxErrorKind::InvalidInput,
149 "MBR image is smaller than one sector",
150 ));
151 }
152
153 let source: Arc<dyn BlockReader> = Arc::new(source);
154 let source_identity = config
155 .source_identity
156 .clone()
157 .unwrap_or_else(|| gibblox_core::block_identity_string(source.as_ref()));
158 let config = config.with_source_identity(source_identity);
159 let byte_reader = AlignedByteReader::new(Arc::clone(&source)).await?;
160
161 let mut mbr_sector = [0u8; MBR_SECTOR_SIZE];
162 byte_reader
163 .read_exact_at(0, &mut mbr_sector, ReadContext::FOREGROUND)
164 .await?;
165 let header = parse_mbr_header(&mbr_sector)?;
166 let disk_signature = mbr_disk_signature(&header);
167
168 let selected = select_partition_entry(&header, &config.selector)?;
169 let partition_offset_bytes = u64::from(selected.first_lba)
170 .checked_mul(MBR_SECTOR_SIZE as u64)
171 .ok_or_else(|| {
172 GibbloxError::with_message(
173 GibbloxErrorKind::OutOfRange,
174 "partition offset overflow",
175 )
176 })?;
177 let partition_size_bytes = u64::from(selected.sector_count)
178 .checked_mul(MBR_SECTOR_SIZE as u64)
179 .ok_or_else(|| {
180 GibbloxError::with_message(GibbloxErrorKind::OutOfRange, "partition size overflow")
181 })?;
182 let partition_end = partition_offset_bytes
183 .checked_add(partition_size_bytes)
184 .ok_or_else(|| {
185 GibbloxError::with_message(GibbloxErrorKind::OutOfRange, "partition range overflow")
186 })?;
187 if partition_end > source_size_bytes {
188 return Err(GibbloxError::with_message(
189 GibbloxErrorKind::OutOfRange,
190 "selected partition exceeds image size",
191 ));
192 }
193
194 let partition_reader = WindowBlockReader::new(
195 Arc::clone(&source),
196 partition_offset_bytes,
197 partition_size_bytes,
198 config.block_size,
199 )
200 .await?;
201
202 let partition_number = u8::try_from(selected.index + 1).map_err(|_| {
203 GibbloxError::with_message(
204 GibbloxErrorKind::OutOfRange,
205 "partition number exceeds addressable range",
206 )
207 })?;
208 let partition_partuuid = format_mbr_partuuid(disk_signature, partition_number);
209
210 info!(
211 partition_index = selected.index,
212 partuuid = %partition_partuuid,
213 first_lba = selected.first_lba,
214 sector_count = selected.sector_count,
215 partition_size_bytes,
216 "resolved MBR partition"
217 );
218
219 Ok(Self {
220 partition_size_bytes,
221 partition_index: selected.index,
222 partition_partuuid,
223 partition_reader,
224 config,
225 })
226 }
227
228 pub fn config(&self) -> &MbrBlockReaderConfig {
229 &self.config
230 }
231
232 pub fn partition_size_bytes(&self) -> u64 {
233 self.partition_size_bytes
234 }
235
236 pub fn partition_index(&self) -> u32 {
237 self.partition_index
238 }
239
240 pub fn partition_partuuid(&self) -> &str {
241 &self.partition_partuuid
242 }
243}
244
245#[async_trait]
246impl BlockReader for MbrBlockReader {
247 fn block_size(&self) -> u32 {
248 self.partition_reader.block_size()
249 }
250
251 async fn total_blocks(&self) -> GibbloxResult<u64> {
252 self.partition_reader.total_blocks().await
253 }
254
255 fn write_identity(&self, out: &mut dyn fmt::Write) -> fmt::Result {
256 self.config.write_identity(out)
257 }
258
259 async fn read_blocks(
260 &self,
261 lba: u64,
262 buf: &mut [u8],
263 ctx: ReadContext,
264 ) -> GibbloxResult<usize> {
265 let read = self.partition_reader.read_blocks(lba, buf, ctx).await?;
266 trace!(
267 lba,
268 requested = buf.len(),
269 read,
270 "reading partition blocks from MBR"
271 );
272 Ok(read)
273 }
274}
275
276#[derive(Clone, Copy, Debug)]
277struct SelectedPartition {
278 index: u32,
279 partition_type: u8,
280 first_lba: u32,
281 sector_count: u32,
282}
283
284impl SelectedPartition {
285 fn is_unused(&self) -> bool {
286 self.partition_type == 0 || self.sector_count == 0
287 }
288
289 fn is_extended(&self) -> bool {
290 matches!(
291 self.partition_type,
292 MBR_EXTENDED_TYPE_CHS | MBR_EXTENDED_TYPE_LBA | MBR_EXTENDED_TYPE_LINUX
293 )
294 }
295}
296
297fn parse_mbr_header(raw: &[u8]) -> GibbloxResult<MasterBootRecord> {
298 if raw.len() < MBR_SECTOR_SIZE {
299 return Err(GibbloxError::with_message(
300 GibbloxErrorKind::InvalidInput,
301 "MBR header block is too short",
302 ));
303 }
304
305 let mbr = *try_from_bytes::<MasterBootRecord>(&raw[..MBR_SECTOR_SIZE]).map_err(|_| {
306 GibbloxError::with_message(
307 GibbloxErrorKind::InvalidInput,
308 "MBR header layout is invalid",
309 )
310 })?;
311 if !mbr.has_valid_signature() {
312 return Err(GibbloxError::with_message(
313 GibbloxErrorKind::InvalidInput,
314 "MBR signature not found at LBA0",
315 ));
316 }
317 let partition_table = mbr.get_partition_table();
318 if !partition_table.is_valid() {
319 return Err(GibbloxError::with_message(
320 GibbloxErrorKind::InvalidInput,
321 "MBR partition entry has invalid boot indicator",
322 ));
323 }
324
325 Ok(mbr)
326}
327
328fn mbr_disk_signature(mbr: &MasterBootRecord) -> u32 {
329 let start = MBR_DISK_SIGNATURE_OFFSET;
330 u32::from_le_bytes([
331 mbr.bootstrap[start],
332 mbr.bootstrap[start + 1],
333 mbr.bootstrap[start + 2],
334 mbr.bootstrap[start + 3],
335 ])
336}
337
338fn select_partition_entry(
339 header: &MasterBootRecord,
340 selector: &MbrPartitionSelector,
341) -> GibbloxResult<SelectedPartition> {
342 let disk_signature = mbr_disk_signature(header);
343 let partition_table = header.get_partition_table();
344
345 let selected = match selector {
346 MbrPartitionSelector::Index(index) => {
347 let index_usize = usize::try_from(*index).map_err(|_| {
348 GibbloxError::with_message(
349 GibbloxErrorKind::InvalidInput,
350 "partition index exceeds addressable range",
351 )
352 })?;
353 if index_usize >= MBR_PRIMARY_PARTITION_COUNT {
354 return Err(GibbloxError::with_message(
355 GibbloxErrorKind::InvalidInput,
356 "MBR partition index not found",
357 ));
358 }
359 let partition = partition_table.partitions[index_usize];
360 SelectedPartition {
361 index: *index,
362 partition_type: partition.part_type,
363 first_lba: partition.start_lba,
364 sector_count: partition.sector_count,
365 }
366 }
367 MbrPartitionSelector::PartUuid(raw_uuid) => {
368 let (target_disk_signature, target_partition_number) =
369 parse_mbr_partuuid_text(raw_uuid)?;
370 if target_partition_number > MBR_PRIMARY_PARTITION_COUNT as u8 {
371 return Err(GibbloxError::with_message(
372 GibbloxErrorKind::Unsupported,
373 "extended/logical MBR partitions are not supported",
374 ));
375 }
376 if disk_signature != target_disk_signature {
377 return Err(GibbloxError::with_message(
378 GibbloxErrorKind::InvalidInput,
379 "MBR partition UUID not found",
380 ));
381 }
382
383 let index = usize::from(target_partition_number - 1);
384 let partition = partition_table.partitions[index];
385 SelectedPartition {
386 index: u32::from(target_partition_number - 1),
387 partition_type: partition.part_type,
388 first_lba: partition.start_lba,
389 sector_count: partition.sector_count,
390 }
391 }
392 };
393
394 if selected.is_unused() {
395 return Err(GibbloxError::with_message(
396 GibbloxErrorKind::InvalidInput,
397 "selected MBR partition entry is unused",
398 ));
399 }
400 if selected.is_extended() {
401 return Err(GibbloxError::with_message(
402 GibbloxErrorKind::Unsupported,
403 "extended/logical MBR partitions are not supported",
404 ));
405 }
406 Ok(selected)
407}
408
409fn parse_mbr_partuuid_text(value: &str) -> GibbloxResult<(u32, u8)> {
410 let trimmed = value.trim();
411 let (disk_signature_text, partition_number_text) = trimmed
412 .split_once('-')
413 .ok_or_else(invalid_mbr_partuuid_error)?;
414 if disk_signature_text.len() != 8 || partition_number_text.len() != 2 {
415 return Err(invalid_mbr_partuuid_error());
416 }
417
418 let disk_signature =
419 u32::from_str_radix(disk_signature_text, 16).map_err(|_| invalid_mbr_partuuid_error())?;
420 let partition_number =
421 u8::from_str_radix(partition_number_text, 16).map_err(|_| invalid_mbr_partuuid_error())?;
422 if partition_number == 0 {
423 return Err(invalid_mbr_partuuid_error());
424 }
425 Ok((disk_signature, partition_number))
426}
427
428fn format_mbr_partuuid(disk_signature: u32, partition_number: u8) -> String {
429 format!("{disk_signature:08x}-{partition_number:02x}")
430}
431
432fn invalid_mbr_partuuid_error() -> GibbloxError {
433 GibbloxError::with_message(
434 GibbloxErrorKind::InvalidInput,
435 "invalid MBR partition UUID text",
436 )
437}
438
439#[cfg(test)]
440mod tests {
441 use super::*;
442 use alloc::vec;
443 use futures::executor::block_on;
444
445 const TEST_BLOCK_SIZE: usize = 512;
446 const TEST_TOTAL_BLOCKS: usize = 256;
447 const TEST_DISK_SIGNATURE: u32 = 0x9439_af65;
448 const TEST_PART1_UUID: &str = "9439af65-01";
449 const TEST_PART2_UUID: &str = "9439af65-02";
450
451 struct FakeReader {
452 block_size: u32,
453 data: alloc::vec::Vec<u8>,
454 }
455
456 #[async_trait]
457 impl BlockReader for FakeReader {
458 fn block_size(&self) -> u32 {
459 self.block_size
460 }
461
462 async fn total_blocks(&self) -> GibbloxResult<u64> {
463 Ok(self.data.len().div_ceil(self.block_size as usize) as u64)
464 }
465
466 fn write_identity(&self, out: &mut dyn fmt::Write) -> fmt::Result {
467 write!(out, "fake-mbr:{}:{}", self.block_size, self.data.len())
468 }
469
470 async fn read_blocks(
471 &self,
472 lba: u64,
473 buf: &mut [u8],
474 _ctx: ReadContext,
475 ) -> GibbloxResult<usize> {
476 if !buf.len().is_multiple_of(self.block_size as usize) {
477 return Err(GibbloxError::with_message(
478 GibbloxErrorKind::InvalidInput,
479 "buffer length must align to block size",
480 ));
481 }
482 let offset = (lba as usize)
483 .checked_mul(self.block_size as usize)
484 .ok_or_else(|| {
485 GibbloxError::with_message(GibbloxErrorKind::OutOfRange, "offset overflow")
486 })?;
487 if offset >= self.data.len() {
488 return Err(GibbloxError::with_message(
489 GibbloxErrorKind::OutOfRange,
490 "requested block out of range",
491 ));
492 }
493
494 let read_len = (self.data.len() - offset).min(buf.len());
495 buf[..read_len].copy_from_slice(&self.data[offset..offset + read_len]);
496 if read_len < buf.len() {
497 buf[read_len..].fill(0);
498 }
499 Ok(buf.len())
500 }
501 }
502
503 #[test]
504 fn mbr_reader_selects_partition_by_partuuid() {
505 let (disk, partition_one_data, _partition_two_data) = build_test_mbr_disk();
506 let reader = FakeReader {
507 block_size: TEST_BLOCK_SIZE as u32,
508 data: disk,
509 };
510
511 let mbr = block_on(MbrBlockReader::new(
512 reader,
513 MbrPartitionSelector::part_uuid(TEST_PART1_UUID),
514 1024,
515 ))
516 .expect("construct MBR partition reader");
517
518 assert_eq!(mbr.partition_size_bytes(), 1536);
519 assert_eq!(mbr.partition_index(), 0);
520 assert_eq!(mbr.partition_partuuid(), TEST_PART1_UUID);
521 assert_eq!(block_on(mbr.total_blocks()).expect("total blocks"), 2);
522
523 let mut first = vec![0u8; 1024];
524 block_on(mbr.read_blocks(0, &mut first, ReadContext::FOREGROUND)).expect("read first");
525 assert_eq!(&first[..], &partition_one_data[..1024]);
526
527 let mut second = vec![0u8; 1024];
528 block_on(mbr.read_blocks(1, &mut second, ReadContext::FOREGROUND)).expect("read second");
529 assert_eq!(&second[..512], &partition_one_data[1024..1536]);
530 assert!(second[512..].iter().all(|byte| *byte == 0));
531 }
532
533 #[test]
534 fn mbr_reader_selects_partition_by_index() {
535 let (disk, _partition_one_data, partition_two_data) = build_test_mbr_disk();
536 let reader = FakeReader {
537 block_size: TEST_BLOCK_SIZE as u32,
538 data: disk,
539 };
540
541 let mbr = block_on(MbrBlockReader::new(
542 reader,
543 MbrPartitionSelector::index(1),
544 TEST_BLOCK_SIZE as u32,
545 ))
546 .expect("construct MBR partition reader");
547
548 assert_eq!(mbr.partition_index(), 1);
549 assert_eq!(mbr.partition_partuuid(), TEST_PART2_UUID);
550 assert_eq!(mbr.partition_size_bytes(), 1024);
551
552 let mut block = vec![0u8; TEST_BLOCK_SIZE];
553 block_on(mbr.read_blocks(0, &mut block, ReadContext::FOREGROUND)).expect("read block");
554 assert_eq!(&block[..], &partition_two_data[..TEST_BLOCK_SIZE]);
555 }
556
557 #[test]
558 fn mbr_reader_reports_missing_partuuid() {
559 let (disk, _partition_one_data, _partition_two_data) = build_test_mbr_disk();
560 let reader = FakeReader {
561 block_size: TEST_BLOCK_SIZE as u32,
562 data: disk,
563 };
564
565 let err = match block_on(MbrBlockReader::new(
566 reader,
567 MbrPartitionSelector::part_uuid("00000000-01"),
568 TEST_BLOCK_SIZE as u32,
569 )) {
570 Ok(_) => panic!("missing uuid should fail"),
571 Err(err) => err,
572 };
573 assert_eq!(err.kind(), GibbloxErrorKind::InvalidInput);
574 }
575
576 #[test]
577 fn mbr_reader_rejects_extended_partition() {
578 let (disk, _partition_one_data, _partition_two_data) = build_test_mbr_disk();
579 let reader = FakeReader {
580 block_size: TEST_BLOCK_SIZE as u32,
581 data: disk,
582 };
583
584 let err = match block_on(MbrBlockReader::new(
585 reader,
586 MbrPartitionSelector::index(2),
587 TEST_BLOCK_SIZE as u32,
588 )) {
589 Ok(_) => panic!("extended partition should be unsupported"),
590 Err(err) => err,
591 };
592 assert_eq!(err.kind(), GibbloxErrorKind::Unsupported);
593 }
594
595 #[test]
596 fn mbr_reader_rejects_linux_extended_partition_type() {
597 let (mut disk, _partition_one_data, _partition_two_data) = build_test_mbr_disk();
598 let entry_offset = MBR_PARTITION_TABLE_OFFSET + (2 * MBR_PARTITION_ENTRY_LEN);
599 disk[entry_offset + 4] = MBR_EXTENDED_TYPE_LINUX;
600
601 let reader = FakeReader {
602 block_size: TEST_BLOCK_SIZE as u32,
603 data: disk,
604 };
605
606 let err = match block_on(MbrBlockReader::new(
607 reader,
608 MbrPartitionSelector::index(2),
609 TEST_BLOCK_SIZE as u32,
610 )) {
611 Ok(_) => panic!("linux extended partition should be unsupported"),
612 Err(err) => err,
613 };
614 assert_eq!(err.kind(), GibbloxErrorKind::Unsupported);
615 }
616
617 fn build_test_mbr_disk() -> (
618 alloc::vec::Vec<u8>,
619 alloc::vec::Vec<u8>,
620 alloc::vec::Vec<u8>,
621 ) {
622 let mut disk = vec![0u8; TEST_BLOCK_SIZE * TEST_TOTAL_BLOCKS];
623
624 disk[MBR_DISK_SIGNATURE_OFFSET..MBR_DISK_SIGNATURE_OFFSET + 4]
625 .copy_from_slice(&TEST_DISK_SIGNATURE.to_le_bytes());
626 disk[MBR_SIGNATURE_OFFSET..MBR_SIGNATURE_OFFSET + 2].copy_from_slice(&MBR_SIGNATURE);
627
628 write_partition_entry(&mut disk, 0, 0x80, 0x83, 2, 3);
629 write_partition_entry(&mut disk, 1, 0x00, 0x83, 10, 2);
630 write_partition_entry(&mut disk, 2, 0x00, MBR_EXTENDED_TYPE_LBA, 20, 64);
631
632 let partition_one_offset = 2 * TEST_BLOCK_SIZE;
633 let partition_one_data: alloc::vec::Vec<u8> =
634 (0..1536).map(|idx| (idx % 251) as u8).collect();
635 disk[partition_one_offset..partition_one_offset + partition_one_data.len()]
636 .copy_from_slice(&partition_one_data);
637
638 let partition_two_offset = 10 * TEST_BLOCK_SIZE;
639 let partition_two_data = vec![0xAA; 1024];
640 disk[partition_two_offset..partition_two_offset + partition_two_data.len()]
641 .copy_from_slice(&partition_two_data);
642
643 (disk, partition_one_data, partition_two_data)
644 }
645
646 fn write_partition_entry(
647 disk: &mut [u8],
648 index: usize,
649 status: u8,
650 partition_type: u8,
651 first_lba: u32,
652 sector_count: u32,
653 ) {
654 let offset = MBR_PARTITION_TABLE_OFFSET + index * MBR_PARTITION_ENTRY_LEN;
655 let entry = &mut disk[offset..offset + MBR_PARTITION_ENTRY_LEN];
656 entry.fill(0);
657 entry[0] = status;
658 entry[4] = partition_type;
659 entry[8..12].copy_from_slice(&first_lba.to_le_bytes());
660 entry[12..16].copy_from_slice(§or_count.to_le_bytes());
661 }
662}