use axum::body::Body;
use axum::http::{header, Request, StatusCode};
use axum::routing::{get, post};
use axum::Router;
use larust_http::csrf;
use larust_http::session::{session_layer as build_session_layer, Session};
use tower::ServiceExt;
async fn show_token(session: Session) -> String {
csrf::token(&session).await
}
async fn submit() -> &'static str {
"ok"
}
async fn shared_pool() -> sqlx::AnyPool {
let dir = tempfile::tempdir().unwrap().keep();
let database_url = format!("sqlite://{}/test.sqlite", dir.display());
let _ = larust_orm::connect(&database_url).await;
larust_orm::pool().unwrap().clone()
}
async fn app() -> Router {
let pool = shared_pool().await;
let session_layer = build_session_layer(&pool, true).await.unwrap();
Router::new()
.route("/token", get(show_token))
.route("/submit", post(submit))
.layer(axum::middleware::from_fn(csrf::verify))
.layer(session_layer)
}
async fn fetch_token(router: &Router) -> (String, String) {
let response = router
.clone()
.oneshot(Request::get("/token").body(Body::empty()).unwrap())
.await
.unwrap();
let cookie = response
.headers()
.get(header::SET_COOKIE)
.expect("session cookie should be set")
.to_str()
.unwrap()
.split(';')
.next()
.unwrap()
.to_string();
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let token = String::from_utf8(body.to_vec()).unwrap();
(token, cookie)
}
fn post_with(cookie: Option<&str>, body: &str) -> Request<Body> {
let mut builder =
Request::post("/submit").header(header::CONTENT_TYPE, "application/x-www-form-urlencoded");
if let Some(cookie) = cookie {
builder = builder.header(header::COOKIE, cookie);
}
builder.body(Body::from(body.to_string())).unwrap()
}
#[tokio::test]
async fn get_requests_are_not_checked() {
let response = app()
.await
.oneshot(Request::get("/token").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn post_with_correct_token_and_matching_session_succeeds() {
let router = app().await;
let (token, cookie) = fetch_token(&router).await;
let response = router
.oneshot(post_with(Some(&cookie), &format!("_csrf_token={token}")))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn post_with_no_token_is_rejected() {
let router = app().await;
let (_token, cookie) = fetch_token(&router).await;
let response = router.oneshot(post_with(Some(&cookie), "")).await.unwrap();
assert_eq!(response.status(), StatusCode::from_u16(419).unwrap());
}
#[tokio::test]
async fn post_with_wrong_token_is_rejected() {
let router = app().await;
let (_token, cookie) = fetch_token(&router).await;
let response = router
.oneshot(post_with(Some(&cookie), "_csrf_token=not-the-real-token"))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::from_u16(419).unwrap());
}
#[tokio::test]
async fn post_with_correct_token_but_different_session_is_rejected() {
let router = app().await;
let (token, _cookie) = fetch_token(&router).await;
let response = router
.oneshot(post_with(None, &format!("_csrf_token={token}")))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::from_u16(419).unwrap());
}
#[tokio::test]
async fn post_with_correct_token_in_header_succeeds_with_no_body_field_at_all() {
let router = app().await;
let (token, cookie) = fetch_token(&router).await;
let response = router
.oneshot(
Request::post("/submit")
.header(header::COOKIE, &cookie)
.header(csrf::HEADER_NAME, &token)
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn a_wrong_header_is_rejected_even_with_a_correct_token_in_the_body() {
let router = app().await;
let (token, cookie) = fetch_token(&router).await;
let response = router
.oneshot(
Request::post("/submit")
.header(header::COOKIE, &cookie)
.header(header::CONTENT_TYPE, "application/x-www-form-urlencoded")
.header(csrf::HEADER_NAME, "not-the-real-token")
.body(Body::from(format!("_csrf_token={token}")))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::from_u16(419).unwrap());
}
#[tokio::test]
async fn a_header_authenticated_request_skips_the_bodys_2mb_read_cap() {
let router = app().await;
let (token, cookie) = fetch_token(&router).await;
let oversized_body = vec![b'a'; 3 * 1024 * 1024];
let response = router
.oneshot(
Request::post("/submit")
.header(header::COOKIE, &cookie)
.header(csrf::HEADER_NAME, &token)
.body(Body::from(oversized_body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}