bb_flasher_sd/
customization.rs

1use crate::{Error, Result};
2use fatfs::FileSystem;
3use fscommon::{BufStream, StreamSlice};
4use std::io::{Read, Seek, SeekFrom, Write};
5
6#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
7pub enum ParitionType {
8    Boot,
9}
10
11impl ParitionType {
12    fn open<T>(&self, dst: T) -> Result<FileSystem<BufStream<StreamSlice<T>>>>
13    where
14        T: Write + Seek + Read + std::fmt::Debug,
15    {
16        match self {
17            Self::Boot => Self::boot_partition(dst),
18        }
19    }
20
21    fn boot_partition<T>(mut dst: T) -> Result<FileSystem<BufStream<StreamSlice<T>>>>
22    where
23        T: Write + Seek + Read + std::fmt::Debug,
24    {
25        // First try GPT partition table. If that fails, try MBR
26        let (start_offset, end_offset) = if let Ok(disk) = gpt::GptConfig::new()
27            .writable(false)
28            .open_from_device(&mut dst)
29        {
30            // FIXME: Add better partition lookup
31            let partition_2 = disk.partitions().get(&2).unwrap();
32
33            let start_offset: u64 = partition_2.first_lba * gpt::disk::DEFAULT_SECTOR_SIZE.as_u64();
34            let end_offset: u64 = partition_2.last_lba * gpt::disk::DEFAULT_SECTOR_SIZE.as_u64();
35
36            (start_offset, end_offset)
37        } else {
38            let mbr =
39                mbrman::MBRHeader::read_from(&mut dst).map_err(|_| Error::InvalidPartitionTable)?;
40
41            let boot_part = mbr.get(1).ok_or(Error::InvalidPartitionTable)?;
42            let start_offset: u64 = (boot_part.starting_lba * 512).into();
43            let end_offset: u64 = start_offset + u64::from(boot_part.sectors) * 512;
44
45            (start_offset, end_offset)
46        };
47        let slice = StreamSlice::new(dst, start_offset, end_offset)
48            .map_err(|_| Error::InvalidPartitionTable)?;
49        let boot_stream = BufStream::new(slice);
50        FileSystem::new(boot_stream, fatfs::FsOptions::new())
51            .map_err(|_| Error::InvalidBootPartition)
52    }
53}
54
55#[derive(Clone, Debug, PartialEq, Eq, Hash)]
56pub enum ContentType {
57    File(Box<std::path::Path>),
58    Data(Box<[u8]>),
59}
60
61impl From<Box<[u8]>> for ContentType {
62    fn from(value: Box<[u8]>) -> Self {
63        Self::Data(value)
64    }
65}
66
67impl From<Box<std::path::Path>> for ContentType {
68    fn from(value: Box<std::path::Path>) -> Self {
69        Self::File(value)
70    }
71}
72
73#[derive(Clone, Debug, PartialEq, Eq, Hash)]
74pub struct Customization {
75    pub partition: ParitionType,
76    pub content: Vec<(Box<str>, ContentType)>,
77}
78
79impl Customization {
80    pub(crate) fn customize(&self, dst: impl Write + Seek + Read + std::fmt::Debug) -> Result<()> {
81        let partition = self.partition.open(dst)?;
82        let root = partition.root_dir();
83
84        for (path, data) in &self.content {
85            let mut f =
86                root.create_file(path)
87                    .map_err(|source| Error::CustomizationFileCreateFail {
88                        source,
89                        file: path.clone(),
90                    })?;
91
92            match data {
93                ContentType::File(path) => {
94                    let mut source = std::fs::File::open(path)?;
95                    std::io::copy(&mut source, &mut f)?;
96                }
97                ContentType::Data(items) => {
98                    f.seek(SeekFrom::End(0))?;
99                    f.write_all(items)?;
100                }
101            }
102        }
103
104        Ok(())
105    }
106}
bb_flasher_sd/customization.rs

bb_flasher_sd/
customization.rs
