Skip to main content

ds_rom/rom/
overlay.rs

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/// An overlay module for ARM9/ARM7.
16#[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
26/// Options for creating an [`Overlay`].
27pub struct OverlayOptions {
28    /// Whether the overlay was originally compressed.
29    pub originally_compressed: bool,
30    /// Overlay info.
31    pub info: OverlayInfo,
32}
33
34/// Errors related to [`Overlay`].
35#[derive(Debug, Snafu)]
36pub enum OverlayError {
37    /// See [`RawHeaderError`].
38    #[snafu(transparent)]
39    RawHeader {
40        /// Source error.
41        source: RawHeaderError,
42    },
43    /// See [`RawFatError`].
44    #[snafu(transparent)]
45    RawFat {
46        /// Source error.
47        source: RawFatError,
48    },
49    /// See [`Arm9OverlaySignaturesError`].
50    #[snafu(transparent)]
51    Arm9OverlaySignatures {
52        /// Source error.
53        source: Arm9OverlaySignaturesError,
54    },
55    /// Occurs when there are no overlay signatures in the ARM9 program.
56    #[snafu(display("no overlay signatures found in ARM9 program:\n{backtrace}"))]
57    NoOverlaySignatures {
58        /// Backtrace to the source of the error.
59        backtrace: Backtrace,
60    },
61    /// Occurs when trying to create a signed ARM7 overlay, but signing ARM7 overlays is not supported.
62    #[snafu(display("signing ARM7 overlays is not supported:\n{backtrace}"))]
63    SignedArm7Overlay {
64        /// Backtrace to the source of the error.
65        backtrace: Backtrace,
66    },
67    /// Occurs when trying to compute the signature but the overlay is not in its originally compressed state.
68    #[snafu(display("cannot compute signature for overlay that is not in its originally compressed state:\n{backtrace}"))]
69    OverlayCompression {
70        /// Backtrace to the source of the error.
71        backtrace: Backtrace,
72    },
73}
74
75impl<'a> Overlay<'a> {
76    /// Creates a new [`Overlay`] from plain data.
77    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    /// Parses an ARM9 [`Overlay`] from a ROM.
85    ///
86    /// # Errors
87    ///
88    /// This function will return an error if the overlay is signed and the ARM9 program does not contain overlay signatures.
89    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    /// Parses an ARM7 [`Overlay`] from a ROM.
113    ///
114    /// # Errors
115    ///
116    /// This function will return an error if the overlay is signed, as ARM7 overlay signatures are not supported.
117    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    /// Builds a raw overlay table entry.
138    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    /// Returns the ID of this [`Overlay`].
159    pub fn id(&self) -> u16 {
160        self.info.id as u16
161    }
162
163    /// Returns the base address of this [`Overlay`].
164    pub fn base_address(&self) -> u32 {
165        self.info.base_address
166    }
167
168    /// Returns the end address of this [`Overlay`].
169    pub fn end_address(&self) -> u32 {
170        self.info.base_address + self.info.code_size + self.info.bss_size
171    }
172
173    /// Returns the size of initialized data in this [`Overlay`].
174    pub fn code_size(&self) -> u32 {
175        self.info.code_size
176    }
177
178    /// Returns the size of uninitialized data in this [`Overlay`].
179    pub fn bss_size(&self) -> u32 {
180        self.info.bss_size
181    }
182
183    /// Returns the offset to the start of the .ctor section.
184    pub fn ctor_start(&self) -> u32 {
185        self.info.ctor_start
186    }
187
188    /// Returns the offset to the end of the .ctor section.
189    pub fn ctor_end(&self) -> u32 {
190        self.info.ctor_end
191    }
192
193    /// Returns the file ID of this [`Overlay`].
194    pub fn file_id(&self) -> u32 {
195        self.info.file_id
196    }
197
198    /// Returns whether this [`Overlay`] is compressed. See [`Self::originally_compressed`] for whether this overlay was
199    /// compressed originally.
200    pub fn is_compressed(&self) -> bool {
201        self.info.compressed
202    }
203
204    /// Returns whether this [`Overlay`] has a signature.
205    pub fn is_signed(&self) -> bool {
206        self.signature.is_some()
207    }
208
209    /// Decompresses this [`Overlay`], but does nothing if already decompressed.
210    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    /// Compresses this [`Overlay`], but does nothing if already compressed.
220    ///
221    /// # Errors
222    ///
223    /// This function will return an error if an I/O operation fails.
224    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    /// Returns a reference to the code of this [`Overlay`].
234    pub fn code(&self) -> &[u8] {
235        &self.data[..self.code_size() as usize]
236    }
237
238    /// Returns a reference to the full data of this [`Overlay`].
239    pub fn full_data(&self) -> &[u8] {
240        &self.data
241    }
242
243    /// Returns a reference to the info of this [`Overlay`].
244    pub fn info(&self) -> &OverlayInfo {
245        &self.info
246    }
247
248    /// Returns whether this [`Overlay`] was compressed originally. See [`Self::is_compressed`] for the current state.
249    pub fn originally_compressed(&self) -> bool {
250        self.originally_compressed
251    }
252
253    /// Computes the signature of this [`Overlay`] using the given HMAC-SHA1 key.
254    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    /// Returns the signature of this [`Overlay`], if it exists.
263    pub fn signature(&self) -> Option<HmacSha1Signature> {
264        self.signature
265    }
266
267    /// Verifies the signature of this [`Overlay`] using the given HMAC-SHA1 key.
268    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    /// Signs this [`Overlay`] using the given HMAC-SHA1 key.
278    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/// Info of an [`Overlay`], similar to an entry in the overlay table.
285#[derive(Serialize, Deserialize, Clone)]
286pub struct OverlayInfo {
287    /// Overlay ID.
288    pub id: u32,
289    /// Base address.
290    pub base_address: u32,
291    /// Initialized size.
292    pub code_size: u32,
293    /// Uninitialized size.
294    pub bss_size: u32,
295    /// Offset to start of .ctor section.
296    pub ctor_start: u32,
297    /// Offset to end of .ctor section.
298    pub ctor_end: u32,
299    /// File ID for the FAT.
300    pub file_id: u32,
301    /// Whether the overlay is compressed.
302    pub compressed: bool,
303}
304
305impl OverlayInfo {
306    /// Creates a new [`OverlayInfo`] from raw data.
307    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}