From 0d9b821953268edeeea8b1c01fe0d33447ba08cb Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 27 Oct 2020 22:53:20 +0000 Subject: [PATCH] Add support for filtering by user ids when requesting guild chunks (#1040) --- src/client/bridge/gateway/mod.rs | 2 +- src/client/bridge/gateway/shard_messenger.rs | 12 ++++---- src/client/bridge/gateway/shard_runner.rs | 4 +-- .../bridge/gateway/shard_runner_message.rs | 25 +++++++++++------ src/gateway/shard.rs | 12 ++++---- src/gateway/ws_client_ext.rs | 28 +++++++++++++------ 6 files changed, 52 insertions(+), 31 deletions(-) diff --git a/src/client/bridge/gateway/mod.rs b/src/client/bridge/gateway/mod.rs index 973b942e5b1..ff0ea5123e0 100644 --- a/src/client/bridge/gateway/mod.rs +++ b/src/client/bridge/gateway/mod.rs @@ -63,7 +63,7 @@ pub use self::shard_manager_monitor::{ShardManagerMonitor, ShardManagerError}; pub use self::shard_messenger::ShardMessenger; pub use self::shard_queuer::ShardQueuer; pub use self::shard_runner::{ShardRunner, ShardRunnerOptions}; -pub use self::shard_runner_message::ShardRunnerMessage; +pub use self::shard_runner_message::{ShardRunnerMessage, ChunkGuildFilter}; pub use self::intents::GatewayIntents; use std::{ diff --git a/src/client/bridge/gateway/shard_messenger.rs b/src/client/bridge/gateway/shard_messenger.rs index b9cc12e8532..a5f81df1d4a 100644 --- a/src/client/bridge/gateway/shard_messenger.rs +++ b/src/client/bridge/gateway/shard_messenger.rs @@ -1,6 +1,6 @@ use crate::gateway::InterMessage; use crate::model::prelude::*; -use super::{ShardClientMessage, ShardRunnerMessage}; +use super::{ShardClientMessage, ShardRunnerMessage, ChunkGuildFilter}; use futures::channel::mpsc::{UnboundedSender as Sender, TrySendError}; use async_tungstenite::tungstenite::Message; #[cfg(feature = "collector")] @@ -53,6 +53,7 @@ impl ShardMessenger { /// /// ```rust,no_run /// # use tokio::sync::Mutex; + /// # use serenity::client::bridge::gateway::ChunkGuildFilter; /// # use serenity::gateway::Shard; /// # use std::sync::Arc; /// # @@ -65,7 +66,7 @@ impl ShardMessenger { /// /// let guild_ids = vec![GuildId(81384788765712384)]; /// - /// shard.chunk_guilds(guild_ids, Some(2000), None, None); + /// shard.chunk_guilds(guild_ids, Some(2000), ChunkGuildFilter::None, None); /// # Ok(()) /// # } /// ``` @@ -75,6 +76,7 @@ impl ShardMessenger { /// /// ```rust,no_run /// # use tokio::sync::Mutex; + /// # use serenity::client::bridge::gateway::ChunkGuildFilter; /// # use serenity::gateway::Shard; /// # use std::sync::Arc; /// # @@ -87,7 +89,7 @@ impl ShardMessenger { /// /// let guild_ids = vec![GuildId(81384788765712384)]; /// - /// shard.chunk_guilds(guild_ids, Some(20), Some("do"), Some("request")); + /// shard.chunk_guilds(guild_ids, Some(20), ChunkGuildFilter::Query("do".to_owned()), Some("request")); /// # Ok(()) /// # } /// ``` @@ -99,7 +101,7 @@ impl ShardMessenger { &self, guild_ids: It, limit: Option, - query: Option, + filter: ChunkGuildFilter, nonce: Option, ) where It: IntoIterator { let guilds = guild_ids.into_iter().collect::>(); @@ -107,7 +109,7 @@ impl ShardMessenger { let _ = self.send_to_shard(ShardRunnerMessage::ChunkGuilds { guild_ids: guilds, limit, - query, + filter, nonce, }); } diff --git a/src/client/bridge/gateway/shard_runner.rs b/src/client/bridge/gateway/shard_runner.rs index 54edc899224..d36c83dbf57 100644 --- a/src/client/bridge/gateway/shard_runner.rs +++ b/src/client/bridge/gateway/shard_runner.rs @@ -362,11 +362,11 @@ impl ShardRunner { true }, - ShardClientMessage::Runner(ShardRunnerMessage::ChunkGuilds { guild_ids, limit, query, nonce }) => { + ShardClientMessage::Runner(ShardRunnerMessage::ChunkGuilds { guild_ids, limit, filter, nonce }) => { self.shard.chunk_guilds( guild_ids, limit, - query.as_deref(), + filter, nonce.as_deref(), ).await.is_ok() }, diff --git a/src/client/bridge/gateway/shard_runner_message.rs b/src/client/bridge/gateway/shard_runner_message.rs index acc354d1ce9..1337dc4d995 100644 --- a/src/client/bridge/gateway/shard_runner_message.rs +++ b/src/client/bridge/gateway/shard_runner_message.rs @@ -1,6 +1,6 @@ use crate::model::{ gateway::Activity, - id::GuildId, + id::{GuildId, UserId}, user::OnlineStatus, }; @@ -8,6 +8,16 @@ use crate::model::{ use crate::collector::{MessageFilter, ReactionFilter}; use async_tungstenite::tungstenite::Message; +#[derive(Clone, Debug)] +pub enum ChunkGuildFilter { + /// Returns all members of the guilds specified. Requires GUILD_MEMBERS intent. + None, + /// A common username prefix filter for the members returned. + Query(String), + /// A set of exact user IDs to query for. + UserIds(Vec), +} + /// A message to send from a shard over a WebSocket. // Once we can use `Box` as part of a pattern, we will reconsider boxing. #[allow(clippy::large_enum_variant)] @@ -15,6 +25,8 @@ use async_tungstenite::tungstenite::Message; pub enum ShardRunnerMessage { /// Indicates that the client is to send a member chunk message. ChunkGuilds { + // FIXME: When intents are enabled, chunking guilds is limited to 1 guild per chunk + // request. Will be mandatory when using API v8. /// The IDs of the [`Guild`]s to chunk. /// /// [`Guild`]: ../../../model/guild/struct.Guild.html @@ -24,14 +36,9 @@ pub enum ShardRunnerMessage { /// /// [`GuildMembersChunkEvent`]: ../../../model/event/struct.GuildMembersChunkEvent.html limit: Option, - /// Text to filter members by. - /// - /// For example, a query of `"s"` will cause only [`Member`]s whose - /// usernames start with `"s"` to be chunked. - /// - /// [`Member`]: ../../../model/guild/struct.Member.html - query: Option, - /// Nonce to identify [`GuildMembersChunkEvent`] responses. + /// A filter to apply to the returned members. + filter: ChunkGuildFilter, + /// Optional nonce to identify [`GuildMembersChunkEvent`] responses. /// /// [`GuildMembersChunkEvent`]: ../../../model/event/struct.GuildMembersChunkEvent.html nonce: Option, diff --git a/src/gateway/shard.rs b/src/gateway/shard.rs index 7c89e7568c1..5394eca7d76 100644 --- a/src/gateway/shard.rs +++ b/src/gateway/shard.rs @@ -7,7 +7,7 @@ use crate::model::{ user::OnlineStatus }; use tokio::sync::Mutex; -use crate::client::bridge::gateway::GatewayIntents; +use crate::client::bridge::gateway::{GatewayIntents, ChunkGuildFilter}; use std::{ sync::Arc, time::{Duration as StdDuration, Instant} @@ -685,6 +685,7 @@ impl Shard { /// /// ```rust,no_run /// # use tokio::sync::Mutex; + /// # use serenity::client::bridge::gateway::ChunkGuildFilter; /// # use serenity::gateway::Shard; /// # use std::sync::Arc; /// # @@ -697,7 +698,7 @@ impl Shard { /// /// let guild_ids = vec![GuildId(81384788765712384)]; /// - /// shard.chunk_guilds(guild_ids, Some(2000), None, None).await?; + /// shard.chunk_guilds(guild_ids, Some(2000), ChunkGuildFilter::None, None).await?; /// # Ok(()) /// # } /// ``` @@ -708,6 +709,7 @@ impl Shard { /// ```rust,no_run /// # use tokio::sync::Mutex; /// # use serenity::gateway::Shard; + /// # use serenity::client::bridge::gateway::ChunkGuildFilter; /// # use std::error::Error; /// # use std::sync::Arc; /// # @@ -720,7 +722,7 @@ impl Shard { /// /// let guild_ids = vec![GuildId(81384788765712384)]; /// - /// shard.chunk_guilds(guild_ids, Some(20), Some("do"), Some("request")).await?; + /// shard.chunk_guilds(guild_ids, Some(20), ChunkGuildFilter::Query("do".to_owned()), Some("request")).await?; /// # Ok(()) /// # } /// ``` @@ -733,7 +735,7 @@ impl Shard { &mut self, guild_ids: It, limit: Option, - query: Option<&str>, + filter: ChunkGuildFilter, nonce: Option<&str>, ) -> Result<()> where It: IntoIterator + Send { debug!("[Shard {:?}] Requesting member chunks", self.shard_info); @@ -742,7 +744,7 @@ impl Shard { guild_ids, &self.shard_info, limit, - query, + filter, nonce, ).await } diff --git a/src/gateway/ws_client_ext.rs b/src/gateway/ws_client_ext.rs index a887c74c2a5..c065f55ebba 100644 --- a/src/gateway/ws_client_ext.rs +++ b/src/gateway/ws_client_ext.rs @@ -1,15 +1,15 @@ +use crate::client::bridge::gateway::{ChunkGuildFilter, GatewayIntents}; use crate::constants::{self, OpCode}; use crate::gateway::{CurrentPresence, WsStream}; -use crate::client::bridge::gateway::GatewayIntents; use crate::internal::prelude::*; use crate::internal::ws_impl::SenderExt; use crate::model::id::GuildId; +use async_trait::async_trait; use serde_json::json; use std::env::consts; use std::time::SystemTime; -use tracing::{debug, trace}; -use async_trait::async_trait; use tracing::instrument; +use tracing::{debug, trace}; #[async_trait] pub trait WebSocketGatewayClientExt { @@ -18,7 +18,7 @@ pub trait WebSocketGatewayClientExt { guild_ids: It, shard_info: &[u64; 2], limit: Option, - query: Option<&str>, + filter: ChunkGuildFilter, nonce: Option<&str>, ) -> Result<()> where It: IntoIterator + Send; @@ -51,20 +51,30 @@ impl WebSocketGatewayClientExt for WsStream { guild_ids: It, shard_info: &[u64; 2], limit: Option, - query: Option<&str>, + filter: ChunkGuildFilter, nonce: Option<&str>, - ) -> Result<()> where It: IntoIterator + Send { + ) -> Result<()> where It: IntoIterator + Send, { debug!("[Shard {:?}] Requesting member chunks", shard_info); - self.send_json(&json!({ + let mut payload = json!({ "op": OpCode::GetGuildMembers.num(), "d": { "guild_id": guild_ids.into_iter().map(|x| x.as_ref().0).collect::>(), "limit": limit.unwrap_or(0), - "query": query.unwrap_or(""), "nonce": nonce.unwrap_or(""), }, - })).await.map_err(From::from) + }); + + match filter { + ChunkGuildFilter::None => {}, + ChunkGuildFilter::Query(query) => payload["d"]["query"] = json!(query), + ChunkGuildFilter::UserIds(user_ids) => { + let ids = user_ids.iter().map(|x| x.as_ref().0).collect::>(); + payload["d"]["user_ids"] = json!(ids) + }, + }; + + self.send_json(&payload).await.map_err(From::from) } #[instrument(skip(self))]