From 391d0c8599a3b71b5eed1498280278aab436095a Mon Sep 17 00:00:00 2001 From: G2 Date: Thu, 15 Aug 2024 00:46:12 -0500 Subject: [PATCH] WASM track upload (#7) * Initial working WASM encryption * Improved encryption implementation --- src/netmd/encryption.rs | 178 +++++++++++++++++++++++++++++++++------- src/netmd/interface.rs | 26 +++--- 2 files changed, 164 insertions(+), 40 deletions(-) diff --git a/src/netmd/encryption.rs b/src/netmd/encryption.rs index fb1b9ab..18c6b07 100644 --- a/src/netmd/encryption.rs +++ b/src/netmd/encryption.rs @@ -9,14 +9,94 @@ use super::interface::DataEncryptorInput; type DesEcbEnc = ecb::Decryptor; type DesCbcEnc = cbc::Encryptor; -pub fn new_thread_encryptor( - _input: DataEncryptorInput, -) -> UnboundedReceiver<(Vec, Vec, Vec)> { - let (tx, rx) = unbounded_channel::<(Vec, Vec, Vec)>(); - let input = Box::from(_input); +pub struct Encryptor { + channel: Option, Vec, Vec)>>, + state: Option, +} - thread::spawn(move || { - let mut iv = [0u8; 8]; +struct EncryptorState { + input_data: Vec, + iv: [u8; 8], + random_key: [u8; 8], + encrypted_random_key: [u8; 8], + default_chunk_size: usize, + current_chunk_size: usize, + offset: usize, + packet_count: usize, + closed: bool, +} + +impl Encryptor { + pub fn new_threaded(input: DataEncryptorInput) -> Self { + let (tx, rx) = unbounded_channel::<(Vec, Vec, Vec)>(); + + thread::spawn(move || { + let mut iv = [0u8; 8]; + + // Create the random key + let mut random_key = [0u8; 8]; + rand::thread_rng().fill_bytes(&mut random_key); + + // Encrypt it with the kek + let mut encrypted_random_key = random_key; + if let Err(x) = DesEcbEnc::new(&input.kek.into()) + .decrypt_padded_mut::(&mut encrypted_random_key) + { + panic!("Cannot create main key {:?}", x) + }; + + let default_chunk_size = match input.chunk_size { + 0 => 0x00100000, + e => e, + }; + + let mut packet_count = 0u32; + let mut current_chunk_size; + + let mut input_data = input.data.clone(); + if (input_data.len() % input.frame_size) != 0 { + let padding_remaining = input.frame_size - (input_data.len() % input.frame_size); + input_data.extend(std::iter::repeat(0).take(padding_remaining)); + } + let input_data_length = input_data.len(); + + let mut offset: usize = 0; + while offset < input_data_length { + if packet_count > 0 { + current_chunk_size = default_chunk_size; + } else { + current_chunk_size = default_chunk_size - 24; + } + + current_chunk_size = std::cmp::min(current_chunk_size, input_data_length - offset); + + let this_data_chunk = &mut input_data[offset..offset + current_chunk_size]; + DesCbcEnc::new(&random_key.into(), &iv.into()) + .encrypt_padded_mut::(this_data_chunk, current_chunk_size) + .unwrap(); + + tx.send(( + encrypted_random_key.to_vec(), + iv.to_vec(), + this_data_chunk.to_vec(), + )) + .unwrap(); + + iv.copy_from_slice(&this_data_chunk[this_data_chunk.len() - 8..]); + + packet_count += 1; + offset += current_chunk_size; + } + }); + + Self { + channel: Some(rx), + state: None + } + } + + pub fn new(input: DataEncryptorInput) -> Self { + let iv = [0u8; 8]; // Create the random key let mut random_key = [0u8; 8]; @@ -35,44 +115,82 @@ pub fn new_thread_encryptor( e => e, }; - let mut packet_count = 0u32; - let mut current_chunk_size; + let packet_count = 0; + let current_chunk_size = 0; let mut input_data = input.data.clone(); if (input_data.len() % input.frame_size) != 0 { let padding_remaining = input.frame_size - (input_data.len() % input.frame_size); input_data.extend(std::iter::repeat(0).take(padding_remaining)); } - let input_data_length = input_data.len(); - let mut offset: usize = 0; - while offset < input_data_length { - if packet_count > 0 { - current_chunk_size = default_chunk_size; - } else { - current_chunk_size = default_chunk_size - 24; + let offset: usize = 0; + + Encryptor { + channel: None, + state: Some(EncryptorState { + input_data, + iv, + random_key, + encrypted_random_key, + current_chunk_size, + offset, + default_chunk_size, + packet_count, + closed: false, + }) + } + } + + /// Get the next encrypted value + pub async fn next(&mut self) -> Option<(Vec, Vec, Vec)> { + let output; + + if let Some(state) = self.state.as_mut() { + if state.closed { + return None } - current_chunk_size = std::cmp::min(current_chunk_size, input_data_length - offset); + if state.packet_count > 0 { + state.current_chunk_size = state.default_chunk_size; + } else { + state.current_chunk_size = state.default_chunk_size - 24; + } - let this_data_chunk = &mut input_data[offset..offset + current_chunk_size]; - DesCbcEnc::new(&random_key.into(), &iv.into()) - .encrypt_padded_mut::(this_data_chunk, current_chunk_size) + state.current_chunk_size = std::cmp::min(state.current_chunk_size, state.input_data.len() - state.offset); + + let this_data_chunk = &mut state.input_data[state.offset..state.offset + state.current_chunk_size]; + DesCbcEnc::new(&state.random_key.into(), &state.iv.into()) + .encrypt_padded_mut::(this_data_chunk, state.current_chunk_size) .unwrap(); - tx.send(( - encrypted_random_key.to_vec(), - iv.to_vec(), + output = Some(( + state.encrypted_random_key.to_vec(), + state.iv.to_vec(), this_data_chunk.to_vec(), - )) - .unwrap(); + )); - iv.copy_from_slice(&this_data_chunk[this_data_chunk.len() - 8..]); + state.iv.copy_from_slice(&this_data_chunk[this_data_chunk.len() - 8..]); - packet_count += 1; - offset += current_chunk_size; + state.packet_count += 1; + state.offset += state.current_chunk_size; + } else if let Some(channel) = self.channel.as_mut() { + output = channel.recv().await + } else { + unreachable!("If you got here, this is bad!"); } - }); - rx + output + } + + /// Call close to return none from subsequent calls + pub fn close(&mut self) { + if let Some(state) = self.state.as_mut() { + state.closed = true; + } else if let Some(channel) = self.channel.as_mut() { + channel.close() + } else { + unreachable!("If you got here, this is bad!"); + } + } } diff --git a/src/netmd/interface.rs b/src/netmd/interface.rs index 6057276..b707ea7 100644 --- a/src/netmd/interface.rs +++ b/src/netmd/interface.rs @@ -15,9 +15,9 @@ use std::collections::HashMap; use std::error::Error; use std::time::Duration; use thiserror::Error; -use tokio::sync::mpsc::UnboundedReceiver; use super::base::NetMD; +use super::encryption::Encryptor; use super::utils::{cross_sleep, to_sjis}; /// An action to take on the player @@ -1690,8 +1690,7 @@ impl NetMDInterface { discformat: u8, frames: u32, pkt_size: u32, - // key, iv, data - mut packets: UnboundedReceiver<(Vec, Vec, Vec)>, + mut packets: Encryptor, hex_session_key: &[u8], progress_callback: F, ) -> Result<(u16, Vec, Vec), InterfaceError> @@ -1730,7 +1729,7 @@ impl NetMDInterface { let mut written_bytes = 0; let mut packet_count = 0; - while let Some((key, iv, data)) = packets.recv().await { + while let Some((key, iv, data)) = packets.next().await { let binpack = if packet_count == 0 { let packed_length: Vec = pkt_size.to_be_bytes().to_vec(); [vec![0, 0, 0, 0], packed_length, key, iv, data].concat() @@ -1863,10 +1862,6 @@ pub struct MDTrack { pub data: Vec, pub chunk_size: usize, pub full_width_title: Option, - - #[allow(clippy::type_complexity)] - pub encrypt_packets_iterator: - Box UnboundedReceiver<(Vec, Vec, Vec)>>, } pub struct DataEncryptorInput { @@ -1921,8 +1916,19 @@ impl MDTrack { [0x14, 0xe3, 0x83, 0x4e, 0xe2, 0xd3, 0xcc, 0xa5] } - pub fn get_encrypting_iterator(&mut self) -> UnboundedReceiver<(Vec, Vec, Vec)> { - (self.encrypt_packets_iterator)(DataEncryptorInput { + #[cfg(not(target_family = "wasm"))] + pub fn get_encrypting_iterator(&mut self) -> Encryptor { + Encryptor::new_threaded(DataEncryptorInput { + kek: self.get_kek(), + frame_size: self.frame_size(), + chunk_size: self.chunk_size(), + data: std::mem::take(&mut self.data), + }) + } + + #[cfg(target_family = "wasm")] + pub fn get_encrypting_iterator(&mut self) -> Encryptor { + Encryptor::new(DataEncryptorInput { kek: self.get_kek(), frame_size: self.frame_size(), chunk_size: self.chunk_size(),