Improve sso auth flow (#6205)

Co-authored-by: Timshel <timshel@users.noreply.github.com>

.env.template

@@ -183,9 +183,9 @@
 ## Defaults to every minute. Set blank to disable this job.
 # DUO_CONTEXT_PURGE_SCHEDULE="30 * * * * *"
 #
-## Cron schedule of the job that cleans sso nonce from incomplete flow
+## Cron schedule of the job that cleans sso auth from incomplete flow
 ## Defaults to daily (20 minutes after midnight). Set blank to disable this job.
-# PURGE_INCOMPLETE_SSO_NONCE="0 20 0 * * *"
+# PURGE_INCOMPLETE_SSO_AUTH="0 20 0 * * *"
 
 ########################
 ### General settings ###

migrations/mysql/2025-08-20-120000_sso_nonce_to_auth/down.sql

@@ -0,0 +1,9 @@
+DROP TABLE IF EXISTS sso_auth;
+
+CREATE TABLE sso_nonce (
+    state               VARCHAR(512) NOT NULL PRIMARY KEY,
+    nonce               TEXT NOT NULL,
+    verifier            TEXT,
+    redirect_uri        TEXT NOT NULL,
+    created_at          TIMESTAMP NOT NULL DEFAULT now()
+);

migrations/mysql/2025-08-20-120000_sso_nonce_to_auth/up.sql

@@ -0,0 +1,12 @@
+DROP TABLE IF EXISTS sso_nonce;
+
+CREATE TABLE sso_auth (
+    state               VARCHAR(512) NOT NULL PRIMARY KEY,
+    client_challenge    TEXT NOT NULL,
+    nonce               TEXT NOT NULL,
+    redirect_uri        TEXT NOT NULL,
+    code_response       TEXT,
+    auth_response       TEXT,
+    created_at          TIMESTAMP NOT NULL DEFAULT now(),
+    updated_at          TIMESTAMP NOT NULL DEFAULT now()
+);

migrations/postgresql/2025-08-20-120000_sso_nonce_to_auth/down.sql

@@ -0,0 +1,9 @@
+DROP TABLE IF EXISTS sso_auth;
+
+CREATE TABLE sso_nonce (
+    state               TEXT NOT NULL PRIMARY KEY,
+    nonce               TEXT NOT NULL,
+    verifier            TEXT,
+    redirect_uri        TEXT NOT NULL,
+    created_at          TIMESTAMP NOT NULL DEFAULT now()
+);

migrations/postgresql/2025-08-20-120000_sso_nonce_to_auth/up.sql

@@ -0,0 +1,12 @@
+DROP TABLE IF EXISTS sso_nonce;
+
+CREATE TABLE sso_auth (
+    state               TEXT NOT NULL PRIMARY KEY,
+    client_challenge    TEXT NOT NULL,
+    nonce               TEXT NOT NULL,
+    redirect_uri        TEXT NOT NULL,
+    code_response       TEXT,
+    auth_response       TEXT,
+    created_at          TIMESTAMP NOT NULL DEFAULT now(),
+    updated_at          TIMESTAMP NOT NULL DEFAULT now()
+);

migrations/sqlite/2025-08-20-120000_sso_nonce_to_auth/down.sql

@@ -0,0 +1,9 @@
+DROP TABLE IF EXISTS sso_auth;
+
+CREATE TABLE sso_nonce (
+  state               TEXT NOT NULL PRIMARY KEY,
+  nonce               TEXT NOT NULL,
+  verifier            TEXT,
+  redirect_uri        TEXT NOT NULL,
+  created_at          DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+);

migrations/sqlite/2025-08-20-120000_sso_nonce_to_auth/up.sql

@@ -0,0 +1,12 @@
+DROP TABLE IF EXISTS sso_nonce;
+
+CREATE TABLE sso_auth (
+    state               TEXT NOT NULL PRIMARY KEY,
+    client_challenge    TEXT NOT NULL,
+    nonce               TEXT NOT NULL,
+    redirect_uri        TEXT NOT NULL,
+    code_response       TEXT,
+    auth_response       TEXT,
+    created_at          DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    updated_at          DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+);

src/api/identity.rs

@@ -24,14 +24,14 @@ use crate::{
     auth::{generate_organization_api_key_login_claims, AuthMethod, ClientHeaders, ClientIp, ClientVersion},
     db::{
         models::{
-            AuthRequest, AuthRequestId, Device, DeviceId, EventType, Invitation, OrganizationApiKey, OrganizationId,
+            AuthRequest, AuthRequestId, Device, DeviceId, EventType, Invitation, OIDCCodeWrapper, OrganizationApiKey,
-            SsoNonce, SsoUser, TwoFactor, TwoFactorIncomplete, TwoFactorType, User, UserId,
+            OrganizationId, SsoAuth, SsoUser, TwoFactor, TwoFactorIncomplete, TwoFactorType, User, UserId,
         },
         DbConn,
     },
     error::MapResult,
     mail, sso,
-    sso::{OIDCCode, OIDCState},
+    sso::{OIDCCode, OIDCCodeChallenge, OIDCCodeVerifier, OIDCState},
     util, CONFIG,
 };
 
@@ -92,6 +92,7 @@ async fn login(
         "authorization_code" if CONFIG.sso_enabled() => {
             _check_is_some(&data.client_id, "client_id cannot be blank")?;
             _check_is_some(&data.code, "code cannot be blank")?;
+            _check_is_some(&data.code_verifier, "code verifier cannot be blank")?;
 
             _check_is_some(&data.device_identifier, "device_identifier cannot be blank")?;
             _check_is_some(&data.device_name, "device_name cannot be blank")?;
@@ -175,17 +176,23 @@ async fn _sso_login(
     // Ratelimit the login
     crate::ratelimit::check_limit_login(&ip.ip)?;
 
-    let code = match data.code.as_ref() {
+    let (state, code_verifier) = match (data.code.as_ref(), data.code_verifier.as_ref()) {
-        None => err!(
+        (None, _) => err!(
             "Got no code in OIDC data",
             ErrorEvent {
                 event: EventType::UserFailedLogIn
             }
         ),
-        Some(code) => code,
+        (_, None) => err!(
+            "Got no code verifier in OIDC data",
+            ErrorEvent {
+                event: EventType::UserFailedLogIn
+            }
+        ),
+        (Some(code), Some(code_verifier)) => (code, code_verifier.clone()),
     };
 
-    let user_infos = sso::exchange_code(code, conn).await?;
+    let (sso_auth, user_infos) = sso::exchange_code(state, code_verifier, conn).await?;
     let user_with_sso = match SsoUser::find_by_identifier(&user_infos.identifier, conn).await {
         None => match SsoUser::find_by_mail(&user_infos.email, conn).await {
             None => None,
@@ -248,7 +255,7 @@ async fn _sso_login(
                 _ => (),
             }
 
-            let mut user = User::new(&user_infos.email, user_infos.user_name);
+            let mut user = User::new(&user_infos.email, user_infos.user_name.clone());
             user.verified_at = Some(now);
             user.save(conn).await?;
 
@@ -272,8 +279,8 @@ async fn _sso_login(
             if user.private_key.is_none() {
                 // User was invited a stub was created
                 user.verified_at = Some(now);
-                if let Some(user_name) = user_infos.user_name {
+                if let Some(ref user_name) = user_infos.user_name {
-                    user.name = user_name;
+                    user.name = user_name.clone();
                 }
 
                 user.save(conn).await?;
@@ -290,28 +297,11 @@ async fn _sso_login(
         }
     };
 
-    // We passed 2FA get full user information
-    let auth_user = sso::redeem(&user_infos.state, conn).await?;
-
-    if sso_user.is_none() {
-        let user_sso = SsoUser {
-            user_uuid: user.uuid.clone(),
-            identifier: user_infos.identifier,
-        };
-        user_sso.save(conn).await?;
-    }
-
     // Set the user_uuid here to be passed back used for event logging.
     *user_id = Some(user.uuid.clone());
 
-    let auth_tokens = sso::create_auth_tokens(
+    // We passed 2FA get auth tokens
-        &device,
+    let auth_tokens = sso::redeem(&device, &user, data.client_id, sso_user, sso_auth, user_infos, conn).await?;
-        &user,
-        data.client_id,
-        auth_user.refresh_token,
-        auth_user.access_token,
-        auth_user.expires_in,
-    )?;
 
     authenticated_response(&user, &mut device, auth_tokens, twofactor_token, &now, conn, ip).await
 }
@@ -997,9 +987,12 @@ struct ConnectData {
     two_factor_remember: Option<i32>,
     #[field(name = uncased("authrequest"))]
     auth_request: Option<AuthRequestId>,
+
     // Needed for authorization code
     #[field(name = uncased("code"))]
-    code: Option<String>,
+    code: Option<OIDCState>,
+    #[field(name = uncased("code_verifier"))]
+    code_verifier: Option<OIDCCodeVerifier>,
 }
 fn _check_is_some<T>(value: &Option<T>, msg: &str) -> EmptyResult {
     if value.is_none() {
@@ -1021,14 +1014,13 @@ fn prevalidate() -> JsonResult {
 }
 
 #[get("/connect/oidc-signin?<code>&<state>", rank = 1)]
-async fn oidcsignin(code: OIDCCode, state: String, conn: DbConn) -> ApiResult<Redirect> {
+async fn oidcsignin(code: OIDCCode, state: String, mut conn: DbConn) -> ApiResult<Redirect> {
-    oidcsignin_redirect(
+    _oidcsignin_redirect(
         state,
-        |decoded_state| sso::OIDCCodeWrapper::Ok {
+        OIDCCodeWrapper::Ok {
-            state: decoded_state,
             code,
         },
-        &conn,
+        &mut conn,
     )
     .await
 }
@@ -1040,42 +1032,44 @@ async fn oidcsignin_error(
     state: String,
     error: String,
     error_description: Option<String>,
-    conn: DbConn,
+    mut conn: DbConn,
 ) -> ApiResult<Redirect> {
-    oidcsignin_redirect(
+    _oidcsignin_redirect(
         state,
-        |decoded_state| sso::OIDCCodeWrapper::Error {
+        OIDCCodeWrapper::Error {
-            state: decoded_state,
             error,
             error_description,
         },
-        &conn,
+        &mut conn,
     )
     .await
 }
 
 // The state was encoded using Base64 to ensure no issue with providers.
 // iss and scope parameters are needed for redirection to work on IOS.
-async fn oidcsignin_redirect(
+// We pass the state as the code to get it back later on.
+async fn _oidcsignin_redirect(
     base64_state: String,
-    wrapper: impl FnOnce(OIDCState) -> sso::OIDCCodeWrapper,
+    code_response: OIDCCodeWrapper,
-    conn: &DbConn,
+    conn: &mut DbConn,
 ) -> ApiResult<Redirect> {
     let state = sso::decode_state(&base64_state)?;
-    let code = sso::encode_code_claims(wrapper(state.clone()));
 
-    let nonce = match SsoNonce::find(&state, conn).await {
+    let mut sso_auth = match SsoAuth::find(&state, conn).await {
-        Some(n) => n,
+        None => err!(format!("Cannot retrieve sso_auth for {state}")),
-        None => err!(format!("Failed to retrieve redirect_uri with {state}")),
+        Some(sso_auth) => sso_auth,
     };
+    sso_auth.code_response = Some(code_response);
+    sso_auth.updated_at = Utc::now().naive_utc();
+    sso_auth.save(conn).await?;
 
-    let mut url = match url::Url::parse(&nonce.redirect_uri) {
+    let mut url = match url::Url::parse(&sso_auth.redirect_uri) {
         Ok(url) => url,
-        Err(err) => err!(format!("Failed to parse redirect uri ({}): {err}", nonce.redirect_uri)),
+        Err(err) => err!(format!("Failed to parse redirect uri ({}): {err}", sso_auth.redirect_uri)),
     };
 
     url.query_pairs_mut()
-        .append_pair("code", &code)
+        .append_pair("code", &state)
         .append_pair("state", &state)
         .append_pair("scope", &AuthMethod::Sso.scope())
         .append_pair("iss", &CONFIG.domain());
@@ -1098,10 +1092,8 @@ struct AuthorizeData {
     #[allow(unused)]
     scope: Option<String>,
     state: OIDCState,
-    #[allow(unused)]
+    code_challenge: OIDCCodeChallenge,
-    code_challenge: Option<String>,
+    code_challenge_method: String,
-    #[allow(unused)]
-    code_challenge_method: Option<String>,
     #[allow(unused)]
     response_mode: Option<String>,
     #[allow(unused)]
@@ -1118,10 +1110,16 @@ async fn authorize(data: AuthorizeData, conn: DbConn) -> ApiResult<Redirect> {
         client_id,
         redirect_uri,
         state,
+        code_challenge,
+        code_challenge_method,
         ..
     } = data;
 
-    let auth_url = sso::authorize_url(state, &client_id, &redirect_uri, conn).await?;
+    if code_challenge_method != "S256" {
+        err!("Unsupported code challenge method");
+    }
+
+    let auth_url = sso::authorize_url(state, code_challenge, &client_id, &redirect_uri, conn).await?;
 
     Ok(Redirect::temporary(String::from(auth_url)))
 }

src/config.rs

@@ -564,9 +564,9 @@ make_config! {
         /// Duo Auth context cleanup schedule |> Cron schedule of the job that cleans expired Duo contexts from the database. Does nothing if Duo MFA is disabled or set to use the legacy iframe prompt.
         /// Defaults to once every minute. Set blank to disable this job.
         duo_context_purge_schedule:   String, false,  def,    "30 * * * * *".to_string();
-        /// Purge incomplete SSO nonce. |> Cron schedule of the job that cleans leftover nonce in db due to incomplete SSO login.
+        /// Purge incomplete SSO auth. |> Cron schedule of the job that cleans leftover auth in db due to incomplete SSO login.
         /// Defaults to daily. Set blank to disable this job.
-        purge_incomplete_sso_nonce: String, false,  def,   "0 20 0 * * *".to_string();
+        purge_incomplete_sso_auth: String, false,  def,   "0 20 0 * * *".to_string();
     },
 
     /// General settings

src/db/mod.rs

@@ -337,6 +337,46 @@ macro_rules! db_run {
     };
 }
 
+// Write all ToSql<Text, DB> and FromSql<Text, DB> given a serializable/deserializable type.
+#[macro_export]
+macro_rules! impl_FromToSqlText {
+    ($name:ty) => {
+        #[cfg(mysql)]
+        impl ToSql<Text, diesel::mysql::Mysql> for $name {
+            fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, diesel::mysql::Mysql>) -> diesel::serialize::Result {
+                serde_json::to_writer(out, self).map(|_| diesel::serialize::IsNull::No).map_err(Into::into)
+            }
+        }
+
+        #[cfg(postgresql)]
+        impl ToSql<Text, diesel::pg::Pg> for $name {
+            fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, diesel::pg::Pg>) -> diesel::serialize::Result {
+                serde_json::to_writer(out, self).map(|_| diesel::serialize::IsNull::No).map_err(Into::into)
+            }
+        }
+
+        #[cfg(sqlite)]
+        impl ToSql<Text, diesel::sqlite::Sqlite> for $name {
+            fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, diesel::sqlite::Sqlite>) -> diesel::serialize::Result {
+                serde_json::to_string(self).map_err(Into::into).map(|str| {
+                    out.set_value(str);
+                    diesel::serialize::IsNull::No
+                })
+            }
+        }
+
+        impl<DB: diesel::backend::Backend> FromSql<Text, DB> for $name
+        where
+            String: FromSql<Text, DB>,
+        {
+            fn from_sql(bytes: DB::RawValue<'_>) -> diesel::deserialize::Result<Self> {
+                <String as FromSql<Text, DB>>::from_sql(bytes)
+                    .and_then(|str| serde_json::from_str(&str).map_err(Into::into))
+            }
+        }
+    };
+}
+
 pub mod schema;
 
 // Reexport the models, needs to be after the macros are defined so it can access them

src/db/models/mod.rs

@@ -11,7 +11,7 @@ mod group;
 mod org_policy;
 mod organization;
 mod send;
-mod sso_nonce;
+mod sso_auth;
 mod two_factor;
 mod two_factor_duo_context;
 mod two_factor_incomplete;
@@ -36,7 +36,7 @@ pub use self::send::{
     id::{SendFileId, SendId},
     Send, SendType,
 };
-pub use self::sso_nonce::SsoNonce;
+pub use self::sso_auth::{OIDCAuthenticatedUser, OIDCCodeWrapper, SsoAuth};
 pub use self::two_factor::{TwoFactor, TwoFactorType};
 pub use self::two_factor_duo_context::TwoFactorDuoContext;
 pub use self::two_factor_incomplete::TwoFactorIncomplete;

src/db/models/sso_auth.rs

@@ -0,0 +1,134 @@
+use chrono::{NaiveDateTime, Utc};
+use std::time::Duration;
+
+use crate::api::EmptyResult;
+use crate::db::schema::sso_auth;
+use crate::db::{DbConn, DbPool};
+use crate::error::MapResult;
+use crate::sso::{OIDCCode, OIDCCodeChallenge, OIDCIdentifier, OIDCState, SSO_AUTH_EXPIRATION};
+
+use diesel::deserialize::FromSql;
+use diesel::expression::AsExpression;
+use diesel::prelude::*;
+use diesel::serialize::{Output, ToSql};
+use diesel::sql_types::Text;
+
+#[derive(AsExpression, Clone, Debug, Serialize, Deserialize, FromSqlRow)]
+#[diesel(sql_type = Text)]
+pub enum OIDCCodeWrapper {
+    Ok {
+        code: OIDCCode,
+    },
+    Error {
+        error: String,
+        error_description: Option<String>,
+    },
+}
+
+impl_FromToSqlText!(OIDCCodeWrapper);
+
+#[derive(AsExpression, Clone, Debug, Serialize, Deserialize, FromSqlRow)]
+#[diesel(sql_type = Text)]
+pub struct OIDCAuthenticatedUser {
+    pub refresh_token: Option<String>,
+    pub access_token: String,
+    pub expires_in: Option<Duration>,
+    pub identifier: OIDCIdentifier,
+    pub email: String,
+    pub email_verified: Option<bool>,
+    pub user_name: Option<String>,
+}
+
+impl_FromToSqlText!(OIDCAuthenticatedUser);
+
+#[derive(Identifiable, Queryable, Insertable, AsChangeset, Selectable)]
+#[diesel(table_name = sso_auth)]
+#[diesel(treat_none_as_null = true)]
+#[diesel(primary_key(state))]
+pub struct SsoAuth {
+    pub state: OIDCState,
+    pub client_challenge: OIDCCodeChallenge,
+    pub nonce: String,
+    pub redirect_uri: String,
+    pub code_response: Option<OIDCCodeWrapper>,
+    pub auth_response: Option<OIDCAuthenticatedUser>,
+    pub created_at: NaiveDateTime,
+    pub updated_at: NaiveDateTime,
+}
+
+/// Local methods
+impl SsoAuth {
+    pub fn new(state: OIDCState, client_challenge: OIDCCodeChallenge, nonce: String, redirect_uri: String) -> Self {
+        let now = Utc::now().naive_utc();
+
+        SsoAuth {
+            state,
+            client_challenge,
+            nonce,
+            redirect_uri,
+            created_at: now,
+            updated_at: now,
+            code_response: None,
+            auth_response: None,
+        }
+    }
+}
+
+/// Database methods
+impl SsoAuth {
+    pub async fn save(&self, conn: &DbConn) -> EmptyResult {
+        db_run! { conn:
+            mysql {
+                diesel::insert_into(sso_auth::table)
+                    .values(self)
+                    .on_conflict(diesel::dsl::DuplicatedKeys)
+                    .do_update()
+                    .set(self)
+                    .execute(conn)
+                    .map_res("Error saving SSO auth")
+            }
+            postgresql, sqlite {
+                diesel::insert_into(sso_auth::table)
+                    .values(self)
+                    .on_conflict(sso_auth::state)
+                    .do_update()
+                    .set(self)
+                    .execute(conn)
+                    .map_res("Error saving SSO auth")
+            }
+        }
+    }
+
+    pub async fn find(state: &OIDCState, conn: &DbConn) -> Option<Self> {
+        let oldest = Utc::now().naive_utc() - *SSO_AUTH_EXPIRATION;
+        db_run! { conn: {
+            sso_auth::table
+                .filter(sso_auth::state.eq(state))
+                .filter(sso_auth::created_at.ge(oldest))
+                .first::<Self>(conn)
+                .ok()
+        }}
+    }
+
+    pub async fn delete(self, conn: &DbConn) -> EmptyResult {
+        db_run! {conn: {
+            diesel::delete(sso_auth::table.filter(sso_auth::state.eq(self.state)))
+                .execute(conn)
+                .map_res("Error deleting sso_auth")
+        }}
+    }
+
+    pub async fn delete_expired(pool: DbPool) -> EmptyResult {
+        debug!("Purging expired sso_auth");
+        if let Ok(conn) = pool.get().await {
+            let oldest = Utc::now().naive_utc() - *SSO_AUTH_EXPIRATION;
+            db_run! { conn: {
+                diesel::delete(sso_auth::table.filter(sso_auth::created_at.lt(oldest)))
+                    .execute(conn)
+                    .map_res("Error deleting expired SSO nonce")
+            }}
+        } else {
+            err!("Failed to get DB connection while purging expired sso_auth")
+        }
+    }
+}

src/db/models/sso_nonce.rs

@@ -1,87 +0,0 @@
-use chrono::{NaiveDateTime, Utc};
-
-use crate::api::EmptyResult;
-use crate::db::schema::sso_nonce;
-use crate::db::{DbConn, DbPool};
-use crate::error::MapResult;
-use crate::sso::{OIDCState, NONCE_EXPIRATION};
-use diesel::prelude::*;
-
-#[derive(Identifiable, Queryable, Insertable)]
-#[diesel(table_name = sso_nonce)]
-#[diesel(primary_key(state))]
-pub struct SsoNonce {
-    pub state: OIDCState,
-    pub nonce: String,
-    pub verifier: Option<String>,
-    pub redirect_uri: String,
-    pub created_at: NaiveDateTime,
-}
-
-/// Local methods
-impl SsoNonce {
-    pub fn new(state: OIDCState, nonce: String, verifier: Option<String>, redirect_uri: String) -> Self {
-        let now = Utc::now().naive_utc();
-
-        SsoNonce {
-            state,
-            nonce,
-            verifier,
-            redirect_uri,
-            created_at: now,
-        }
-    }
-}
-
-/// Database methods
-impl SsoNonce {
-    pub async fn save(&self, conn: &DbConn) -> EmptyResult {
-        db_run! { conn:
-            sqlite, mysql {
-                diesel::replace_into(sso_nonce::table)
-                    .values(self)
-                    .execute(conn)
-                    .map_res("Error saving SSO nonce")
-            }
-            postgresql {
-                diesel::insert_into(sso_nonce::table)
-                    .values(self)
-                    .execute(conn)
-                    .map_res("Error saving SSO nonce")
-            }
-        }
-    }
-
-    pub async fn delete(state: &OIDCState, conn: &DbConn) -> EmptyResult {
-        db_run! { conn: {
-            diesel::delete(sso_nonce::table.filter(sso_nonce::state.eq(state)))
-                .execute(conn)
-                .map_res("Error deleting SSO nonce")
-        }}
-    }
-
-    pub async fn find(state: &OIDCState, conn: &DbConn) -> Option<Self> {
-        let oldest = Utc::now().naive_utc() - *NONCE_EXPIRATION;
-        db_run! { conn: {
-            sso_nonce::table
-                .filter(sso_nonce::state.eq(state))
-                .filter(sso_nonce::created_at.ge(oldest))
-                .first::<Self>(conn)
-                .ok()
-        }}
-    }
-
-    pub async fn delete_expired(pool: DbPool) -> EmptyResult {
-        debug!("Purging expired sso_nonce");
-        if let Ok(conn) = pool.get().await {
-            let oldest = Utc::now().naive_utc() - *NONCE_EXPIRATION;
-            db_run! { conn: {
-                diesel::delete(sso_nonce::table.filter(sso_nonce::created_at.lt(oldest)))
-                    .execute(conn)
-                    .map_res("Error deleting expired SSO nonce")
-            }}
-        } else {
-            err!("Failed to get DB connection while purging expired sso_nonce")
-        }
-    }
-}

src/db/schema.rs

@@ -256,12 +256,15 @@ table! {
 }
 
 table! {
-    sso_nonce (state) {
+    sso_auth (state) {
         state -> Text,
+        client_challenge -> Text,
         nonce -> Text,
-        verifier -> Nullable<Text>,
         redirect_uri -> Text,
+        code_response -> Nullable<Text>,
+        auth_response -> Nullable<Text>,
         created_at -> Timestamp,
+        updated_at -> Timestamp,
     }
 }
 

src/main.rs

@@ -699,10 +699,10 @@ fn schedule_jobs(pool: db::DbPool) {
                 }));
             }
 
-            // Purge sso nonce from incomplete flow (default to daily at 00h20).
+            // Purge sso auth from incomplete flow (default to daily at 00h20).
-            if !CONFIG.purge_incomplete_sso_nonce().is_empty() {
+            if !CONFIG.purge_incomplete_sso_auth().is_empty() {
-                sched.add(Job::new(CONFIG.purge_incomplete_sso_nonce().parse().unwrap(), || {
+                sched.add(Job::new(CONFIG.purge_incomplete_sso_auth().parse().unwrap(), || {
-                    runtime.spawn(db::models::SsoNonce::delete_expired(pool.clone()));
+                    runtime.spawn(db::models::SsoAuth::delete_expired(pool.clone()));
                 }));
             }
 

src/sso.rs

@@ -1,8 +1,7 @@
 use std::{sync::LazyLock, time::Duration};
 
 use chrono::Utc;
-use derive_more::{AsRef, Deref, Display, From};
+use derive_more::{AsRef, Deref, Display, From, Into};
-use mini_moka::sync::Cache;
 use regex::Regex;
 use url::Url;
 
@@ -11,7 +10,7 @@ use crate::{
     auth,
     auth::{AuthMethod, AuthTokens, TokenWrapper, BW_EXPIRATION, DEFAULT_REFRESH_VALIDITY},
     db::{
-        models::{Device, SsoNonce, User},
+        models::{Device, OIDCAuthenticatedUser, OIDCCodeWrapper, SsoAuth, SsoUser, User},
         DbConn,
     },
     sso_client::Client,
@@ -20,12 +19,10 @@ use crate::{
 
 pub static FAKE_IDENTIFIER: &str = "VW_DUMMY_IDENTIFIER_FOR_OIDC";
 
-static AC_CACHE: LazyLock<Cache<OIDCState, AuthenticatedUser>> =
-    LazyLock::new(|| Cache::builder().max_capacity(1000).time_to_live(Duration::from_secs(10 * 60)).build());
-
 static SSO_JWT_ISSUER: LazyLock<String> = LazyLock::new(|| format!("{}|sso", CONFIG.domain_origin()));
 
-pub static NONCE_EXPIRATION: LazyLock<chrono::Duration> = LazyLock::new(|| chrono::TimeDelta::try_minutes(10).unwrap());
+pub static SSO_AUTH_EXPIRATION: LazyLock<chrono::Duration> =
+    LazyLock::new(|| chrono::TimeDelta::try_minutes(10).unwrap());
 
 #[derive(
     Clone,
@@ -47,6 +44,47 @@ pub static NONCE_EXPIRATION: LazyLock<chrono::Duration> = LazyLock::new(|| chron
 #[from(forward)]
 pub struct OIDCCode(String);
 
+#[derive(
+    Clone,
+    Debug,
+    Default,
+    DieselNewType,
+    FromForm,
+    PartialEq,
+    Eq,
+    Hash,
+    Serialize,
+    Deserialize,
+    AsRef,
+    Deref,
+    Display,
+    From,
+    Into,
+)]
+#[deref(forward)]
+#[into(owned)]
+pub struct OIDCCodeChallenge(String);
+
+#[derive(
+    Clone,
+    Debug,
+    Default,
+    DieselNewType,
+    FromForm,
+    PartialEq,
+    Eq,
+    Hash,
+    Serialize,
+    Deserialize,
+    AsRef,
+    Deref,
+    Display,
+    Into,
+)]
+#[deref(forward)]
+#[into(owned)]
+pub struct OIDCCodeVerifier(String);
+
 #[derive(
     Clone,
     Debug,
@@ -91,40 +129,6 @@ pub fn encode_ssotoken_claims() -> String {
     auth::encode_jwt(&claims)
 }
 
-#[derive(Debug, Serialize, Deserialize)]
-pub enum OIDCCodeWrapper {
-    Ok {
-        state: OIDCState,
-        code: OIDCCode,
-    },
-    Error {
-        state: OIDCState,
-        error: String,
-        error_description: Option<String>,
-    },
-}
-
-#[derive(Debug, Serialize, Deserialize)]
-struct OIDCCodeClaims {
-    // Expiration time
-    pub exp: i64,
-    // Issuer
-    pub iss: String,
-
-    pub code: OIDCCodeWrapper,
-}
-
-pub fn encode_code_claims(code: OIDCCodeWrapper) -> String {
-    let time_now = Utc::now();
-    let claims = OIDCCodeClaims {
-        exp: (time_now + chrono::TimeDelta::try_minutes(5).unwrap()).timestamp(),
-        iss: SSO_JWT_ISSUER.to_string(),
-        code,
-    };
-
-    auth::encode_jwt(&claims)
-}
-
 #[derive(Clone, Debug, Serialize, Deserialize)]
 struct BasicTokenClaims {
     iat: Option<i64>,
@@ -178,9 +182,14 @@ pub fn decode_state(base64_state: &str) -> ApiResult<OIDCState> {
     Ok(state)
 }
 
-// The `nonce` allow to protect against replay attacks
 // redirect_uri from: https://github.com/bitwarden/server/blob/main/src/Identity/IdentityServer/ApiClient.cs
-pub async fn authorize_url(state: OIDCState, client_id: &str, raw_redirect_uri: &str, conn: DbConn) -> ApiResult<Url> {
+pub async fn authorize_url(
+    state: OIDCState,
+    client_challenge: OIDCCodeChallenge,
+    client_id: &str,
+    raw_redirect_uri: &str,
+    conn: DbConn,
+) -> ApiResult<Url> {
     let redirect_uri = match client_id {
         "web" | "browser" => format!("{}/sso-connector.html", CONFIG.domain()),
         "desktop" | "mobile" => "bitwarden://sso-callback".to_string(),
@@ -194,8 +203,8 @@ pub async fn authorize_url(state: OIDCState, client_id: &str, raw_redirect_uri:
         _ => err!(format!("Unsupported client {client_id}")),
     };
 
-    let (auth_url, nonce) = Client::authorize_url(state, redirect_uri).await?;
+    let (auth_url, sso_auth) = Client::authorize_url(state, client_challenge, redirect_uri).await?;
-    nonce.save(&conn).await?;
+    sso_auth.save(&conn).await?;
     Ok(auth_url)
 }
 
@@ -225,78 +234,45 @@ impl OIDCIdentifier {
     }
 }
 
-#[derive(Clone, Debug)]
-pub struct AuthenticatedUser {
-    pub refresh_token: Option<String>,
-    pub access_token: String,
-    pub expires_in: Option<Duration>,
-    pub identifier: OIDCIdentifier,
-    pub email: String,
-    pub email_verified: Option<bool>,
-    pub user_name: Option<String>,
-}
-
-#[derive(Clone, Debug)]
-pub struct UserInformation {
-    pub state: OIDCState,
-    pub identifier: OIDCIdentifier,
-    pub email: String,
-    pub email_verified: Option<bool>,
-    pub user_name: Option<String>,
-}
-
-async fn decode_code_claims(code: &str, conn: &DbConn) -> ApiResult<(OIDCCode, OIDCState)> {
-    match auth::decode_jwt::<OIDCCodeClaims>(code, SSO_JWT_ISSUER.to_string()) {
-        Ok(code_claims) => match code_claims.code {
-            OIDCCodeWrapper::Ok {
-                state,
-                code,
-            } => Ok((code, state)),
-            OIDCCodeWrapper::Error {
-                state,
-                error,
-                error_description,
-            } => {
-                if let Err(err) = SsoNonce::delete(&state, conn).await {
-                    error!("Failed to delete database sso_nonce using {state}: {err}")
-                }
-                err!(format!(
-                    "SSO authorization failed: {error}, {}",
-                    error_description.as_ref().unwrap_or(&String::new())
-                ))
-            }
-        },
-        Err(err) => err!(format!("Failed to decode code wrapper: {err}")),
-    }
-}
-
 // During the 2FA flow we will
 //  - retrieve the user information and then only discover he needs 2FA.
-//  - second time we will rely on the `AC_CACHE` since the `code` has already been exchanged.
+//  - second time we will rely on `SsoAuth.auth_response` since the `code` has already been exchanged.
-// The `nonce` will ensure that the user is authorized only once.
+// The `SsoAuth` will ensure that the user is authorized only once.
-// We return only the `UserInformation` to force calling `redeem` to obtain the `refresh_token`.
+pub async fn exchange_code(
-pub async fn exchange_code(wrapped_code: &str, conn: &DbConn) -> ApiResult<UserInformation> {
+    state: &OIDCState,
+    client_verifier: OIDCCodeVerifier,
+    conn: &DbConn,
+) -> ApiResult<(SsoAuth, OIDCAuthenticatedUser)> {
     use openidconnect::OAuth2TokenResponse;
 
-    let (code, state) = decode_code_claims(wrapped_code, conn).await?;
+    let mut sso_auth = match SsoAuth::find(state, conn).await {
+        None => err!(format!("Invalid state cannot retrieve sso auth")),
+        Some(sso_auth) => sso_auth,
+    };
 
-    if let Some(authenticated_user) = AC_CACHE.get(&state) {
+    if let Some(authenticated_user) = sso_auth.auth_response.clone() {
-        return Ok(UserInformation {
+        return Ok((sso_auth, authenticated_user));
-            state,
-            identifier: authenticated_user.identifier,
-            email: authenticated_user.email,
-            email_verified: authenticated_user.email_verified,
-            user_name: authenticated_user.user_name,
-        });
     }
 
-    let nonce = match SsoNonce::find(&state, conn).await {
+    let code = match sso_auth.code_response.clone() {
-        None => err!(format!("Invalid state cannot retrieve nonce")),
+        Some(OIDCCodeWrapper::Ok {
-        Some(nonce) => nonce,
+            code,
+        }) => code.clone(),
+        Some(OIDCCodeWrapper::Error {
+            error,
+            error_description,
+        }) => {
+            sso_auth.delete(conn).await?;
+            err!(format!("SSO authorization failed: {error}, {}", error_description.as_ref().unwrap_or(&String::new())))
+        }
+        None => {
+            sso_auth.delete(conn).await?;
+            err!("Missing authorization provider return");
+        }
     };
 
     let client = Client::cached().await?;
-    let (token_response, id_claims) = client.exchange_code(code, nonce).await?;
+    let (token_response, id_claims) = client.exchange_code(code, client_verifier, &sso_auth).await?;
 
     let user_info = client.user_info(token_response.access_token().to_owned()).await?;
 
@@ -316,7 +292,7 @@ pub async fn exchange_code(wrapped_code: &str, conn: &DbConn) -> ApiResult<UserI
 
     let identifier = OIDCIdentifier::new(id_claims.issuer(), id_claims.subject());
 
-    let authenticated_user = AuthenticatedUser {
+    let authenticated_user = OIDCAuthenticatedUser {
         refresh_token: refresh_token.cloned(),
         access_token: token_response.access_token().secret().clone(),
         expires_in: token_response.expires_in(),
@@ -327,29 +303,49 @@ pub async fn exchange_code(wrapped_code: &str, conn: &DbConn) -> ApiResult<UserI
     };
 
     debug!("Authenticated user {authenticated_user:?}");
+    sso_auth.auth_response = Some(authenticated_user.clone());
+    sso_auth.updated_at = Utc::now().naive_utc();
+    sso_auth.save(conn).await?;
 
-    AC_CACHE.insert(state.clone(), authenticated_user);
+    Ok((sso_auth, authenticated_user))
-
-    Ok(UserInformation {
-        state,
-        identifier,
-        email,
-        email_verified,
-        user_name,
-    })
 }
 
-// User has passed 2FA flow we can delete `nonce` and clear the cache.
+// User has passed 2FA flow we can delete auth info from database
-pub async fn redeem(state: &OIDCState, conn: &DbConn) -> ApiResult<AuthenticatedUser> {
+pub async fn redeem(
-    if let Err(err) = SsoNonce::delete(state, conn).await {
+    device: &Device,
-        error!("Failed to delete database sso_nonce using {state}: {err}")
+    user: &User,
+    client_id: Option<String>,
+    sso_user: Option<SsoUser>,
+    sso_auth: SsoAuth,
+    auth_user: OIDCAuthenticatedUser,
+    conn: &DbConn,
+) -> ApiResult<AuthTokens> {
+    sso_auth.delete(conn).await?;
+
+    if sso_user.is_none() {
+        let user_sso = SsoUser {
+            user_uuid: user.uuid.clone(),
+            identifier: auth_user.identifier.clone(),
+        };
+        user_sso.save(conn).await?;
     }
 
-    if let Some(au) = AC_CACHE.get(state) {
+    if !CONFIG.sso_auth_only_not_session() {
-        AC_CACHE.invalidate(state);
+        let now = Utc::now();
-        Ok(au)
+
+        let (ap_nbf, ap_exp) =
+            match (decode_token_claims("access_token", &auth_user.access_token), auth_user.expires_in) {
+                (Ok(ap), _) => (ap.nbf(), ap.exp),
+                (Err(_), Some(exp)) => (now.timestamp(), (now + exp).timestamp()),
+                _ => err!("Non jwt access_token and empty expires_in"),
+            };
+
+        let access_claims =
+            auth::LoginJwtClaims::new(device, user, ap_nbf, ap_exp, AuthMethod::Sso.scope_vec(), client_id, now);
+
+        _create_auth_tokens(device, auth_user.refresh_token, access_claims, auth_user.access_token)
     } else {
-        err!("Failed to retrieve user info from sso cache")
+        Ok(AuthTokens::new(device, user, AuthMethod::Sso, client_id))
     }
 }
 

src/sso_client.rs

@@ -7,8 +7,8 @@ use url::Url;
 
 use crate::{
     api::{ApiResult, EmptyResult},
-    db::models::SsoNonce,
+    db::models::SsoAuth,
-    sso::{OIDCCode, OIDCState},
+    sso::{OIDCCode, OIDCCodeChallenge, OIDCCodeVerifier, OIDCState},
     CONFIG,
 };
 
@@ -107,7 +107,11 @@ impl Client {
     }
 
     // The `state` is encoded using base64 to ensure no issue with providers (It contains the Organization identifier).
-    pub async fn authorize_url(state: OIDCState, redirect_uri: String) -> ApiResult<(Url, SsoNonce)> {
+    pub async fn authorize_url(
+        state: OIDCState,
+        client_challenge: OIDCCodeChallenge,
+        redirect_uri: String,
+    ) -> ApiResult<(Url, SsoAuth)> {
         let scopes = CONFIG.sso_scopes_vec().into_iter().map(Scope::new);
         let base64_state = data_encoding::BASE64.encode(state.to_string().as_bytes());
 
@@ -122,22 +126,21 @@ impl Client {
             .add_scopes(scopes)
             .add_extra_params(CONFIG.sso_authorize_extra_params_vec());
 
-        let verifier = if CONFIG.sso_pkce() {
+        if CONFIG.sso_pkce() {
-            let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256();
+            auth_req = auth_req
-            auth_req = auth_req.set_pkce_challenge(pkce_challenge);
+                .add_extra_param::<&str, String>("code_challenge", client_challenge.clone().into())
-            Some(pkce_verifier.into_secret())
+                .add_extra_param("code_challenge_method", "S256");
-        } else {
+        }
-            None
-        };
 
         let (auth_url, _, nonce) = auth_req.url();
-        Ok((auth_url, SsoNonce::new(state, nonce.secret().clone(), verifier, redirect_uri)))
+        Ok((auth_url, SsoAuth::new(state, client_challenge, nonce.secret().clone(), redirect_uri)))
     }
 
     pub async fn exchange_code(
         &self,
         code: OIDCCode,
-        nonce: SsoNonce,
+        client_verifier: OIDCCodeVerifier,
+        sso_auth: &SsoAuth,
     ) -> ApiResult<(
         StandardTokenResponse<
             IdTokenFields<
@@ -155,17 +158,21 @@ impl Client {
 
         let mut exchange = self.core_client.exchange_code(oidc_code);
 
+        let verifier = PkceCodeVerifier::new(client_verifier.into());
         if CONFIG.sso_pkce() {
-            match nonce.verifier {
+            exchange = exchange.set_pkce_verifier(verifier);
-                None => err!(format!("Missing verifier in the DB nonce table")),
+        } else {
-                Some(secret) => exchange = exchange.set_pkce_verifier(PkceCodeVerifier::new(secret)),
+            let challenge = PkceCodeChallenge::from_code_verifier_sha256(&verifier);
+            if challenge.as_str() != String::from(sso_auth.client_challenge.clone()) {
+                err!(format!("PKCE client challenge failed"))
+                // Might need to notify admin ? how ?
             }
         }
 
         match exchange.request_async(&self.http_client).await {
             Err(err) => err!(format!("Failed to contact token endpoint: {:?}", err)),
             Ok(token_response) => {
-                let oidc_nonce = Nonce::new(nonce.nonce);
+                let oidc_nonce = Nonce::new(sso_auth.nonce.clone());
 
                 let id_token = match token_response.extra_fields().id_token() {
                     None => err!("Token response did not contain an id_token"),