Skip to content

Commit

Permalink
fix: roundtrip works. get_file test does not.
Browse files Browse the repository at this point in the history
  • Loading branch information
@tuddman
tuddman committed Apr 3, 2024
1 parent d849075 commit 9ea77be
Showing 1 changed file with 90 additions and 59 deletions.
149 changes: 90 additions & 59 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,10 @@ async fn upload_file(
mut payload: web::Payload,
data: web::Data<AppState>,
) -> 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();
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -108,24 +112,18 @@ 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,
Err(_) => return Err(HttpResponse::BadRequest().body("Invalid HMAC format")),
};

// 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);
Expand Down Expand Up @@ -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::<Sha256>::new_from_slice(secret_key).expect("HMAC can take key of any size");
Expand All @@ -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"
);
}
}

0 comments on commit 9ea77be

Please sign in to comment.