Use middleware instead of manual auth checks

backend/sync_server/src/server.rs

@@ -11,7 +11,7 @@ mod responses;
 mod update_document;
 mod websocket;
 
-use std::{ffi::OsString, sync::Arc};
+use std::{ffi::OsString, sync::Arc, time::Duration};
 
 use aide::{
     axum::{
@@ -23,10 +23,12 @@ use aide::{
     transform::TransformOpenApi,
 };
 use anyhow::{Context as _, Result, anyhow};
+use auth::auth_middleware;
 use axum::{
     Extension, Json,
     extract::{DefaultBodyLimit, Request},
     http::{self, HeaderValue, Method},
+    middleware,
     response::IntoResponse,
     routing::IntoMakeService,
 };
@@ -36,6 +38,7 @@ use tower_http::{
     LatencyUnit,
     cors::CorsLayer,
     limit::RequestBodyLimitLayer,
+    timeout::TimeoutLayer,
     trace::{
         DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, DefaultOnRequest, DefaultOnResponse,
         TraceLayer,
@@ -46,7 +49,7 @@ use tracing::{Level, info_span};
 use crate::{
     app_state::AppState,
     config::server_config::ServerConfig,
-    errors::{SerializedError, not_found_error},
+    errors::{SerializedError, client_error, not_found_error},
 };
 
 pub async fn create_server(config_path: Option<OsString>) -> Result<()> {
@@ -61,12 +64,85 @@ pub async fn create_server(config_path: Option<OsString>) -> Result<()> {
 
     let mut api = create_open_api();
     let app = ApiRouter::new()
+        .nest("/", get_authed_routes(app_state.clone()))
         .api_route("/vaults/:vault_id/ping", get(ping::ping))
+        .route("/vaults/:vault_id/ws", get(websocket::websocket_handler))
+        .route("/", Scalar::new("/api.json").axum_route())
+        .route("/api.json", axum::routing::get(serve_api))
+        .layer(DefaultBodyLimit::disable())
+        .layer(RequestBodyLimitLayer::new(
+            app_state.config.server.max_body_size_mb * 1024 * 1024,
+        ))
+        .layer(TimeoutLayer::new(Duration::from_secs(
+            server_config.response_timeout_seconds,
+        )))
+        .layer(
+            CorsLayer::new()
+                .allow_origin("*".parse::<HeaderValue>().expect("Failed to parse origin"))
+                .allow_headers([http::header::CONTENT_TYPE, http::header::AUTHORIZATION])
+                .allow_methods([Method::GET, Method::POST, Method::PUT, Method::DELETE]),
+        )
+        .layer(
+            TraceLayer::new_for_http()
+                .make_span_with(|request: &Request<_>| {
+                    info_span!(
+                        "http_request",
+                        method = ?request.method(),
+                        uri = ?request.uri(),
+                    )
+                })
+                .on_request(DefaultOnRequest::new().level(Level::INFO))
+                .on_response(
+                    DefaultOnResponse::new()
+                        .level(Level::INFO)
+                        .latency_unit(LatencyUnit::Millis),
+                )
+                .on_body_chunk(DefaultOnBodyChunk::new())
+                .on_eos(DefaultOnEos::new())
+                .on_failure(DefaultOnFailure::new().level(Level::ERROR)),
+        )
+        .with_state(app_state)
+        .finish_api_with(&mut api, add_api_docs_error_example)
+        .layer(Extension(Arc::new(api))) // https://github.com/tamasfe/aide/blob/507f4a8822bc0c13cbda0f589da1e0f4cbcdb812/examples/example-axum/src/main.rs#L39
+        .fallback(handle_404)
+        .fallback(handle_405)
+        .into_make_service();
+
+    start_server(app, &server_config).await
+}
+
+async fn serve_api(Extension(api): Extension<Arc<OpenApi>>) -> impl IntoResponse { Json(api) }
+
+fn create_open_api() -> OpenApi {
+    OpenApi {
+        info: Info {
+            title: "VaultLink sync server".to_owned(),
+            summary: Some(
+                "Simple API for syncing documents between concurrent clients.".to_owned(),
+            ),
+            description: Some(include_str!("../README.md").to_owned()),
+            version: env!("CARGO_PKG_VERSION").to_owned(),
+            ..Info::default()
+        },
+        ..OpenApi::default()
+    }
+}
+
+fn add_api_docs_error_example(api: TransformOpenApi<'_>) -> TransformOpenApi<'_> {
+    api.default_response_with::<Json<SerializedError>, _>(|res| {
+        res.example(SerializedError {
+            message: "An error has occurred".to_owned(),
+            causes: vec![],
+        })
+    })
+}
+
+fn get_authed_routes(app_state: AppState) -> ApiRouter<AppState> {
+    ApiRouter::new()
         .api_route(
             "/vaults/:vault_id/documents",
             get(fetch_latest_documents::fetch_latest_documents),
         )
-        .route("/vaults/:vault_id/ws", get(websocket::websocket_handler))
         .api_route(
             "/vaults/:vault_id/documents",
             post(create_document::create_document_multipart),
@@ -99,70 +175,7 @@ pub async fn create_server(config_path: Option<OsString>) -> Result<()> {
             "/vaults/:vault_id/documents/:document_id",
             delete(delete_document::delete_document),
         )
-        .route("/", Scalar::new("/api.json").axum_route())
+        .layer(middleware::from_fn_with_state(app_state, auth_middleware))
-        .route("/api.json", axum::routing::get(serve_api))
-        .layer(
-            TraceLayer::new_for_http()
-                .make_span_with(|request: &Request<_>| {
-                    info_span!(
-                        "http_request",
-                        method = ?request.method(),
-                        uri = ?request.uri(),
-                    )
-                })
-                .on_request(DefaultOnRequest::new().level(Level::INFO))
-                .on_response(
-                    DefaultOnResponse::new()
-                        .level(Level::INFO)
-                        .latency_unit(LatencyUnit::Millis),
-                )
-                .on_body_chunk(DefaultOnBodyChunk::new())
-                .on_eos(DefaultOnEos::new())
-                .on_failure(DefaultOnFailure::new().level(Level::ERROR)),
-        )
-        .layer(DefaultBodyLimit::disable())
-        .layer(RequestBodyLimitLayer::new(
-            app_state.config.server.max_body_size_mb * 1024 * 1024,
-        ))
-        .layer(
-            CorsLayer::new()
-                .allow_origin("*".parse::<HeaderValue>().expect("Failed to parse origin"))
-                .allow_headers([http::header::CONTENT_TYPE, http::header::AUTHORIZATION])
-                .allow_methods([Method::GET, Method::POST, Method::PUT, Method::DELETE]),
-        )
-        .with_state(app_state)
-        .finish_api_with(&mut api, add_api_docs_error_example)
-        .layer(Extension(Arc::new(api))) // https://github.com/tamasfe/aide/blob/507f4a8822bc0c13cbda0f589da1e0f4cbcdb812/examples/example-axum/src/main.rs#L39
-        .fallback(handler_404)
-        .into_make_service();
-
-    start_server(app, &server_config).await
-}
-
-async fn serve_api(Extension(api): Extension<Arc<OpenApi>>) -> impl IntoResponse { Json(api) }
-
-fn create_open_api() -> OpenApi {
-    OpenApi {
-        info: Info {
-            title: "VaultLink sync server".to_owned(),
-            summary: Some(
-                "Simple API for syncing documents between concurrent clients.".to_owned(),
-            ),
-            description: Some(include_str!("../README.md").to_owned()),
-            version: env!("CARGO_PKG_VERSION").to_owned(),
-            ..Info::default()
-        },
-        ..OpenApi::default()
-    }
-}
-
-fn add_api_docs_error_example(api: TransformOpenApi<'_>) -> TransformOpenApi<'_> {
-    api.default_response_with::<Json<SerializedError>, _>(|res| {
-        res.example(SerializedError {
-            message: "An error has occurred".to_owned(),
-            causes: vec![],
-        })
-    })
 }
 
 async fn start_server(app: IntoMakeService<axum::Router>, config: &ServerConfig) -> Result<()> {
@@ -209,4 +222,6 @@ async fn shutdown_signal() {
     }
 }
 
-async fn handler_404() -> impl IntoResponse { not_found_error(anyhow!("Page not found")) }
+async fn handle_404() -> impl IntoResponse { not_found_error(anyhow!("Page not found")) }
+
+async fn handle_405() -> impl IntoResponse { client_error(anyhow!("Method not allowed")) }

backend/sync_server/src/server/create_document.rs

@@ -1,19 +1,12 @@
 use aide_axum_typed_multipart::TypedMultipart;
 use anyhow::Context as _;
 use axum::extract::{Path, State};
-use axum_extra::{
-    TypedHeader,
-    headers::{Authorization, authorization::Bearer},
-};
 use axum_jsonschema::Json;
 use schemars::JsonSchema;
 use serde::Deserialize;
 use sync_lib::base64_to_bytes;
 
-use super::{
+use super::requests::{CreateDocumentVersion, CreateDocumentVersionMultipart};
-    auth::auth,
-    requests::{CreateDocumentVersion, CreateDocumentVersionMultipart},
-};
 use crate::{
     app_state::{
         AppState,
@@ -36,7 +29,6 @@ pub struct CreateDocumentPathParams {
 /// with their content merged.
 #[axum::debug_handler]
 pub async fn create_document_multipart(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(CreateDocumentPathParams { vault_id }): Path<CreateDocumentPathParams>,
     State(state): State<AppState>,
     TypedMultipart(axum_typed_multipart::TypedMultipart(request)): TypedMultipart<
@@ -44,7 +36,6 @@ pub async fn create_document_multipart(
     >,
 ) -> Result<Json<DocumentVersionWithoutContent>, SyncServerError> {
     internal_create_document(
-        auth_header,
         state,
         vault_id,
         request.document_id,
@@ -59,7 +50,6 @@ pub async fn create_document_multipart(
 /// with their content merged.
 #[axum::debug_handler]
 pub async fn create_document_json(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(CreateDocumentPathParams { vault_id }): Path<CreateDocumentPathParams>,
     State(state): State<AppState>,
     Json(request): Json<CreateDocumentVersion>,
@@ -69,7 +59,6 @@ pub async fn create_document_json(
         .map_err(client_error)?;
 
     internal_create_document(
-        auth_header,
         state,
         vault_id,
         request.document_id,
@@ -80,15 +69,12 @@ pub async fn create_document_json(
 }
 
 async fn internal_create_document(
-    auth_header: Authorization<Bearer>,
     state: AppState,
     vault_id: VaultId,
     document_id: Option<DocumentId>,
     relative_path: String,
     content: Vec<u8>,
 ) -> Result<Json<DocumentVersionWithoutContent>, SyncServerError> {
-    auth(&state, auth_header.token(), &vault_id)?;
-
     let mut transaction = state
         .database
         .create_write_transaction(&vault_id)

backend/sync_server/src/server/delete_document.rs

@@ -1,14 +1,10 @@
 use anyhow::Context as _;
 use axum::extract::{Path, State};
-use axum_extra::{
-    TypedHeader,
-    headers::{Authorization, authorization::Bearer},
-};
 use axum_jsonschema::Json;
 use schemars::JsonSchema;
 use serde::Deserialize;
 
-use super::{auth::auth, requests::DeleteDocumentVersion};
+use super::requests::DeleteDocumentVersion;
 use crate::{
     app_state::{
         AppState,
@@ -29,7 +25,6 @@ pub struct DeleteDocumentPathParams {
 
 #[axum::debug_handler]
 pub async fn delete_document(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(DeleteDocumentPathParams {
         vault_id,
         document_id,
@@ -37,8 +32,6 @@ pub async fn delete_document(
     State(state): State<AppState>,
     Json(request): Json<DeleteDocumentVersion>,
 ) -> Result<Json<DocumentVersionWithoutContent>, SyncServerError> {
-    auth(&state, auth_header.token(), &vault_id)?;
-
     let mut transaction = state
         .database
         .create_write_transaction(&vault_id)

backend/sync_server/src/server/fetch_document_version.rs

@@ -1,14 +1,9 @@
 use anyhow::anyhow;
 use axum::extract::{Path, State};
-use axum_extra::{
-    TypedHeader,
-    headers::{Authorization, authorization::Bearer},
-};
 use axum_jsonschema::Json;
 use schemars::JsonSchema;
 use serde::Deserialize;
 
-use super::auth::auth;
 use crate::{
     app_state::{
         AppState,
@@ -27,7 +22,6 @@ pub struct FetchDocumentVersionPathParams {
 
 #[axum::debug_handler]
 pub async fn fetch_document_version(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(FetchDocumentVersionPathParams {
         vault_id,
         document_id,
@@ -35,8 +29,6 @@ pub async fn fetch_document_version(
     }): Path<FetchDocumentVersionPathParams>,
     State(state): State<AppState>,
 ) -> Result<Json<DocumentVersion>, SyncServerError> {
-    auth(&state, auth_header.token(), &vault_id)?;
-
     let result = state
         .database
         .get_document_version(&vault_id, vault_update_id, None)

backend/sync_server/src/server/fetch_document_version_content.rs

@@ -3,14 +3,9 @@ use axum::{
     body::Bytes,
     extract::{Path, State},
 };
-use axum_extra::{
-    TypedHeader,
-    headers::{Authorization, authorization::Bearer},
-};
 use schemars::JsonSchema;
 use serde::Deserialize;
 
-use super::auth::auth;
 use crate::{
     app_state::{
         AppState,
@@ -29,7 +24,6 @@ pub struct FetchDocumentVersionContentPathParams {
 
 #[axum::debug_handler]
 pub async fn fetch_document_version_content(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(FetchDocumentVersionContentPathParams {
         vault_id,
         document_id,
@@ -37,8 +31,6 @@ pub async fn fetch_document_version_content(
     }): Path<FetchDocumentVersionContentPathParams>,
     State(state): State<AppState>,
 ) -> Result<Bytes, SyncServerError> {
-    auth(&state, auth_header.token(), &vault_id)?;
-
     let result = state
         .database
         .get_document_version(&vault_id, vault_update_id, None)

backend/sync_server/src/server/fetch_latest_document_version.rs

@@ -1,14 +1,9 @@
 use anyhow::anyhow;
 use axum::extract::{Path, State};
-use axum_extra::{
-    TypedHeader,
-    headers::{Authorization, authorization::Bearer},
-};
 use axum_jsonschema::Json;
 use schemars::JsonSchema;
 use serde::Deserialize;
 
-use super::auth::auth;
 use crate::{
     app_state::{
         AppState,
@@ -26,15 +21,12 @@ pub struct FetchLatestDocumentVersionPathParams {
 
 #[axum::debug_handler]
 pub async fn fetch_latest_document_version(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(FetchLatestDocumentVersionPathParams {
         vault_id,
         document_id,
     }): Path<FetchLatestDocumentVersionPathParams>,
     State(state): State<AppState>,
 ) -> Result<Json<DocumentVersion>, SyncServerError> {
-    auth(&state, auth_header.token(), &vault_id)?;
-
     let latest_version = state
         .database
         .get_latest_document(&vault_id, &document_id, None)

backend/sync_server/src/server/fetch_latest_documents.rs

@@ -1,13 +1,9 @@
 use axum::extract::{Path, Query, State};
-use axum_extra::{
-    TypedHeader,
-    headers::{Authorization, authorization::Bearer},
-};
 use axum_jsonschema::Json;
 use schemars::JsonSchema;
 use serde::Deserialize;
 
-use super::{auth::auth, responses::FetchLatestDocumentsResponse};
+use super::responses::FetchLatestDocumentsResponse;
 use crate::{
     app_state::{
         AppState,
@@ -30,13 +26,10 @@ pub struct QueryParams {
 
 #[axum::debug_handler]
 pub async fn fetch_latest_documents(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(FetchLatestDocumentsPathParams { vault_id }): Path<FetchLatestDocumentsPathParams>,
     Query(QueryParams { since_update_id }): Query<QueryParams>,
     State(state): State<AppState>,
 ) -> Result<Json<FetchLatestDocumentsResponse>, SyncServerError> {
-    auth(&state, auth_header.token(), &vault_id)?;
-
     let documents = if let Some(since_update_id) = since_update_id {
         state
             .database

backend/sync_server/src/server/update_document.rs

@@ -1,10 +1,6 @@
 use aide_axum_typed_multipart::TypedMultipart;
 use anyhow::{Context as _, anyhow};
 use axum::extract::{Path, State};
-use axum_extra::{
-    TypedHeader,
-    headers::{Authorization, authorization::Bearer},
-};
 use axum_jsonschema::Json;
 use log::info;
 use schemars::JsonSchema;
@@ -12,7 +8,6 @@ use serde::Deserialize;
 use sync_lib::{base64_to_bytes, is_file_type_mergable, merge};
 
 use super::{
-    auth::auth,
     requests::{UpdateDocumentVersion, UpdateDocumentVersionMultipart},
     responses::DocumentUpdateResponse,
 };
@@ -34,7 +29,6 @@ pub struct UpdateDocumentPathParams {
 
 #[axum::debug_handler]
 pub async fn update_document_multipart(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(UpdateDocumentPathParams {
         vault_id,
         document_id,
@@ -45,7 +39,6 @@ pub async fn update_document_multipart(
     >,
 ) -> Result<Json<DocumentUpdateResponse>, SyncServerError> {
     internal_update_document(
-        auth_header,
         state,
         vault_id,
         document_id,
@@ -58,7 +51,6 @@ pub async fn update_document_multipart(
 
 #[axum::debug_handler]
 pub async fn update_document_json(
-    TypedHeader(auth_header): TypedHeader<Authorization<Bearer>>,
     Path(UpdateDocumentPathParams {
         vault_id,
         document_id,
@@ -71,7 +63,6 @@ pub async fn update_document_json(
         .map_err(client_error)?;
 
     internal_update_document(
-        auth_header,
         state,
         vault_id,
         document_id,
@@ -84,7 +75,6 @@ pub async fn update_document_json(
 
 #[allow(clippy::too_many_arguments, clippy::too_many_lines)]
 async fn internal_update_document(
-    auth_header: Authorization<Bearer>,
     state: AppState,
     vault_id: VaultId,
     document_id: DocumentId,
@@ -92,8 +82,6 @@ async fn internal_update_document(
     relative_path: String,
     content: Vec<u8>,
 ) -> Result<Json<DocumentUpdateResponse>, SyncServerError> {
-    auth(&state, auth_header.token(), &vault_id)?;
-
     // No need for a transaction as document versions are immutable
     let parent_document = state
         .database