ds_rom/rom/
arm9.rs

1use std::{borrow::Cow, io, mem::replace, ops::Range};
2
3use serde::{Deserialize, Serialize};
4use snafu::{Backtrace, Snafu};
5
6use super::{
7    raw::{AutoloadInfo, AutoloadKind, BuildInfo, RawAutoloadInfoError, RawBuildInfoError},
8    Autoload,
9};
10use crate::{
11    compress::lz77::{Lz77, Lz77DecompressError},
12    crc::CRC_16_MODBUS,
13    crypto::blowfish::{Blowfish, BlowfishError, BlowfishKey, BlowfishLevel},
14};
15
16/// ARM9 program.
17#[derive(Clone)]
18pub struct Arm9<'a> {
19    data: Cow<'a, [u8]>,
20    offsets: Arm9Offsets,
21    originally_compressed: bool,
22    originally_encrypted: bool,
23}
24
25/// Offsets in the ARM9 program.
26#[derive(Serialize, Deserialize, Clone, Copy)]
27pub struct Arm9Offsets {
28    /// Base address.
29    pub base_address: u32,
30    /// Entrypoint function address.
31    pub entry_function: u32,
32    /// Build info offset.
33    pub build_info: u32,
34    /// Autoload callback address.
35    pub autoload_callback: u32,
36}
37
38const SECURE_AREA_ID: [u8; 8] = [0xff, 0xde, 0xff, 0xe7, 0xff, 0xde, 0xff, 0xe7];
39const SECURE_AREA_ENCRY_OBJ: &[u8] = "encryObj".as_bytes();
40
41const LZ77: Lz77 = Lz77 {};
42
43const COMPRESSION_START: usize = 0x4000;
44
45/// Errors related to [`Arm9`].
46#[derive(Debug, Snafu)]
47pub enum Arm9Error {
48    /// Occurs when the program is too small to contain a secure area.
49    #[snafu(display("expected {expected:#x} bytes for secure area but had only {actual:#x}:\n{backtrace}"))]
50    DataTooSmall {
51        /// Expected minimum size.
52        expected: usize,
53        /// Actual size.
54        actual: usize,
55        /// Backtrace to the source of the error.
56        backtrace: Backtrace,
57    },
58    /// See [`BlowfishError`].
59    #[snafu(transparent)]
60    Blowfish {
61        /// Source error.
62        source: BlowfishError,
63    },
64    /// Occurs when the string "encryObj" is not found when de/encrypting the secure area.
65    #[snafu(display("invalid encryption, 'encryObj' not found"))]
66    NotEncryObj {
67        /// Backtrace to the source of the error.
68        backtrace: Backtrace,
69    },
70    /// See [`RawBuildInfoError`].
71    #[snafu(transparent)]
72    RawBuildInfo {
73        /// Source error.
74        source: RawBuildInfoError,
75    },
76    /// See [`Lz77DecompressError`].
77    #[snafu(transparent)]
78    Lz77Decompress {
79        /// Source error.
80        source: Lz77DecompressError,
81    },
82    /// See [`io::Error`].
83    #[snafu(transparent)]
84    Io {
85        /// Source error.
86        source: io::Error,
87    },
88}
89
90/// Errors related to ARM9 autoloads.
91#[derive(Debug, Snafu)]
92pub enum Arm9AutoloadError {
93    /// See [`RawBuildInfoError`].
94    #[snafu(transparent)]
95    RawBuildInfo {
96        /// Source error.
97        source: RawBuildInfoError,
98    },
99    /// See [`RawAutoloadInfoError`].
100    #[snafu(transparent)]
101    RawAutoloadInfo {
102        /// Source error.
103        source: RawAutoloadInfoError,
104    },
105    /// Occurs when trying to access autoload blocks while the ARM9 program is compressed.
106    #[snafu(display("ARM9 program must be decompressed before accessing autoload blocks:\n{backtrace}"))]
107    Compressed {
108        /// Backtrace to the source of the error.
109        backtrace: Backtrace,
110    },
111    /// Occurs when trying to access a kind of autoload block which doesn't exist in the ARM9 program.
112    #[snafu(display("autoload block {kind} could not be found:\n{backtrace}"))]
113    NotFound {
114        /// Kind of autoload.
115        kind: AutoloadKind,
116        /// Backtrace to the source of the error.
117        backtrace: Backtrace,
118    },
119}
120
121/// Options for [`Arm9::with_two_tcms`].
122pub struct Arm9WithTcmsOptions {
123    /// Whether the program was compressed originally.
124    pub originally_compressed: bool,
125    /// Whether the program was encrypted originally.
126    pub originally_encrypted: bool,
127}
128
129impl<'a> Arm9<'a> {
130    /// Creates a new ARM9 program from raw data.
131    pub fn new<T: Into<Cow<'a, [u8]>>>(data: T, offsets: Arm9Offsets) -> Result<Self, RawBuildInfoError> {
132        let mut arm9 = Arm9 { data: data.into(), offsets, originally_compressed: false, originally_encrypted: false };
133        arm9.originally_compressed = arm9.is_compressed()?;
134        arm9.originally_encrypted = arm9.is_encrypted();
135        Ok(arm9)
136    }
137
138    /// Creates a new ARM9 program with raw data and two autoloads (ITCM and DTCM).
139    ///
140    /// # Errors
141    ///
142    /// See [`Self::build_info_mut`].
143    pub fn with_two_tcms(
144        mut data: Vec<u8>,
145        itcm: Autoload,
146        dtcm: Autoload,
147        offsets: Arm9Offsets,
148        options: Arm9WithTcmsOptions,
149    ) -> Result<Self, RawBuildInfoError> {
150        let autoload_infos = [itcm.info().clone(), dtcm.info().clone()];
151
152        let autoload_blocks = data.len() as u32 + offsets.base_address;
153        data.extend(itcm.into_data().iter());
154        data.extend(dtcm.into_data().iter());
155        let autoload_infos_start = data.len() as u32 + offsets.base_address;
156        data.extend(bytemuck::bytes_of(&autoload_infos));
157        let autoload_infos_end = data.len() as u32 + offsets.base_address;
158
159        let Arm9WithTcmsOptions { originally_compressed, originally_encrypted } = options;
160        let mut arm9 = Self { data: data.into(), offsets, originally_compressed, originally_encrypted };
161
162        let build_info = arm9.build_info_mut()?;
163        build_info.autoload_blocks = autoload_blocks;
164        build_info.autoload_infos_start = autoload_infos_start;
165        build_info.autoload_infos_end = autoload_infos_end;
166
167        Ok(arm9)
168    }
169
170    /// Creates a new ARM9 program with raw data and a list of autoloads.
171    ///
172    /// # Errors
173    ///
174    /// See [`Self::build_info_mut`].
175    pub fn with_autoloads(
176        mut data: Vec<u8>,
177        autoloads: &[Autoload],
178        offsets: Arm9Offsets,
179        options: Arm9WithTcmsOptions,
180    ) -> Result<Self, RawBuildInfoError> {
181        let autoload_blocks = data.len() as u32 + offsets.base_address;
182
183        for autoload in autoloads {
184            data.extend(autoload.full_data());
185        }
186
187        let autoload_infos_start = data.len() as u32 + offsets.base_address;
188        for autoload in autoloads {
189            data.extend(bytemuck::bytes_of(autoload.info()));
190        }
191        let autoload_infos_end = data.len() as u32 + offsets.base_address;
192
193        let Arm9WithTcmsOptions { originally_compressed, originally_encrypted } = options;
194        let mut arm9 = Self { data: data.into(), offsets, originally_compressed, originally_encrypted };
195
196        let build_info = arm9.build_info_mut()?;
197        build_info.autoload_blocks = autoload_blocks;
198        build_info.autoload_infos_start = autoload_infos_start;
199        build_info.autoload_infos_end = autoload_infos_end;
200
201        Ok(arm9)
202    }
203
204    /// Returns whether the secure area is encrypted. See [`Self::originally_encrypted`] for whether the secure area was
205    /// encrypted originally.
206    pub fn is_encrypted(&self) -> bool {
207        self.data.len() < 8 || self.data[0..8] != SECURE_AREA_ID
208    }
209
210    /// Decrypts the secure area. Does nothing if already decrypted.
211    ///
212    /// # Errors
213    ///
214    /// This function will return an error if the program is too small to contain a secure area, [`Blowfish::decrypt`] fails or
215    /// "encryObj" was not found.
216    pub fn decrypt(&mut self, key: &BlowfishKey, gamecode: u32) -> Result<(), Arm9Error> {
217        if !self.is_encrypted() {
218            return Ok(());
219        }
220
221        if self.data.len() < 0x4000 {
222            DataTooSmallSnafu { expected: 0x4000usize, actual: self.data.len() }.fail()?;
223        }
224
225        let mut secure_area = [0u8; 0x4000];
226        secure_area.clone_from_slice(&self.data[0..0x4000]);
227
228        let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level2);
229        blowfish.decrypt(&mut secure_area[0..8])?;
230
231        let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level3);
232        blowfish.decrypt(&mut secure_area[0..0x800])?;
233
234        if &secure_area[0..8] != SECURE_AREA_ENCRY_OBJ {
235            NotEncryObjSnafu {}.fail()?;
236        }
237
238        secure_area[0..8].copy_from_slice(&SECURE_AREA_ID);
239        self.data.to_mut()[0..0x4000].copy_from_slice(&secure_area);
240        Ok(())
241    }
242
243    /// Encrypts the secure area. Does nothing if already encrypted.
244    ///
245    /// # Errors
246    ///
247    /// This function will return an error if the program is too small to contain a secure area, or the secure area ID was not
248    /// found.
249    pub fn encrypt(&mut self, key: &BlowfishKey, gamecode: u32) -> Result<(), Arm9Error> {
250        if self.is_encrypted() {
251            return Ok(());
252        }
253
254        if self.data.len() < 0x4000 {
255            DataTooSmallSnafu { expected: 0x4000usize, actual: self.data.len() }.fail()?;
256        }
257
258        if self.data[0..8] != SECURE_AREA_ID {
259            NotEncryObjSnafu {}.fail()?;
260        }
261
262        let secure_area = self.encrypted_secure_area(key, gamecode);
263        self.data.to_mut()[0..0x4000].copy_from_slice(&secure_area);
264        Ok(())
265    }
266
267    /// Returns an encrypted copy of the secure area.
268    pub fn encrypted_secure_area(&self, key: &BlowfishKey, gamecode: u32) -> [u8; 0x4000] {
269        let mut secure_area = [0u8; 0x4000];
270        secure_area.copy_from_slice(&self.data[0..0x4000]);
271        if self.is_encrypted() {
272            return secure_area;
273        }
274
275        secure_area[0..8].copy_from_slice(SECURE_AREA_ENCRY_OBJ);
276
277        let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level3);
278        blowfish.encrypt(&mut secure_area[0..0x800]).unwrap();
279
280        let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level2);
281        blowfish.encrypt(&mut secure_area[0..8]).unwrap();
282
283        secure_area
284    }
285
286    /// Returns a CRC checksum of the encrypted secure area.
287    pub fn secure_area_crc(&self, key: &BlowfishKey, gamecode: u32) -> u16 {
288        let secure_area = self.encrypted_secure_area(key, gamecode);
289        let checksum = CRC_16_MODBUS.checksum(&secure_area);
290        checksum
291    }
292
293    /// Returns a reference to the build info.
294    ///
295    /// # Errors
296    ///
297    /// See [`BuildInfo::borrow_from_slice`].
298    pub fn build_info(&self) -> Result<&BuildInfo, RawBuildInfoError> {
299        BuildInfo::borrow_from_slice(&self.data[self.offsets.build_info as usize..])
300    }
301
302    /// Returns a mutable reference to the build info.
303    ///
304    /// # Errors
305    ///
306    /// See [`BuildInfo::borrow_from_slice_mut`].
307    pub fn build_info_mut(&mut self) -> Result<&mut BuildInfo, RawBuildInfoError> {
308        BuildInfo::borrow_from_slice_mut(&mut self.data.to_mut()[self.offsets.build_info as usize..])
309    }
310
311    /// Returns whether this ARM9 program is compressed. See [`Self::originally_compressed`] for whether the program was
312    /// compressed originally.
313    ///
314    /// # Errors
315    ///
316    /// See [`Self::build_info`].
317    pub fn is_compressed(&self) -> Result<bool, RawBuildInfoError> {
318        Ok(self.build_info()?.is_compressed())
319    }
320
321    /// Decompresses this ARM9 program. Does nothing if already decompressed.
322    ///
323    /// # Errors
324    ///
325    /// See [`Self::is_compressed`] and [`Self::build_info_mut`].
326    pub fn decompress(&mut self) -> Result<(), Arm9Error> {
327        if !self.is_compressed()? {
328            return Ok(());
329        }
330
331        let data: Cow<[u8]> = LZ77.decompress(&self.data)?.into_vec().into();
332        let old_data = replace(&mut self.data, data);
333        let build_info = match self.build_info_mut() {
334            Ok(build_info) => build_info,
335            Err(e) => {
336                self.data = old_data;
337                return Err(e.into());
338            }
339        };
340        build_info.compressed_code_end = 0;
341        Ok(())
342    }
343
344    /// Compresses this ARM9 program. Does nothing if already compressed.
345    ///
346    /// # Errors
347    ///
348    /// See [`Self::is_compressed`], [`Lz77::compress`] and [`Self::build_info_mut`].
349    pub fn compress(&mut self) -> Result<(), Arm9Error> {
350        if self.is_compressed()? {
351            return Ok(());
352        }
353
354        let data: Cow<[u8]> = LZ77.compress(&self.data, COMPRESSION_START)?.into_vec().into();
355        let length = data.len();
356        let old_data = replace(&mut self.data, data);
357        let base_address = self.base_address();
358        let build_info = match self.build_info_mut() {
359            Ok(build_info) => build_info,
360            Err(e) => {
361                self.data = old_data;
362                return Err(e.into());
363            }
364        };
365        build_info.compressed_code_end = base_address + length as u32;
366        Ok(())
367    }
368
369    fn get_autoload_infos(&self, build_info: &BuildInfo) -> Result<&[AutoloadInfo], Arm9AutoloadError> {
370        let start = (build_info.autoload_infos_start - self.base_address()) as usize;
371        let end = (build_info.autoload_infos_end - self.base_address()) as usize;
372        let autoload_info = AutoloadInfo::borrow_from_slice(&self.data[start..end])?;
373        Ok(autoload_info)
374    }
375
376    /// Returns the autoload infos of this [`Arm9`].
377    ///
378    /// # Errors
379    ///
380    /// This function will return an error if [`Self::build_info`] or [`Self::get_autoload_infos`] fails or this ARM9 program
381    /// is compressed.
382    pub fn autoload_infos(&self) -> Result<&[AutoloadInfo], Arm9AutoloadError> {
383        let build_info: &BuildInfo = self.build_info()?;
384        if build_info.is_compressed() {
385            CompressedSnafu {}.fail()?;
386        }
387        self.get_autoload_infos(build_info)
388    }
389
390    /// Returns the autoloads of this [`Arm9`].
391    ///
392    /// # Errors
393    ///
394    /// This function will return an error if [`Self::build_info`] or [`Self::get_autoload_infos`] fails or this ARM9 program
395    /// is compressed.
396    pub fn autoloads(&self) -> Result<Box<[Autoload]>, Arm9AutoloadError> {
397        let build_info = self.build_info()?;
398        if build_info.is_compressed() {
399            CompressedSnafu {}.fail()?;
400        }
401        let autoload_infos = self.get_autoload_infos(build_info)?;
402
403        let mut autoloads = vec![];
404        let mut load_offset = build_info.autoload_blocks - self.base_address();
405        for autoload_info in autoload_infos {
406            let start = load_offset as usize;
407            let end = start + autoload_info.code_size as usize;
408            let data = &self.data[start..end];
409            autoloads.push(Autoload::new(data, *autoload_info));
410            load_offset += autoload_info.code_size;
411        }
412
413        Ok(autoloads.into_boxed_slice())
414    }
415
416    /// Returns the number of unknown autoloads of this [`Arm9`].
417    ///
418    /// # Errors
419    ///
420    /// See [`Self::autoloads`].
421    pub fn num_unknown_autoloads(&self) -> Result<usize, Arm9AutoloadError> {
422        Ok(self.autoloads()?.iter().filter(|a| matches!(a.kind(), AutoloadKind::Unknown(_))).count())
423    }
424
425    /// Returns the code of this ARM9 program.
426    ///
427    /// # Errors
428    ///
429    /// See [`Self::build_info`].
430    pub fn code(&self) -> Result<&[u8], RawBuildInfoError> {
431        let build_info = self.build_info()?;
432        Ok(&self.data[..(build_info.bss_start - self.base_address()) as usize])
433    }
434
435    /// Returns a reference to the full data.
436    pub fn full_data(&self) -> &[u8] {
437        &self.data
438    }
439
440    /// Returns the base address.
441    pub fn base_address(&self) -> u32 {
442        self.offsets.base_address
443    }
444
445    /// Returns the end address.
446    pub fn end_address(&self) -> Result<u32, RawBuildInfoError> {
447        let build_info = self.build_info()?;
448        Ok(build_info.bss_end)
449    }
450
451    /// Returns the entry function address.
452    pub fn entry_function(&self) -> u32 {
453        self.offsets.entry_function
454    }
455
456    /// Returns the build info offset.
457    pub fn build_info_offset(&self) -> u32 {
458        self.offsets.build_info
459    }
460
461    /// Returns the autoload callback address.
462    pub fn autoload_callback(&self) -> u32 {
463        self.offsets.autoload_callback
464    }
465
466    /// Returns the [`Range`] of uninitialized data in this ARM9 program.
467    ///
468    /// # Errors
469    ///
470    /// See [`Self::build_info`].
471    pub fn bss(&self) -> Result<Range<u32>, RawBuildInfoError> {
472        let build_info = self.build_info()?;
473        Ok(build_info.bss_start..build_info.bss_end)
474    }
475
476    /// Returns a reference to the ARM9 offsets.
477    pub fn offsets(&self) -> &Arm9Offsets {
478        &self.offsets
479    }
480
481    /// Returns whether the ARM9 program was compressed originally. See [`Self::is_compressed`] for the current state.
482    pub fn originally_compressed(&self) -> bool {
483        self.originally_compressed
484    }
485
486    /// Returns whether the ARM9 program was encrypted originally. See [`Self::is_encrypted`] for the current state.
487    pub fn originally_encrypted(&self) -> bool {
488        self.originally_encrypted
489    }
490}
491
492impl<'a> AsRef<[u8]> for Arm9<'a> {
493    fn as_ref(&self) -> &[u8] {
494        &self.data
495    }
496}
ds_rom/rom/arm9.rs

ds_rom/rom/
arm9.rs
