From 3127eca67cbc9a26d0d2250048760750940057de Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Wed, 4 Mar 2026 20:34:55 +0000
Subject: [PATCH] Add conditional UIAA flows for SSO and password. (#314)

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/api/router/auth/uiaa.rs | 60 +++++++++++++++++++++++++++----------
 src/service/uiaa/mod.rs     | 17 +++++++++++
 2 files changed, 61 insertions(+), 16 deletions(-)

diff --git a/src/api/router/auth/uiaa.rs b/src/api/router/auth/uiaa.rs
index 4f7306fa..3e355acd 100644
--- a/src/api/router/auth/uiaa.rs
+++ b/src/api/router/auth/uiaa.rs
@@ -5,7 +5,13 @@ use ruma::{
 		client::uiaa::{AuthData, AuthFlow, AuthType, Jwt, UiaaInfo},
 	},
 };
-use tuwunel_core::{Err, Error, Result, err, is_equal_to, utils};
+use tuwunel_core::{
+	Err, Error, Result, err, utils,
+	utils::{
+		OptionExt,
+		future::{OptionFutureExt, TryExtExt},
+	},
+};
 use tuwunel_service::{Services, uiaa::SESSION_ID_LENGTH};
 
 use crate::{Ruma, client::jwt};
@@ -15,14 +21,46 @@ where
 	T: IncomingRequest + Send + Sync,
 {
 	let sender_device = body.sender_device()?;
+	let sender_user = body.sender_user.as_deref();
 
-	let flows = [
-		AuthFlow::new([AuthType::Password].into()),
-		AuthFlow::new([AuthType::Jwt].into()),
-	];
+	let password_flow = [AuthType::Password];
+	let has_password = sender_user
+		.map_async(|sender_user| {
+			services
+				.users
+				.has_password(sender_user)
+				.unwrap_or(false)
+		})
+		.unwrap_or(false)
+		.await;
+
+	//TODO: UIAA for SSO.
+	let sso_flow = [AuthType::Sso];
+	let has_sso = false;
+	let _has_sso = sender_user
+		.map_async(|sender_user| {
+			services
+				.oauth
+				.sessions
+				.exists_for_user(sender_user)
+		})
+		.unwrap_or(false)
+		.await;
+
+	//NOTE: Not implemented as a fallback/web for now.
+	let has_jwt = false;
+	let jwt_flow = [AuthType::Jwt];
 
 	let mut uiaainfo = UiaaInfo {
-		flows: flows.into(),
+		flows: has_password
+			.then_some(password_flow)
+			.into_iter()
+			.chain(has_sso.then_some(sso_flow))
+			.chain(has_jwt.then_some(jwt_flow))
+			.map(Vec::from)
+			.map(AuthFlow::new)
+			.collect(),
+
 		..Default::default()
 	};
 
@@ -70,16 +108,6 @@ where
 					.as_deref()
 					.ok_or_else(|| err!(Request(MissingToken("Missing access token."))))?;
 
-				// Skip UIAA for SSO/OIDC users.
-				if services
-					.users
-					.origin(sender_user)
-					.await
-					.is_ok_and(is_equal_to!("sso"))
-				{
-					return Ok(sender_user.to_owned());
-				}
-
 				uiaainfo.session = Some(utils::random_string(SESSION_ID_LENGTH));
 				services
 					.uiaa
diff --git a/src/service/uiaa/mod.rs b/src/service/uiaa/mod.rs
index c28bd587..51d866e6 100644
--- a/src/service/uiaa/mod.rs
+++ b/src/service/uiaa/mod.rs
@@ -16,6 +16,8 @@ use tuwunel_core::{
 };
 use tuwunel_database::{Deserialized, Json, Map};
 
+use crate::users::PASSWORD_SENTINEL;
+
 pub struct Service {
 	userdevicesessionid_uiaarequest: RwLock<RequestMap>,
 	db: Data,
@@ -107,9 +109,11 @@ pub async fn try_auth(
 			// Check if password is correct
 			let user_id = user_id_from_username;
 			let mut password_verified = false;
+			let mut password_sentinel = false;
 
 			// First try local password hash verification
 			if let Ok(hash) = self.services.users.password_hash(&user_id).await {
+				password_sentinel = hash == PASSWORD_SENTINEL;
 				password_verified = hash::verify_password(password, &hash).is_ok();
 			}
 
@@ -130,6 +134,19 @@ pub async fn try_auth(
 				}
 			}
 
+			// For SSO users that have never set a password, allow.
+			if !password_verified
+				&& password_sentinel
+				&& self
+					.services
+					.oauth
+					.sessions
+					.exists_for_user(&user_id)
+					.await
+			{
+				return Ok((true, uiaainfo));
+			}
+
 			if !password_verified {
 				uiaainfo.auth_error = Some(StandardErrorBody {
 					kind: ErrorKind::forbidden(),