simplify OIDC infinite redirect prevention logic

This update introduces a new function, `handle_oidc_callback_error`, to streamline error handling during OIDC callback processing. It enhances the management of redirect counts and separates the logic for handling maximum redirect limits into `handle_max_redirect_count_reached`. Additionally, the `build_auth_provider_redirect_response` function is updated to accept the redirect count, ensuring accurate tracking of redirects. This refactor aims to prevent infinite redirect loops and improve code clarity.

Author: lovasoa

src/webserver/oidc.rs

@@ -403,7 +403,9 @@ async fn handle_unauthenticated_request(
     log::debug!("Redirecting to OIDC provider");
 
     let initial_url = request.uri().to_string();
-    let response = build_auth_provider_redirect_response(oidc_state, &initial_url).await;
+    let redirect_count = get_redirect_count(&request);
+    let response =
+        build_auth_provider_redirect_response(oidc_state, &initial_url, redirect_count).await;
     MiddlewareResponse::Respond(request.into_response(response))
 }
 
@@ -413,25 +415,40 @@ async fn handle_oidc_callback(oidc_state: &OidcState, request: ServiceRequest) -
             clear_redirect_count_cookie(&mut response);
             request.into_response(response)
         }
-        Err(e) => {
-            let redirect_count = get_redirect_count(&request);
-            if redirect_count >= MAX_OIDC_REDIRECTS {
-                log::error!(
-                    "Failed to process OIDC callback after {redirect_count} attempts. \
-                     Stopping to avoid infinite redirections: {e:#}"
-                );
-                let resp = build_oidc_error_response(&request, &e);
-                return request.into_response(resp);
-            }
-            log::error!("Failed to process OIDC callback (attempt {redirect_count}). Refreshing oidc provider metadata, then redirecting to home page: {e:#}");
-            oidc_state.refresh_on_error(&request).await;
-            let mut resp = build_auth_provider_redirect_response(oidc_state, "/").await;
-            set_redirect_count_cookie(&mut resp, redirect_count + 1);
-            request.into_response(resp)
-        }
+        Err(e) => handle_oidc_callback_error(oidc_state, request, e).await,
     }
 }
 
+async fn handle_oidc_callback_error(
+    oidc_state: &OidcState,
+    request: ServiceRequest,
+    e: anyhow::Error,
+) -> ServiceResponse {
+    let redirect_count = get_redirect_count(&request);
+    if redirect_count >= MAX_OIDC_REDIRECTS {
+        return handle_max_redirect_count_reached(request, &e, redirect_count);
+    }
+    log::error!(
+        "Failed to process OIDC callback (attempt {redirect_count}). Refreshing oidc provider metadata, then redirecting to home page: {e:#}"
+    );
+    oidc_state.refresh_on_error(&request).await;
+    let resp = build_auth_provider_redirect_response(oidc_state, "/", redirect_count).await;
+    request.into_response(resp)
+}
+
+fn handle_max_redirect_count_reached(
+    request: ServiceRequest,
+    e: &anyhow::Error,
+    redirect_count: u8,
+) -> ServiceResponse {
+    log::error!(
+        "Failed to process OIDC callback after {redirect_count} attempts. \
+         Stopping to avoid infinite redirections: {e:#}"
+    );
+    let resp = build_oidc_error_response(&request, e);
+    request.into_response(resp)
+}
+
 async fn handle_oidc_logout(oidc_state: &OidcState, request: ServiceRequest) -> ServiceResponse {
     match process_oidc_logout(oidc_state, &request).await {
         Ok(response) => request.into_response(response),
@@ -684,12 +701,23 @@ fn set_auth_cookie(response: &mut HttpResponse, id_token: &OidcToken) {
 async fn build_auth_provider_redirect_response(
     oidc_state: &OidcState,
     initial_url: &str,
+    redirect_count: u8,
 ) -> HttpResponse {
     let AuthUrl { url, params } = build_auth_url(oidc_state).await;
     let tmp_login_flow_state_cookie = create_tmp_login_flow_state_cookie(&params, initial_url);
+    let redirect_count_cookie = Cookie::build(
+        SQLPAGE_OIDC_REDIRECT_COUNT_COOKIE,
+        (redirect_count + 1).to_string(),
+    )
+    .path("/")
+    .http_only(true)
+    .same_site(actix_web::cookie::SameSite::Lax)
+    .max_age(LOGIN_FLOW_STATE_COOKIE_EXPIRATION)
+    .finish();
     HttpResponse::SeeOther()
         .append_header((header::LOCATION, url.to_string()))
         .cookie(tmp_login_flow_state_cookie)
+        .cookie(redirect_count_cookie)
         .body("Redirecting...")
 }
 
@@ -706,16 +734,6 @@ fn get_redirect_count(request: &ServiceRequest) -> u8 {
         .unwrap_or(0)
 }
 
-fn set_redirect_count_cookie(response: &mut HttpResponse, count: u8) {
-    let cookie = Cookie::build(SQLPAGE_OIDC_REDIRECT_COUNT_COOKIE, count.to_string())
-        .path("/")
-        .http_only(true)
-        .same_site(actix_web::cookie::SameSite::Lax)
-        .max_age(LOGIN_FLOW_STATE_COOKIE_EXPIRATION)
-        .finish();
-    response.add_cookie(&cookie).ok();
-}
-
 fn clear_redirect_count_cookie(response: &mut HttpResponse) {
     let cookie = Cookie::build(SQLPAGE_OIDC_REDIRECT_COUNT_COOKIE, "")
         .path("/")