diff --git a/src/connector.rs b/src/connector.rs index 01ff126..8cafe64 100755 --- a/src/connector.rs +++ b/src/connector.rs @@ -185,27 +185,47 @@ impl Connector { loop { let rx_msg = rx.recv().await; - match rx_msg.unwrap() { - Message::ClientReq(name) => { - let req = Request::new(name, metadata).unwrap(); // TODO: handle - self.request(conn, req).await?; - } - Message::Shutdown => { - let msg = b"DISCONNECT".to_vec(); - comms::send( - &mut conn.writer, - Some(&mut conn.cipher), - Some(&mut conn.rng), - &msg, - ) - .await?; - - break; - } - _ => continue, + if rx_msg.is_none() { + continue; } + + match self.msg_handler(rx_msg.unwrap(), conn, metadata).await? { + true => continue, + false => break, + }; } Ok(()) } + + async fn msg_handler( + &self, + msg: Message, + conn: &mut Connection<'_>, + metadata: &HashMap, + ) -> Result> { + match msg { + Message::ClientReq(name) => { + let req = match Request::new(name, metadata) { + Some(req) => req, + None => return Ok(true), + }; + self.request(conn, req).await?; + Ok(true) + } + Message::Shutdown => { + let msg = b"DISCONNECT".to_vec(); + comms::send( + &mut conn.writer, + Some(&mut conn.cipher), + Some(&mut conn.rng), + &msg, + ) + .await?; + + Ok(false) + } + _ => Ok(true), + } + } } diff --git a/src/crypto.rs b/src/crypto.rs index 538e282..daf33a6 100755 --- a/src/crypto.rs +++ b/src/crypto.rs @@ -62,7 +62,10 @@ pub fn aes_encrypt( rng: &mut OsRng, ) -> Result, Box> { let nonce = generate_nonce(rng); - let encrypted = cipher.encrypt(&nonce, data.as_ref()).unwrap(); // TODO: handle errors + let encrypted = match cipher.encrypt(&nonce, data.as_ref()) { + Ok(d) => d, + Err(_) => return Err("AES encryption failed".into()), + }; let mut data = nonce.to_vec(); data.extend_from_slice(&encrypted); @@ -74,9 +77,10 @@ pub fn aes_decrypt( cipher: &mut AesGcm, ) -> Result, Box> { let (nonce_bytes, data) = data.split_at(AES_NONCE_SIZE); - let decrypted = cipher - .decrypt(Nonce::from_slice(nonce_bytes), data.as_ref()) - .unwrap(); // TODO: handle errors + let decrypted = match cipher.decrypt(Nonce::from_slice(nonce_bytes), data.as_ref()) { + Ok(d) => d, + Err(_) => return Err("AES decryption failed".into()), + }; Ok(decrypted) } diff --git a/src/error.rs b/src/error.rs deleted file mode 100644 index ff7bd09..0000000 --- a/src/error.rs +++ /dev/null @@ -1 +0,0 @@ -// placeholder