Add missing fields to GuildMembersChunkEvent and add nonce to `Shar…

…dMessenger::chunk_guilds` (#1018)
serenity-rs · Oct 15, 2020 · 8c30b40 · 8c30b40
1 parent b1187ce
commit 8c30b40
Show file tree

Hide file tree

Showing 8 changed files with 36 additions and 10 deletions.
diff --git a/src/client/bridge/gateway/shard_messenger.rs b/src/client/bridge/gateway/shard_messenger.rs
@@ -65,13 +65,13 @@ impl ShardMessenger {
     ///
     /// let guild_ids = vec![GuildId(81384788765712384)];
     ///
-    /// shard.chunk_guilds(guild_ids, Some(2000), None);
+    /// shard.chunk_guilds(guild_ids, Some(2000), None, None);
     /// #     Ok(())
     /// # }
     /// ```
     ///
-    /// Chunk a single guild by Id, limiting to 20 members, and specifying a
-    /// query parameter of `"do"`:
+    /// Chunk a single guild by Id, limiting to 20 members, specifying a
+    /// query parameter of `"do"` and a nonce of `"request"`:
     ///
     /// ```rust,no_run
     /// # use tokio::sync::Mutex;
@@ -87,7 +87,7 @@ impl ShardMessenger {
     ///
     /// let guild_ids = vec![GuildId(81384788765712384)];
     ///
-    /// shard.chunk_guilds(guild_ids, Some(20), Some("do"));
+    /// shard.chunk_guilds(guild_ids, Some(20), Some("do"), Some("request"));
     /// #     Ok(())
     /// # }
     /// ```
@@ -100,13 +100,15 @@ impl ShardMessenger {
         guild_ids: It,
         limit: Option<u16>,
         query: Option<String>,
+        nonce: Option<String>,
     ) where It: IntoIterator<Item=GuildId> {
         let guilds = guild_ids.into_iter().collect::<Vec<GuildId>>();
 
         let _ = self.send_to_shard(ShardRunnerMessage::ChunkGuilds {
             guild_ids: guilds,
             limit,
             query,
+            nonce,
         });
     }
 

diff --git a/src/client/bridge/gateway/shard_runner.rs b/src/client/bridge/gateway/shard_runner.rs
@@ -362,11 +362,12 @@ impl ShardRunner {
 
                         true
                     },
-                ShardClientMessage::Runner(ShardRunnerMessage::ChunkGuilds { guild_ids, limit, query }) => {
+                ShardClientMessage::Runner(ShardRunnerMessage::ChunkGuilds { guild_ids, limit, query, nonce }) => {
                     self.shard.chunk_guilds(
                         guild_ids,
                         limit,
                         query.as_deref(),
+                        nonce.as_deref(),
                     ).await.is_ok()
                 },
                 ShardClientMessage::Runner(ShardRunnerMessage::Close(code, reason)) => {

diff --git a/src/client/bridge/gateway/shard_runner_message.rs b/src/client/bridge/gateway/shard_runner_message.rs
@@ -31,6 +31,10 @@ pub enum ShardRunnerMessage {
         ///
         /// [`Member`]: ../../../model/guild/struct.Member.html
         query: Option<String>,
+        /// Nonce to identify [`GuildMembersChunkEvent`] responses.
+        ///
+        /// [`GuildMembersChunkEvent`]: ../../../model/event/struct.GuildMembersChunkEvent.html
+        nonce: Option<String>,
     },
     /// Indicates that the client is to close with the given status code and
     /// reason.

diff --git a/src/client/dispatch.rs b/src/client/dispatch.rs
@@ -530,7 +530,7 @@ async fn handle_event(
             let event_handler = Arc::clone(event_handler);
 
             tokio::spawn(async move {
-                event_handler.guild_members_chunk(context, event.guild_id, event.members, event.nonce).await;
+                event_handler.guild_members_chunk(context, event).await;
             });
         },
         DispatchEvent::Model(Event::GuildRoleCreate(mut event)) => {

diff --git a/src/client/event_handler.rs b/src/client/event_handler.rs
@@ -152,7 +152,7 @@ pub trait EventHandler: Send + Sync {
     /// Dispatched when the data for offline members was requested.
     ///
     /// Provides the guild's id and the data.
-    async fn guild_members_chunk(&self, _ctx: Context, _guild_id: GuildId, _offline_members: HashMap<UserId, Member>, _nonce: Option<String>) {}
+    async fn guild_members_chunk(&self, _ctx: Context, _chunk: GuildMembersChunkEvent) {}
 
     /// Dispatched when a role is created.
     ///

diff --git a/src/gateway/shard.rs b/src/gateway/shard.rs
@@ -697,13 +697,13 @@ impl Shard {
     ///
     /// let guild_ids = vec![GuildId(81384788765712384)];
     ///
-    /// shard.chunk_guilds(guild_ids, Some(2000), None).await?;
+    /// shard.chunk_guilds(guild_ids, Some(2000), None, None).await?;
     /// #     Ok(())
     /// # }
     /// ```
     ///
     /// Chunk a single guild by Id, limiting to 20 members, and specifying a
-    /// query parameter of `"do"`:
+    /// query parameter of `"do"` and a nonce of `"request"`:
     ///
     /// ```rust,no_run
     /// # use tokio::sync::Mutex;
@@ -720,7 +720,7 @@ impl Shard {
     ///
     /// let guild_ids = vec![GuildId(81384788765712384)];
     ///
-    /// shard.chunk_guilds(guild_ids, Some(20), Some("do")).await?;
+    /// shard.chunk_guilds(guild_ids, Some(20), Some("do"), Some("request")).await?;
     /// #     Ok(())
     /// # }
     /// ```
@@ -734,6 +734,7 @@ impl Shard {
         guild_ids: It,
         limit: Option<u16>,
         query: Option<&str>,
+        nonce: Option<&str>,
     ) -> Result<()> where It: IntoIterator<Item=GuildId> + Send {
         debug!("[Shard {:?}] Requesting member chunks", self.shard_info);
 
@@ -742,6 +743,7 @@ impl Shard {
             &self.shard_info,
             limit,
             query,
+            nonce,
         ).await
     }
 

diff --git a/src/gateway/ws_client_ext.rs b/src/gateway/ws_client_ext.rs
@@ -19,6 +19,7 @@ pub trait WebSocketGatewayClientExt {
         shard_info: &[u64; 2],
         limit: Option<u16>,
         query: Option<&str>,
+        nonce: Option<&str>,
     ) -> Result<()> where It: IntoIterator<Item=GuildId> + Send;
 
     async fn send_heartbeat(&mut self, shard_info: &[u64; 2], seq: Option<u64>)
@@ -51,6 +52,7 @@ impl WebSocketGatewayClientExt for WsStream {
         shard_info: &[u64; 2],
         limit: Option<u16>,
         query: Option<&str>,
+        nonce: Option<&str>,
     ) -> Result<()> where It: IntoIterator<Item=GuildId> + Send {
         debug!("[Shard {:?}] Requesting member chunks", shard_info);
 
@@ -60,6 +62,7 @@ impl WebSocketGatewayClientExt for WsStream {
                 "guild_id": guild_ids.into_iter().map(|x| x.as_ref().0).collect::<Vec<u64>>(),
                 "limit": limit.unwrap_or(0),
                 "query": query.unwrap_or(""),
+                "nonce": nonce.unwrap_or(""),
             },
         })).await.map_err(From::from)
     }

diff --git a/src/model/event.rs b/src/model/event.rs
@@ -557,6 +557,8 @@ impl CacheUpdate for GuildMemberUpdateEvent {
 pub struct GuildMembersChunkEvent {
     pub guild_id: GuildId,
     pub members: HashMap<UserId, Member>,
+    pub chunk_index: u32,
+    pub chunk_count: u32,
     pub nonce: Option<String>,
     #[serde(skip)]
     pub(crate) _nonexhaustive: (),
@@ -592,6 +594,16 @@ impl<'de> Deserialize<'de> for GuildMembersChunkEvent {
         let mut members = map.remove("members")
             .ok_or_else(|| DeError::custom("missing member chunk members"))?;
 
+        let chunk_index = map.get("chunk_index")
+            .ok_or_else(|| DeError::custom("missing member chunk index"))
+            .and_then(u32::deserialize)
+            .map_err(DeError::custom)?;
+
+        let chunk_count = map.get("chunk_count")
+            .ok_or_else(|| DeError::custom("missing member chunk count"))
+            .and_then(u32::deserialize)
+            .map_err(DeError::custom)?;
+
         if let Some(members) = members.as_array_mut() {
             let num = Value::Number(Number::from(guild_id.0));
 
@@ -621,6 +633,8 @@ impl<'de> Deserialize<'de> for GuildMembersChunkEvent {
         Ok(GuildMembersChunkEvent {
             guild_id,
             members,
+            chunk_index,
+            chunk_count,
             nonce,
             _nonexhaustive: (),
         })