feat: Add user locking and unlocking commands and functionality

Also corrects the response code returned by UserSuspended

src/admin/user/commands.rs

@@ -974,3 +974,34 @@ pub(super) async fn force_leave_remote_room(
 	self.write_str(&format!("{user_id} successfully left {room_id} via remote server."))
 		.await
 }
+
+#[admin_command]
+pub(super) async fn lock(&self, user_id: String) -> Result {
+	let user_id = parse_local_user_id(self.services, &user_id)?;
+	assert!(
+		self.services.globals.user_is_local(&user_id),
+		"Parsed user_id must be a local user"
+	);
+
+	self.services
+		.users
+		.lock_account(&user_id, self.sender_or_service_user())
+		.await;
+
+	self.write_str(&format!("User {user_id} has been locked."))
+		.await
+}
+
+#[admin_command]
+pub(super) async fn unlock(&self, user_id: String) -> Result {
+	let user_id = parse_local_user_id(self.services, &user_id)?;
+	assert!(
+		self.services.globals.user_is_local(&user_id),
+		"Parsed user_id must be a local user"
+	);
+
+	self.services.users.unlock_account(&user_id).await;
+
+	self.write_str(&format!("User {user_id} has been unlocked."))
+		.await
+}

src/admin/user/mod.rs

@@ -81,6 +81,26 @@ pub enum UserCommand {
 		user_id: String,
 	},
 
+	/// - Lock a user
+	///
+	/// Locked users are unable to use their accounts beyond logging out. This
+	/// is akin to a temporary deactivation that does not change the user's
+	/// password. This can be used to quickly prevent a user from accessing
+	/// their account.
+	Lock {
+		/// Username of the user to lock
+		user_id: String,
+	},
+
+	/// - Unlock a user
+	///
+	/// Reverses the effects of the `lock` command, allowing the user to use
+	/// their account again.
+	Unlock {
+		/// Username of the user to unlock
+		user_id: String,
+	},
+
 	/// - List local users in the database
 	#[clap(alias = "list")]
 	ListUsers,

src/api/router/auth.rs

@@ -137,12 +137,30 @@ pub(super) async fn auth(
 		| (
 			AuthScheme::AccessToken | AuthScheme::AccessTokenOptional | AuthScheme::None,
 			Token::User((user_id, device_id)),
-		) => Ok(Auth {
-			origin: None,
-			sender_user: Some(user_id),
-			sender_device: Some(device_id),
-			appservice_info: None,
-		}),
+		) => {
+			let is_locked = services.users.is_locked(&user_id).await.map_err(|e| {
+				err!(Request(Forbidden(warn!("Failed to check user lock status: {e}"))))
+			})?;
+			if is_locked {
+				// Only /logout and /logout/all are allowed for locked users
+				if !matches!(
+					metadata,
+					&ruma::api::client::session::logout::v3::Request::METADATA
+						| &ruma::api::client::session::logout_all::v3::Request::METADATA
+				) {
+					return Err(Error::BadRequest(
+						ErrorKind::UserLocked,
+						"This account has been locked.",
+					));
+				}
+			}
+			Ok(Auth {
+				origin: None,
+				sender_user: Some(user_id),
+				sender_device: Some(device_id),
+				appservice_info: None,
+			})
+		},
 		| (AuthScheme::ServerSignatures, Token::None) =>
 			Ok(auth_server(services, request, json_body).await?),
 		| (

src/core/error/response.rs

@@ -75,10 +75,12 @@ pub(super) fn bad_request_code(kind: &ErrorKind) -> StatusCode {
 		| ThreepidDenied
 		| InviteBlocked
 		| WrongRoomKeysVersion { .. }
+		| UserSuspended
 		| Forbidden { .. } => StatusCode::FORBIDDEN,
 
 		// 401
-		| UnknownToken { .. } | MissingToken | Unauthorized => StatusCode::UNAUTHORIZED,
+		| UnknownToken { .. } | MissingToken | Unauthorized | UserLocked =>
+			StatusCode::UNAUTHORIZED,
 
 		// 400
 		| _ => StatusCode::BAD_REQUEST,

src/database/maps.rs

@@ -386,6 +386,10 @@ pub(super) static MAPS: &[Descriptor] = &[
 		name: "userid_suspension",
 		..descriptor::RANDOM_SMALL
 	},
+	Descriptor {
+		name: "userid_lock",
+		..descriptor::RANDOM_SMALL
+	},
 	Descriptor {
 		name: "userid_presenceid",
 		..descriptor::RANDOM_SMALL

src/service/users/mod.rs

@@ -77,6 +77,7 @@ struct Data {
 	userid_origin: Arc<Map>,
 	userid_password: Arc<Map>,
 	userid_suspension: Arc<Map>,
+	userid_lock: Arc<Map>,
 	userid_selfsigningkeyid: Arc<Map>,
 	userid_usersigningkeyid: Arc<Map>,
 	useridprofilekey_value: Arc<Map>,
@@ -115,6 +116,7 @@ impl crate::Service for Service {
 				userid_origin: args.db["userid_origin"].clone(),
 				userid_password: args.db["userid_password"].clone(),
 				userid_suspension: args.db["userid_suspension"].clone(),
+				userid_lock: args.db["userid_lock"].clone(),
 				userid_selfsigningkeyid: args.db["userid_selfsigningkeyid"].clone(),
 				userid_usersigningkeyid: args.db["userid_usersigningkeyid"].clone(),
 				useridprofilekey_value: args.db["useridprofilekey_value"].clone(),
@@ -220,6 +222,26 @@ impl Service {
 		self.db.userid_suspension.remove(user_id);
 	}
 
+	pub async fn lock_account(&self, user_id: &UserId, locking_user: &UserId) {
+		// NOTE: Locking is basically just suspension with a more severe effect,
+		// so we'll just re-use the suspension data structure to store the lock state.
+		let suspension = self
+			.db
+			.userid_lock
+			.get(user_id)
+			.await
+			.deserialized::<UserSuspension>()
+			.unwrap_or_else(|_| UserSuspension {
+				suspended: true,
+				suspended_at: MilliSecondsSinceUnixEpoch::now().get().into(),
+				suspended_by: locking_user.to_string(),
+			});
+
+		self.db.userid_lock.raw_put(user_id, Json(suspension));
+	}
+
+	pub async fn unlock_account(&self, user_id: &UserId) { self.db.userid_lock.remove(user_id); }
+
 	/// Check if a user has an account on this homeserver.
 	#[inline]
 	pub async fn exists(&self, user_id: &UserId) -> bool {
@@ -255,6 +277,24 @@ impl Service {
 		}
 	}
 
+	pub async fn is_locked(&self, user_id: &UserId) -> Result<bool> {
+		match self
+			.db
+			.userid_lock
+			.get(user_id)
+			.await
+			.deserialized::<UserSuspension>()
+		{
+			| Ok(s) => Ok(s.suspended),
+			| Err(e) =>
+				if e.is_not_found() {
+					Ok(false)
+				} else {
+					Err(e)
+				},
+		}
+	}
+
 	/// Check if account is active, infallible
 	pub async fn is_active(&self, user_id: &UserId) -> bool {
 		!self.is_deactivated(user_id).await.unwrap_or(true)