change player subscription handling a bit

Author: Wyvest

src/api/websocket/endpoint.rs

@@ -33,8 +33,8 @@ use crate::api::{
 
 /// Max UUIDs in a single `SubscribePlayers` or `GetActiveCosmetics` message.
 const MAX_PLAYERS_PER_REQUEST: usize = 64;
-/// Max distinct players a connection may subscribe to at once (render distance).
-const MAX_PLAYER_SUBSCRIPTIONS: usize = 128;
+const MAX_PLAYER_SUBSCRIPTIONS: usize = 512;
+const REAL_PLAYER_UUID_VERSION: usize = 4;
 
 fn enforce_max_players_per_request(players: &[Uuid]) -> Result<(), WebsocketError> {
 	if players.len() > MAX_PLAYERS_PER_REQUEST {
@@ -45,6 +45,21 @@ fn enforce_max_players_per_request(players: &[Uuid]) -> Result<(), WebsocketErro
 	Ok(())
 }
 
+fn is_fake_player(player: &Uuid) -> bool {
+	player.get_version_num() != REAL_PLAYER_UUID_VERSION
+}
+
+fn empty_snapshot(rejected: Vec<Uuid>, request_id: Option<u64>) -> ClientBoundPacket {
+	ClientBoundPacket::SubscriptionSnapshot {
+		equipped: HashMap::new(),
+		active_emotes: HashMap::new(),
+		particle_colors: HashMap::new(),
+		users: Vec::new(),
+		rejected,
+		request_id,
+	}
+}
+
 pub(super) fn router() -> ApiRouter<ApiState> {
 	ApiRouter::new()
 		.route("/websocket", get(self::endpoint))
@@ -377,42 +392,49 @@ async fn subscribe(
 	state: &ApiState,
 	connection_id: ConnectionId,
 	players: Vec<Uuid>,
+	request_id: Option<u64>,
 ) -> Result<ClientBoundPacket, WebsocketError> {
-	let requested = players.into_iter().collect::<HashSet<_>>();
-	let newly_subscribed = {
+	let mut seen = HashSet::new();
+	let requested = players
+		.into_iter()
+		.filter(|player| seen.insert(*player))
+		.collect::<Vec<_>>();
+
+	let (newly_subscribed, rejected) = {
 		let mut connections = state.realtime.connections.write().await;
 		let Some(connection) = connections.get_mut(&connection_id) else {
-			return Ok(ClientBoundPacket::SubscriptionSnapshot {
-				equipped: HashMap::new(),
-				active_emotes: HashMap::new(),
-				particle_colors: HashMap::new(),
-				users: Vec::new(),
-			});
+			return Ok(empty_snapshot(Vec::new(), request_id));
 		};
 
-		let pending = requested
-			.iter()
-			.filter(|player| !connection.subscriptions.contains(player))
-			.count();
-		if connection.subscriptions.len() + pending > MAX_PLAYER_SUBSCRIPTIONS {
-			return Err(WebsocketError::SubscriptionLimitExceeded {
-				limit: MAX_PLAYER_SUBSCRIPTIONS,
-			});
+		let mut newly_subscribed = Vec::new();
+		let mut rejected = Vec::new();
+		for player in requested {
+			if is_fake_player(&player) {
+				rejected.push(player);
+				continue;
+			}
+			if connection.subscriptions.contains(&player) {
+				continue;
+			}
+			if connection.subscriptions.len() >= MAX_PLAYER_SUBSCRIPTIONS {
+				rejected.push(player);
+				continue;
+			}
+			connection.subscriptions.insert(player);
+			newly_subscribed.push(player);
 		}
-
-		requested
-			.into_iter()
-			.filter(|player| connection.subscriptions.insert(*player))
-			.collect::<Vec<_>>()
+		(newly_subscribed, rejected)
 	};
 
+	if !rejected.is_empty() {
+		warn!(
+			"Connection {connection_id} rejected {} subscription(s) (cap {MAX_PLAYER_SUBSCRIPTIONS})",
+			rejected.len()
+		);
+	}
+
 	if newly_subscribed.is_empty() {
-		return Ok(ClientBoundPacket::SubscriptionSnapshot {
-			equipped: HashMap::new(),
-			active_emotes: HashMap::new(),
-			particle_colors: HashMap::new(),
-			users: Vec::new(),
-		});
+		return Ok(empty_snapshot(rejected, request_id));
 	}
 
 	{
@@ -480,6 +502,8 @@ async fn subscribe(
 		active_emotes,
 		particle_colors,
 		users,
+		rejected,
+		request_id,
 	})
 }
 
@@ -530,14 +554,28 @@ async fn broadcast_to_watchers(
 	}
 }
 
+struct RequestError {
+	error: WebsocketError,
+	request_id: Option<u64>,
+}
+
+impl From<WebsocketError> for RequestError {
+	fn from(error: WebsocketError) -> Self {
+		Self {
+			error,
+			request_id: None,
+		}
+	}
+}
+
 async fn handle_msg(
 	socket: &mut WebSocket,
 	state: &ApiState,
 	player: &entities::user::Model,
 	connection_id: ConnectionId,
 	msg: Result<Message, axum::Error>,
-) -> Result<(), WebsocketError> {
-	let msg = msg?;
+) -> Result<(), RequestError> {
+	let msg = msg.map_err(WebsocketError::from)?;
 
 	// Ignore control/keepalive frames. Ping/Pong carry an opaque payload (Ktor
 	// sends a Ping every pingInterval) that is not a serializable request, and
@@ -548,7 +586,20 @@ async fn handle_msg(
 
 	let parsed = serde_json::from_slice::<ServerBoundPacket>(&msg.into_data())
 		.map_err(WebsocketError::Deserialization)?;
+	let request_id = parsed.request_id();
 
+	handle_packet(socket, state, player, connection_id, parsed)
+		.await
+		.map_err(|error| RequestError { error, request_id })
+}
+
+async fn handle_packet(
+	socket: &mut WebSocket,
+	state: &ApiState,
+	player: &entities::user::Model,
+	connection_id: ConnectionId,
+	parsed: ServerBoundPacket,
+) -> Result<(), WebsocketError> {
 	match parsed {
 		ServerBoundPacket::GetActiveCosmetics { players } => {
 			enforce_max_players_per_request(&players)?;
@@ -560,9 +611,12 @@ async fn handle_msg(
 			)
 			.await?;
 		}
-		ServerBoundPacket::SubscribePlayers { players } => {
+		ServerBoundPacket::SubscribePlayers {
+			players,
+			request_id,
+		} => {
 			enforce_max_players_per_request(&players)?;
-			let snapshot = subscribe(state, connection_id, players).await?;
+			let snapshot = subscribe(state, connection_id, players, request_id).await?;
 			send_packet(socket, snapshot).await?;
 		}
 		ServerBoundPacket::UnsubscribePlayers { players } => {
@@ -670,8 +724,14 @@ async fn endpoint(
 		let equipped = match load_equipped(&state, player.id).await {
 			Ok(equipped) => equipped,
 			Err(error) => {
-				let _ =
-					send_packet(&mut socket, ClientBoundPacket::Error { error }).await;
+				let _ = send_packet(
+					&mut socket,
+					ClientBoundPacket::Error {
+						error,
+						request_id: None,
+					},
+				)
+				.await;
 				return;
 			}
 		};
@@ -697,15 +757,18 @@ async fn endpoint(
 					let Some(packet) = packet else {
 						break;
 					};
-					send_packet(&mut socket, packet).await
+					send_packet(&mut socket, packet).await.map_err(RequestError::from)
 				}
 			};
 
 			match result {
 				Ok(_) => continue,
-				Err(WebsocketError::Fatal(_)) => break,
-				Err(e) => {
-					let e = ClientBoundPacket::Error { error: e };
+				Err(RequestError {
+					error: WebsocketError::Fatal(_),
+					..
+				}) => break,
+				Err(RequestError { error, request_id }) => {
+					let e = ClientBoundPacket::Error { error, request_id };
 					if send_packet(&mut socket, e).await.is_err() {
 						break;
 					};

src/api/websocket/structs.rs

@@ -23,8 +23,6 @@ pub enum WebsocketError {
 	UnownedEmote(i32),
 	#[error("Too many players in one request (max {limit})")]
 	TooManyPlayersInRequest { limit: usize },
-	#[error("Too many player subscriptions (max {limit})")]
-	SubscriptionLimitExceeded { limit: usize },
 }
 
 impl WebsocketError {
@@ -37,8 +35,7 @@ impl WebsocketError {
 			Self::DatabaseQuery(_) | Self::Serialization(_) => Self::ERROR_CODES[1],
 			Self::Deserialization(_)
 			| Self::InvalidSlot { .. }
-			| Self::TooManyPlayersInRequest { .. }
-			| Self::SubscriptionLimitExceeded { .. } => Self::ERROR_CODES[2],
+			| Self::TooManyPlayersInRequest { .. } => Self::ERROR_CODES[2],
 			Self::UnownedCosmetic(_) | Self::UnownedEmote(_) => Self::ERROR_CODES[3],
 		}
 	}
@@ -92,6 +89,8 @@ pub enum ServerBoundPacket {
 	SubscribePlayers {
 		/// Player UUIDs to watch. Capped per request and per connection total.
 		players: Vec<Uuid>,
+		#[serde(default)]
+		request_id: Option<u64>,
 	},
 	UnsubscribePlayers {
 		players: Vec<Uuid>,
@@ -133,6 +132,8 @@ pub enum ClientBoundPacket {
 		/// The subset of subscribed players that currently have a live PolyPlus
 		/// session connected. Used to render a "uses PolyPlus" indicator.
 		users: Vec<Uuid>,
+		rejected: Vec<Uuid>,
+		request_id: Option<u64>,
 	},
 	/// A subscribed player's PolyPlus session came online or went offline.
 	PlayerPresence {
@@ -166,9 +167,19 @@ pub enum ClientBoundPacket {
 	Error {
 		#[serde(flatten)]
 		error: WebsocketError,
+		request_id: Option<u64>,
 	},
 }
 
+impl ServerBoundPacket {
+	pub fn request_id(&self) -> Option<u64> {
+		match self {
+			Self::SubscribePlayers { request_id, .. } => *request_id,
+			_ => None,
+		}
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use std::collections::HashMap;
@@ -203,13 +214,28 @@ mod tests {
 		.expect("packet should parse");
 
 		match packet {
-			ServerBoundPacket::SubscribePlayers { players } => {
+			ServerBoundPacket::SubscribePlayers {
+				players,
+				request_id,
+			} => {
 				assert_eq!(players, vec![player]);
+				assert_eq!(request_id, None);
 			}
 			_ => panic!("unexpected packet variant"),
 		}
 	}
 
+	#[test]
+	fn parses_player_subscription_with_request_id() {
+		let player = Uuid::nil();
+		let packet: ServerBoundPacket = serde_json::from_str(&format!(
+			r#"{{"type":"SubscribePlayers","players":["{player}"],"request_id":7}}"#
+		))
+		.expect("packet should parse");
+
+		assert_eq!(packet.request_id(), Some(7));
+	}
+
 	#[test]
 	fn serializes_subscription_snapshot_packet() {
 		let player = Uuid::nil();
@@ -218,12 +244,29 @@ mod tests {
 			active_emotes: HashMap::from([(player, 6)]),
 			particle_colors: HashMap::from([(player, 0xFF_0000)]),
 			users: vec![player],
+			rejected: Vec::new(),
+			request_id: Some(3),
 		};
 
 		let serialized = serde_json::to_value(packet).expect("packet should serialize");
 		assert_eq!(serialized["type"], "SubscriptionSnapshot");
 		assert_eq!(serialized["equipped"][player.to_string()]["cape"], 1);
 		assert_eq!(serialized["users"][0], player.to_string());
+		assert_eq!(serialized["rejected"].as_array().map(Vec::len), Some(0));
+		assert_eq!(serialized["request_id"], 3);
+	}
+
+	#[test]
+	fn serializes_error_packet_with_request_id() {
+		let packet = ClientBoundPacket::Error {
+			error: super::WebsocketError::TooManyPlayersInRequest { limit: 64 },
+			request_id: Some(9),
+		};
+
+		let serialized = serde_json::to_value(packet).expect("packet should serialize");
+		assert_eq!(serialized["type"], "Error");
+		assert_eq!(serialized["error_code"], "bad_request");
+		assert_eq!(serialized["request_id"], 9);
 	}
 
 	#[test]