From 8a767c4b1087f48edbfdb8a48c33b364a7637ec2 Mon Sep 17 00:00:00 2001
From: strawberry <strawberry@puppygock.gay>
Date: Sun, 21 Apr 2024 23:36:20 -0400
Subject: [PATCH] on room alias joins, attempt to find the room ID through
 *more* servers if available

Signed-off-by: strawberry <strawberry@puppygock.gay>
---
 src/api/client_server/alias.rs                | 130 ++++++++++++------
 src/api/client_server/membership.rs           |   2 +-
 .../admin/room/room_moderation_commands.rs    |   6 +-
 3 files changed, 91 insertions(+), 47 deletions(-)

diff --git a/src/api/client_server/alias.rs b/src/api/client_server/alias.rs
index 3fafee72..1345ef0c 100644
--- a/src/api/client_server/alias.rs
+++ b/src/api/client_server/alias.rs
@@ -10,8 +10,9 @@ use ruma::{
 	},
 	OwnedRoomAliasId, OwnedServerName,
 };
+use tracing::{debug, info};
 
-use crate::{services, Error, Result, Ruma};
+use crate::{debug_info, services, Error, Result, Ruma};
 
 /// # `PUT /_matrix/client/v3/directory/room/{roomAlias}`
 ///
@@ -118,12 +119,19 @@ pub async fn delete_alias_route(body: Ruma<delete_alias::v3::Request>) -> Result
 ///
 /// Resolve an alias locally or over federation.
 pub async fn get_alias_route(body: Ruma<get_alias::v3::Request>) -> Result<get_alias::v3::Response> {
-	get_alias_helper(body.body.room_alias).await
+	get_alias_helper(body.body.room_alias, None).await
 }
 
-pub(crate) async fn get_alias_helper(room_alias: OwnedRoomAliasId) -> Result<get_alias::v3::Response> {
-	if room_alias.server_name() != services().globals.server_name() {
-		let response = services()
+pub(crate) async fn get_alias_helper(
+	room_alias: OwnedRoomAliasId, servers: Option<Vec<OwnedServerName>>,
+) -> Result<get_alias::v3::Response> {
+	if room_alias.server_name() != services().globals.server_name()
+		&& (!servers
+			.as_ref()
+			.is_some_and(|servers| servers.contains(&services().globals.server_name().to_owned()))
+			|| servers.as_ref().is_none())
+	{
+		let mut response = services()
 			.sending
 			.send_federation_request(
 				room_alias.server_name(),
@@ -131,47 +139,83 @@ pub(crate) async fn get_alias_helper(room_alias: OwnedRoomAliasId) -> Result<get
 					room_alias: room_alias.clone(),
 				},
 			)
-			.await?;
+			.await;
 
-		let room_id = response.room_id;
-
-		let mut servers = response.servers;
-
-		// since the room alias server_name responded, insert it into the list
-		servers.push(room_alias.server_name().into());
-
-		// find active servers in room state cache to suggest
-		servers.extend(
-			services()
-				.rooms
-				.state_cache
-				.room_servers(&room_id)
-				.filter_map(Result::ok),
-		);
-
-		servers.sort_unstable();
-		servers.dedup();
-
-		// shuffle list of servers randomly after sort and dedupe
-		servers.shuffle(&mut rand::thread_rng());
-
-		// prefer the very first server to be ourselves if available, else prefer the
-		// room alias server first
-		if let Some(server_index) = servers
-			.iter()
-			.position(|server| server == services().globals.server_name())
-		{
-			servers.remove(server_index);
-			servers.insert(0, services().globals.server_name().to_owned());
-		} else if let Some(alias_server_index) = servers
-			.iter()
-			.position(|server| server == room_alias.server_name())
-		{
-			servers.remove(alias_server_index);
-			servers.insert(0, room_alias.server_name().into());
+		if let Err(ref e) = response {
+			debug_info!(
+				"Server {} of the original room alias failed to assist in resolving room alias: {e}",
+				room_alias.server_name()
+			);
 		}
 
-		return Ok(get_alias::v3::Response::new(room_id, servers));
+		if let Some(servers) = servers {
+			if !servers.is_empty() {
+				for server in servers {
+					response = services()
+						.sending
+						.send_federation_request(
+							&server,
+							federation::query::get_room_information::v1::Request {
+								room_alias: room_alias.clone(),
+							},
+						)
+						.await;
+
+					if response.is_err() {
+						continue;
+					}
+
+					break;
+				}
+			}
+		}
+
+		if let Ok(response) = response {
+			let room_id = response.room_id;
+
+			let mut servers = response.servers;
+
+			// since the room alias server_name responded, insert it into the list
+			servers.push(room_alias.server_name().into());
+
+			// find active servers in room state cache to suggest
+			servers.extend(
+				services()
+					.rooms
+					.state_cache
+					.room_servers(&room_id)
+					.filter_map(Result::ok),
+			);
+
+			servers.sort_unstable();
+			servers.dedup();
+
+			// shuffle list of servers randomly after sort and dedupe
+			servers.shuffle(&mut rand::thread_rng());
+
+			// prefer the very first server to be ourselves if available, else prefer the
+			// room alias server first
+			if let Some(server_index) = servers
+				.iter()
+				.position(|server| server == services().globals.server_name())
+			{
+				servers.remove(server_index);
+				servers.insert(0, services().globals.server_name().to_owned());
+			} else if let Some(alias_server_index) = servers
+				.iter()
+				.position(|server| server == room_alias.server_name())
+			{
+				servers.remove(alias_server_index);
+				servers.insert(0, room_alias.server_name().into());
+			}
+
+			return Ok(get_alias::v3::Response::new(room_id, servers));
+		}
+
+		return Err(Error::BadRequest(
+			ErrorKind::Unknown,
+			"No servers could assist in resolving the room alias",
+		));
 	}
 
 	let mut room_id = None;
diff --git a/src/api/client_server/membership.rs b/src/api/client_server/membership.rs
index 1c289d66..0c5744c3 100644
--- a/src/api/client_server/membership.rs
+++ b/src/api/client_server/membership.rs
@@ -198,7 +198,7 @@ pub async fn join_room_by_id_or_alias_route(
 				));
 			}
 
-			let response = get_alias_helper(room_alias.clone()).await?;
+			let response = get_alias_helper(room_alias.clone(), Some(body.server_name.clone())).await?;
 
 			if services().rooms.metadata.is_banned(&response.room_id)? && !services().users.is_admin(sender_user)? {
 				return Err(Error::BadRequest(
diff --git a/src/service/admin/room/room_moderation_commands.rs b/src/service/admin/room/room_moderation_commands.rs
index c02d8d16..fe870d87 100644
--- a/src/service/admin/room/room_moderation_commands.rs
+++ b/src/service/admin/room/room_moderation_commands.rs
@@ -71,7 +71,7 @@ pub(crate) async fn process(command: RoomModerationCommand, body: Vec<&str>) ->
 						 federation"
 					);
 
-					match get_alias_helper(room_alias).await {
+					match get_alias_helper(room_alias, None).await {
 						Ok(response) => {
 							debug!("Got federation response fetching room ID for room {room}: {:?}", response);
 							response.room_id
@@ -233,7 +233,7 @@ pub(crate) async fn process(command: RoomModerationCommand, body: Vec<&str>) ->
 												 fetch room ID over federation"
 											);
 
-											match get_alias_helper(room_alias).await {
+											match get_alias_helper(room_alias, None).await {
 												Ok(response) => {
 													debug!(
 														"Got federation response fetching room ID for room {room}: \
@@ -432,7 +432,7 @@ pub(crate) async fn process(command: RoomModerationCommand, body: Vec<&str>) ->
 						 federation"
 					);
 
-					match get_alias_helper(room_alias).await {
+					match get_alias_helper(room_alias, None).await {
 						Ok(response) => {
 							debug!("Got federation response fetching room ID for room {room}: {:?}", response);
 							response.room_id