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(), *dtcm.info()];
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        CRC_16_MODBUS.checksum(&secure_area)
290    }
291
292    /// Returns a reference to the build info.
293    ///
294    /// # Errors
295    ///
296    /// See [`BuildInfo::borrow_from_slice`].
297    pub fn build_info(&self) -> Result<&BuildInfo, RawBuildInfoError> {
298        BuildInfo::borrow_from_slice(&self.data[self.offsets.build_info as usize..])
299    }
300
301    /// Returns a mutable reference to the build info.
302    ///
303    /// # Errors
304    ///
305    /// See [`BuildInfo::borrow_from_slice_mut`].
306    pub fn build_info_mut(&mut self) -> Result<&mut BuildInfo, RawBuildInfoError> {
307        BuildInfo::borrow_from_slice_mut(&mut self.data.to_mut()[self.offsets.build_info as usize..])
308    }
309
310    /// Returns whether this ARM9 program is compressed. See [`Self::originally_compressed`] for whether the program was
311    /// compressed originally.
312    ///
313    /// # Errors
314    ///
315    /// See [`Self::build_info`].
316    pub fn is_compressed(&self) -> Result<bool, RawBuildInfoError> {
317        Ok(self.build_info()?.is_compressed())
318    }
319
320    /// Decompresses this ARM9 program. Does nothing if already decompressed.
321    ///
322    /// # Errors
323    ///
324    /// See [`Self::is_compressed`] and [`Self::build_info_mut`].
325    pub fn decompress(&mut self) -> Result<(), Arm9Error> {
326        if !self.is_compressed()? {
327            return Ok(());
328        }
329
330        let data: Cow<[u8]> = LZ77.decompress(&self.data)?.into_vec().into();
331        let old_data = replace(&mut self.data, data);
332        let build_info = match self.build_info_mut() {
333            Ok(build_info) => build_info,
334            Err(e) => {
335                self.data = old_data;
336                return Err(e.into());
337            }
338        };
339        build_info.compressed_code_end = 0;
340        Ok(())
341    }
342
343    /// Compresses this ARM9 program. Does nothing if already compressed.
344    ///
345    /// # Errors
346    ///
347    /// See [`Self::is_compressed`], [`Lz77::compress`] and [`Self::build_info_mut`].
348    pub fn compress(&mut self) -> Result<(), Arm9Error> {
349        if self.is_compressed()? {
350            return Ok(());
351        }
352
353        let data: Cow<[u8]> = LZ77.compress(&self.data, COMPRESSION_START)?.into_vec().into();
354        let length = data.len();
355        let old_data = replace(&mut self.data, data);
356        let base_address = self.base_address();
357        let build_info = match self.build_info_mut() {
358            Ok(build_info) => build_info,
359            Err(e) => {
360                self.data = old_data;
361                return Err(e.into());
362            }
363        };
364        build_info.compressed_code_end = base_address + length as u32;
365        Ok(())
366    }
367
368    fn get_autoload_infos(&self, build_info: &BuildInfo) -> Result<&[AutoloadInfo], Arm9AutoloadError> {
369        let start = (build_info.autoload_infos_start - self.base_address()) as usize;
370        let end = (build_info.autoload_infos_end - self.base_address()) as usize;
371        let autoload_info = AutoloadInfo::borrow_from_slice(&self.data[start..end])?;
372        Ok(autoload_info)
373    }
374
375    /// Returns the autoload infos of this [`Arm9`].
376    ///
377    /// # Errors
378    ///
379    /// This function will return an error if [`Self::build_info`] or [`Self::get_autoload_infos`] fails or this ARM9 program
380    /// is compressed.
381    pub fn autoload_infos(&self) -> Result<&[AutoloadInfo], Arm9AutoloadError> {
382        let build_info: &BuildInfo = self.build_info()?;
383        if build_info.is_compressed() {
384            CompressedSnafu {}.fail()?;
385        }
386        self.get_autoload_infos(build_info)
387    }
388
389    /// Returns the autoloads of this [`Arm9`].
390    ///
391    /// # Errors
392    ///
393    /// This function will return an error if [`Self::build_info`] or [`Self::get_autoload_infos`] fails or this ARM9 program
394    /// is compressed.
395    pub fn autoloads(&self) -> Result<Box<[Autoload]>, Arm9AutoloadError> {
396        let build_info = self.build_info()?;
397        if build_info.is_compressed() {
398            CompressedSnafu {}.fail()?;
399        }
400        let autoload_infos = self.get_autoload_infos(build_info)?;
401
402        let mut autoloads = vec![];
403        let mut load_offset = build_info.autoload_blocks - self.base_address();
404        for autoload_info in autoload_infos {
405            let start = load_offset as usize;
406            let end = start + autoload_info.code_size as usize;
407            let data = &self.data[start..end];
408            autoloads.push(Autoload::new(data, *autoload_info));
409            load_offset += autoload_info.code_size;
410        }
411
412        Ok(autoloads.into_boxed_slice())
413    }
414
415    /// Returns the number of unknown autoloads of this [`Arm9`].
416    ///
417    /// # Errors
418    ///
419    /// See [`Self::autoloads`].
420    pub fn num_unknown_autoloads(&self) -> Result<usize, Arm9AutoloadError> {
421        Ok(self.autoloads()?.iter().filter(|a| matches!(a.kind(), AutoloadKind::Unknown(_))).count())
422    }
423
424    /// Returns the code of this ARM9 program.
425    ///
426    /// # Errors
427    ///
428    /// See [`Self::build_info`].
429    pub fn code(&self) -> Result<&[u8], RawBuildInfoError> {
430        let build_info = self.build_info()?;
431        Ok(&self.data[..(build_info.bss_start - self.base_address()) as usize])
432    }
433
434    /// Returns a reference to the full data.
435    pub fn full_data(&self) -> &[u8] {
436        &self.data
437    }
438
439    /// Returns the base address.
440    pub fn base_address(&self) -> u32 {
441        self.offsets.base_address
442    }
443
444    /// Returns the end address.
445    pub fn end_address(&self) -> Result<u32, RawBuildInfoError> {
446        let build_info = self.build_info()?;
447        Ok(build_info.bss_end)
448    }
449
450    /// Returns the entry function address.
451    pub fn entry_function(&self) -> u32 {
452        self.offsets.entry_function
453    }
454
455    /// Returns the build info offset.
456    pub fn build_info_offset(&self) -> u32 {
457        self.offsets.build_info
458    }
459
460    /// Returns the autoload callback address.
461    pub fn autoload_callback(&self) -> u32 {
462        self.offsets.autoload_callback
463    }
464
465    /// Returns the [`Range`] of uninitialized data in this ARM9 program.
466    ///
467    /// # Errors
468    ///
469    /// See [`Self::build_info`].
470    pub fn bss(&self) -> Result<Range<u32>, RawBuildInfoError> {
471        let build_info = self.build_info()?;
472        Ok(build_info.bss_start..build_info.bss_end)
473    }
474
475    /// Returns a reference to the ARM9 offsets.
476    pub fn offsets(&self) -> &Arm9Offsets {
477        &self.offsets
478    }
479
480    /// Returns whether the ARM9 program was compressed originally. See [`Self::is_compressed`] for the current state.
481    pub fn originally_compressed(&self) -> bool {
482        self.originally_compressed
483    }
484
485    /// Returns whether the ARM9 program was encrypted originally. See [`Self::is_encrypted`] for the current state.
486    pub fn originally_encrypted(&self) -> bool {
487        self.originally_encrypted
488    }
489}
490
491impl AsRef<[u8]> for Arm9<'_> {
492    fn as_ref(&self) -> &[u8] {
493        &self.data
494    }
495}
ds_rom/rom/arm9.rs

ds_rom/rom/
arm9.rs
