feat(!783): Add admin commands for managing tokens

conduwuit-example.toml

@@ -421,7 +421,7 @@
 # `yes_i_am_very_very_sure_i_want_an_open_registration_server_prone_to_abuse`
 #
 # If you would like registration only via token reg, please configure
-# `registration_token` or `registration_token_file`.
+# `registration_token`.
 #
 #allow_registration = false
 
@@ -458,16 +458,6 @@
 #
 #registration_token =
 
-# Path to a file on the system that gets read for additional registration
-# tokens. Multiple tokens can be added if you separate them with
-# whitespace
-#
-# continuwuity must be able to access the file, and it must not be empty
-#
-# example: "/etc/continuwuity/.reg_token"
-#
-#registration_token_file =
-
 # The public site key for reCaptcha. If this is provided, reCaptcha
 # becomes required during registration. If both captcha *and*
 # registration token are enabled, both will be required during

src/admin/admin.rs

@@ -2,10 +2,17 @@ use clap::Parser;
 use conduwuit::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)]
@@ -19,6 +26,10 @@ pub enum AdminCommand {
 	/// - Commands for managing local users
 	Users(UserCommand),
 
+	#[command(subcommand)]
+	/// - Commands for managing registration tokens
+	Tokens(TokenCommand),
+
 	#[command(subcommand)]
 	/// - Commands for managing rooms
 	Rooms(RoomCommand),
@@ -64,6 +75,11 @@ pub(super) async fn process(command: AdminCommand, context: &Context<'_>) -> Res
 			context.bail_restricted()?;
 			user::process(command, context).await
 		},
+		| Tokens(command) => {
+			// token commands are all restricted
+			context.bail_restricted()?;
+			token::process(command, context).await
+		},
 		| Rooms(command) => room::process(command, context).await,
 		| Federation(command) => federation::process(command, context).await,
 		| Server(command) => server::process(command, context).await,

src/admin/mod.rs

@@ -17,6 +17,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;
 
 extern crate conduwuit_api as api;

src/admin/token/commands.rs

@@ -0,0 +1,74 @@
+use conduwuit::{Err, Result, utils};
+use conduwuit_macros::admin_command;
+use futures::StreamExt;
+use service::registration_tokens::TokenExpires;
+
+#[admin_command]
+pub(super) async fn issue_token(&self, expires: super::TokenExpires) -> Result {
+	let expires = {
+		if expires.immortal {
+			None
+		} else if let Some(max_uses) = expires.max_uses {
+			Some(TokenExpires::AfterUses(max_uses))
+		} else if let Some(max_age) = expires
+			.max_age
+			.as_deref()
+			.map(|max_age| utils::time::timepoint_from_now(utils::time::parse_duration(max_age)?))
+			.transpose()?
+		{
+			Some(TokenExpires::AfterTime(max_age))
+		} else {
+			unreachable!();
+		}
+	};
+
+	let (token, info) = self
+		.services
+		.registration_tokens
+		.issue_token(self.sender_or_service_user().into(), expires);
+
+	self.write_str(&format!(
+		"New registration token issued: `{token}`. {}.",
+		if let Some(expires) = info.expires {
+			format!("{expires}")
+		} else {
+			"Never expires".to_owned()
+		}
+	))
+	.await
+}
+
+#[admin_command]
+pub(super) async fn revoke_token(&self, token: String) -> Result {
+	let Some(token) = self
+		.services
+		.registration_tokens
+		.validate_token(token)
+		.await
+	else {
+		return Err!("This token does not exist or has already expired.");
+	};
+
+	self.services.registration_tokens.revoke_token(token)?;
+
+	self.write_str("Token revoked successfully.").await
+}
+
+#[admin_command]
+pub(super) async fn list_tokens(&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,47 @@
+mod commands;
+
+use clap::{Args, Subcommand};
+use conduwuit::Result;
+
+use crate::admin_command_dispatch;
+
+#[admin_command_dispatch]
+#[derive(Debug, Subcommand)]
+pub enum TokenCommand {
+	/// - Issue a new registration token
+	#[clap(name = "issue")]
+	IssueToken {
+		/// When this token will expire.
+		#[command(flatten)]
+		expires: TokenExpires,
+	},
+
+	/// - Revoke a registration token
+	#[clap(name = "revoke")]
+	RevokeToken {
+		/// The token to revoke.
+		token: String,
+	},
+
+	/// - List all registration tokens
+	#[clap(name = "list")]
+	ListTokens,
+}
+
+#[derive(Debug, Args)]
+#[group(required = true, multiple = false)]
+pub struct TokenExpires {
+	/// 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>,
+
+	/// This token will never expire.
+	#[arg(long)]
+	immortal: bool,
+}

src/api/client/account.rs

@@ -864,7 +864,7 @@ pub(crate) async fn check_registration_token_validity(
 
 	let valid = services
 		.registration_tokens
-		.validate_token(&body.token)
+		.validate_token(body.token.clone())
 		.await
 		.is_some();
 

src/database/maps.rs

@@ -141,6 +141,10 @@ pub(super) static MAPS: &[Descriptor] = &[
 		name: "referencedevents",
 		..descriptor::RANDOM
 	},
+	Descriptor {
+		name: "registrationtoken_info",
+		..descriptor::RANDOM_SMALL
+	},
 	Descriptor {
 		name: "roomid_invitedcount",
 		..descriptor::RANDOM_SMALL

src/service/registration_tokens/data.rs

@@ -1,6 +1,9 @@
 use std::{sync::Arc, time::SystemTime};
 
-use conduwuit::utils::stream::{ReadyExt, TryIgnore};
+use conduwuit::utils::{
+	self,
+	stream::{ReadyExt, TryIgnore},
+};
 use database::{Database, Deserialized, Json, Map};
 use futures::Stream;
 use ruma::OwnedUserId;
@@ -32,7 +35,7 @@ impl DatabaseTokenInfo {
 	#[must_use]
 	pub fn is_valid(&self) -> bool {
 		match self.expires {
-			| Some(TokenExpires::AfterUses(max_uses)) => self.uses <= max_uses,
+			| Some(TokenExpires::AfterUses(max_uses)) => self.uses < max_uses,
 			| Some(TokenExpires::AfterTime(expiry_time)) => {
 				let now = SystemTime::now();
 
@@ -43,12 +46,46 @@ impl DatabaseTokenInfo {
 	}
 }
 
+impl std::fmt::Display for DatabaseTokenInfo {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "Token created by {} and used {} time. ", &self.creator, self.uses)?;
+		if let Some(expires) = &self.expires {
+			write!(f, "{expires}.")?;
+		} else {
+			write!(f, "Never expires.")?;
+		}
+
+		Ok(())
+	}
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub enum TokenExpires {
 	AfterUses(u64),
 	AfterTime(SystemTime),
 }
 
+impl std::fmt::Display for TokenExpires {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		match self {
+			| Self::AfterUses(max_uses) => write!(f, "Expires after {max_uses} uses"),
+			| Self::AfterTime(max_age) => {
+				let now = SystemTime::now();
+				let formatted_expiry = utils::time::format(*max_age, "%+");
+
+				match max_age.duration_since(now) {
+					| Ok(duration) => write!(
+						f,
+						"Expires in {} ({formatted_expiry})",
+						utils::time::pretty(duration)
+					),
+					| Err(_) => write!(f, "Expired at {formatted_expiry}"),
+				}
+			},
+		}
+	}
+}
+
 impl Data {
 	pub(super) fn new(db: &Arc<Database>) -> Self {
 		Self {
@@ -58,16 +95,16 @@ impl Data {
 
 	/// Associate a registration token with its metadata in the database.
 	pub(super) fn save_token(&self, token: &str, info: &DatabaseTokenInfo) {
-		self.registrationtoken_info.put(token, Json(info));
+		self.registrationtoken_info.raw_put(token, Json(info));
 	}
 
 	/// Delete a registration token.
-	pub(super) fn revoke_token(&self, token: &str) { self.registrationtoken_info.del(token); }
+	pub(super) fn revoke_token(&self, token: &str) { self.registrationtoken_info.remove(token); }
 
 	/// Look up a registration token's metadata.
 	pub(super) async fn lookup_token_info(&self, token: &str) -> Option<DatabaseTokenInfo> {
 		self.registrationtoken_info
-			.qry(token)
+			.get(token)
 			.await
 			.deserialized()
 			.ok()
@@ -80,12 +117,12 @@ impl Data {
 		self.registrationtoken_info
 			.stream()
 			.ignore_err()
-			.ready_filter(|item: &(&str, DatabaseTokenInfo)| {
+			.ready_filter_map(|item: (&str, DatabaseTokenInfo)| {
 				if item.1.is_valid() {
-					true
+					Some(item)
 				} else {
-					self.registrationtoken_info.del(item.0);
-					false
+					self.registrationtoken_info.remove(item.0);
+					None
 				}
 			})
 	}

src/service/registration_tokens/mod.rs

@@ -10,7 +10,7 @@ use ruma::OwnedUserId;
 
 use crate::{Dep, config};
 
-const RANDOM_TOKEN_LENGTH: usize = 64;
+const RANDOM_TOKEN_LENGTH: usize = 16;
 
 pub struct Service {
 	db: Data,
@@ -22,16 +22,24 @@ struct Services {
 }
 
 /// A validated registration token which may be used to create an account.
-pub struct ValidToken<'token> {
-	pub token: &'token str,
+#[derive(Debug)]
+pub struct ValidToken {
+	pub token: String,
 	pub source: ValidTokenSource,
 }
 
-impl PartialEq<str> for ValidToken<'_> {
+impl std::fmt::Display for ValidToken {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "`{}` --- {}", self.token, &self.source)
+	}
+}
+
+impl PartialEq<str> for ValidToken {
 	fn eq(&self, other: &str) -> bool { self.token == other }
 }
 
 /// The source of a valid database token.
+#[derive(Debug)]
 pub enum ValidTokenSource {
 	/// The static token set in the homeserver's config file, which is
 	/// always valid.
@@ -40,6 +48,15 @@ pub enum ValidTokenSource {
 	Database(DatabaseTokenInfo),
 }
 
+impl std::fmt::Display for ValidTokenSource {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		match self {
+			| Self::ConfigFile => write!(f, "Token defined in config."),
+			| Self::Database(info) => info.fmt(f),
+		}
+	}
+}
+
 impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
 		Ok(Arc::new(Self {
@@ -68,11 +85,11 @@ impl Service {
 	}
 
 	/// Get the registration token set in the config file, if it exists.
-	pub fn get_config_file_token(&self) -> Option<ValidToken<'_>> {
+	pub fn get_config_file_token(&self) -> Option<ValidToken> {
 		self.services
 			.config
 			.registration_token
-			.as_deref()
+			.clone()
 			.map(|token| ValidToken {
 				token,
 				source: ValidTokenSource::ConfigFile,
@@ -80,7 +97,7 @@ impl Service {
 	}
 
 	/// Validate a registration token.
-	pub async fn validate_token<'token>(&self, token: &'token str) -> Option<ValidToken<'token>> {
+	pub async fn validate_token(&self, token: String) -> Option<ValidToken> {
 		// Check the registration token in the config first
 		if self
 			.get_config_file_token()
@@ -93,7 +110,7 @@ impl Service {
 		}
 
 		// Now check the database
-		if let Some(token_info) = self.db.lookup_token_info(token).await
+		if let Some(token_info) = self.db.lookup_token_info(&token).await
 			&& token_info.is_valid()
 		{
 			return Some(ValidToken {
@@ -107,7 +124,7 @@ impl Service {
 	}
 
 	/// Mark a valid token as having been used to create a new account.
-	pub fn mark_token_as_used(&self, ValidToken { token, source }: ValidToken<'_>) {
+	pub fn mark_token_as_used(&self, ValidToken { token, source }: ValidToken) {
 		match source {
 			| ValidTokenSource::ConfigFile => {
 				// we don't track uses of the config file token, do nothing
@@ -115,7 +132,7 @@ impl Service {
 			| ValidTokenSource::Database(mut info) => {
 				info.uses = info.uses.saturating_add(1);
 
-				self.db.save_token(token, &info);
+				self.db.save_token(&token, &info);
 			},
 		}
 	}
@@ -124,7 +141,7 @@ impl Service {
 	///
 	/// Note that some tokens (like the one set in the config file) cannot be
 	/// revoked.
-	pub fn revoke_token(&self, ValidToken { token, source }: ValidToken<'_>) -> Result {
+	pub fn revoke_token(&self, ValidToken { token, source }: ValidToken) -> Result {
 		match source {
 			| ValidTokenSource::ConfigFile => {
 				// the config file token cannot be revoked
@@ -134,19 +151,22 @@ impl Service {
 				)
 			},
 			| ValidTokenSource::Database(_) => {
-				self.db.revoke_token(token);
+				self.db.revoke_token(&token);
 				Ok(())
 			},
 		}
 	}
 
 	/// Iterate over all valid registration tokens.
-	pub fn iterate_tokens(&self) -> impl Stream<Item = ValidToken<'_>> + Send + '_ {
-		stream::iter(self.get_config_file_token()).chain(self.db.iterate_and_clean_tokens().map(
-			|(token, info)| ValidToken {
-				token,
+	pub fn iterate_tokens(&self) -> impl Stream<Item = ValidToken> + Send + '_ {
+		let db_tokens = self
+			.db
+			.iterate_and_clean_tokens()
+			.map(|(token, info)| ValidToken {
+				token: token.to_owned(),
 				source: ValidTokenSource::Database(info),
-			},
-		))
+			});
+
+		stream::iter(self.get_config_file_token()).chain(db_tokens)
 	}
 }

src/service/uiaa/mod.rs

@@ -209,7 +209,7 @@ pub async fn try_auth(
 			}
 		},
 		| AuthData::RegistrationToken(t) => {
-			let token = t.token.trim();
+			let token = t.token.trim().to_owned();
 
 			if let Some(valid_token) = self
 				.services