From 0b1b30baf1d97c6e508e70b3ca6d49ae0dd0928a Mon Sep 17 00:00:00 2001
From: mdecimus <mauro@stalw.art>
Date: Tue, 2 Jul 2024 22:30:41 +0200
Subject: [PATCH] Fixed OTP validation order

---
 .../directory/src/backend/internal/manage.rs  |  7 +-
 crates/directory/src/core/secret.rs           | 69 ++++++++++---------
 2 files changed, 44 insertions(+), 32 deletions(-)

diff --git a/crates/directory/src/backend/internal/manage.rs b/crates/directory/src/backend/internal/manage.rs
index 85eb4408..10598ef6 100644
--- a/crates/directory/src/backend/internal/manage.rs
+++ b/crates/directory/src/backend/internal/manage.rs
@@ -420,7 +420,12 @@ impl ManageDirectory for Store {
                     PrincipalValue::String(secret),
                 ) => {
                     if !principal.inner.secrets.contains(&secret) {
-                        principal.inner.secrets.push(secret);
+                        if secret.is_otp_auth() && !principal.inner.secrets.is_empty() {
+                            // Add OTP Auth URLs to the beginning of the list
+                            principal.inner.secrets.insert(0, secret);
+                        } else {
+                            principal.inner.secrets.push(secret);
+                        }
                     }
                 }
                 (
diff --git a/crates/directory/src/core/secret.rs b/crates/directory/src/core/secret.rs
index 39fe487a..09d9c36a 100644
--- a/crates/directory/src/core/secret.rs
+++ b/crates/directory/src/core/secret.rs
@@ -36,40 +36,38 @@ impl<T: serde::Serialize + serde::de::DeserializeOwned> Principal<T> {
                 // Account is disabled, no need to check further
 
                 return Ok(false);
-            } else if secret.is_otp_auth() && !is_totp_verified && !is_totp_token_missing {
-                is_totp_required = true;
+            } else if secret.is_otp_auth() {
+                if !is_totp_verified && !is_totp_token_missing {
+                    is_totp_required = true;
 
-                let totp_token = if let Some(totp_token) = totp_token {
-                    totp_token
-                } else if let Some((_code, _totp_token)) = code
-                    .rsplit_once('$')
-                    .filter(|(c, t)| !c.is_empty() && !t.is_empty())
+                    let totp_token = if let Some(totp_token) = totp_token {
+                        totp_token
+                    } else if let Some((_code, _totp_token)) = code
+                        .rsplit_once('$')
+                        .filter(|(c, t)| !c.is_empty() && !t.is_empty())
+                    {
+                        totp_token = Some(_totp_token);
+                        code = _code;
+                        _totp_token
+                    } else {
+                        is_totp_token_missing = true;
+                        continue;
+                    };
+
+                    // Token needs to validate with at least one of the TOTP secrets
+                    is_totp_verified = TOTP::from_url(secret)
+                        .map_err(DirectoryError::InvalidTotpUrl)?
+                        .check_current(totp_token)
+                        .unwrap_or(false);
+                }
+            } else if !is_authenticated && !is_app_authenticated {
+                if let Some((_, app_secret)) =
+                    secret.strip_prefix("$app$").and_then(|s| s.split_once('$'))
                 {
-                    totp_token = Some(_totp_token);
-                    code = _code;
-                    _totp_token
+                    is_app_authenticated = verify_secret_hash(app_secret, code).await;
                 } else {
-                    is_totp_token_missing = true;
-                    continue;
-                };
-
-                // Token needs to validate with at least one of the TOPT secrets
-                is_totp_verified = TOTP::from_url(secret)
-                    .map_err(DirectoryError::InvalidTotpUrl)?
-                    .check_current(totp_token)
-                    .unwrap_or(false);
-            }
-
-            if is_app_authenticated || is_authenticated {
-                continue;
-            }
-
-            if let Some((_, app_secret)) =
-                secret.strip_prefix("$app$").and_then(|s| s.split_once('$'))
-            {
-                is_app_authenticated = verify_secret_hash(app_secret, code).await;
-            } else {
-                is_authenticated = verify_secret_hash(secret, code).await;
+                    is_authenticated = verify_secret_hash(secret, code).await;
+                }
             }
         }
 
@@ -93,6 +91,15 @@ impl<T: serde::Serialize + serde::de::DeserializeOwned> Principal<T> {
 
             Ok(true)
         } else {
+            if is_totp_verified {
+                // TOTP URL appeared after password hash in secrets list
+                for secret in &self.secrets {
+                    if secret.is_password() && verify_secret_hash(secret, code).await {
+                        return Ok(true);
+                    }
+                }
+            }
+
             Ok(false)
         }
     }