Limited use registration token support

Co-authored-by: Ginger <ginger@gingershaped.computer>
Signed-off-by: Jason Volk <jason@zemos.net>

src/admin/admin.rs

@@ -2,10 +2,17 @@ use clap::Parser;
 use tuwunel_core::Result;
 
 use crate::{
-	appservice, appservice::AppserviceCommand, check, check::CheckCommand, context::Context,
-	debug, debug::DebugCommand, federation, federation::FederationCommand, media,
-	media::MediaCommand, query, query::QueryCommand, room, room::RoomCommand, server,
-	server::ServerCommand, user, user::UserCommand,
+	appservice::{self, AppserviceCommand},
+	check::{self, CheckCommand},
+	context::Context,
+	debug::{self, DebugCommand},
+	federation::{self, FederationCommand},
+	media::{self, MediaCommand},
+	query::{self, QueryCommand},
+	room::{self, RoomCommand},
+	server::{self, ServerCommand},
+	token::{self, TokenCommand},
+	user::{self, UserCommand},
 };
 
 #[derive(Debug, Parser)]
@@ -46,6 +53,10 @@ pub(super) enum AdminCommand {
 	#[command(subcommand)]
 	/// - Low-level queries for database getters and iterators
 	Query(QueryCommand),
+
+	#[command(subcommand)]
+	/// - Commands for managing registration tokens
+	Token(TokenCommand),
 }
 
 #[tracing::instrument(skip_all, name = "command")]
@@ -62,5 +73,6 @@ pub(super) async fn process(command: AdminCommand, context: &Context<'_>) -> Res
 		| Debug(command) => debug::process(command, context).await,
 		| Query(command) => query::process(command, context).await,
 		| Check(command) => check::process(command, context).await,
+		| Token(command) => token::process(command, context).await,
 	}
 }

src/admin/mod.rs

@@ -15,6 +15,7 @@ pub(crate) mod media;
 pub(crate) mod query;
 pub(crate) mod room;
 pub(crate) mod server;
+pub(crate) mod token;
 pub(crate) mod user;
 
 pub(crate) use tuwunel_macros::{admin_command, admin_command_dispatch};

src/admin/token/commands.rs

@@ -0,0 +1,61 @@
+use futures::StreamExt;
+use tuwunel_core::{Result, utils};
+use tuwunel_macros::admin_command;
+use tuwunel_service::registration_tokens::TokenExpires;
+
+#[admin_command]
+pub(super) async fn issue(
+	&self,
+	max_uses: Option<u64>,
+	max_age: Option<String>,
+	once: bool,
+) -> Result {
+	let expires = TokenExpires {
+		max_uses: max_uses.or_else(|| once.then_some(1)),
+		max_age: max_age
+			.map(|max_age| {
+				let duration = utils::time::parse_duration(&max_age)?;
+				utils::time::timepoint_from_now(duration)
+			})
+			.transpose()?,
+	};
+
+	let (token, info) = self
+		.services
+		.registration_tokens
+		.issue_token(expires)
+		.await?;
+
+	self.write_str(&format!("New registration token issued: `{token}` - {info}",))
+		.await
+}
+
+#[admin_command]
+pub(super) async fn revoke(&self, token: String) -> Result {
+	self.services
+		.registration_tokens
+		.revoke_token(&token)
+		.await?;
+
+	self.write_str("Token revoked successfully.")
+		.await
+}
+
+#[admin_command]
+pub(super) async fn list(&self) -> Result {
+	let tokens: Vec<_> = self
+		.services
+		.registration_tokens
+		.iterate_tokens()
+		.collect()
+		.await;
+
+	self.write_str(&format!("Found {} registration tokens:\n", tokens.len()))
+		.await?;
+
+	for token in tokens {
+		self.write_str(&format!("- {token}\n")).await?;
+	}
+
+	Ok(())
+}

src/admin/token/mod.rs

@@ -0,0 +1,36 @@
+mod commands;
+
+use clap::Subcommand;
+use tuwunel_core::Result;
+
+use crate::admin_command_dispatch;
+
+#[admin_command_dispatch]
+#[derive(Debug, Subcommand)]
+pub(crate) enum TokenCommand {
+	/// - Issue a new registration token
+	Issue {
+		/// The maximum number of times this token is allowed to be used before
+		/// it expires.
+		#[arg(long)]
+		max_uses: Option<u64>,
+
+		/// The maximum age of this token (e.g. 30s, 5m, 7d). It will expire
+		/// after this much time has passed.
+		#[arg(long)]
+		max_age: Option<String>,
+
+		/// A shortcut for `--max-uses 1`.
+		#[arg(long)]
+		once: bool,
+	},
+
+	/// - Revoke a registration token
+	Revoke {
+		/// The token to revoke.
+		token: String,
+	},
+
+	/// - List all registration tokens
+	List,
+}