refactor: address code review feedback for auth and pagination improvements

- Extract duplicated thread/message pagination functions to shared utils module
- Refactor pagination token parsing to use Option combinators instead of defaults
- Split access token generation from assignment for clearer error handling
- Add appservice token collision detection at startup and registration
- Allow appservice re-registration with same token (for config updates)
- Simplify thread relation chunk building using iterator chaining
- Fix saturating_inc edge case in relation queries with explicit filtering
- Add concise comments explaining non-obvious behaviour choices

src/service/appservice/mod.rs

@@ -4,7 +4,7 @@ mod registration_info;
 use std::{collections::BTreeMap, iter::IntoIterator, sync::Arc};
 
 use async_trait::async_trait;
-use conduwuit::{Result, err, utils::stream::IterStream};
+use conduwuit::{Err, Result, err, utils::stream::IterStream};
 use database::Map;
 use futures::{Future, FutureExt, Stream, TryStreamExt};
 use ruma::{RoomAliasId, RoomId, UserId, api::appservice::Registration};
@@ -48,36 +48,50 @@ impl crate::Service for Service {
 	}
 
 	async fn worker(self: Arc<Self>) -> Result {
-		self.iter_db_ids()
-			.try_for_each(async |appservice| {
-				let (id, registration) = appservice;
+		// First, collect all appservices to check for token conflicts
+		let appservices: Vec<(String, Registration)> = self.iter_db_ids().try_collect().await?;
 
-				// During startup, resolve any token collisions in favour of appservices
-				// by logging out conflicting user devices
-				if let Ok((user_id, device_id)) = self
-					.services
-					.users
-					.find_from_token(&registration.as_token)
-					.await
-				{
-					conduwuit::warn!(
-						"Token collision detected during startup: Appservice '{}' token was \
-						 also used by user '{}' device '{}'. Logging out the user device to \
-						 resolve conflict.",
-						id,
-						user_id.localpart(),
-						device_id
-					);
-
-					self.services
-						.users
-						.remove_device(&user_id, &device_id)
-						.await;
+		// Check for appservice-to-appservice token conflicts
+		for i in 0..appservices.len() {
+			for j in i.saturating_add(1)..appservices.len() {
+				if appservices[i].1.as_token == appservices[j].1.as_token {
+					return Err!(Database(error!(
+						"Token collision detected: Appservices '{}' and '{}' have the same token",
+						appservices[i].0, appservices[j].0
+					)));
 				}
+			}
+		}
 
-				self.start_appservice(id, registration).await
-			})
-			.await
+		// Process each appservice
+		for (id, registration) in appservices {
+			// During startup, resolve any token collisions in favour of appservices
+			// by logging out conflicting user devices
+			if let Ok((user_id, device_id)) = self
+				.services
+				.users
+				.find_from_token(&registration.as_token)
+				.await
+			{
+				conduwuit::warn!(
+					"Token collision detected during startup: Appservice '{}' token was also \
+					 used by user '{}' device '{}'. Logging out the user device to resolve \
+					 conflict.",
+					id,
+					user_id.localpart(),
+					device_id
+				);
+
+				self.services
+					.users
+					.remove_device(&user_id, &device_id)
+					.await;
+			}
+
+			self.start_appservice(id, registration).await?;
+		}
+
+		Ok(())
 	}
 
 	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
@@ -125,6 +139,18 @@ impl Service {
 	) -> Result {
 		//TODO: Check for collisions between exclusive appservice namespaces
 
+		// Check for token collision with other appservices (allow re-registration of
+		// same appservice)
+		if let Ok(existing) = self.find_from_token(&registration.as_token).await {
+			if existing.registration.id != registration.id {
+				return Err(err!(Request(InvalidParam(
+					"Cannot register appservice: Token is already used by appservice '{}'. \
+					 Please generate a different token.",
+					existing.registration.id
+				))));
+			}
+		}
+
 		// Prevent token collision with existing user tokens
 		if self
 			.services
@@ -182,6 +208,7 @@ impl Service {
 			.map(|info| info.registration)
 	}
 
+	/// Returns Result to match users::find_from_token for select_ok usage
 	pub async fn find_from_token(&self, token: &str) -> Result<RegistrationInfo> {
 		self.read()
 			.await

src/service/rooms/pdu_metadata/data.rs

@@ -61,6 +61,8 @@ impl Data {
 		from: PduCount,
 		dir: Direction,
 	) -> impl Stream<Item = (PduCount, impl Event)> + Send + '_ {
+		// Query from exact position then filter excludes it (saturating_inc could skip
+		// events at min/max boundaries)
 		let from_unsigned = from.into_unsigned();
 		let mut current = ArrayVec::<u8, 16>::new();
 		current.extend(target.to_be_bytes());

src/service/users/mod.rs

@@ -393,6 +393,31 @@ impl Service {
 		self.db.userdeviceid_token.qry(&key).await.deserialized()
 	}
 
+	/// Generate a unique access token that doesn't collide with existing tokens
+	pub async fn generate_unique_token(&self) -> String {
+		loop {
+			let token = utils::random_string(32);
+
+			// Check for collision with appservice tokens
+			if self
+				.services
+				.appservice
+				.find_from_token(&token)
+				.await
+				.is_ok()
+			{
+				continue;
+			}
+
+			// Check for collision with user tokens
+			if self.db.token_userdeviceid.get(&token).await.is_ok() {
+				continue;
+			}
+
+			return token;
+		}
+	}
+
 	/// Replaces the access token of one device.
 	pub async fn set_token(
 		&self,
@@ -409,25 +434,18 @@ impl Service {
 			)));
 		}
 
-		// Prevent token collisions with appservice tokens
-		let final_token = if self
+		// Check for token collision with appservices
+		if self
 			.services
 			.appservice
 			.find_from_token(token)
 			.await
 			.is_ok()
 		{
-			let new_token = utils::random_string(32);
-			conduwuit::debug_warn!(
-				"Token collision prevented: Generated new token for user '{}' device '{}' \
-				 (original token conflicted with an appservice)",
-				user_id.localpart(),
-				device_id
-			);
-			new_token
-		} else {
-			token.to_owned()
-		};
+			return Err!(Request(InvalidParam(
+				"Token conflicts with an existing appservice token"
+			)));
+		}
 
 		// Remove old token
 		if let Ok(old_token) = self.db.userdeviceid_token.qry(&key).await {
@@ -436,8 +454,8 @@ impl Service {
 		}
 
 		// Assign token to user device combination
-		self.db.userdeviceid_token.put_raw(key, &final_token);
-		self.db.token_userdeviceid.raw_put(&final_token, key);
+		self.db.userdeviceid_token.put_raw(key, token);
+		self.db.token_userdeviceid.raw_put(token, key);
 
 		Ok(())
 	}