1use std::{backtrace::Backtrace, borrow::Cow, io};
2
3use serde::{Deserialize, Serialize};
4use snafu::Snafu;
5
6use super::{
7 raw::{self, HmacSha1Signature, OverlayFlags, RawFatError, RawHeaderError},
8 Arm9, Arm9OverlaySignaturesError,
9};
10use crate::{
11 compress::lz77::{Lz77, Lz77DecompressError},
12 crypto::hmac_sha1::HmacSha1,
13};
14
15#[derive(Clone)]
17pub struct Overlay<'a> {
18 originally_compressed: bool,
19 info: OverlayInfo,
20 signature: Option<HmacSha1Signature>,
21 data: Cow<'a, [u8]>,
22}
23
24const LZ77: Lz77 = Lz77 {};
25
26pub struct OverlayOptions {
28 pub originally_compressed: bool,
30 pub info: OverlayInfo,
32}
33
34#[derive(Debug, Snafu)]
36pub enum OverlayError {
37 #[snafu(transparent)]
39 RawHeader {
40 source: RawHeaderError,
42 },
43 #[snafu(transparent)]
45 RawFat {
46 source: RawFatError,
48 },
49 #[snafu(transparent)]
51 Arm9OverlaySignatures {
52 source: Arm9OverlaySignaturesError,
54 },
55 #[snafu(display("no overlay signatures found in ARM9 program:\n{backtrace}"))]
57 NoOverlaySignatures {
58 backtrace: Backtrace,
60 },
61 #[snafu(display("signing ARM7 overlays is not supported:\n{backtrace}"))]
63 SignedArm7Overlay {
64 backtrace: Backtrace,
66 },
67 #[snafu(display("cannot compute signature for overlay that is not in its originally compressed state:\n{backtrace}"))]
69 OverlayCompression {
70 backtrace: Backtrace,
72 },
73}
74
75impl<'a> Overlay<'a> {
76 pub fn new<T: Into<Cow<'a, [u8]>>>(data: T, options: OverlayOptions) -> Result<Self, OverlayError> {
78 let OverlayOptions { originally_compressed, info } = options;
79 let data = data.into();
80
81 Ok(Self { originally_compressed, info, signature: None, data })
82 }
83
84 pub fn parse_arm9(overlay: &raw::Overlay, rom: &'a raw::Rom, arm9: &Arm9) -> Result<Self, OverlayError> {
90 let fat = rom.fat()?;
91
92 let alloc = fat[overlay.file_id as usize];
93 let data = &rom.data()[alloc.range()];
94
95 let mut signature = None;
96 if overlay.flags.is_signed() {
97 let num_overlays = rom.num_arm9_overlays()?;
98 let signatures = arm9.overlay_signatures(num_overlays)?;
99 signature = Some(signatures.map(|s| s[overlay.id as usize]).ok_or_else(|| NoOverlaySignaturesSnafu {}.build())?);
100 }
101
102 let overlay = Self {
103 originally_compressed: overlay.flags.is_compressed(),
104 info: OverlayInfo::new(overlay),
105 signature,
106 data: Cow::Borrowed(data),
107 };
108
109 Ok(overlay)
110 }
111
112 pub fn parse_arm7(overlay: &raw::Overlay, rom: &'a raw::Rom) -> Result<Self, OverlayError> {
118 let fat = rom.fat()?;
119
120 let alloc = fat[overlay.file_id as usize];
121 let data = &rom.data()[alloc.range()];
122
123 if overlay.flags.is_signed() {
124 return SignedArm7OverlaySnafu {}.fail();
125 }
126
127 let overlay = Self {
128 originally_compressed: overlay.flags.is_compressed(),
129 info: OverlayInfo::new(overlay),
130 signature: None,
131 data: Cow::Borrowed(data),
132 };
133
134 Ok(overlay)
135 }
136
137 pub fn build(&self) -> raw::Overlay {
139 let mut flags = OverlayFlags::new();
140 flags.set_is_compressed(self.is_compressed());
141 flags.set_is_signed(self.is_signed());
142 if self.is_compressed() {
143 flags.set_size(self.data.len());
144 }
145
146 raw::Overlay {
147 id: self.id() as u32,
148 base_addr: self.base_address(),
149 code_size: self.code_size(),
150 bss_size: self.bss_size(),
151 ctor_start: self.ctor_start(),
152 ctor_end: self.ctor_end(),
153 file_id: self.file_id(),
154 flags,
155 }
156 }
157
158 pub fn id(&self) -> u16 {
160 self.info.id as u16
161 }
162
163 pub fn base_address(&self) -> u32 {
165 self.info.base_address
166 }
167
168 pub fn end_address(&self) -> u32 {
170 self.info.base_address + self.info.code_size + self.info.bss_size
171 }
172
173 pub fn code_size(&self) -> u32 {
175 self.info.code_size
176 }
177
178 pub fn bss_size(&self) -> u32 {
180 self.info.bss_size
181 }
182
183 pub fn ctor_start(&self) -> u32 {
185 self.info.ctor_start
186 }
187
188 pub fn ctor_end(&self) -> u32 {
190 self.info.ctor_end
191 }
192
193 pub fn file_id(&self) -> u32 {
195 self.info.file_id
196 }
197
198 pub fn is_compressed(&self) -> bool {
201 self.info.compressed
202 }
203
204 pub fn is_signed(&self) -> bool {
206 self.signature.is_some()
207 }
208
209 pub fn decompress(&mut self) -> Result<(), Lz77DecompressError> {
211 if !self.is_compressed() {
212 return Ok(());
213 }
214 self.data = LZ77.decompress(&self.data)?.into_vec().into();
215 self.info.compressed = false;
216 Ok(())
217 }
218
219 pub fn compress(&mut self) -> Result<(), io::Error> {
225 if self.is_compressed() {
226 return Ok(());
227 }
228 self.data = LZ77.compress(&self.data, 0)?.into_vec().into();
229 self.info.compressed = true;
230 Ok(())
231 }
232
233 pub fn code(&self) -> &[u8] {
235 &self.data[..self.code_size() as usize]
236 }
237
238 pub fn full_data(&self) -> &[u8] {
240 &self.data
241 }
242
243 pub fn info(&self) -> &OverlayInfo {
245 &self.info
246 }
247
248 pub fn originally_compressed(&self) -> bool {
250 self.originally_compressed
251 }
252
253 pub fn compute_signature(&self, hmac_sha1: &HmacSha1) -> Result<HmacSha1Signature, OverlayError> {
255 if self.is_compressed() != self.originally_compressed {
256 OverlayCompressionSnafu {}.fail()?;
257 }
258
259 Ok(HmacSha1Signature::from_hmac_sha1(hmac_sha1, self.data.as_ref()))
260 }
261
262 pub fn signature(&self) -> Option<HmacSha1Signature> {
264 self.signature
265 }
266
267 pub fn verify_signature(&self, hmac_sha1: &HmacSha1) -> Result<bool, OverlayError> {
269 let Some(signature) = self.signature() else {
270 return Ok(true);
271 };
272
273 let computed_signature = self.compute_signature(hmac_sha1)?;
274 Ok(computed_signature == signature)
275 }
276
277 pub fn sign(&mut self, hmac_sha1: &HmacSha1) -> Result<(), OverlayError> {
279 self.signature = Some(self.compute_signature(hmac_sha1)?);
280 Ok(())
281 }
282}
283
284#[derive(Serialize, Deserialize, Clone)]
286pub struct OverlayInfo {
287 pub id: u32,
289 pub base_address: u32,
291 pub code_size: u32,
293 pub bss_size: u32,
295 pub ctor_start: u32,
297 pub ctor_end: u32,
299 pub file_id: u32,
301 pub compressed: bool,
303}
304
305impl OverlayInfo {
306 pub fn new(overlay: &raw::Overlay) -> Self {
308 Self {
309 id: overlay.id,
310 base_address: overlay.base_addr,
311 code_size: overlay.code_size,
312 bss_size: overlay.bss_size,
313 ctor_start: overlay.ctor_start,
314 ctor_end: overlay.ctor_end,
315 file_id: overlay.file_id,
316 compressed: overlay.flags.is_compressed(),
317 }
318 }
319}