Refactor OIDC state management to use separate nonce and redirect cookies

Co-authored-by: contact <[email protected]>

Author: cursoragent

src/webserver/oidc.rs

@@ -210,11 +210,11 @@ impl OidcState {
     async fn get_token_claims(
         &self,
         id_token: OidcToken,
-        state: Option<&OidcLoginState>,
+        expected_nonce: Option<&Nonce>,
     ) -> anyhow::Result<OidcClaims> {
         let client = &self.get_client().await;
         let verifier = self.config.create_id_token_verifier(client);
-        let nonce_verifier = |nonce: Option<&Nonce>| check_nonce(nonce, state);
+        let nonce_verifier = |nonce: Option<&Nonce>| check_nonce(nonce, expected_nonce);
         let claims: OidcClaims = id_token
             .into_claims(&verifier, nonce_verifier)
             .map_err(|e| anyhow::anyhow!("Could not verify the ID token: {}", e))?;
@@ -379,11 +379,8 @@ async fn handle_oidc_callback(oidc_state: &OidcState, request: ServiceRequest) -
             // Try to get redirect URL from query params if available
             let redirect_url =
                 if let Ok(params) = Query::<OidcCallbackParams>::from_query(query_string) {
-                    if let Ok(state) = get_state_from_cookies(&request, &params.state) {
-                        state.initial_url
-                    } else {
-                        "/".to_string()
-                    }
+                    get_redirect_url_from_cookie(&request, &params.state)
+                        .unwrap_or_else(|_| "/".to_string())
                 } else {
                     "/".to_string()
                 };
@@ -398,16 +395,13 @@ async fn handle_oidc_callback(oidc_state: &OidcState, request: ServiceRequest) -
 /// When an user has already authenticated (potentially in another tab), we ignore the callback and redirect to the initial URL.
 fn handle_authenticated_oidc_callback(request: ServiceRequest) -> ServiceResponse {
     // Try to get redirect URL from query params if available
-    let redirect_url =
-        if let Ok(params) = Query::<OidcCallbackParams>::from_query(request.query_string()) {
-            if let Ok(state) = get_state_from_cookies(&request, &params.state) {
-                state.initial_url
-            } else {
-                "/".to_string()
-            }
-        } else {
-            "/".to_string()
-        };
+    let redirect_url = if let Ok(params) =
+        Query::<OidcCallbackParams>::from_query(request.query_string())
+    {
+        get_redirect_url_from_cookie(&request, &params.state).unwrap_or_else(|_| "/".to_string())
+    } else {
+        "/".to_string()
+    };
     log::debug!("OIDC callback received for authenticated user. Redirecting to {redirect_url}");
     request.into_response(build_redirect_response(redirect_url))
 }
@@ -450,34 +444,23 @@ async fn process_oidc_callback(
         })?
         .into_inner();
 
-    let state = get_state_from_cookies(request, &params.state)
-        .context("Failed to read oidc state cookies")?;
+    let redirect_url = get_redirect_url_from_cookie(request, &params.state)
+        .context("Failed to read redirect URL from cookie")?;
 
     let client = oidc_state.get_client().await;
     log::debug!("Processing OIDC callback with params: {params:?}. Requesting token...");
-    let token_response = exchange_code_for_token(&client, http_client, params).await?;
+    let token_response = exchange_code_for_token(&client, http_client, params.clone()).await?;
     log::debug!("Received token response: {token_response:?}");
 
-    let redirect_target = validate_redirect_url(state.initial_url);
+    let redirect_target = validate_redirect_url(redirect_url);
     log::info!("Redirecting to {redirect_target} after a successful login");
     let mut response = build_redirect_response(redirect_target);
     set_auth_cookie(&mut response, &token_response, oidc_state)
         .await
         .context("Failed to set auth cookie")?;
 
     // Clean up the state-specific cookie after successful authentication
-    let state_cookie_name = format!(
-        "{}{}",
-        SQLPAGE_STATE_COOKIE_PREFIX,
-        state.csrf_token.secret()
-    );
-    let cleanup_cookie = Cookie::build(state_cookie_name, "")
-        .secure(true)
-        .http_only(true)
-        .same_site(actix_web::cookie::SameSite::Lax)
-        .path("/")
-        .max_age(actix_web::cookie::time::Duration::seconds(0)) // Expire immediately
-        .finish();
+    let cleanup_cookie = create_cleanup_cookie(&params.state);
     response.add_cookie(&cleanup_cookie).unwrap();
 
     Ok(response)
@@ -539,8 +522,8 @@ async fn build_auth_provider_redirect_response(
     initial_url: String,
 ) -> HttpResponse {
     let AuthUrl { url, params } = build_auth_url(oidc_state).await;
-    let state = OidcLoginState::new(initial_url, params);
-    let (nonce_cookie, redirect_cookie) = create_state_cookies(&state);
+    let nonce_cookie = create_nonce_cookie(&params.nonce);
+    let redirect_cookie = create_redirect_cookie(&params.csrf_token, &initial_url);
     HttpResponse::TemporaryRedirect()
         .append_header(("Location", url.to_string()))
         .cookie(nonce_cookie)
@@ -566,10 +549,10 @@ async fn get_authenticated_user_info(
     let id_token = OidcToken::from_str(&cookie_value)
         .with_context(|| format!("Invalid SQLPage auth cookie: {cookie_value:?}"))?;
 
-    // Try to get state from cookies if this is a callback request
-    let state = if request.path() == SQLPAGE_REDIRECT_URI {
-        if let Ok(params) = Query::<OidcCallbackParams>::from_query(request.query_string()) {
-            get_state_from_cookies(request, &params.state).ok()
+    // Try to get nonce from cookies if this is a callback request
+    let nonce = if request.path() == SQLPAGE_REDIRECT_URI {
+        if let Ok(_params) = Query::<OidcCallbackParams>::from_query(request.query_string()) {
+            get_nonce_from_cookie(request).ok()
         } else {
             None
         }
@@ -579,7 +562,7 @@ async fn get_authenticated_user_info(
 
     log::debug!("Verifying id token: {id_token:?}");
     let claims = oidc_state
-        .get_token_claims(id_token, state.as_ref())
+        .get_token_claims(id_token, nonce.as_ref())
         .await?;
     log::debug!("The current user is: {claims:?}");
     Ok(Some(claims))
@@ -731,7 +714,7 @@ fn make_oidc_client(
     Ok(client)
 }
 
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Deserialize, Clone)]
 struct OidcCallbackParams {
     code: String,
     state: CsrfToken,
@@ -783,13 +766,13 @@ fn hash_nonce(nonce: &Nonce) -> String {
 
 fn check_nonce(
     id_token_nonce: Option<&Nonce>,
-    login_state: Option<&OidcLoginState>,
+    expected_nonce: Option<&Nonce>,
 ) -> Result<(), String> {
-    let Some(state) = login_state else {
-        return Ok(()); // No login state, no nonce to check
+    let Some(expected) = expected_nonce else {
+        return Ok(()); // No expected nonce, no validation needed
     };
     match id_token_nonce {
-        Some(id_token_nonce) => nonce_matches(id_token_nonce, &state.nonce),
+        Some(id_token_nonce) => nonce_matches(id_token_nonce, expected),
         None => Err("No nonce found in the ID token".to_string()),
     }
 }
@@ -816,107 +799,78 @@ fn nonce_matches(id_token_nonce: &Nonce, state_nonce: &Nonce) -> Result<(), Stri
     Ok(())
 }
 
-#[derive(Debug, Serialize, Deserialize)]
-struct OidcLoginState {
-    /// The URL to redirect to after the login process is complete.
-    #[serde(rename = "u")]
-    initial_url: String,
-    /// The CSRF token to use for the login process.
-    #[serde(rename = "c")]
-    csrf_token: CsrfToken,
-    /// The source nonce to use for the login process. It must be checked against the hash
-    /// stored in the ID token.
-    #[serde(rename = "n")]
-    nonce: Nonce,
-}
-
 #[derive(Debug, Serialize, Deserialize)]
 struct OidcNonceState {
-    /// The source nonce to use for the login process. It must be checked against the hash
-    /// stored in the ID token.
     #[serde(rename = "n")]
     nonce: Nonce,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 struct OidcRedirectState {
-    /// The URL to redirect to after the login process is complete.
     #[serde(rename = "u")]
     initial_url: String,
 }
 
-impl OidcLoginState {
-    fn new(initial_url: String, auth_url: AuthUrlParams) -> Self {
-        Self {
-            initial_url,
-            csrf_token: auth_url.csrf_token,
-            nonce: auth_url.nonce,
-        }
-    }
-}
-
-fn create_state_cookies(login_state: &OidcLoginState) -> (Cookie<'_>, Cookie<'_>) {
-    // Create nonce cookie (longer-lived for security)
+fn create_nonce_cookie(nonce: &Nonce) -> Cookie<'_> {
     let nonce_state = OidcNonceState {
-        nonce: login_state.nonce.clone(),
+        nonce: nonce.clone(),
     };
     let nonce_json = serde_json::to_string(&nonce_state).unwrap();
-    let nonce_cookie = Cookie::build(SQLPAGE_NONCE_COOKIE_NAME, nonce_json)
+    Cookie::build(SQLPAGE_NONCE_COOKIE_NAME, nonce_json)
         .secure(true)
         .http_only(true)
         .same_site(actix_web::cookie::SameSite::Lax)
         .path("/")
-        .max_age(actix_web::cookie::time::Duration::minutes(10)) // 10 minutes should be enough for login flow
-        .finish();
+        .max_age(actix_web::cookie::time::Duration::minutes(10))
+        .finish()
+}
 
-    // Create state-specific redirect cookie (short-lived)
+fn create_redirect_cookie(csrf_token: &CsrfToken, initial_url: &str) -> Cookie<'static> {
     let redirect_state = OidcRedirectState {
-        initial_url: login_state.initial_url.clone(),
+        initial_url: initial_url.to_string(),
     };
     let redirect_json = serde_json::to_string(&redirect_state).unwrap();
-    let state_cookie_name = format!(
-        "{}{}",
-        SQLPAGE_STATE_COOKIE_PREFIX,
-        login_state.csrf_token.secret()
-    );
-    let redirect_cookie = Cookie::build(state_cookie_name, redirect_json)
+    let cookie_name = format!("{}{}", SQLPAGE_STATE_COOKIE_PREFIX, csrf_token.secret());
+    Cookie::build(cookie_name, redirect_json)
         .secure(true)
         .http_only(true)
         .same_site(actix_web::cookie::SameSite::Lax)
         .path("/")
-        .max_age(actix_web::cookie::time::Duration::minutes(5)) // Short-lived, only needed during login flow
-        .finish();
+        .max_age(actix_web::cookie::time::Duration::minutes(5))
+        .finish()
+}
 
-    (nonce_cookie, redirect_cookie)
+fn create_cleanup_cookie(csrf_token: &CsrfToken) -> Cookie<'_> {
+    let cookie_name = format!("{}{}", SQLPAGE_STATE_COOKIE_PREFIX, csrf_token.secret());
+    Cookie::build(cookie_name, "")
+        .secure(true)
+        .http_only(true)
+        .same_site(actix_web::cookie::SameSite::Lax)
+        .path("/")
+        .max_age(actix_web::cookie::time::Duration::seconds(0))
+        .finish()
+}
+
+fn get_nonce_from_cookie(request: &ServiceRequest) -> anyhow::Result<Nonce> {
+    let cookie = request
+        .cookie(SQLPAGE_NONCE_COOKIE_NAME)
+        .with_context(|| format!("No {SQLPAGE_NONCE_COOKIE_NAME} cookie found"))?;
+    let nonce_state: OidcNonceState = serde_json::from_str(cookie.value())
+        .with_context(|| format!("Failed to parse nonce from cookie: {cookie}"))?;
+    Ok(nonce_state.nonce)
 }
 
-fn get_state_from_cookies(
+fn get_redirect_url_from_cookie(
     request: &ServiceRequest,
     csrf_token: &CsrfToken,
-) -> anyhow::Result<OidcLoginState> {
-    // Get nonce from the nonce cookie
-    let nonce_cookie = request.cookie(SQLPAGE_NONCE_COOKIE_NAME).with_context(|| {
-        format!("No {SQLPAGE_NONCE_COOKIE_NAME} cookie found for {SQLPAGE_REDIRECT_URI}")
-    })?;
-    let nonce_state: OidcNonceState = serde_json::from_str(nonce_cookie.value())
-        .with_context(|| format!("Failed to parse OIDC nonce from cookie: {nonce_cookie}"))?;
-
-    // Get redirect URL from the state-specific cookie
-    let state_cookie_name = format!("{}{}", SQLPAGE_STATE_COOKIE_PREFIX, csrf_token.secret());
-    let redirect_cookie = request.cookie(&state_cookie_name).with_context(|| {
-        format!("No {state_cookie_name} cookie found for {SQLPAGE_REDIRECT_URI}")
-    })?;
-    let redirect_state: OidcRedirectState = serde_json::from_str(redirect_cookie.value())
-        .with_context(|| {
-            format!("Failed to parse OIDC redirect state from cookie: {redirect_cookie}")
-        })?;
-
-    // Reconstruct the full login state
-    Ok(OidcLoginState {
-        initial_url: redirect_state.initial_url,
-        csrf_token: csrf_token.clone(),
-        nonce: nonce_state.nonce,
-    })
+) -> anyhow::Result<String> {
+    let cookie_name = format!("{}{}", SQLPAGE_STATE_COOKIE_PREFIX, csrf_token.secret());
+    let cookie = request
+        .cookie(&cookie_name)
+        .with_context(|| format!("No {cookie_name} cookie found"))?;
+    let redirect_state: OidcRedirectState = serde_json::from_str(cookie.value())
+        .with_context(|| format!("Failed to parse redirect state from cookie: {cookie}"))?;
+    Ok(redirect_state.initial_url)
 }
 
 /// Given an audience, verify if it is trusted. The `client_id` is always trusted, independently of this function.