add resolve_with_servers() to alias service; simplify api

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

src/admin/user/commands.rs

@@ -381,13 +381,18 @@ pub(super) async fn force_join_room(
 	&self, user_id: String, room_id: OwnedRoomOrAliasId,
 ) -> Result<RoomMessageEventContent> {
 	let user_id = parse_local_user_id(self.services, &user_id)?;
-	let room_id = self.services.rooms.alias.resolve(&room_id).await?;
+	let (room_id, servers) = self
+		.services
+		.rooms
+		.alias
+		.resolve_with_servers(&room_id, None)
+		.await?;
 
 	assert!(
 		self.services.globals.user_is_local(&user_id),
 		"Parsed user_id must be a local user"
 	);
-	join_room_by_id_helper(self.services, &user_id, &room_id, None, &[], None, &None).await?;
+	join_room_by_id_helper(self.services, &user_id, &room_id, None, &servers, None, &None).await?;
 
 	Ok(RoomMessageEventContent::notice_markdown(format!(
 		"{user_id} has been joined to {room_id}.",

src/api/client/alias.rs

@@ -86,25 +86,19 @@ pub(crate) async fn get_alias_route(
 	State(services): State<crate::State>, body: Ruma<get_alias::v3::Request>,
 ) -> Result<get_alias::v3::Response> {
 	let room_alias = body.body.room_alias;
-	let servers = None;
 
-	let Ok((room_id, pre_servers)) = services
-		.rooms
-		.alias
-		.resolve_alias(&room_alias, servers.as_ref())
-		.await
-	else {
+	let Ok((room_id, servers)) = services.rooms.alias.resolve_alias(&room_alias, None).await else {
 		return Err!(Request(NotFound("Room with alias not found.")));
 	};
 
-	let servers = room_available_servers(&services, &room_id, &room_alias, &pre_servers).await;
+	let servers = room_available_servers(&services, &room_id, &room_alias, servers).await;
 	debug!(?room_alias, ?room_id, "available servers: {servers:?}");
 
 	Ok(get_alias::v3::Response::new(room_id, servers))
 }
 
 async fn room_available_servers(
-	services: &Services, room_id: &RoomId, room_alias: &RoomAliasId, pre_servers: &Option<Vec<OwnedServerName>>,
+	services: &Services, room_id: &RoomId, room_alias: &RoomAliasId, pre_servers: Vec<OwnedServerName>,
 ) -> Vec<OwnedServerName> {
 	// find active servers in room state cache to suggest
 	let mut servers: Vec<OwnedServerName> = services
@@ -117,9 +111,7 @@ async fn room_available_servers(
 
 	// push any servers we want in the list already (e.g. responded remote alias
 	// servers, room alias server itself)
-	if let Some(pre_servers) = pre_servers {
-		servers.extend(pre_servers.clone());
-	};
+	servers.extend(pre_servers);
 
 	servers.sort_unstable();
 	servers.dedup();

src/api/client/membership.rs

@@ -9,8 +9,9 @@ use axum_client_ip::InsecureClientIp;
 use conduit::{
 	debug, debug_info, debug_warn, err, error, info, pdu,
 	pdu::{gen_event_id_canonical_json, PduBuilder},
+	result::FlatOk,
 	trace, utils,
-	utils::{IterStream, ReadyExt},
+	utils::{shuffle, IterStream, ReadyExt},
 	warn, Err, Error, PduEvent, Result,
 };
 use futures::{FutureExt, StreamExt};
@@ -188,6 +189,10 @@ pub(crate) async fn join_room_by_id_route(
 		servers.push(server.into());
 	}
 
+	servers.sort_unstable();
+	servers.dedup();
+	shuffle(&mut servers);
+
 	join_room_by_id_helper(
 		&services,
 		sender_user,
@@ -251,45 +256,48 @@ pub(crate) async fn join_room_by_id_or_alias_route(
 				servers.push(server.to_owned());
 			}
 
+			servers.sort_unstable();
+			servers.dedup();
+			shuffle(&mut servers);
+
 			(servers, room_id)
 		},
 		Err(room_alias) => {
-			let response = services
+			let (room_id, mut servers) = services
 				.rooms
 				.alias
-				.resolve_alias(&room_alias, Some(&body.via.clone()))
+				.resolve_alias(&room_alias, Some(body.via.clone()))
 				.await?;
-			let (room_id, mut pre_servers) = response;
 
 			banned_room_check(&services, sender_user, Some(&room_id), Some(room_alias.server_name()), client).await?;
 
-			let mut servers = body.via;
-			if let Some(pre_servers) = &mut pre_servers {
-				servers.append(pre_servers);
-			}
+			let addl_via_servers = services
+				.rooms
+				.state_cache
+				.servers_invite_via(&room_id)
+				.map(ToOwned::to_owned);
 
-			servers.extend(
-				services
-					.rooms
-					.state_cache
-					.servers_invite_via(&room_id)
-					.map(ToOwned::to_owned)
-					.collect::<Vec<_>>()
-					.await,
-			);
+			let addl_state_servers = services
+				.rooms
+				.state_cache
+				.invite_state(sender_user, &room_id)
+				.await
+				.unwrap_or_default();
 
-			servers.extend(
-				services
-					.rooms
-					.state_cache
-					.invite_state(sender_user, &room_id)
-					.await
-					.unwrap_or_default()
-					.iter()
-					.filter_map(|event| event.get_field("sender").ok().flatten())
-					.filter_map(|sender: &str| UserId::parse(sender).ok())
-					.map(|user| user.server_name().to_owned()),
-			);
+			let mut addl_servers: Vec<_> = addl_state_servers
+				.iter()
+				.map(|event| event.get_field("sender"))
+				.filter_map(FlatOk::flat_ok)
+				.map(|user: &UserId| user.server_name().to_owned())
+				.stream()
+				.chain(addl_via_servers)
+				.collect()
+				.await;
+
+			addl_servers.sort_unstable();
+			addl_servers.dedup();
+			shuffle(&mut addl_servers);
+			servers.append(&mut addl_servers);
 
 			(servers, room_id)
 		},

src/service/rooms/alias/mod.rs

@@ -112,40 +112,51 @@ impl Service {
 		Ok(())
 	}
 
+	#[inline]
 	pub async fn resolve(&self, room: &RoomOrAliasId) -> Result<OwnedRoomId> {
+		self.resolve_with_servers(room, None)
+			.await
+			.map(|(room_id, _)| room_id)
+	}
+
+	pub async fn resolve_with_servers(
+		&self, room: &RoomOrAliasId, servers: Option<Vec<OwnedServerName>>,
+	) -> Result<(OwnedRoomId, Vec<OwnedServerName>)> {
 		if room.is_room_id() {
-			let room_id: &RoomId = &RoomId::parse(room).expect("valid RoomId");
-			Ok(room_id.to_owned())
+			let room_id = RoomId::parse(room).expect("valid RoomId");
+			Ok((room_id, servers.unwrap_or_default()))
 		} else {
-			let alias: &RoomAliasId = &RoomAliasId::parse(room).expect("valid RoomAliasId");
-			Ok(self.resolve_alias(alias, None).await?.0)
+			let alias = &RoomAliasId::parse(room).expect("valid RoomAliasId");
+			self.resolve_alias(alias, servers).await
 		}
 	}
 
 	#[tracing::instrument(skip(self), name = "resolve")]
 	pub async fn resolve_alias(
-		&self, room_alias: &RoomAliasId, servers: Option<&Vec<OwnedServerName>>,
-	) -> Result<(OwnedRoomId, Option<Vec<OwnedServerName>>)> {
-		if !self
-			.services
-			.globals
-			.server_is_ours(room_alias.server_name())
-			&& (!servers
+		&self, room_alias: &RoomAliasId, servers: Option<Vec<OwnedServerName>>,
+	) -> Result<(OwnedRoomId, Vec<OwnedServerName>)> {
+		let server_name = room_alias.server_name();
+		let server_is_ours = self.services.globals.server_is_ours(server_name);
+		let servers_contains_ours = || {
+			servers
 				.as_ref()
-				.is_some_and(|servers| servers.contains(&self.services.globals.server_name().to_owned()))
-				|| servers.as_ref().is_none())
-		{
-			return self.remote_resolve(room_alias, servers).await;
+				.is_some_and(|servers| servers.contains(&self.services.globals.config.server_name))
+		};
+
+		if !server_is_ours && !servers_contains_ours() {
+			return self
+				.remote_resolve(room_alias, servers.unwrap_or_default())
+				.await;
 		}
 
-		let room_id: Option<OwnedRoomId> = match self.resolve_local_alias(room_alias).await {
+		let room_id = match self.resolve_local_alias(room_alias).await {
 			Ok(r) => Some(r),
 			Err(_) => self.resolve_appservice_alias(room_alias).await?,
 		};
 
 		room_id.map_or_else(
-			|| Err(Error::BadRequest(ErrorKind::NotFound, "Room with alias not found.")),
-			|room_id| Ok((room_id, None)),
+			|| Err!(Request(NotFound("Room with alias not found."))),
+			|room_id| Ok((room_id, Vec::new())),
 		)
 	}
 

src/service/rooms/alias/remote.rs

@@ -1,75 +1,67 @@
-use conduit::{debug, debug_warn, Error, Result};
-use ruma::{
-	api::{client::error::ErrorKind, federation},
-	OwnedRoomId, OwnedServerName, RoomAliasId,
-};
+use std::iter::once;
 
-impl super::Service {
-	pub(super) async fn remote_resolve(
-		&self, room_alias: &RoomAliasId, servers: Option<&Vec<OwnedServerName>>,
-	) -> Result<(OwnedRoomId, Option<Vec<OwnedServerName>>)> {
-		debug!(?room_alias, ?servers, "resolve");
+use conduit::{debug, debug_error, err, implement, Result};
+use federation::query::get_room_information::v1::Response;
+use ruma::{api::federation, OwnedRoomId, OwnedServerName, RoomAliasId, ServerName};
 
-		let mut response = self
-			.services
-			.sending
-			.send_federation_request(
-				room_alias.server_name(),
-				federation::query::get_room_information::v1::Request {
-					room_alias: room_alias.to_owned(),
-				},
-			)
-			.await;
+#[implement(super::Service)]
+pub(super) async fn remote_resolve(
+	&self, room_alias: &RoomAliasId, servers: Vec<OwnedServerName>,
+) -> Result<(OwnedRoomId, Vec<OwnedServerName>)> {
+	debug!(?room_alias, servers = ?servers, "resolve");
+	let servers = once(room_alias.server_name())
+		.map(ToOwned::to_owned)
+		.chain(servers.into_iter());
 
-		debug!("room alias server_name get_alias_helper response: {response:?}");
+	let mut resolved_servers = Vec::new();
+	let mut resolved_room_id: Option<OwnedRoomId> = None;
+	for server in servers {
+		match self.remote_request(room_alias, &server).await {
+			Err(e) => debug_error!("Failed to query for {room_alias:?} from {server}: {e}"),
+			Ok(Response {
+				room_id,
+				servers,
+			}) => {
+				debug!("Server {server} answered with {room_id:?} for {room_alias:?} servers: {servers:?}");
 
-		if let Err(ref e) = response {
-			debug_warn!(
-				"Server {} of the original room alias failed to assist in resolving room alias: {e}",
-				room_alias.server_name(),
-			);
-		}
+				resolved_room_id.get_or_insert(room_id);
+				add_server(&mut resolved_servers, server);
 
-		if response.as_ref().is_ok_and(|resp| resp.servers.is_empty()) || response.as_ref().is_err() {
-			if let Some(servers) = servers {
-				for server in servers {
-					response = self
-						.services
-						.sending
-						.send_federation_request(
-							server,
-							federation::query::get_room_information::v1::Request {
-								room_alias: room_alias.to_owned(),
-							},
-						)
-						.await;
-					debug!("Got response from server {server} for room aliases: {response:?}");
-
-					if let Ok(ref response) = response {
-						if !response.servers.is_empty() {
-							break;
-						}
-						debug_warn!(
-							"Server {server} responded with room aliases, but was empty? Response: {response:?}"
-						);
-					}
+				if !servers.is_empty() {
+					add_servers(&mut resolved_servers, servers);
+					break;
 				}
-			}
+			},
 		}
+	}
 
-		if let Ok(response) = response {
-			let room_id = response.room_id;
+	resolved_room_id
+		.map(|room_id| (room_id, resolved_servers))
+		.ok_or_else(|| err!(Request(NotFound("No servers could assist in resolving the room alias"))))
+}
 
-			let mut pre_servers = response.servers;
-			// since the room alis server responded, insert it into the list
-			pre_servers.push(room_alias.server_name().into());
+#[implement(super::Service)]
+async fn remote_request(&self, room_alias: &RoomAliasId, server: &ServerName) -> Result<Response> {
+	use federation::query::get_room_information::v1::Request;
 
-			return Ok((room_id, Some(pre_servers)));
-		}
+	let request = Request {
+		room_alias: room_alias.to_owned(),
+	};
 
-		Err(Error::BadRequest(
-			ErrorKind::NotFound,
-			"No servers could assist in resolving the room alias",
-		))
+	self.services
+		.sending
+		.send_federation_request(server, request)
+		.await
+}
+
+fn add_servers(servers: &mut Vec<OwnedServerName>, new: Vec<OwnedServerName>) {
+	for server in new {
+		add_server(servers, server);
+	}
+}
+
+fn add_server(servers: &mut Vec<OwnedServerName>, server: OwnedServerName) {
+	if !servers.contains(&server) {
+		servers.push(server);
 	}
 }