1use std::collections::HashMap;
22use std::sync::Arc;
23
24use sha2::{Digest as _, Sha256};
25
26use crate::image::{detect_chip, ChipError};
27use crate::{layout, Chip};
28
29pub type Digest = [u8; 32];
31
32const MASKED: [layout::Range; 5] = [
35 layout::BOOT_PAGE,
36 layout::UHD_REVOCATION,
37 layout::BD_REVOCATION,
38 layout::SETTINGS,
39 layout::CALIBRATION,
40];
41
42pub fn content_key(image: &[u8]) -> Option<Digest> {
47 if image.len() != layout::IMAGE_SIZE {
48 return None;
49 }
50 let mut masked: Vec<layout::Range> = MASKED.to_vec();
51 masked.sort_by_key(|r| r.start);
52 let mut hasher = Sha256::new();
53 let mut at = 0;
54 for range in masked {
55 hasher.update(&image[at..range.start]);
56 at = range.end();
57 }
58 hasher.update(&image[at..]);
59 Some(hasher.finalize().into())
60}
61
62pub fn sha256(bytes: &[u8]) -> Digest {
64 Sha256::digest(bytes).into()
65}
66
67#[derive(Clone, Debug, PartialEq, Eq)]
69pub struct Factory {
70 boot_page: Arc<[u8]>,
71 regions: [Arc<[u8]>; 4],
72}
73
74impl Factory {
75 pub fn new(boot_page: Arc<[u8]>, regions: [Arc<[u8]>; 4]) -> Result<Self, RebuildError> {
78 if boot_page.len() != layout::BOOT_PAGE.len {
79 return Err(RebuildError::Factory(layout::BOOT_PAGE.name));
80 }
81 for (range, bytes) in layout::DRIVE_WRITTEN.iter().zip(®ions) {
82 if bytes.len() != range.len {
83 return Err(RebuildError::Factory(range.name));
84 }
85 }
86 Ok(Self { boot_page, regions })
87 }
88
89 pub fn from_image(image: &[u8]) -> Result<Self, RebuildError> {
92 if image.len() != layout::IMAGE_SIZE {
93 return Err(RebuildError::Size(image.len()));
94 }
95 if crate::image::is_drive_read(image) {
96 return Err(RebuildError::Factory(layout::BOOT_PAGE.name));
97 }
98 let take =
99 |range: layout::Range| -> Arc<[u8]> { Arc::from(&image[range.start..range.end()]) };
100 Self::new(take(layout::BOOT_PAGE), layout::DRIVE_WRITTEN.map(take))
101 }
102
103 pub fn boot_page(&self) -> &[u8] {
105 &self.boot_page
106 }
107
108 pub fn regions(&self) -> &[Arc<[u8]>; 4] {
110 &self.regions
111 }
112}
113
114#[derive(Clone, Debug, PartialEq, Eq)]
116pub struct Build {
117 pub model: String,
119 pub revision: String,
121 pub factory: Factory,
123}
124
125#[derive(Clone, Debug, Default)]
127pub struct Catalog {
128 exact: HashMap<Digest, Build>,
129 defaults: Vec<(Chip, Factory)>,
130}
131
132impl Catalog {
133 pub fn new() -> Self {
135 Self::default()
136 }
137
138 pub fn insert(&mut self, key: Digest, build: Build) {
140 self.exact.insert(key, build);
141 }
142
143 pub fn set_default(&mut self, chip: Chip, factory: Factory) {
145 self.defaults.retain(|(c, _)| *c != chip);
146 self.defaults.push((chip, factory));
147 }
148
149 pub fn build(&self, key: &Digest) -> Option<&Build> {
151 self.exact.get(key)
152 }
153
154 pub fn default_for(&self, chip: Chip) -> Option<&Factory> {
156 self.defaults
157 .iter()
158 .find(|(c, _)| *c == chip)
159 .map(|(_, f)| f)
160 }
161
162 pub fn knows_boot_page(&self, page: &[u8]) -> bool {
165 self.defaults.iter().any(|(_, f)| f.boot_page() == page)
166 || self.exact.values().any(|b| b.factory.boot_page() == page)
167 }
168
169 pub fn len(&self) -> usize {
171 self.exact.len()
172 }
173
174 pub fn is_empty(&self) -> bool {
176 self.exact.is_empty()
177 }
178}
179
180#[derive(Clone, Debug, PartialEq, Eq)]
182pub enum Source {
183 Exact {
185 model: String,
187 revision: String,
189 },
190 ChipDefault,
192}
193
194#[derive(Clone, Copy, Debug, PartialEq, Eq)]
196pub enum BootPage {
197 Stored,
199 Restored,
201}
202
203#[derive(Clone, Debug, PartialEq, Eq)]
205pub struct Provenance {
206 pub chip: Chip,
208 pub source: Source,
210 pub boot_page: BootPage,
212 pub regions_reset: Vec<(&'static str, usize)>,
214}
215
216impl Provenance {
217 pub fn is_exact(&self) -> bool {
219 matches!(self.source, Source::Exact { .. })
220 }
221}
222
223#[derive(Clone, Debug, PartialEq, Eq)]
225pub struct Rebuild {
226 pub image: Vec<u8>,
228 pub provenance: Provenance,
230}
231
232#[derive(Clone, Debug, PartialEq, Eq)]
234#[non_exhaustive]
235pub enum RebuildError {
236 Size(usize),
238 Chip(ChipError),
240 NoFactory(Chip),
242 UnknownBootPage(Digest),
245 Factory(&'static str),
247}
248
249impl core::fmt::Display for RebuildError {
250 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
251 match self {
252 Self::Size(n) => write!(
253 f,
254 "image is {n} bytes; an OEM image is exactly {} bytes",
255 layout::IMAGE_SIZE
256 ),
257 Self::Chip(e) => write!(f, "{e}"),
258 Self::NoFactory(chip) => write!(f, "no factory contents are known for {chip}"),
259 Self::UnknownBootPage(sha) => {
260 f.write_str(
261 "boot page is neither a stored page nor a drive read's mirror (sha256 ",
262 )?;
263 sha.iter().try_for_each(|b| write!(f, "{b:02x}"))?;
264 f.write_str(")")
265 }
266 Self::Factory(name) => write!(f, "factory {name} has the wrong length or form"),
267 }
268 }
269}
270
271impl std::error::Error for RebuildError {
272 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
273 match self {
274 Self::Chip(source) => Some(source),
275 _ => None,
276 }
277 }
278}
279
280pub fn rebuild(read: &[u8], catalog: &Catalog) -> Result<Rebuild, RebuildError> {
288 let key = content_key(read).ok_or(RebuildError::Size(read.len()))?;
289 let chip = detect_chip(read).map_err(RebuildError::Chip)?.family;
290 let (factory, source, exact) = match catalog.build(&key) {
291 Some(build) => (
292 &build.factory,
293 Source::Exact {
294 model: build.model.clone(),
295 revision: build.revision.clone(),
296 },
297 true,
298 ),
299 None => (
300 catalog
301 .default_for(chip)
302 .ok_or(RebuildError::NoFactory(chip))?,
303 Source::ChipDefault,
304 false,
305 ),
306 };
307
308 let mut image = read.to_vec();
309 let boot = &read[..layout::BOOT_PAGE.len];
310 let boot_page = if boot == factory.boot_page() {
311 BootPage::Stored
312 } else if crate::image::is_drive_read(read) {
313 image[..layout::BOOT_PAGE.len].copy_from_slice(factory.boot_page());
314 BootPage::Restored
315 } else {
316 return Err(RebuildError::UnknownBootPage(sha256(boot)));
317 };
318
319 let mut regions_reset = Vec::new();
320 for (range, contents) in layout::DRIVE_WRITTEN.iter().zip(factory.regions()) {
321 let current = &read[range.start..range.end()];
322 if !exact && layout::is_erased(current) {
323 continue;
324 }
325 let changed = current
326 .iter()
327 .zip(contents.iter())
328 .filter(|(a, b)| a != b)
329 .count();
330 if changed > 0 {
331 image[range.start..range.end()].copy_from_slice(contents);
332 regions_reset.push((range.name, changed));
333 }
334 }
335
336 Ok(Rebuild {
337 image,
338 provenance: Provenance {
339 chip,
340 source,
341 boot_page,
342 regions_reset,
343 },
344 })
345}
346
347#[cfg(test)]
348#[path = "oem_tests.rs"]
349mod tests;