Make the compat_router also recover from errors in human-facing routes

Author: sandhose

crates/cli/src/server.rs

@@ -269,7 +269,7 @@ pub fn build_router(
             }
             mas_config::HttpResource::OAuth => router.merge(mas_handlers::api_router::<AppState>()),
             mas_config::HttpResource::Compat => {
-                router.merge(mas_handlers::compat_router::<AppState>())
+                router.merge(mas_handlers::compat_router::<AppState>(templates.clone()))
             }
             mas_config::HttpResource::AdminApi => {
                 let (_, api_router) = mas_handlers::admin_api_router::<AppState>();

crates/handlers/src/lib.rs

@@ -257,7 +257,7 @@ where
 }
 
 #[allow(clippy::trait_duplication_in_bounds)]
-pub fn compat_router<S>() -> Router<S>
+pub fn compat_router<S>(templates: Templates) -> Router<S>
 where
     S: Clone + Send + Sync + 'static,
     UrlBuilder: FromRef<S>,
@@ -272,7 +272,28 @@ where
     BoxClock: FromRequestParts<S>,
     BoxRng: FromRequestParts<S>,
 {
-    Router::new()
+    // A sub-router for human-facing routes with error handling
+    let human_router = Router::new()
+        .route(
+            mas_router::CompatLoginSsoRedirect::route(),
+            get(self::compat::login_sso_redirect::get),
+        )
+        .route(
+            mas_router::CompatLoginSsoRedirectIdp::route(),
+            get(self::compat::login_sso_redirect::get),
+        )
+        .route(
+            mas_router::CompatLoginSsoRedirectSlash::route(),
+            get(self::compat::login_sso_redirect::get),
+        )
+        .layer(AndThenLayer::new(
+            async move |response: axum::response::Response| {
+                Ok::<_, Infallible>(recover_error(&templates, response))
+            },
+        ));
+
+    // A sub-router for API-facing routes with CORS
+    let api_router = Router::new()
         .route(
             mas_router::CompatLogin::route(),
             get(self::compat::login::get).post(self::compat::login::post),
@@ -289,18 +310,6 @@ where
             mas_router::CompatRefresh::route(),
             post(self::compat::refresh::post),
         )
-        .route(
-            mas_router::CompatLoginSsoRedirect::route(),
-            get(self::compat::login_sso_redirect::get),
-        )
-        .route(
-            mas_router::CompatLoginSsoRedirectIdp::route(),
-            get(self::compat::login_sso_redirect::get),
-        )
-        .route(
-            mas_router::CompatLoginSsoRedirectSlash::route(),
-            get(self::compat::login_sso_redirect::get),
-        )
         .layer(
             CorsLayer::new()
                 .allow_origin(Any)
@@ -314,7 +323,9 @@ where
                     HeaderName::from_static("x-requested-with"),
                 ])
                 .max_age(Duration::from_secs(60 * 60)),
-        )
+        );
+
+    Router::new().merge(human_router).merge(api_router)
 }
 
 #[allow(clippy::too_many_lines)]
@@ -454,22 +465,29 @@ where
         )
         .layer(AndThenLayer::new(
             async move |response: axum::response::Response| {
-                // Error responses should have an ErrorContext attached to them
-                let ext = response.extensions().get::<ErrorContext>();
-                if let Some(ctx) = ext {
-                    if let Ok(res) = templates.render_error(ctx) {
-                        let (mut parts, _original_body) = response.into_parts();
-                        parts.headers.remove(CONTENT_TYPE);
-                        parts.headers.remove(CONTENT_LENGTH);
-                        return Ok((parts, Html(res)).into_response());
-                    }
-                }
-
-                Ok::<_, Infallible>(response)
+                Ok::<_, Infallible>(recover_error(&templates, response))
             },
         ))
 }
 
+fn recover_error(
+    templates: &Templates,
+    response: axum::response::Response,
+) -> axum::response::Response {
+    // Error responses should have an ErrorContext attached to them
+    let ext = response.extensions().get::<ErrorContext>();
+    if let Some(ctx) = ext {
+        if let Ok(res) = templates.render_error(ctx) {
+            let (mut parts, _original_body) = response.into_parts();
+            parts.headers.remove(CONTENT_TYPE);
+            parts.headers.remove(CONTENT_LENGTH);
+            return (parts, Html(res)).into_response();
+        }
+    }
+
+    response
+}
+
 /// The fallback handler for all routes that don't match anything else.
 ///
 /// # Errors

crates/handlers/src/test_utils.rs

@@ -319,7 +319,7 @@ impl TestState {
         let app = crate::healthcheck_router()
             .merge(crate::discovery_router())
             .merge(crate::api_router())
-            .merge(crate::compat_router())
+            .merge(crate::compat_router(self.templates.clone()))
             .merge(crate::human_router(self.templates.clone()))
             // We enable undocumented_oauth2_access for the tests, as it is easier to query the API
             // with it