ape_mbr/
lib.rs

1//! [![MIT License](https://img.shields.io/badge/license-MIT-blue.svg)](./LICENSE.txt)
2//! [![Crates.io](https://img.shields.io/crates/v/ape-mbr)](https://crates.io/crates/ape-mbr)
3//! [![Documentation](https://docs.rs/ape-mbr/badge.svg)](https://docs.rs/ape-mbr)
4//! [![APE](https://img.shields.io/badge/-APE-%2359118e)](https://openapeshop.org/)
5//! ## *simple crate to interface between a disk and it's partitions*
6//!
7//! This crate is especially designed to provide an interface between a disk
8//! and a file system library, where both are able to implement embedded_io.
9//!
10//! Dead simple, as it should be.
11//!
12//! # Usage
13//!
14//! This crate can be used by adding `ape-mbr` to the dependencies in your
15//! project's `Cargo.toml`.
16//!
17//! ```toml
18//! [dependencies]
19//! ape-mbr = "0.1.0"
20//! ```
21//!
22//! # Examples
23//!
24//! Here's `ape-mbr` being coupled with `ape-fatfs`
25//!
26//! ```rust
27//! use std::io::prelude::*;
28//! use ape_fatfs::{
29//!     fs::{
30//!         FsOptions,
31//!         FileSystem,
32//!     },
33//!     io::{
34//!         StdIoWrapper
35//!     }
36//! };
37//!
38//! use ape_mbr::{
39//!     PartitionId,
40//!     MBR,
41//! };
42//!
43//! fn main() {
44//!     # std::fs::copy("resources/test2.img", "test.img").unwrap();
45//!     // Initialize the MBR
46//!     let img_file = std::fs::OpenOptions::new().read(true).write(true)
47//!         .open("test.img").unwrap();
48//!
49//!     let img_file = StdIoWrapper::new(img_file);
50//!    
51//!     let mut mbr = MBR::new(img_file).unwrap();
52//!     let mut p1 = mbr.get_partition(PartitionId::One).unwrap();
53//!     
54//!     let fs = FileSystem::new(p1, FsOptions::new()).unwrap();
55//!     let root_dir = fs.root_dir();
56//!
57//!     // Write a file
58//!     root_dir.create_dir("foo").unwrap();
59//!     let mut file = root_dir.create_file("foo/hello.txt").unwrap();
60//!     file.truncate().unwrap();
61//!     file.write_all(b"Hello World!").unwrap();
62//!
63//!     // Read a directory
64//!     let dir = root_dir.open_dir("foo").unwrap();
65//!     for r in dir.iter() {
66//!         let entry = r.unwrap();
67//!         println!("{}", entry.file_name());
68//!     }
69//!     # std::fs::remove_file("test.img").unwrap();
70//! }
71//! ```
72#![cfg_attr(not(test), no_std)]
73
74use core::cmp;
75use embedded_io::{
76    blocking::{Read, Seek, Write},
77    Io, SeekFrom,
78};
79use types::PartitionType;
80
81pub mod types;
82
83/// Length of each record in bytes
84pub const RECORD_LEN: usize = 16;
85/// Number of record in MBR
86pub const RECORD_COUNT: usize = 4;
87/// Size of blocks in bytes
88pub const BLOCK_SIZE: u64 = 512;
89/// Offset to the start of the partition records
90pub const RECORDS_START: u64 = 0x1be;
91/// Offset of the relative sector field in a partition record
92pub const RELATIVE_SECTOR_OFFSET: usize = 8;
93/// Offset of the total sectors field in a partition record
94pub const TOTAL_SECTORS_OFFSET: usize = 12;
95/// Offset of the system id field in a partition record
96pub const SYSTEM_ID_OFFSET: usize = 4;
97/// Offset of the boot indicator flag in a partition record
98pub const BOOT_FLAG_OFFSET: usize = 0;
99
100/// ID of each partition
101#[repr(usize)]
102pub enum PartitionId {
103    One = 0,
104    Two = 1,
105    Three = 2,
106    Four = 3,
107}
108
109#[inline]
110/// Convert an LBA address to a u64
111pub fn lba_to_u64(lba: u32) -> u64 {
112    (lba as u64) * BLOCK_SIZE
113}
114
115/// Used to interface with partitions
116pub struct Partition<'a, IO> {
117    start_pos: u64,
118    end_pos: u64,
119    pos: u64,
120    io: &'a mut IO,
121}
122
123impl<'a, IO: Io + Seek> Partition<'a, IO> {
124    /// Create a new partition given the start and end position
125    pub fn new(start_pos: u64, end_pos: u64, io: &'a mut IO) -> Result<Self, <Self as Io>::Error> {
126        // Seek to the start of the partition
127        io.seek(SeekFrom::Start(start_pos))?;
128
129        Ok(Self {
130            start_pos,
131            end_pos,
132            pos: 0,
133            io,
134        })
135    }
136}
137
138impl<'a, IO> Partition<'a, IO> {
139    #[inline]
140    /// Get the length of the partition in bytes
141    pub fn len(&self) -> u64 {
142        self.end_pos - self.start_pos
143    }
144}
145
146impl<'a, IO: Io> Io for Partition<'a, IO> {
147    type Error = IO::Error;
148}
149
150impl<'a, IO: Read> Read for Partition<'a, IO> {
151    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
152        // Limit the amount of data available to read to the size of the partition
153        let available = self.len() - self.pos;
154
155        let buf_slice = match (buf.len() as u64) < available {
156            true => buf,
157            false => &mut buf[..available as usize],
158        };
159
160        self.pos += buf_slice.len() as u64;
161
162        self.io.read(buf_slice)
163    }
164}
165
166impl<'a, IO: Write> Write for Partition<'a, IO> {
167    fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
168        // Limit the amount of data available to write to the size of the partition
169        let available = self.len() - self.pos;
170
171        let buf_slice = match (buf.len() as u64) < available {
172            true => buf,
173            false => &buf[..available as usize],
174        };
175
176        self.pos += buf_slice.len() as u64;
177
178        self.io.write(buf_slice)
179    }
180
181    #[inline]
182    fn flush(&mut self) -> Result<(), Self::Error> {
183        self.io.flush()
184    }
185}
186
187impl<'a, IO: Seek> Seek for Partition<'a, IO> {
188    fn seek(&mut self, pos: embedded_io::SeekFrom) -> Result<u64, Self::Error> {
189        self.pos = match pos {
190            SeekFrom::Start(pos) => {
191                // Ensure that we don't go past the partition boundries
192                cmp::min(pos, self.len())
193            }
194            SeekFrom::Current(pos) => {
195                // Ensure that we don't go past the partition boundries
196                cmp::max(cmp::min((pos as i64) + pos, self.len() as i64), 0) as u64
197            }
198            SeekFrom::End(pos) => {
199                // Ensure that we don't go past the partition boundries
200                cmp::max(cmp::min((self.len() as i64) + pos, self.len() as i64), 0) as u64
201            }
202        };
203
204        self.io.seek(SeekFrom::Start(self.start_pos + self.pos))?;
205
206        Ok(self.pos)
207    }
208}
209
210/// Used to store data about partitions in the MBR
211#[derive(Debug, Copy, Clone, Default)]
212pub struct PartitionRecord {
213    relative_sector: u32,
214    total_sectors: u32,
215    partition_type: PartitionType,
216    boot_flag: bool,
217}
218
219impl PartitionRecord {
220    /// Create a partition record from bytes
221    pub fn from_bytes(bytes: &[u8; RECORD_LEN]) -> Self {
222        let relative_sector_array: [u8; 4] = bytes[RELATIVE_SECTOR_OFFSET..TOTAL_SECTORS_OFFSET]
223            .try_into()
224            .unwrap();
225        let total_sectors_array: [u8; 4] =
226            bytes[TOTAL_SECTORS_OFFSET..RECORD_LEN].try_into().unwrap();
227
228        let relative_sector = u32::from_le_bytes(relative_sector_array);
229        let total_sectors = u32::from_le_bytes(total_sectors_array);
230        
231        let system_id: u8 = bytes[SYSTEM_ID_OFFSET];
232        let boot_flag: bool = bytes[BOOT_FLAG_OFFSET] == 0x80;
233
234        Self {
235            relative_sector,
236            total_sectors,
237            partition_type: system_id.try_into().unwrap(),
238            boot_flag,
239        }
240    }
241
242    #[inline]
243    /// Get the starting position of a partition
244    pub fn get_start_pos(&self) -> u64 {
245        lba_to_u64(self.relative_sector)
246    }
247
248    #[inline]
249    /// Get the end position of a partition
250    pub fn get_end_pos(&self) -> u64 {
251        lba_to_u64(self.relative_sector) + lba_to_u64(self.total_sectors)
252    }
253
254    #[inline]
255    /// Get the type of a partition
256    pub fn get_partition_type(&self) -> PartitionType {
257        self.partition_type
258    }
259
260    #[inline]
261    /// Check to see if the partition's boot flag is set
262    pub fn is_bootable(&self) -> bool {
263        self.boot_flag
264    }
265}
266
267/// Used to grab partitions from the MBR
268pub struct MBR<IO: Read + Seek> {
269    partitions: [PartitionRecord; RECORD_COUNT],
270    io: IO,
271}
272
273impl<IO: Read + Seek> MBR<IO> {
274    /// Create a new MBR from anything that implements embedded_io
275    pub fn new(mut io: IO) -> Result<Self, <IO as Io>::Error> {
276        let mut partitions: [PartitionRecord; RECORD_COUNT] =
277            [PartitionRecord::default(); RECORD_COUNT];
278        let mut buffer: [u8; RECORD_LEN * RECORD_COUNT] = [0; RECORD_LEN * RECORD_COUNT];
279
280        io.seek(SeekFrom::Start(RECORDS_START))?;
281        io.read(&mut buffer)?;
282
283        for i in 0..RECORD_COUNT {
284            let buffer_i = i * RECORD_LEN;
285
286            let record_slice = &buffer[buffer_i..buffer_i + RECORD_LEN];
287
288            partitions[i] = PartitionRecord::from_bytes(record_slice.try_into().unwrap());
289        }
290
291        Ok(Self { partitions, io })
292    }
293
294    #[inline]
295    /// Get a partition from the MBR
296    pub fn get_partition(&mut self, id: PartitionId) -> Result<Partition<IO>, IO::Error> {
297        let record = self.partitions[id as usize];
298
299        Partition::new(record.get_start_pos(), record.get_end_pos(), &mut self.io)
300    }
301
302    #[inline]
303    /// Get the partition type from the MBR
304    pub fn get_partition_type(&self, id: PartitionId) -> PartitionType {
305        let record = self.partitions[id as usize];
306
307        record.get_partition_type()
308    }
309
310    #[inline]
311    /// Check if a partition is bootable in the MBR
312    pub fn is_partition_bootable(&self, id: PartitionId) -> bool {
313        let record = self.partitions[id as usize];
314
315        record.is_bootable()
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use core::panic::AssertUnwindSafe;
322    use std::{io::Cursor, panic};
323
324    use ape_fatfs::{fs::{FileSystem, FsOptions, FatType}, io::StdIoWrapper};
325    use embedded_io::{
326        adapters::FromStd,
327        blocking::{Read, Seek, Write},
328    };
329
330    use crate::*;
331
332    static TEST_IMG_1: &[u8] = include_bytes!("../resources/test1.img");
333    static TEST_IMG_2: &[u8] = include_bytes!("../resources/test2.img");
334    static TEST_STR_1: [u8; 10] = *b"Partition1";
335    static TEST_STR_2: [u8; 10] = *b"Partition2";
336    static TEST_STR_3: [u8; 10] = *b"Partition3";
337    static TEST_STR_4: [u8; 10] = *b"Partition4";
338
339    #[test]
340    /// The dummy image is a four partition image with "Partition" witten to the
341    /// start of each partition and the partition number written to the end
342    /// each partition
343    ///
344    /// Partitions sizes in sectors:
345    ///     First partition size: 17
346    ///     Second partition size: 33
347    ///     Third partition size: 65
348    ///     Fourth partition size: 84
349    fn test_dummy_img() {
350        let img = FromStd::new(Cursor::new(TEST_IMG_1.to_vec()));
351
352        let mut mbr = MBR::new(img).unwrap();
353
354        // Test partition 1
355        let mut partition_1 = mbr.get_partition(PartitionId::One).unwrap();
356
357        let mut buf: [u8; 10] = [0; 10];
358
359        partition_1.read_exact(&mut buf[..9]).unwrap();
360        partition_1.seek(embedded_io::SeekFrom::End(-1)).unwrap();
361        partition_1.read_exact(&mut buf[9..]).unwrap();
362        assert_eq!(partition_1.len(), 17 * BLOCK_SIZE);
363        assert_eq!(buf, TEST_STR_1);
364        drop(partition_1);
365
366        // Test partition 2
367        let mut partition_2 = mbr.get_partition(PartitionId::Two).unwrap();
368
369        partition_2.read_exact(&mut buf[..9]).unwrap();
370        partition_2.seek(embedded_io::SeekFrom::End(-1)).unwrap();
371        partition_2.read_exact(&mut buf[9..]).unwrap();
372        assert_eq!(partition_2.len(), 33 * BLOCK_SIZE);
373        assert_eq!(buf, TEST_STR_2);
374        drop(partition_2);
375
376        // Test partition 3
377        let mut partition_3 = mbr.get_partition(PartitionId::Three).unwrap();
378
379        partition_3.read_exact(&mut buf[..9]).unwrap();
380        partition_3.seek(embedded_io::SeekFrom::End(-1)).unwrap();
381        partition_3.read_exact(&mut buf[9..]).unwrap();
382        assert_eq!(partition_3.len(), 65 * BLOCK_SIZE);
383        assert_eq!(buf, TEST_STR_3);
384        drop(partition_3);
385
386        // Test partition 2
387        let mut partition_4 = mbr.get_partition(PartitionId::Four).unwrap();
388
389        partition_4.read_exact(&mut buf[..9]).unwrap();
390        partition_4.seek(embedded_io::SeekFrom::End(-1)).unwrap();
391        partition_4.read_exact(&mut buf[9..]).unwrap();
392        assert_eq!(partition_4.len(), 84 * BLOCK_SIZE);
393        assert_eq!(buf, TEST_STR_4);
394    }
395
396    #[test]
397    /// The "real" image is a three partition image designed to simulate a real
398    /// drive
399    ///
400    /// Partition layout is as follows:
401    ///     fat12, 2000 sectors, boot flag set
402    ///     fat16, 5000 sectors
403    ///     fat32, 68000 sectors
404    fn test_real_img() {
405        let img = StdIoWrapper::new(Cursor::new(TEST_IMG_2.to_vec()));
406
407        let mut mbr = MBR::new(img).unwrap();
408
409        // Test partition 1
410        assert_eq!(mbr.get_partition_type(PartitionId::One), PartitionType::Fat12);
411        assert!(mbr.is_partition_bootable(PartitionId::One));
412
413        {
414            let partition = mbr.get_partition(PartitionId::One).unwrap();
415
416            let fs = FileSystem::new(partition, FsOptions::new()).unwrap();
417
418            assert_eq!(fs.fat_type(), FatType::Fat12);
419        }
420
421        // Test partition 2
422        assert_eq!(mbr.get_partition_type(PartitionId::Two), PartitionType::Fat16);
423        assert!(!mbr.is_partition_bootable(PartitionId::Two));
424
425        {
426            let partition = mbr.get_partition(PartitionId::Two).unwrap();
427
428            let fs = FileSystem::new(partition, FsOptions::new()).unwrap();
429
430            assert_eq!(fs.fat_type(), FatType::Fat16);
431        }
432
433        // Test partition 3
434        assert_eq!(mbr.get_partition_type(PartitionId::Three), PartitionType::W95Fat32);
435        assert!(!mbr.is_partition_bootable(PartitionId::Three));
436
437        {
438            let partition = mbr.get_partition(PartitionId::Three).unwrap();
439
440            let fs = FileSystem::new(partition, FsOptions::new()).unwrap();
441
442            assert_eq!(fs.fat_type(), FatType::Fat32);
443        }
444    }
445
446    #[test]
447    /// Ensure that we cannot read or write past the end of the partition
448    fn test_bounds() {
449        let img = FromStd::new(Cursor::new(TEST_IMG_1.to_vec()));
450
451        let mut mbr = MBR::new(img).unwrap();
452
453        let mut buf: [u8; 10] = [0; 10];
454        let mut partition_1 = mbr.get_partition(PartitionId::One).unwrap();
455
456        partition_1.seek(embedded_io::SeekFrom::End(0)).unwrap();
457
458        partition_1.read_exact(&mut buf).unwrap_err();
459        // write_all panics on failure
460        assert!(panic::catch_unwind(AssertUnwindSafe(|| {
461            partition_1.write_all(&buf).unwrap_err();
462        }))
463        .is_err());
464    }
465}
ape_mbr/lib.rs

ape_mbr/
lib.rs
