diff --git a/src/clean.rs b/src/clean.rs index e7798fb..a3c45b6 100644 --- a/src/clean.rs +++ b/src/clean.rs @@ -51,8 +51,12 @@ pub fn clean( warn!("Could not relocate: {e}") } } - let cleaned = new_me.fpt_area.to_vec(); - let size = cleaned.len(); - data[me.base..me.base + size].copy_from_slice(&cleaned); - Ok(data.to_vec()) + match new_me.fpt_area.to_vec() { + Ok(cleaned) => { + let size = cleaned.len(); + data[me.base..me.base + size].copy_from_slice(&cleaned); + Ok(data.to_vec()) + } + Err(e) => Err(e), + } } diff --git a/src/me.rs b/src/me.rs index 11fd9d8..f6df1f1 100644 --- a/src/me.rs +++ b/src/me.rs @@ -128,14 +128,17 @@ impl FPTArea { } /// Clear out fully removable partitions and adjust FPT - pub fn to_vec(&self) -> Vec { + pub fn to_vec(&self) -> Result, String> { let debug = true; if debug { println!("Recreate ME region from components"); } - let mut res = self.partitions.to_vec(); + let mut res = match self.partitions.to_vec() { + Ok(r) => r, + Err(e) => return Err(e), + }; if debug { println!(" Minimum size: {:08x}", res.len()); } @@ -168,7 +171,7 @@ impl FPTArea { } res[..s].copy_from_slice(&raw_fpt); - res + Ok(res) } } diff --git a/src/part/partitions.rs b/src/part/partitions.rs index b21605c..3f9e64e 100644 --- a/src/part/partitions.rs +++ b/src/part/partitions.rs @@ -185,7 +185,28 @@ impl Partitions { Ok(()) } - pub fn to_vec(&self) -> Vec { + /// Get the last FPT entry based on offsets, i.e., the one with the highest + /// offset. Use its offset and size and round to next 4K to calculate the + /// smallest possible ME region. + pub fn last_entry(&self) -> Option { + let sorted_parts = &self.get_sorted(); + // NOTE: We need to filter out NVRAM partitions, which have an offset of + // 0xffff_ffff. + if let Some(last) = &sorted_parts + .into_iter() + .filter(|p| { + let f = p.entry().flags; + f.kind() != PartitionKind::NVRAM + }) + .last() + { + Some(last.entry().clone()) + } else { + None + } + } + + pub fn to_vec(&self) -> Result, String> { use log::debug; fn copy_parts(parts: &Vec<&dyn Partition>, data: &mut Vec) { for p in parts { @@ -208,40 +229,31 @@ impl Partitions { } } - let sorted_parts = &self.get_sorted(); - // This gets us the smallest possible slice to copy into. - // NOTE: We need to filter out NVRAM partitions, which have an offset of - // 0xffff_ffff. - let last = &sorted_parts - .into_iter() - .filter(|p| { - let f = p.entry().flags; - f.kind() != PartitionKind::NVRAM - }) - .last() - .unwrap() - .entry(); - let o = last.offset(); - let size = o + last.size(); - debug!("Last partition @ {o:08x}; final size: {size:08x}"); - let mut data = vec![EMPTY; size]; - - match sorted_parts { - Partitions::Gen2(parts) => copy_parts( - &parts.iter().map(|p| p as &dyn Partition).collect(), - &mut data, - ), - Partitions::Gen3(parts) => copy_parts( - &parts.iter().map(|p| p as &dyn Partition).collect(), - &mut data, - ), - Partitions::Unknown(parts) => copy_parts( - &parts.iter().map(|p| p as &dyn Partition).collect(), - &mut data, - ), - }; + if let Some(last) = self.last_entry() { + let o = last.offset(); + let size = o + last.size(); + debug!("Last partition @ {o:08x}; final size: {size:08x}"); + let mut data = vec![EMPTY; size]; - data + match &self { + Partitions::Gen2(parts) => copy_parts( + &parts.iter().map(|p| p as &dyn Partition).collect(), + &mut data, + ), + Partitions::Gen3(parts) => copy_parts( + &parts.iter().map(|p| p as &dyn Partition).collect(), + &mut data, + ), + Partitions::Unknown(parts) => copy_parts( + &parts.iter().map(|p| p as &dyn Partition).collect(), + &mut data, + ), + }; + + Ok(data) + } else { + Err("no partition entries found".into()) + } } }