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())
-        .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![],
-        })
-    })
+        .layer(middleware::from_fn_with_state(app_state, auth_middleware))
 }
 
 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")) }