Decrypt error handling (#10)

Co-authored-by: ssolit <sol@seismic.systems>
SeismicSystems · Oct 25, 2024 · 83d629c · 83d629c
1 parent c5d8b01
commit 83d629c
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 4 deletions.
diff --git 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<Body>) -> Result<Response<Body>,
         .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
@@ -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<T: Encodable>(key: &Key<Aes256Gcm>, plaintext: &T, nonce: u64
 ///
 /// # Panics
 /// This function will panic if decryption or decoding fails.
-pub fn aes_decrypt<T>(key: &Key<Aes256Gcm>, ciphertext: &[u8], nonce: u64) -> T
+pub fn aes_decrypt<T>(
+    key: &Key<Aes256Gcm>,
+    ciphertext: &[u8],
+    nonce: u64,
+) -> Result<T, anyhow::Error>
 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
@@ -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(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<Body> {
+    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()
+}