Add privileged support for SSO account associations. (#252)

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

src/admin/query/oauth.rs

@@ -1,7 +1,7 @@
 use clap::Subcommand;
 use futures::{StreamExt, TryStreamExt};
 use ruma::OwnedUserId;
-use tuwunel_core::{Err, Result, utils::stream::IterStream};
+use tuwunel_core::{Err, Result, apply, err, itertools::Itertools, utils::stream::IterStream};
 use tuwunel_service::{
 	Services,
 	oauth::{Provider, Session},
@@ -13,6 +13,19 @@ use crate::{admin_command, admin_command_dispatch};
 #[derive(Debug, Subcommand)]
 /// Query OAuth service state
 pub(crate) enum OauthCommand {
+	/// Associate existing user with future authorization claims.
+	Associate {
+		/// ID of configured provider to listen on.
+		provider: String,
+
+		/// MXID of local user to associate.
+		user_id: OwnedUserId,
+
+		/// List of claims to match in key=value format.
+		#[arg(long, required = true)]
+		claim: Vec<String>,
+	},
+
 	/// List configured OAuth providers.
 	ListProviders,
 
@@ -51,6 +64,53 @@ pub(crate) enum OauthCommand {
 	},
 }
 
+#[admin_command]
+pub(super) async fn oauth_associate(
+	&self,
+	provider: String,
+	user_id: OwnedUserId,
+	claim: Vec<String>,
+) -> Result {
+	if !self.services.globals.user_is_local(&user_id) {
+		return Err!(Request(NotFound("User {user_id:?} does not belong to this server.")));
+	}
+
+	if !self.services.users.exists(&user_id).await {
+		return Err!(Request(NotFound("User {user_id:?} is not registered")));
+	}
+
+	let provider = self
+		.services
+		.oauth
+		.providers
+		.get(&provider)
+		.await?;
+
+	let claim = claim
+		.iter()
+		.map(|kv| {
+			let (key, val) = kv
+				.split_once('=')
+				.ok_or_else(|| err!("Missing '=' in --claim {kv}=???"))?;
+
+			if !key.is_empty() && !val.is_empty() {
+				Ok((key, val))
+			} else {
+				Err!("Missing key or value in --claim=key=value argument")
+			}
+		})
+		.map_ok(apply!(2, ToOwned::to_owned))
+		.collect::<Result<_>>()?;
+
+	let _replaced = self
+		.services
+		.oauth
+		.sessions
+		.set_user_association_pending(provider.id(), &user_id, claim);
+
+	Ok(())
+}
+
 #[admin_command]
 pub(super) async fn oauth_list_providers(&self) -> Result {
 	self.services