feat: support chunked data

.gitignore

@@ -1 +1,2 @@
 /target
+/pkg

src/lib.rs

@@ -28,12 +28,12 @@ use hmac::{Hmac, Mac};
 use js_sys::Function;
 use prost::Message;
 use sha2::{Digest, Sha256, Sha512};
+use std::collections::HashMap;
 use std::io::{self, Read};
 use wasm_bindgen::prelude::*;
 
 type HmacSha256 = Hmac<Sha256>;
 
-// Keep the original protobuf module
 pub mod signal {
     include!(concat!(env!("OUT_DIR"), "/signal.rs"));
 }
@@ -63,19 +63,22 @@ impl DecryptionResult {
     }
 }
 
-// Helper struct to read from byte slice
+struct ByteReader {
-struct ByteReader<'a> {
+    data: Vec<u8>,
-    data: &'a [u8],
     position: usize,
 }
 
-impl<'a> ByteReader<'a> {
+impl ByteReader {
-    fn new(data: &'a [u8]) -> Self {
+    fn new(data: Vec<u8>) -> Self {
         ByteReader { data, position: 0 }
     }
+
+    fn remaining_data(&self) -> &[u8] {
+        &self.data[self.position..]
+    }
 }
 
-impl<'a> Read for ByteReader<'a> {
+impl Read for ByteReader {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         let available = self.data.len() - self.position;
         let amount = buf.len().min(available);
@@ -90,7 +93,6 @@ impl<'a> Read for ByteReader<'a> {
     }
 }
 
-// Keep the original structs but without Debug derive
 struct HeaderData {
     initialisation_vector: Vec<u8>,
     salt: Vec<u8>,
@@ -230,125 +232,113 @@ fn decrypt_frame(
         .map_err(|e| JsValue::from_str(&format!("Failed to decode frame: {}", e)))
 }
 
-// this would be used for attachments, stickers, avatars, which we might need later
+#[wasm_bindgen]
-fn decrypt_frame_payload(
+pub struct BackupDecryptor {
-    reader: &mut ByteReader,
+    reader: ByteReader,
-    length: usize,
+    keys: Option<Keys>,
-    hmac_key: &[u8],
+    header_data: Option<HeaderData>,
-    cipher_key: &[u8],
+    initialisation_vector: Option<Vec<u8>>,
-    initialisation_vector: &[u8],
+    database_bytes: Vec<u8>,
-    chunk_size: usize,
+    preferences: HashMap<String, HashMap<String, HashMap<String, serde_json::Value>>>,
-) -> Result<Vec<u8>, JsValue> {
+    key_values: HashMap<String, HashMap<String, serde_json::Value>>,
-    let mut hmac = <HmacSha256 as Mac>::new_from_slice(hmac_key)
+    ciphertext_buf: Vec<u8>,
-        .map_err(|_| JsValue::from_str("Invalid HMAC key"))?;
+    plaintext_buf: Vec<u8>,
-    Mac::update(&mut hmac, initialisation_vector);
+    total_bytes_processed: usize,
-
+    is_initialized: bool,
-    let mut ctr =
-        <Ctr32BE<Aes256> as KeyIvInit>::new_from_slices(cipher_key, initialisation_vector)
-            .map_err(|_| JsValue::from_str("Invalid CTR parameters"))?;
-
-    let mut decrypted_data = Vec::new();
-    let mut remaining_length = length;
-
-    while remaining_length > 0 {
-        let this_chunk_length = remaining_length.min(chunk_size);
-        remaining_length -= this_chunk_length;
-
-        let mut ciphertext = vec![0u8; this_chunk_length];
-        reader
-            .read_exact(&mut ciphertext)
-            .map_err(|e| JsValue::from_str(&format!("Failed to read chunk: {}", e)))?;
-        Mac::update(&mut hmac, &ciphertext);
-
-        let mut decrypted_chunk = ciphertext;
-        ctr.apply_keystream(&mut decrypted_chunk);
-        decrypted_data.extend(decrypted_chunk);
-    }
-
-    let mut their_mac = [0u8; 10];
-    reader
-        .read_exact(&mut their_mac)
-        .map_err(|e| JsValue::from_str(&format!("Failed to read MAC: {}", e)))?;
-    let our_mac = hmac.finalize().into_bytes();
-
-    if &their_mac != &our_mac[..10] {
-        return Err(JsValue::from_str(
-            "Bad MAC found. Passphrase may be incorrect or file corrupted or incompatible.",
-        ));
-    }
-
-    Ok(decrypted_data)
 }
 
 #[wasm_bindgen]
-pub fn decrypt_backup(
+impl BackupDecryptor {
-    backup_data: &[u8],
+    #[wasm_bindgen(constructor)]
+    pub fn new() -> Self {
+        Self {
+            reader: ByteReader::new(Vec::new()),
+            keys: None,
+            header_data: None,
+            initialisation_vector: None,
+            database_bytes: Vec::new(),
+            preferences: HashMap::new(),
+            key_values: HashMap::new(),
+            ciphertext_buf: Vec::with_capacity(1024 * 1024),
+            plaintext_buf: Vec::with_capacity(1024 * 1024),
+            total_bytes_processed: 0,
+            is_initialized: false,
+        }
+    }
+
+    #[wasm_bindgen]
+    pub fn feed_data(&mut self, chunk: &[u8]) {
+        let mut new_data = self.reader.remaining_data().to_vec();
+        new_data.extend_from_slice(chunk);
+        self.total_bytes_processed += chunk.len();
+        self.reader = ByteReader::new(new_data);
+    }
+
+    #[wasm_bindgen]
+    pub fn process_chunk(
+        &mut self,
         passphrase: &str,
         progress_callback: &Function,
-) -> Result<DecryptionResult, JsValue> {
+    ) -> Result<bool, JsValue> {
-    let mut reader = ByteReader::new(backup_data);
+        if !self.is_initialized {
-    let total_size = backup_data.len();
+            // Initialize on first chunk
-    let mut last_percentage = 0;
+            self.header_data = Some(read_backup_header(&mut self.reader)?);
+            let header_data = self.header_data.as_ref().unwrap();
+            self.keys = Some(derive_keys(passphrase, &header_data.salt)?);
+            self.initialisation_vector = Some(header_data.initialisation_vector.clone());
+            self.is_initialized = true;
+        }
 
-    // Set up collections for results
+        // Calculate progress based on current position
-    let mut database_bytes = Vec::new();
+        let current_position = self.reader.position;
-    let mut preferences: std::collections::HashMap<
+        let percentage =
-        String,
+            ((current_position as f64 / self.total_bytes_processed as f64) * 100.0) as u32;
-        std::collections::HashMap<String, std::collections::HashMap<String, serde_json::Value>>,
-    > = std::collections::HashMap::new();
-    let mut key_values: std::collections::HashMap<
-        String,
-        std::collections::HashMap<String, serde_json::Value>,
-    > = std::collections::HashMap::new();
 
-    // Read header and derive keys
+        // Report progress
-    let header_data = read_backup_header(&mut reader)?;
-    let keys = derive_keys(passphrase, &header_data.salt)?;
-    let mut initialisation_vector = header_data.initialisation_vector.clone();
-
-    // Pre-allocate buffers for frame decryption
-    let mut ciphertext: Vec<u8> = Vec::with_capacity(1024 * 1024);
-    let mut plaintext: Vec<u8> = Vec::with_capacity(1024 * 1024);
-
-    // Main decryption loop
-    loop {
-        // Update progress
-        let current_position = reader.position;
-        let percentage = ((current_position as f64 / total_size as f64) * 100.0) as u32;
-        if percentage != last_percentage {
         progress_callback
             .call1(&JsValue::NULL, &JsValue::from_f64(percentage as f64))
             .map_err(|e| JsValue::from_str(&format!("Failed to report progress: {:?}", e)))?;
-            last_percentage = percentage;
-        }
 
-        let backup_frame = decrypt_frame(
+        let keys = self.keys.as_ref().unwrap();
-            &mut reader,
+        let header_data = self.header_data.as_ref().unwrap();
+        let iv = self.initialisation_vector.as_ref().unwrap();
+
+        match decrypt_frame(
+            &mut self.reader,
             &keys.hmac_key,
             &keys.cipher_key,
-            &initialisation_vector,
+            iv,
             header_data.version,
-            &mut ciphertext,
+            &mut self.ciphertext_buf,
-            &mut plaintext,
+            &mut self.plaintext_buf,
-        )?;
+        ) {
-
+            Ok(backup_frame) => {
-        initialisation_vector = increment_initialisation_vector(&initialisation_vector);
+                self.initialisation_vector = Some(increment_initialisation_vector(iv));
 
                 if backup_frame.end.unwrap_or(false) {
-            break;
+                    // Report 100% completion when done
-        } else if let Some(statement) = backup_frame.statement {
+                    progress_callback
+                        .call1(&JsValue::NULL, &JsValue::from_f64(100.0))
+                        .map_err(|e| {
+                            JsValue::from_str(&format!("Failed to report final progress: {:?}", e))
+                        })?;
+                    return Ok(true);
+                }
+
+                // Process frame contents
+                if let Some(statement) = backup_frame.statement {
                     if let Some(sql) = statement.statement {
                         if !sql.to_lowercase().starts_with("create table sqlite_")
                             && !sql.contains("sms_fts_")
                             && !sql.contains("mms_fts_")
                         {
-                    // Store SQL statements and parameters for database reconstruction
+                            self.database_bytes.extend_from_slice(sql.as_bytes());
-                    database_bytes.extend_from_slice(sql.as_bytes());
+                            self.database_bytes.push(b';');
-                    database_bytes.push(b';');
                         }
                     }
                 } else if let Some(preference) = backup_frame.preference {
-            let value_dict = preferences
+                    let value_dict = self
+                        .preferences
                         .entry(preference.file.unwrap_or_default())
                         .or_default()
                         .entry(preference.key.unwrap_or_default())
@@ -376,9 +366,8 @@ pub fn decrypt_backup(
                         );
                     }
                 } else if let Some(key_value) = backup_frame.key_value {
-            let value_dict = key_values
+                    let key = key_value.key.clone().unwrap_or_default();
-                .entry(key_value.key.unwrap_or_default())
+                    let value_dict = self.key_values.entry(key).or_default();
-                .or_default();
 
                     if let Some(boolean_value) = key_value.boolean_value {
                         value_dict.insert(
@@ -422,19 +411,82 @@ pub fn decrypt_backup(
                         );
                     }
                 }
-        // Note: We're skipping attachments, stickers, and avatars for now
+
+                Ok(false)
+            }
+            Err(e) => {
+                if e.as_string()
+                    .unwrap_or_default()
+                    .contains("unexpected end of file")
+                {
+                    Ok(false)
+                } else {
+                    Err(e)
+                }
+            }
+        }
     }
 
-    // Final progress update
+    #[wasm_bindgen]
-    progress_callback
+    pub fn finish(self) -> Result<DecryptionResult, JsValue> {
-        .call1(&JsValue::NULL, &JsValue::from_f64(100.0))
-        .map_err(|e| JsValue::from_str(&format!("Failed to report final progress: {:?}", e)))?;
-
         Ok(DecryptionResult {
-        database_bytes,
+            database_bytes: self.database_bytes,
-        preferences: serde_json::to_string(&preferences)
+            preferences: serde_json::to_string(&self.preferences).map_err(|e| {
-            .map_err(|e| JsValue::from_str(&format!("Failed to serialize preferences: {}", e)))?,
+                JsValue::from_str(&format!("Failed to serialize preferences: {}", e))
-        key_values: serde_json::to_string(&key_values)
+            })?,
-            .map_err(|e| JsValue::from_str(&format!("Failed to serialize key_values: {}", e)))?,
+            key_values: serde_json::to_string(&self.key_values).map_err(|e| {
+                JsValue::from_str(&format!("Failed to serialize key_values: {}", e))
+            })?,
         })
     }
+}
+
+// this would be used for attachments, stickers, avatars, which we might need later
+// fn decrypt_frame_payload(
+//     reader: &mut ByteReader,
+//     length: usize,
+//     hmac_key: &[u8],
+//     cipher_key: &[u8],
+//     initialisation_vector: &[u8],
+//     chunk_size: usize,
+// ) -> Result<Vec<u8>, JsValue> {
+//     let mut hmac = <HmacSha256 as Mac>::new_from_slice(hmac_key)
+//         .map_err(|_| JsValue::from_str("Invalid HMAC key"))?;
+//     Mac::update(&mut hmac, initialisation_vector);
+
+//     let mut ctr =
+//         <Ctr32BE<Aes256> as KeyIvInit>::new_from_slices(cipher_key, initialisation_vector)
+//             .map_err(|_| JsValue::from_str("Invalid CTR parameters"))?;
+
+//     let mut decrypted_data = Vec::new();
+//     let mut remaining_length = length;
+
+//     while remaining_length > 0 {
+//         let this_chunk_length = remaining_length.min(chunk_size);
+//         remaining_length -= this_chunk_length;
+
+//         let mut ciphertext = vec![0u8; this_chunk_length];
+//         reader
+//             .read_exact(&mut ciphertext)
+//             .map_err(|e| JsValue::from_str(&format!("Failed to read chunk: {}", e)))?;
+//         Mac::update(&mut hmac, &ciphertext);
+
+//         let mut decrypted_chunk = ciphertext;
+//         ctr.apply_keystream(&mut decrypted_chunk);
+//         decrypted_data.extend(decrypted_chunk);
+//     }
+
+//     let mut their_mac = [0u8; 10];
+//     reader
+//         .read_exact(&mut their_mac)
+//         .map_err(|e| JsValue::from_str(&format!("Failed to read MAC: {}", e)))?;
+//     let our_mac = hmac.finalize().into_bytes();
+
+//     if &their_mac != &our_mac[..10] {
+//         return Err(JsValue::from_str(
+//             "Bad MAC found. Passphrase may be incorrect or file corrupted or incompatible.",
+//         ));
+//     }
+
+//     Ok(decrypted_data)
+// }