From 83d629cf8b012b8492ee9b8484ece7e0bac70985 Mon Sep 17 00:00:00 2001 From: ssolit <49683577+ssolit@users.noreply.github.com> Date: Fri, 25 Oct 2024 09:25:25 -0400 Subject: [PATCH] Decrypt error handling (#10) Co-authored-by: ssolit --- src/tx_io/handlers.rs | 37 ++++++++++++++++++++++++++++++++++++- src/utils/crypto_utils.rs | 14 +++++++++++--- src/utils/response_utils.rs | 13 +++++++++++++ 3 files changed, 60 insertions(+), 4 deletions(-) diff --git a/src/tx_io/handlers.rs b/src/tx_io/handlers.rs index 33a087e1..1fd9ba73 100644 --- a/src/tx_io/handlers.rs +++ b/src/tx_io/handlers.rs @@ -6,7 +6,9 @@ use secp256k1::SecretKey; use super::structs::*; use crate::utils::crypto_utils::*; -use crate::utils::response_utils::{invalid_json_body_resp, invalid_req_body_resp}; +use crate::utils::response_utils::{ + invalid_ciphertext_resp, invalid_json_body_resp, invalid_req_body_resp, +}; /// Handles an IO encryption request, encrypting the provided data using AES. /// @@ -87,6 +89,13 @@ pub async fn tx_io_decrypt_handler(req: Request) -> Result, .unwrap(); let decrypted_data = aes_decrypt(&aes_key, &decryption_request.data, decryption_request.nonce); + let decrypted_data = match decrypted_data { + Ok(data) => data, + Err(e) => { + return Ok(invalid_ciphertext_resp(e)); + } + }; + let response_body = IoDecryptionResponse { decrypted_data }; let response_json = serde_json::to_string(&response_body).unwrap(); @@ -182,6 +191,8 @@ mod tests { let enc_response: IoEncryptionResponse = serde_json::from_slice(&body).unwrap(); assert!(!enc_response.encrypted_data.is_empty()); + println!("Encrypted data: {:?}", enc_response.encrypted_data); + // check that decryption returns the original data // Prepare decrypt request body let decryption_request = IoDecryptionRequest { @@ -207,4 +218,28 @@ mod tests { let dec_response: IoDecryptionResponse = serde_json::from_slice(&body).unwrap(); assert_eq!(dec_response.decrypted_data, data_to_encrypt); } + + #[tokio::test] + async fn test_decrypt_invalid_ciphertext() { + let base_url = "http://localhost:7878"; + let bad_ciphertext = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + let decryption_request = IoDecryptionRequest { + msg_sender: PublicKey::from_str( + "03e31e68908a6404a128904579c677534d19d0e5db80c7d9cf4de6b4b7fe0518bd", + ) + .unwrap(), + data: bad_ciphertext, + nonce: 12345678, + }; + let payload_json = serde_json::to_string(&decryption_request).unwrap(); + let req = Request::builder() + .method("POST") + .uri(format!("{}/tx_io/decrypt", base_url)) + .header("Content-Type", "application/json") + .body(Body::from(payload_json)) + .unwrap(); + + let res = tx_io_decrypt_handler(req).await.unwrap(); + assert_eq!(res.status(), 422); + } } diff --git a/src/utils/crypto_utils.rs b/src/utils/crypto_utils.rs index f50bab36..26fe9e4c 100644 --- a/src/utils/crypto_utils.rs +++ b/src/utils/crypto_utils.rs @@ -3,6 +3,7 @@ use aes_gcm::{ Aes256Gcm, Key, }; use alloy_rlp::{Decodable, Encodable}; +use anyhow::anyhow; use hkdf::Hkdf; use secp256k1::{ecdh::SharedSecret, ecdsa::Signature, Message, PublicKey, Secp256k1, SecretKey}; use serde::{Deserialize, Serialize}; @@ -80,7 +81,11 @@ pub fn aes_encrypt(key: &Key, plaintext: &T, nonce: u64 /// /// # Panics /// This function will panic if decryption or decoding fails. -pub fn aes_decrypt(key: &Key, ciphertext: &[u8], nonce: u64) -> T +pub fn aes_decrypt( + key: &Key, + ciphertext: &[u8], + nonce: u64, +) -> Result where T: Decodable, { @@ -90,10 +95,13 @@ where // recover the plaintext byte encoding of the object let buf = cipher .decrypt(&nonce, ciphertext.as_ref()) - .expect("AES decryption failed"); + .map_err(|e| anyhow!("AES decryption failed: {:?}", e))?; // recover the object from the byte encoding - T::decode(&mut &buf[..]).unwrap_or_else(|err| panic!("Failed to decode: {:?}", err)) + let plaintext = + T::decode(&mut &buf[..]).map_err(|e| anyhow!("Failed to decode plaintext: {:?}", e))?; + + Ok(plaintext) } /// Derives an AES key from a shared secret using HKDF and SHA-256. diff --git a/src/utils/response_utils.rs b/src/utils/response_utils.rs index 8f8eab9d..6a65f6f8 100644 --- a/src/utils/response_utils.rs +++ b/src/utils/response_utils.rs @@ -1,3 +1,4 @@ +use anyhow::Error; use hyper::{Body, Response, StatusCode}; use serde_json::json; @@ -20,3 +21,15 @@ pub fn invalid_json_body_resp() -> Response { .body(Body::from(error_response)) .unwrap() } + +// Returns 422 Unprocessable Entity +// Meant to be used if decrypting the ciphertext fails +pub fn invalid_ciphertext_resp(e: Error) -> Response { + let error_message = format!("Invalid ciphertext: {}", e); // Use error's Display trait + let error_response = json!({ "error": error_message }).to_string(); + + Response::builder() + .status(StatusCode::UNPROCESSABLE_ENTITY) + .body(Body::from(error_response)) + .unwrap() +}