From 9ea77be1b1b14afe15f362eddfb0901b8351b8a1 Mon Sep 17 00:00:00 2001 From: tuddman Date: Wed, 3 Apr 2024 14:37:21 +0200 Subject: [PATCH] fix: roundtrip works. get_file test does not. --- src/main.rs | 149 +++++++++++++++++++++++++++++++--------------------- 1 file changed, 90 insertions(+), 59 deletions(-) diff --git a/src/main.rs b/src/main.rs index a16f8c1..625dfc4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,7 +24,10 @@ async fn upload_file( mut payload: web::Payload, data: web::Data, ) -> impl Responder { - let secret_key = data.secret_key.as_bytes(); + let hmac_header = match req.headers().get("X-HMAC") { + Some(value) => value.to_str().unwrap_or_default(), + None => return HttpResponse::BadRequest().body("Missing X-HMAC header"), + }; // Create a new UUID for the file. let file_id = Uuid::new_v4(); @@ -52,7 +55,8 @@ async fn upload_file( } // Verify HMAC - if let Err(response) = verify_hmac(&req, &buffer, secret_key) { + let secret_key = data.secret_key.as_bytes(); + if let Err(response) = verify_hmac(hmac_header, &buffer, secret_key) { let _ = std::fs::remove_file(file_name); return response; } @@ -80,7 +84,7 @@ async fn get_file( // Read the file's content let mut file = match File::open(file_name) { Ok(f) => f, - Err(_) => return HttpResponse::InternalServerError().body("File not found"), + Err(_) => return HttpResponse::NotFound().body("File not found"), }; let mut buffer = Vec::new(); if file.read_to_end(&mut buffer).is_err() { @@ -108,16 +112,10 @@ async fn get_file( /// Verifies the HMAC of the request. fn verify_hmac( - req: &HttpRequest, + hmac_header: &str, file_bytes: &[u8], secret_key: &[u8], ) -> Result<(), HttpResponse> { - // Extract the HMAC from the header - let hmac_header = match req.headers().get("X-HMAC") { - Some(h) => h.to_str().unwrap_or_default(), - None => return Err(HttpResponse::Unauthorized().body("Missing HMAC header")), - }; - // Decode the hex HMAC let received_hmac = match decode(hmac_header) { Ok(hmac) => hmac, @@ -125,7 +123,7 @@ fn verify_hmac( }; // Create an instance of the HMAC-SHA256 - let mut mac = HmacSha256::new_from_slice(secret_key).expect("HMAC can take key of any size"); + let mut mac = HmacSha256::new_from_slice(secret_key).expect("Insufficient HMAC key size"); // Input the data to the HMAC instance mac.update(file_bytes); @@ -161,14 +159,12 @@ async fn main() -> std::io::Result<()> { #[cfg(test)] mod tests { use super::*; - use actix_web::dev::ServiceResponse; - use actix_web::{test, web, App}; - use hex::encode; - use sha2::Sha256; + use actix_web::test; + use tempfile::tempdir; const SECRET_KEY: &[u8] = b"TEST_SECRET_KEY"; - // Helper function to create a HMAC signature + // Test helper function to create a HMAC signature fn create_hmac_signature(secret_key: &[u8], data: &[u8]) -> String { let mut mac = Hmac::::new_from_slice(secret_key).expect("HMAC can take key of any size"); @@ -177,83 +173,118 @@ mod tests { } // Tests the HMAC verification logic - #[actix_rt::test] + #[test] async fn test_hmac_verification() { - let app = test::init_service(App::new().route( - "/test", - web::post().to(|req: HttpRequest, body: web::Bytes| async move { - // Attempt to verify the HMAC - let req = req.clone(); - let body = body.clone(); - match verify_hmac(&req, &body, SECRET_KEY) { - Ok(_) => HttpResponse::Ok().finish(), - Err(err) => err, - } - }), - )) - .await; - + // test valid HMAC passes as Ok(()) let correct_payload = b"correct payload"; + let correct_hmac = create_hmac_signature(SECRET_KEY, correct_payload); + let verify_correct = verify_hmac(&correct_hmac, correct_payload, SECRET_KEY); + + assert!(verify_correct.is_ok(), "Should succeed with correct HMAC"); + + // test invvalid HMAC returns an HttpResponse()) let incorrect_payload = b"incorrect payload"; + let verify_incorrect = verify_hmac(&correct_hmac, incorrect_payload, SECRET_KEY); + assert!(verify_incorrect.is_err()); + } + + #[actix_rt::test] + async fn test_upload_file() { + // set up application state + let data = web::Data::new(AppState { + secret_key: std::str::from_utf8(SECRET_KEY).unwrap().to_string(), + file_map: Mutex::new(HashMap::new()), + }); + + let app = test::init_service( + App::new() + .app_data(data.clone()) + .route("/upload", web::post().to(upload_file)), + ) + .await; - // Create a correct HMAC for the correct payload + let correct_payload = b"correct payload"; let correct_hmac = create_hmac_signature(SECRET_KEY, correct_payload); - // Simulate sending a request with the correct HMAC let req = test::TestRequest::post() - .uri("/test") + .uri("/upload") .insert_header(("X-HMAC", correct_hmac)) .set_payload(correct_payload.to_vec()) .to_request(); - let resp: ServiceResponse = test::call_service(&app, req).await; - assert!( - resp.status().is_success(), - "Should succeed with correct HMAC" - ); - - // Create an incorrect HMAC for the purpose of testing - let incorrect_hmac = create_hmac_signature(SECRET_KEY, incorrect_payload); - - // Simulate sending a request with the incorrect HMAC - let req = test::TestRequest::post() - .uri("/test") - .insert_header(("X-HMAC", incorrect_hmac)) - .set_payload(correct_payload.to_vec()) - .to_request(); + let resp = test::call_service(&app, req).await; + assert!(resp.status().is_success()); - let resp: ServiceResponse = test::call_service(&app, req).await; + // Verify that the file was created and the HMAC was correctly verified + let file_map = data.file_map.lock().unwrap(); assert!( - resp.status().is_client_error(), - "Should fail with incorrect HMAC" + !file_map.is_empty(), + "File map should not be empty after upload" ); - } + // Cleanup: TempDir will automatically remove the temporary directory once it goes out of scope + } #[actix_rt::test] async fn test_get_file() { + let data = web::Data::new(AppState { + secret_key: std::str::from_utf8(SECRET_KEY).unwrap().to_string(), + file_map: Mutex::new(HashMap::new()), + }); + let file_contents = b"this too shall pass!"; // Simulated file contents let correct_hmac = create_hmac_signature(SECRET_KEY, file_contents); let incorrect_hmac = "nope".to_string(); + let temp_dir = tempdir().unwrap(); + let upload_path = temp_dir.path().join("uploads"); + std::fs::create_dir_all(&upload_path).unwrap(); + + let file_id = Uuid::new_v4(); + let file_name = upload_path.as_path().join(file_id.to_string()); + let mut file = File::create(&file_name).unwrap(); + let _ = file.write_all(file_contents); + + // Insert the file ID and name into the state map. + let _ = data + .file_map + .lock() + .unwrap() + .insert(file_id, file_name.display().to_string()); + + dbg!(&data.file_map); + let app = test::init_service( - App::new().route("/file", web::get().to(get_file)) - ).await; + App::new() + .app_data(data) + .route(&format!("/files/{file_id}"), web::get().to(get_file)), + ) + .await; // Test with correct HMAC let req = test::TestRequest::get() - .uri("/file") + .uri(&format!("/files/{file_id}")) .insert_header(("X-HMAC", correct_hmac)) .to_request(); + dbg!(&req); let resp = test::call_service(&app, req).await; - assert!(resp.status().is_success(), "Should succeed with correct HMAC"); + dbg!(&file_id); + dbg!(&file_name); + dbg!(&resp); + assert!( + resp.status().is_success(), + "Should succeed with correct HMAC" + ); // Test with incorrect HMAC let req = test::TestRequest::get() - .uri("/file") + .uri(&format!("/files/{file_id}")) .insert_header(("X-HMAC", incorrect_hmac)) .to_request(); let resp = test::call_service(&app, req).await; - assert!(resp.status().is_client_error(), "Should fail with incorrect HMAC"); + assert!( + resp.status().is_client_error(), + "Should fail with incorrect HMAC" + ); } }