use std::{collections::HashSet, sync::Arc}; use gloo_console::log; use image::GenericImageView; use js_sys::{Array, Uint8Array}; use omegaupload_common::{ crypto::{open_in_place, Key, Nonce}, Expiration, }; use wasm_bindgen::JsCast; use web_sys::Blob; use yew::worker::{Agent, AgentLink, Context, HandlerId}; use yew::{html::Scope, worker::Public}; use crate::{DecryptedData, Paste, PasteCompleteConstructionError}; #[derive(Clone)] pub struct DecryptionAgent { link: AgentLink, } impl Agent for DecryptionAgent { type Reach = Public; type Message = (); type Input = DecryptionAgentMessage; type Output = Result<(DecryptedData, PasteContext), PasteCompleteConstructionError>; fn create(link: AgentLink) -> Self { Self { link } } fn update(&mut self, _: Self::Message) {} fn handle_input( &mut self, DecryptionAgentMessage { context, params }: Self::Input, id: HandlerId, ) { let DecryptionParams { data, key, nonce, maybe_password, } = params; self.link.respond( id, decrypt(data, key, nonce, maybe_password).map(|res| (res, context)), ) } } pub struct DecryptionAgentMessage { context: PasteContext, params: DecryptionParams, } impl DecryptionAgentMessage { pub fn new(context: PasteContext, params: DecryptionParams) -> Self { Self { context, params } } } pub struct PasteContext { pub link: Scope, pub expires: Option, } impl PasteContext { pub fn new(link: Scope, expires: Option) -> Self { Self { link, expires } } } pub struct DecryptionParams { data: Vec, key: Key, nonce: Nonce, maybe_password: Option, } impl DecryptionParams { pub fn new(data: Vec, key: Key, nonce: Nonce, maybe_password: Option) -> Self { Self { data, key, nonce, maybe_password, } } } fn decrypt( mut container: Vec, key: Key, nonce: Nonce, maybe_password: Option, ) -> Result { let container = &mut container; log!("stage 1 decryption start"); if let Some(password) = maybe_password { open_in_place(container, &nonce.increment(), &password) .map_err(|_| PasteCompleteConstructionError::StageOneFailure)?; } log!("stage 2 decryption start"); open_in_place(container, &nonce, &key) .map_err(|_| PasteCompleteConstructionError::StageTwoFailure)?; log!("stage 2 decryption end"); if let Ok(decrypted) = std::str::from_utf8(&container) { Ok(DecryptedData::String(Arc::new(decrypted.to_owned()))) } else { log!("blob conversion start"); let blob_chunks = Array::new_with_length(container.chunks(65536).len().try_into().unwrap()); for (i, chunk) in container.chunks(65536).enumerate() { let array = Uint8Array::new_with_length(chunk.len().try_into().unwrap()); array.copy_from(&chunk); blob_chunks.set(i.try_into().unwrap(), array.dyn_into().unwrap()); } let blob = Arc::new(Blob::new_with_u8_array_sequence(blob_chunks.dyn_ref().unwrap()).unwrap()); log!("blob conversion end"); if let Ok(image) = image::load_from_memory(&container) { Ok(DecryptedData::Image( blob, image.dimensions(), container.len(), )) } else { Ok(DecryptedData::Blob(blob)) } } }