Use Rc<RefCell<T>>> for content_records so that they sync in a Title

When using a Title, the content_records value stored in the TMD and ContentRegion instances will now point to the same data, meaning that they stay in sync. Previously, you had to manually sync the content records between them as they were modified, and not doing so would cause problems when editing a WAD.

src/title/content.rs

@@ -3,7 +3,9 @@
 //
 // Implements content parsing and editing.
 
+use std::cell::RefCell;
 use std::io::{Cursor, Read, Seek, SeekFrom, Write};
+use std::rc::Rc;
 use sha1::{Sha1, Digest};
 use thiserror::Error;
 use crate::title::tmd::{ContentRecord, ContentType};
@@ -31,9 +33,8 @@ pub enum ContentError {
 #[derive(Debug)]
 /// A structure that represents the block of data containing the content of a digital Wii title.
 pub struct ContentRegion {
-    pub content_records: Vec<ContentRecord>,
+    pub content_records: Rc<RefCell<Vec<ContentRecord>>>,
     pub content_region_size: u32,
-    pub num_contents: u16,
     pub content_start_offsets: Vec<u64>,
     pub contents: Vec<Vec<u8>>,
 }
@@ -41,32 +42,31 @@ pub struct ContentRegion {
 impl ContentRegion {
     /// Creates a ContentRegion instance that can be used to parse and edit content stored in a 
     /// digital Wii title from the content area of a WAD and the ContentRecords from a TMD.
-    pub fn from_bytes(data: &[u8], content_records: Vec<ContentRecord>) -> Result<Self, ContentError> {
+    pub fn from_bytes(data: &[u8], content_records: Rc<RefCell<Vec<ContentRecord>>>) -> Result<Self, ContentError> {
         let content_region_size = data.len() as u32;
-        let num_contents = content_records.len() as u16;
+        let num_contents = content_records.borrow().len() as u16;
         // Calculate the starting offsets of each content.
         let content_start_offsets: Vec<u64> = std::iter::once(0)
-            .chain(content_records.iter().scan(0, |offset, record| {
+            .chain(content_records.borrow().iter().scan(0, |offset, record| {
                 *offset += record.content_size;
                 if record.content_size % 64 != 0 {
                     *offset += 64 - (record.content_size % 64);
                 }
                 Some(*offset)
-            })).take(content_records.len()).collect(); // Trims the extra final entry.
+            })).take(content_records.borrow().len()).collect(); // Trims the extra final entry.
         // Parse the content blob and create a vector of vectors from it.
         let mut contents: Vec<Vec<u8>> = Vec::with_capacity(num_contents as usize);
         let mut buf = Cursor::new(data);
         for i in 0..num_contents {
             buf.seek(SeekFrom::Start(content_start_offsets[i as usize]))?;
-            let size = (content_records[i as usize].content_size + 15) & !15;
+            let size = (content_records.borrow()[i as usize].content_size + 15) & !15;
             let mut content = vec![0u8; size as usize];
             buf.read_exact(&mut content)?;
             contents.push(content);
         }
         Ok(ContentRegion {
-            content_records,
+            content_records: Rc::clone(&content_records),
             content_region_size,
-            num_contents,
             content_start_offsets,
             contents,
         })
@@ -74,29 +74,29 @@ impl ContentRegion {
 
     /// Creates a ContentRegion instance that can be used to parse and edit content stored in a 
     /// digital Wii title from a vector of contents and the ContentRecords from a TMD.
-    pub fn from_contents(contents: Vec<Vec<u8>>, content_records: Vec<ContentRecord>) -> Result<Self, ContentError> {
-        if contents.len() != content_records.len() {
-            return Err(ContentError::MissingContents { required: content_records.len(), found: contents.len()});
+    pub fn from_contents(contents: Vec<Vec<u8>>, content_records: Rc<RefCell<Vec<ContentRecord>>>) -> Result<Self, ContentError> {
+        if contents.len() != content_records.borrow().len() {
+            return Err(ContentError::MissingContents { required: content_records.borrow().len(), found: contents.len()});
         }
-        let mut content_region = Self::new(content_records)?;
+        let mut content_region = Self::new(Rc::clone(&content_records))?;
         for i in 0..contents.len() {
-            content_region.load_enc_content(&contents[i], content_region.content_records[i].index as usize)?;
+            let target_index = content_region.content_records.borrow()[i].index;
+            content_region.load_enc_content(&contents[i], target_index as usize)?;
         }
         Ok(content_region)
     }
     
     /// Creates a ContentRegion instance from the ContentRecords of a TMD that contains no actual
     /// content. This can be used to load existing content from files.
-    pub fn new(content_records: Vec<ContentRecord>) -> Result<Self, ContentError> {
-        let content_region_size: u64 = content_records.iter().map(|x| (x.content_size + 63) & !63).sum();
+    pub fn new(content_records: Rc<RefCell<Vec<ContentRecord>>>) -> Result<Self, ContentError> {
+        let content_region_size: u64 = content_records.borrow().iter().map(|x| (x.content_size + 63) & !63).sum();
         let content_region_size = content_region_size as u32;
-        let num_contents = content_records.len() as u16;
+        let num_contents = content_records.borrow().len() as u16;
         let content_start_offsets: Vec<u64> = vec![0; num_contents as usize];
         let contents: Vec<Vec<u8>> = vec![Vec::new(); num_contents as usize];
         Ok(ContentRegion {
-            content_records,
+            content_records:  Rc::clone(&content_records),
             content_region_size,
-            num_contents,
             content_start_offsets,
             contents,
         })
@@ -105,7 +105,7 @@ impl ContentRegion {
     /// Dumps the entire ContentRegion back into binary data that can be written to a file.
     pub fn to_bytes(&self) -> Result<Vec<u8>, std::io::Error> {
         let mut buf: Vec<u8> = Vec::new();
-        for i in 0..self.num_contents {
+        for i in 0..self.content_records.borrow().len() {
             let mut content = self.contents[i as usize].clone();
             // Round up size to nearest 64 to add appropriate padding.
             content.resize((content.len() + 63) & !63, 0);
@@ -118,7 +118,7 @@ impl ContentRegion {
     pub fn get_index_from_cid(&self, cid: u32) -> Result<usize, ContentError> {
         // Use fancy Rust find and map methods to find the index matching the provided CID. Take
         // that libWiiPy!
-        let content_index = self.content_records.iter()
+        let content_index = self.content_records.borrow().iter()
             .find(|record| record.content_id == cid)
             .map(|record| record.index);
         if let Some(index) = content_index {
@@ -130,7 +130,7 @@ impl ContentRegion {
 
     /// Gets the encrypted content file from the ContentRegion at the specified index.
     pub fn get_enc_content_by_index(&self, index: usize) -> Result<Vec<u8>, ContentError> {
-        let content = self.contents.get(index).ok_or(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 })?;
+        let content = self.contents.get(index).ok_or(ContentError::IndexOutOfRange { index, max: self.content_records.borrow().len() - 1 })?;
         Ok(content.clone())
     }
 
@@ -138,20 +138,20 @@ impl ContentRegion {
     pub fn get_content_by_index(&self, index: usize, title_key: [u8; 16]) -> Result<Vec<u8>, ContentError> {
         let content = self.get_enc_content_by_index(index)?;
         // Verify the hash of the decrypted content against its record.
-        let mut content_dec = crypto::decrypt_content(&content, title_key, self.content_records[index].index);
-        content_dec.resize(self.content_records[index].content_size as usize, 0);
+        let mut content_dec = crypto::decrypt_content(&content, title_key, self.content_records.borrow()[index].index);
+        content_dec.resize(self.content_records.borrow()[index].content_size as usize, 0);
         let mut hasher = Sha1::new();
         hasher.update(content_dec.clone());
         let result = hasher.finalize();
-        if result[..] != self.content_records[index].content_hash {
-            return Err(ContentError::BadHash { hash: hex::encode(result), expected: hex::encode(self.content_records[index].content_hash) });
+        if result[..] != self.content_records.borrow()[index].content_hash {
+            return Err(ContentError::BadHash { hash: hex::encode(result), expected: hex::encode(self.content_records.borrow()[index].content_hash) });
         }
         Ok(content_dec)
     }
 
     /// Gets the encrypted content file from the ContentRegion with the specified Content ID.
     pub fn get_enc_content_by_cid(&self, cid: u32) -> Result<Vec<u8>, ContentError> {
-        let index = self.content_records.iter().position(|x| x.content_id == cid);
+        let index = self.content_records.borrow().iter().position(|x| x.content_id == cid);
         if let Some(index) = index {
             let content = self.get_enc_content_by_index(index).map_err(|_| ContentError::CIDNotFound(cid))?;
             Ok(content)
@@ -162,7 +162,7 @@ impl ContentRegion {
 
     /// Gets the decrypted content file from the ContentRegion with the specified Content ID.
     pub fn get_content_by_cid(&self, cid: u32, title_key: [u8; 16]) -> Result<Vec<u8>, ContentError> {
-        let index = self.content_records.iter().position(|x| x.content_id == cid);
+        let index = self.content_records.borrow().iter().position(|x| x.content_id == cid);
         if let Some(index) = index {
             let content_dec = self.get_content_by_index(index, title_key)?;
             Ok(content_dec)
@@ -174,8 +174,8 @@ impl ContentRegion {
     /// Loads existing content into the specified index of a ContentRegion instance. This content 
     /// must be encrypted.
     pub fn load_enc_content(&mut self, content: &[u8], index: usize) -> Result<(), ContentError> {
-        if index >= self.content_records.len() {
-            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 });
+        if index >= self.content_records.borrow().len() {
+            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.borrow().len() - 1 });
         }
         self.contents[index] = content.to_vec();
         Ok(())
@@ -186,20 +186,20 @@ impl ContentRegion {
     /// values can be set in the corresponding content record. Optionally, a new Content ID or
     /// content type can be provided, with the existing values being preserved by default.
     pub fn set_enc_content(&mut self, content: &[u8], index: usize, content_size: u64, content_hash: [u8; 20], cid: Option<u32>, content_type: Option<ContentType>) -> Result<(), ContentError> {
-        if index >= self.content_records.len() {
-            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 });
+        if index >= self.content_records.borrow().len() {
+            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.borrow().len() - 1 });
         }
-        self.content_records[index].content_size = content_size;
-        self.content_records[index].content_hash = content_hash;
+        self.content_records.borrow_mut()[index].content_size = content_size;
+        self.content_records.borrow_mut()[index].content_hash = content_hash;
         if cid.is_some() {
             // Make sure that the new CID isn't already in use.
-            if self.content_records.iter().any(|record| record.content_id == cid.unwrap()) {
+            if self.content_records.borrow().iter().any(|record| record.content_id == cid.unwrap()) {
                 return Err(ContentError::CIDAlreadyExists(cid.unwrap()));
             }
-            self.content_records[index].content_id = cid.unwrap();
+            self.content_records.borrow_mut()[index].content_id = cid.unwrap();
         }
         if content_type.is_some() {
-            self.content_records[index].content_type = content_type.unwrap();
+            self.content_records.borrow_mut()[index].content_type = content_type.unwrap();
         }
         self.contents[index] = content.to_vec();
         Ok(())
@@ -209,18 +209,18 @@ impl ContentRegion {
     /// must be decrypted and needs to match the size and hash listed in the content record at that
     /// index.
     pub fn load_content(&mut self, content: &[u8], index: usize, title_key: [u8; 16]) -> Result<(), ContentError> {
-        if index >= self.content_records.len() {
-            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 });
+        if index >= self.content_records.borrow().len() {
+            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.borrow().len() - 1 });
         }
         // Hash the content we're trying to load to ensure it matches the hash expected in the
         // matching record.
         let mut hasher = Sha1::new();
         hasher.update(content);
         let result = hasher.finalize();
-        if result[..] != self.content_records[index].content_hash {
-            return Err(ContentError::BadHash { hash: hex::encode(result), expected: hex::encode(self.content_records[index].content_hash) });
+        if result[..] != self.content_records.borrow()[index].content_hash {
+            return Err(ContentError::BadHash { hash: hex::encode(result), expected: hex::encode(self.content_records.borrow()[index].content_hash) });
         }
-        let content_enc = encrypt_content(content, title_key, self.content_records[index].index, self.content_records[index].content_size);
+        let content_enc = encrypt_content(content, title_key, self.content_records.borrow()[index].index, self.content_records.borrow()[index].content_size);
         self.contents[index] = content_enc;
         Ok(())
     }
@@ -243,12 +243,11 @@ impl ContentRegion {
     /// may leave a gap in the indexes recorded in the content records, but this should not cause
     /// issues on the Wii or with correctly implemented WAD parsers.
     pub fn remove_content(&mut self, index: usize) -> Result<(), ContentError> {
-        if self.contents.get(index).is_none() || self.content_records.get(index).is_none() {
-            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 });
+        if self.contents.get(index).is_none() || self.content_records.borrow().get(index).is_none() {
+            return Err(ContentError::IndexOutOfRange { index, max: self.content_records.borrow().len() - 1 });
         }
         self.contents.remove(index);
-        self.content_records.remove(index);
-        self.num_contents -= 1;
+        self.content_records.borrow_mut().remove(index);
         Ok(())
     }
 
@@ -256,15 +255,14 @@ impl ContentRegion {
     /// Content ID, type, index, and decrypted hash will be added to the record.
     pub fn add_enc_content(&mut self, content: &[u8], index: u16, cid: u32, content_type: ContentType, content_size: u64, content_hash: [u8; 20]) -> Result<(), ContentError> {
         // Return an error if the specified index or CID already exist in the records.
-        if self.content_records.iter().any(|record| record.index == index) {
+        if self.content_records.borrow().iter().any(|record| record.index == index) {
             return Err(ContentError::IndexAlreadyExists(index));
         }
-        if self.content_records.iter().any(|record| record.content_id == cid) {
+        if self.content_records.borrow().iter().any(|record| record.content_id == cid) {
             return Err(ContentError::CIDAlreadyExists(cid));
         }
         self.contents.push(content.to_vec());
-        self.content_records.push(ContentRecord { content_id: cid, index, content_type, content_size, content_hash });
-        self.num_contents += 1;
+        self.content_records.borrow_mut().push(ContentRecord { content_id: cid, index, content_type, content_size, content_hash });
         Ok(())
     }
     
@@ -273,7 +271,7 @@ impl ContentRegion {
     /// index will be automatically assigned based on the highest index currently recorded in the
     /// content records.
     pub fn add_content(&mut self, content: &[u8], cid: u32, content_type: ContentType, title_key: [u8; 16]) -> Result<(), ContentError> {
-        let max_index = self.content_records.iter()
+        let max_index = self.content_records.borrow().iter()
             .max_by_key(|record| record.index)
             .map(|record| record.index)
             .unwrap_or(0); // This should be impossible, but I guess 0 is a safe value just in case?

src/title/mod.rs

@@ -13,6 +13,7 @@ pub mod tmd;
 pub mod versions;
 pub mod wad;
 
+use std::rc::Rc;
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -52,7 +53,7 @@ impl Title {
         let cert_chain = cert::CertificateChain::from_bytes(&wad.cert_chain()).map_err(TitleError::CertificateError)?;
         let ticket = ticket::Ticket::from_bytes(&wad.ticket()).map_err(TitleError::Ticket)?;
         let tmd = tmd::TMD::from_bytes(&wad.tmd()).map_err(TitleError::TMD)?;
-        let content = content::ContentRegion::from_bytes(&wad.content(), tmd.content_records.clone()).map_err(TitleError::Content)?;
+        let content = content::ContentRegion::from_bytes(&wad.content(), Rc::clone(&tmd.content_records)).map_err(TitleError::Content)?;
         Ok(Title {
             cert_chain,
             crl: wad.crl(),
@@ -137,7 +138,6 @@ impl Title {
     /// content type can be provided, with the existing values being preserved by default.
     pub fn set_content(&mut self, content: &[u8], index: usize, cid: Option<u32>, content_type: Option<tmd::ContentType>) -> Result<(), TitleError> {
         self.content.set_content(content, index, cid, content_type, self.ticket.dec_title_key())?;
-        self.tmd.content_records = self.content.content_records.clone();
         Ok(())
     }
 
@@ -147,7 +147,6 @@ impl Title {
     /// content records.
     pub fn add_content(&mut self, content: &[u8], cid: u32, content_type: tmd::ContentType) -> Result<(), TitleError> {
         self.content.add_content(content, cid, content_type, self.ticket.dec_title_key())?;
-        self.tmd.content_records = self.content.content_records.clone();
         Ok(())
     }
     
@@ -159,7 +158,7 @@ impl Title {
         // accurate results.
         title_size += self.tmd.to_bytes().map_err(|x| TitleError::TMD(tmd::TMDError::IO(x)))?.len();
         title_size += self.ticket.to_bytes().map_err(|x| TitleError::Ticket(ticket::TicketError::IO(x)))?.len();
-        for record in &self.tmd.content_records {
+        for record in self.tmd.content_records.borrow().iter() {
             if matches!(record.content_type, tmd::ContentType::Shared) {
                 if absolute == Some(true) {
                     title_size += record.content_size as usize;

src/title/nus.rs

@@ -80,7 +80,7 @@ pub fn download_content(title_id: [u8; 8], content_id: u32, wiiu_endpoint: bool)
 
 /// Downloads all contents from the specified title from the NUS.
 pub fn download_contents(tmd: &tmd::TMD, wiiu_endpoint: bool) -> Result<Vec<Vec<u8>>, NUSError> {
-    let content_ids: Vec<u32> = tmd.content_records.iter().map(|record| { record.content_id }).collect();
+    let content_ids: Vec<u32> = tmd.content_records.borrow().iter().map(|record| { record.content_id }).collect();
     let mut contents: Vec<Vec<u8>> = Vec::new();
     for id in content_ids {
         contents.push(download_content(tmd.title_id, id, wiiu_endpoint)?);

src/title/tmd.rs

@@ -3,9 +3,11 @@
 //
 // Implements the structures and methods required for TMD parsing and editing.
 
+use std::cell::RefCell;
 use std::fmt;
 use std::io::{Cursor, Read, Write};
 use std::ops::Index;
+use std::rc::Rc;
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use sha1::{Sha1, Digest};
 use thiserror::Error;
@@ -110,7 +112,7 @@ pub struct TMD {
     pub num_contents: u16,
     pub boot_index: u16,
     pub minor_version: u16, // Normally unused, but good for fakesigning!
-    pub content_records: Vec<ContentRecord>,
+    pub content_records: Rc<RefCell<Vec<ContentRecord>>>,
 }
 
 impl TMD {
@@ -204,7 +206,7 @@ impl TMD {
             num_contents,
             boot_index,
             minor_version,
-            content_records,
+            content_records: Rc::new(RefCell::new(content_records)),
         })
     }
     
@@ -231,11 +233,11 @@ impl TMD {
         buf.write_all(&self.reserved2)?;
         buf.write_u32::<BigEndian>(self.access_rights)?;
         buf.write_u16::<BigEndian>(self.title_version)?;
-        buf.write_u16::<BigEndian>(self.num_contents)?;
+        buf.write_u16::<BigEndian>(self.content_records.borrow().len() as u16)?;
         buf.write_u16::<BigEndian>(self.boot_index)?;
         buf.write_u16::<BigEndian>(self.minor_version)?;
         // Iterate over content records and write out content record data.
-        for content in &self.content_records {
+        for content in self.content_records.borrow().iter() {
             buf.write_u32::<BigEndian>(content.content_id)?;
             buf.write_u16::<BigEndian>(content.index)?;
             match content.content_type {
@@ -317,11 +319,11 @@ impl TMD {
         // Find possible content indices, because the provided one could exist while the indices
         // are out of order, which could cause problems finding the content.
         let mut content_indices = Vec::new();
-        for record in &self.content_records {
+        for record in self.content_records.borrow().iter() {
             content_indices.push(record.index);
         }
         let target_index = content_indices.index(index);
-        match self.content_records[*target_index as usize].content_type {
+        match self.content_records.borrow()[*target_index as usize].content_type {
             ContentType::Normal => ContentType::Normal,
             ContentType::Development => ContentType::Development,
             ContentType::HashTree => ContentType::HashTree,