Skip to main content

webdav_server/api/middleware/
auth_guard.rs

1use crate::{
2    api::Error::{BadRequest, Unauthorized},
3    app::State as AppState,
4    model::session::{Session, TokenHash},
5};
6use axum::{
7    extract::{Request, State},
8    middleware::Next,
9    response::Response,
10};
11use blake3::Hasher;
12
13use crate::Result;
14
15/// Middleware that guards protected routes by validating the session token.
16///
17/// Extracts the Bearer token from the `Authorization` header, computes its hash,
18/// and checks the in-memory cache. If missing from the cache, it verifies the
19/// token against the database and caches the result for future requests.
20///
21/// # Errors
22/// - Returns an `Unauthorized` if the `Authorization` header is missing, malformed, or contains an invalid/expired token.
23/// - Returns a `BadRequest` if the token string cannot be serialized.
24/// - Returns an internal error if a database query fails.
25pub async fn auth_guard(
26    State(state): State<AppState>,
27    mut request: Request,
28    next: Next,
29) -> Result<Response> {
30    let header = request
31        .headers()
32        .get("Authorization")
33        .ok_or_else(|| Unauthorized("Missing Header".into()))?;
34
35    let token = header
36        .to_str()
37        .map_err(|_| BadRequest("Auth failed to serialize".into()))?
38        .strip_prefix("Bearer ")
39        .ok_or_else(|| {
40            Unauthorized("Tokens must start with 'Bearer'".into())
41        })?;
42
43    let token_hash: TokenHash = Hasher::new()
44        .update(token.as_bytes())
45        .finalize()
46        .as_bytes()
47        .into();
48
49    if let Some(user_id) = state.session_cache.get(&token_hash).await {
50        request.extensions_mut().insert(user_id);
51        return Ok(next.run(request).await);
52    }
53
54    let session = Session::verify(&state.db, token_hash.as_ref())
55        .await?
56        .ok_or_else(|| Unauthorized("Invalid or expired session".into()))?;
57
58    state
59        .session_cache
60        .insert(token_hash, session.user_id)
61        .await;
62
63    request.extensions_mut().insert(session.user_id);
64
65    Ok(next.run(request).await)
66}
67
68#[cfg(test)]
69mod tests {
70    use crate::app::AppStateBuilder;
71    use axum::{Router, body::Body, http::Request, routing::get};
72    use blake3::hash;
73    use sqlx::sqlite::SqlitePoolOptions;
74    use tower::ServiceExt;
75
76    use super::*;
77
78    // This test checks, that in-memory cache is being used first, instead of
79    // querrying the database first.
80    #[tokio::test]
81    async fn auth_guard_bypasses_db_on_cache_hit() {
82        // Even though we simulate establishing a connection to db,
83        // but accessing this db will itself result in error.
84        // and Ofcourse, this error will be bypassed if AppState::session_cache
85        // returns the user_id, which is exactly what we want to know.
86        let pool = SqlitePoolOptions::new()
87            .connect("sqlite::memory:")
88            .await
89            .unwrap();
90
91        let state = AppStateBuilder::new().vault_path("/tmp").db(pool).build();
92
93        let token = "top_secret";
94        let token_hash = TokenHash::from(hash(token.as_bytes()).as_bytes());
95
96        state.session_cache.insert(token_hash, 4).await;
97
98        let app = Router::new()
99            .route("/", get(|| async { "Success!" }))
100            .route_layer(axum::middleware::from_fn_with_state(
101                state.clone(),
102                auth_guard,
103            ))
104            .with_state(state);
105
106        let req = Request::builder()
107            .header("Authorization", format!("Bearer {token}"))
108            .body(Body::empty())
109            .unwrap();
110
111        let response = app.oneshot(req).await.unwrap();
112
113        // Test: We bypassed db querry for token check?
114        assert_eq!(response.status(), 200);
115    }
116}