Implement associated multi-provider single-sign-on flow support. (#252)

Add experimental note for multi-provider flow. (#252)

Signed-off-by: Jason Volk <jason@zemos.net>

src/admin/query/oauth.rs

@@ -1,11 +1,14 @@
 use clap::Subcommand;
 use futures::{StreamExt, TryStreamExt};
 use ruma::OwnedUserId;
-use tuwunel_core::{Err, Result, apply, err, itertools::Itertools, utils::stream::IterStream};
-use tuwunel_service::{
-	Services,
-	oauth::{Provider, Session},
+use tuwunel_core::{
+	Err, Result, apply,
+	either::{Either, Left, Right},
+	err,
+	itertools::Itertools,
+	utils::stream::{IterStream, ReadyExt},
 };
+use tuwunel_service::oauth::{Provider, ProviderId, SessionId};
 
 use crate::{admin_command, admin_command_dispatch};
 
@@ -29,12 +32,18 @@ pub(crate) enum OauthCommand {
 	/// List configured OAuth providers.
 	ListProviders,
 
-	/// List users associated with an OAuth provider
+	/// List users associated with any OAuth session
 	ListUsers,
 
+	/// List session ID's
+	ListSessions {
+		#[arg(long)]
+		user: Option<OwnedUserId>,
+	},
+
 	/// Show active configuration of a provider.
 	ShowProvider {
-		id: String,
+		id: ProviderId,
 
 		#[arg(long)]
 		config: bool,
@@ -42,28 +51,43 @@ pub(crate) enum OauthCommand {
 
 	/// Show session state
 	ShowSession {
-		id: String,
+		id: SessionId,
+	},
+
+	/// Show user sessions
+	ShowUser {
+		user_id: OwnedUserId,
 	},
 
 	/// Token introspection request to provider.
 	TokenInfo {
-		id: String,
+		id: SessionId,
 	},
 
 	/// Revoke token for user_id or sess_id.
 	Revoke {
-		id: String,
+		#[arg(value_parser = session_or_user_id)]
+		id: Either<SessionId, OwnedUserId>,
 	},
 
 	/// Remove oauth state (DANGER!)
-	Remove {
-		id: String,
+	Delete {
+		#[arg(value_parser = session_or_user_id)]
+		id: Either<SessionId, OwnedUserId>,
 
 		#[arg(long)]
 		force: bool,
 	},
 }
 
+type SessionOrUserId = Either<SessionId, OwnedUserId>;
+
+fn session_or_user_id(input: &str) -> Result<SessionOrUserId> {
+	OwnedUserId::parse(input)
+		.map(Right)
+		.or_else(|_| Ok(Left(input.to_owned())))
+}
+
 #[admin_command]
 pub(super) async fn oauth_associate(
 	&self,
@@ -137,7 +161,34 @@ pub(super) async fn oauth_list_users(&self) -> Result {
 }
 
 #[admin_command]
-pub(super) async fn oauth_show_provider(&self, id: String, config: bool) -> Result {
+pub(super) async fn oauth_list_sessions(&self, user_id: Option<OwnedUserId>) -> Result {
+	if let Some(user_id) = user_id.as_deref() {
+		return self
+			.services
+			.oauth
+			.sessions
+			.get_sess_id_by_user(user_id)
+			.map_ok(|id| format!("{id}\n"))
+			.try_for_each(async |id: String| self.write_str(&id).await)
+			.await;
+	}
+
+	self.services
+		.oauth
+		.sessions
+		.stream()
+		.ready_filter_map(|sess| sess.sess_id)
+		.map(|sess_id| format!("{sess_id:?}\n"))
+		.for_each(async |id: String| {
+			self.write_str(&id).await.ok();
+		})
+		.await;
+
+	Ok(())
+}
+
+#[admin_command]
+pub(super) async fn oauth_show_provider(&self, id: ProviderId, config: bool) -> Result {
 	if config {
 		let config = self.services.oauth.providers.get_config(&id)?;
 
@@ -151,15 +202,29 @@ pub(super) async fn oauth_show_provider(&self, id: String, config: bool) -> Resu
 }
 
 #[admin_command]
-pub(super) async fn oauth_show_session(&self, id: String) -> Result {
-	let session = find_session(self.services, &id).await?;
+pub(super) async fn oauth_show_session(&self, id: SessionId) -> Result {
+	let session = self.services.oauth.sessions.get(&id).await?;
 
 	self.write_str(&format!("{session:#?}\n")).await
 }
 
 #[admin_command]
-pub(super) async fn oauth_token_info(&self, id: String) -> Result {
-	let session = find_session(self.services, &id).await?;
+pub(super) async fn oauth_show_user(&self, user_id: OwnedUserId) -> Result {
+	self.services
+		.oauth
+		.sessions
+		.get_sess_id_by_user(&user_id)
+		.try_for_each(async |id| {
+			let session = self.services.oauth.sessions.get(&id).await?;
+
+			self.write_str(&format!("{session:#?}\n")).await
+		})
+		.await
+}
+
+#[admin_command]
+pub(super) async fn oauth_token_info(&self, id: SessionId) -> Result {
+	let session = self.services.oauth.sessions.get(&id).await?;
 
 	let provider = self
 		.services
@@ -172,61 +237,64 @@ pub(super) async fn oauth_token_info(&self, id: String) -> Result {
 		.services
 		.oauth
 		.request_tokeninfo((&provider, &session))
-		.await;
+		.await?;
 
 	self.write_str(&format!("{tokeninfo:#?}\n")).await
 }
 
 #[admin_command]
-pub(super) async fn oauth_revoke(&self, id: String) -> Result {
-	let session = find_session(self.services, &id).await?;
+pub(super) async fn oauth_revoke(&self, id: SessionOrUserId) -> Result {
+	match id {
+		| Left(sess_id) => {
+			let session = self.services.oauth.sessions.get(&sess_id).await?;
 
-	let provider = self
-		.services
-		.oauth
-		.sessions
-		.provider(&session)
-		.await?;
+			let provider = self
+				.services
+				.oauth
+				.sessions
+				.provider(&session)
+				.await?;
 
-	self.services
-		.oauth
-		.revoke_token((&provider, &session))
-		.await?;
+			self.services
+				.oauth
+				.revoke_token((&provider, &session))
+				.await
+				.ok();
+		},
+		| Right(user_id) =>
+			self.services
+				.oauth
+				.revoke_user_tokens(&user_id)
+				.await,
+	}
 
-	self.write_str("done").await
+	self.write_str("revoked").await
 }
 
 #[admin_command]
-pub(super) async fn oauth_remove(&self, id: String, force: bool) -> Result {
-	let session = find_session(self.services, &id).await?;
-
-	let Some(sess_id) = session.sess_id else {
-		return Err!("Missing sess_id in oauth Session state");
-	};
-
+pub(super) async fn oauth_delete(&self, id: SessionOrUserId, force: bool) -> Result {
 	if !force {
 		return Err!(
 			"Deleting these records can cause registration conflicts. Use --force to be sure."
 		);
 	}
 
-	self.services
-		.oauth
-		.sessions
-		.delete(&sess_id)
-		.await;
-
-	self.write_str("done").await
-}
-
-async fn find_session(services: &Services, id: &str) -> Result<Session> {
-	if let Ok(user_id) = OwnedUserId::parse(id) {
-		services
-			.oauth
-			.sessions
-			.get_by_user(&user_id)
-			.await
-	} else {
-		services.oauth.sessions.get(id).await
+	match id {
+		| Left(sess_id) => {
+			self.services
+				.oauth
+				.sessions
+				.delete(&sess_id)
+				.await;
+		},
+		| Right(user_id) => {
+			self.services
+				.oauth
+				.delete_user_sessions(&user_id)
+				.await;
+		},
 	}
+
+	self.write_str("deleted any oauth state for {id}")
+		.await
 }