Add Function to neutralise Mentions (#414)

4 files changed, 579 insertions, 7 deletions

diff --git a/examples/05_command_framework/src/main.rs b/examples/05_command_framework/src/main.rs
index e69934c..843fb43 100644
--- a/examples/05_command_framework/src/main.rs
+++ b/examples/05_command_framework/src/main.rs
@@ -20,6 +20,7 @@ use serenity::model::gateway::Ready;
 use serenity::model::Permissions;
 use serenity::prelude::Mutex;
 use serenity::prelude::*;
+use serenity::utils::{content_safe, ContentSafeOptions};
 use std::collections::HashMap;
 use std::env;
 use std::fmt::Write;
@@ -175,6 +176,10 @@ fn main() {
             // Make this command use the "complicated" bucket.
             .bucket("complicated")
             .cmd(commands))
+        // Command that will repeat passed arguments and remove user and
+        // role mentions with safe alternative.
+        .command("say", |c| c
+            .cmd(say))
         .group("Emoji", |g| g
             // Sets multiple prefixes for a group.
             // This requires us to call commands in this group
@@ -250,6 +255,32 @@ command!(commands(ctx, msg, _args) {
     }
 });
 
+// Repeats what the user passed as argument but ensures that user and role
+// mentions are replaced with a safe textual alternative.
+// In this example channel mentions are excluded via the `ContentSafeOptions`.
+command!(say(_ctx, msg, args) {
+    let mut settings = if let Some(guild_id) = msg.guild_id {
+       // By default roles, users, and channel mentions are cleaned.
+       ContentSafeOptions::default()
+            // We do not want to clean channal mentions as they
+            // do not ping users.
+            .clean_channel(false)
+            // If it's a guild channel, we want mentioned users to be displayed
+            // as their display name.
+            .display_as_member_from(guild_id)
+    } else {
+        ContentSafeOptions::default()
+            .clean_channel(false)
+            .clean_role(false)
+    };
+
+    let mut content = content_safe(&args.full(), &settings);
+
+    if let Err(why) = msg.channel_id.say(&content) {
+        println!("Error sending message: {:?}", why);
+    }
+});
+
 // A function which acts as a "check", to determine whether to call a command.
 //
 // In this case, this command checks to ensure you are the owner of the message
diff --git a/src/model/channel/channel_id.rs b/src/model/channel/channel_id.rs
index a897002..3f52055 100644
--- a/src/model/channel/channel_id.rs
+++ b/src/model/channel/channel_id.rs
@@ -14,6 +14,8 @@ use builder::{
 };
 #[cfg(all(feature = "cache", feature = "model"))]
 use CACHE;
+#[cfg(all(feature = "cache", feature = "model"))]
+use Cache;
 #[cfg(feature = "model")]
 use http::{self, AttachmentType};
 #[cfg(feature = "model")]
@@ -293,7 +295,17 @@ impl ChannelId {
     /// [`Channel`]: ../channel/enum.Channel.html
     #[cfg(feature = "cache")]
     #[inline]
-    pub fn to_channel_cached(self) -> Option<Channel> { CACHE.read().channel(self) }
+    pub fn to_channel_cached(self) -> Option<Channel> {
+        self._to_channel_cached(&CACHE)
+    }
+
+    /// To allow testing pass their own cache instead of using the globale one.
+    #[cfg(feature = "cache")]
+    #[inline]
+    pub(crate) fn _to_channel_cached(self, cache: &RwLock<Cache>) -> Option<Channel> {
+        cache.read().channel(self)
+    }
+
 
     /// Search the cache for the channel. If it can't be found, the channel is
     /// requested over REST.
diff --git a/src/model/guild/role.rs b/src/model/guild/role.rs
index ac969c0..53239a8 100644
--- a/src/model/guild/role.rs
+++ b/src/model/guild/role.rs
@@ -6,7 +6,7 @@ use builder::EditRole;
 #[cfg(all(feature = "cache", feature = "model"))]
 use internal::prelude::*;
 #[cfg(all(feature = "cache", feature = "model"))]
-use {CACHE, http};
+use {CACHE, Cache, http};
 
 #[cfg(all(feature = "cache", feature = "model", feature = "utils"))]
 use std::str::FromStr;
@@ -182,9 +182,12 @@ impl RoleId {
     /// [`Role`]: ../guild/struct.Role.html
     #[cfg(feature = "cache")]
     pub fn to_role_cached(self) -> Option<Role> {
-        let cache = CACHE.read();
+        self._to_role_cached(&CACHE)
+    }
 
-        for guild in cache.guilds.values() {
+    #[cfg(feature = "cache")]
+    pub(crate) fn _to_role_cached(self, cache: &RwLock<Cache>) -> Option<Role> {
+        for guild in cache.read().guilds.values() {
             let guild = guild.read();
 
             if !guild.roles.contains_key(&self) {
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 7cd60cf..0696aad 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -13,15 +13,26 @@ pub use self::{
 
 use base64;
 use internal::prelude::*;
-use model::id::EmojiId;
-use model::misc::EmojiIdentifier;
+use prelude::RwLock;
+use model::{
+    channel::Channel,
+    misc::EmojiIdentifier,
+    id::{
+        ChannelId,
+        GuildId,
+        EmojiId,
+        RoleId,
+        UserId,
+    },
+};
 use std::{
     collections::HashMap,
     ffi::OsStr,
     fs::File,
     hash::{BuildHasher, Hash},
     io::Read,
-    path::Path
+    path::Path,
+    str::FromStr,
 };
 
 #[cfg(feature = "cache")]
@@ -527,6 +538,328 @@ pub fn with_cache_mut<T, F>(mut f: F) -> T
     f(&mut cache)
 }
 
+/// Struct that allows to alter [`content_safe`]'s behaviour.
+///
+/// [`content_safe`]: fn.content_safe.html
+#[cfg(feature = "cache")]
+#[derive(Clone, Debug)]
+pub struct ContentSafeOptions {
+    clean_role: bool,
+    clean_user: bool,
+    clean_channel: bool,
+    clean_here: bool,
+    clean_everyone: bool,
+    show_discriminator: bool,
+    guild_reference: Option<GuildId>,
+}
+
+#[cfg(feature = "cache")]
+impl ContentSafeOptions {
+    pub fn new() -> Self {
+        ContentSafeOptions::default()
+    }
+
+    /// [`content_safe`] will replace role mentions (`<@&{id}>`) with its name
+    /// prefixed with `@` (`@rolename`) or with `@deleted-role` if the
+    /// identifier is invalid.
+    ///
+    /// [`content_safe`]: fn.content_safe.html
+    pub fn clean_role(mut self, b: bool) -> Self {
+        self.clean_role = b;
+
+        self
+    }
+
+    /// If set to true, [`content_safe`] will replace user mentions
+    /// (`<@!{id}>` or `<@{id}>`) with the user's name prefixed with `@`
+    /// (`@username`) or with `@invalid-user` if the identifier is invalid.
+    ///
+    /// [`content_safe`]: fn.content_safe.html
+    pub fn clean_user(mut self, b: bool) -> Self {
+        self.clean_user = b;
+
+        self
+    }
+
+    /// If set to true, [`content_safe`] will replace channel mentions
+    /// (`<#{id}>`) with the channel's name prefixed with `#`
+    /// (`#channelname`) or with `#deleted-channel` if the identifier is
+    /// invalid.
+    ///
+    /// [`content_safe`]: fn.content_safe.html
+    pub fn clean_channel(mut self, b: bool) -> Self {
+        self.clean_channel = b;
+
+        self
+    }
+
+    /// If set to true, if [`content_safe`] replaces a user mention it will
+    /// add their four digit discriminator with a preceeding `#`,
+    /// turning `@username` to `@username#discriminator`.
+    ///
+    /// [`content_safe`]: fn.content_safe.html
+    pub fn show_discriminator(mut self, b: bool) -> Self {
+        self.show_discriminator = b;
+
+        self
+    }
+
+    /// If set, [`content_safe`] will replace a user mention with the user's
+    /// display name in passed `guild`.
+    ///
+    /// [`content_safe`]: fn.content_safe.html
+    pub fn display_as_member_from<G: Into<GuildId>>(mut self, guild: G) -> Self {
+        self.guild_reference = Some(guild.into());
+
+        self
+    }
+
+    /// If set, [`content_safe`] will replace `@here` with a non-pinging
+    /// alternative.
+    ///
+    /// [`content_safe`]: fn.content_safe.html
+    pub fn clean_here(mut self, b: bool) -> Self {
+        self.clean_here = b;
+
+        self
+    }
+
+    /// If set, [`content_safe`] will replace `@everyone` with a non-pinging
+    /// alternative.
+    ///
+    /// [`content_safe`]: fn.content_safe.html
+    pub fn clean_everyone(mut self, b: bool) -> Self {
+        self.clean_everyone = b;
+
+        self
+    }
+}
+
+#[cfg(feature = "cache")]
+impl Default for ContentSafeOptions {
+    /// Instantiates with all options set to `true`.
+    fn default() -> Self {
+        ContentSafeOptions {
+            clean_role: true,
+            clean_user: true,
+            clean_channel: true,
+            clean_here: true,
+            clean_everyone: true,
+            show_discriminator: true,
+            guild_reference: None,
+        }
+    }
+}
+
+#[cfg(feature = "cache")]
+#[inline]
+fn clean_roles(cache: &RwLock<Cache>, s: &mut String) {
+    let mut progress = 0;
+
+    while let Some(mut mention_start) = s[progress..].find("<@&") {
+        mention_start += progress;
+
+        if let Some(mut mention_end) = s[mention_start..].find(">") {
+            mention_end += mention_start;
+            mention_start += "<@&".len();
+
+            if let Ok(id) = RoleId::from_str(&s[mention_start..mention_end]) {
+                let to_replace = format!("<@&{}>", &s[mention_start..mention_end]);
+
+                *s = if let Some(role) = id._to_role_cached(&cache) {
+                    s.replace(&to_replace, &format!("@{}", &role.name))
+                } else {
+                    s.replace(&to_replace, &"@deleted-role")
+                };
+            } else  {
+                let id = &s[mention_start..mention_end].to_string();
+
+                if !id.is_empty() && id.as_bytes().iter()
+                    .all(u8::is_ascii_digit){
+                    let to_replace = format!("<@&{}>", id);
+
+                    *s = s.replace(&to_replace, &"@deleted-role");
+                } else {
+                    progress = mention_end;
+                }
+            }
+        } else {
+            break;
+        }
+    }
+}
+
+#[cfg(feature = "cache")]
+#[inline]
+fn clean_channels(cache: &RwLock<Cache>, s: &mut String) {
+    let mut progress = 0;
+
+    while let Some(mut mention_start) = s[progress..].find("<#") {
+        mention_start += progress;
+
+        if let Some(mut mention_end) = s[mention_start..].find(">") {
+            mention_end += mention_start;
+            mention_start += "<#".len();
+
+            if let Ok(id) = ChannelId::from_str(&s[mention_start..mention_end]) {
+                let to_replace = format!("<#{}>", &s[mention_start..mention_end]);
+
+                *s = if let Some(Channel::Guild(channel)) = id._to_channel_cached(&cache) {
+                    let replacement = format!("#{}", &channel.read().name);
+                    s.replace(&to_replace, &replacement)
+                } else {
+                    s.replace(&to_replace, &"#deleted-channel")
+                };
+            } else {
+                let id = &s[mention_start..mention_end].to_string();
+
+                if !id.is_empty() && id.as_bytes().iter()
+                    .all(u8::is_ascii_digit) {
+                    let to_replace = format!("<#{}>", id);
+
+                    *s = s.replace(&to_replace, &"#deleted-channel");
+                } else {
+                    progress = mention_end;
+                }
+            }
+        } else {
+            break;
+        }
+    }
+}
+
+#[cfg(feature = "cache")]
+#[inline]
+fn clean_users(cache: &RwLock<Cache>, s: &mut String, show_discriminator: bool, guild: Option<GuildId>) {
+    let mut progress = 0;
+
+    while let Some(mut mention_start) = s[progress..].find("<@") {
+        mention_start += progress;
+
+        if let Some(mut mention_end) = s[mention_start..].find(">") {
+            mention_end += mention_start;
+            mention_start += "<@".len();
+            let mut has_exclamation = false;
+
+            if s[mention_start..].as_bytes().get(0).map_or(false, |c| *c == b'!') {
+                mention_start += "!".len();
+                has_exclamation = true;
+            }
+
+            if let Ok(id) = UserId::from_str(&s[mention_start..mention_end]) {
+                let mut replacement = if let Some(guild) = guild {
+
+                    if let Some(guild) = cache.read().guild(&guild) {
+
+                        if let Some(member) = guild.read().members.get(&id) {
+
+                            if show_discriminator {
+                                format!("@{}", member.distinct())
+                            } else {
+                                format!("@{}", member.display_name())
+                            }
+                        } else {
+                            "@invalid-user".to_string()
+                        }
+                    } else {
+                        "@invalid-user".to_string()
+                    }
+                } else {
+                    let user = cache.read().users.get(&id).map(|user| user.clone());
+
+                    if let Some(user) = user {
+                        let user = user.read();
+
+                        if show_discriminator {
+                            format!("@{}#{:04}", user.name, user.discriminator)
+                        } else {
+                            format!("@{}", user.name)
+                        }
+                    } else {
+                        "@invalid-user".to_string()
+                    }
+                };
+
+                let code_start = if has_exclamation { "<@!" } else { "<@" };
+                let to_replace = format!("{}{}>", code_start, &s[mention_start..mention_end]);
+
+                *s = s.replace(&to_replace, &replacement)
+            } else {
+                let id = &s[mention_start..mention_end].to_string();
+
+                if !id.is_empty() && id.as_bytes().iter().all(u8::is_ascii_digit) {
+                    let code_start = if has_exclamation { "<@!" } else { "<@" };
+                    let to_replace = format!("{}{}>", code_start, id);
+
+                    *s = s.replace(&to_replace, &"@invalid-user");
+                } else {
+                    progress = mention_end;
+                }
+            }
+        } else {
+            break;
+        }
+    }
+}
+
+/// Transforms role, channel, user, `@everyone` and `@here` mentions
+/// into raw text by using the [`Cache`] only.
+///
+/// [`ContentSafeOptions`] decides what kind of mentions should be filtered
+/// and how the raw-text will be displayed.
+///
+/// # Examples
+///
+/// Sanitise an `@everyone` mention.
+///
+/// ```rust
+/// use serenity::utils::{
+///     content_safe,
+///     ContentSafeOptions,
+/// };
+///
+/// let with_mention = "@everyone";
+/// let without_mention = content_safe(&with_mention, &ContentSafeOptions::default());
+///
+/// assert_eq!("@\u{200B}everyone".to_string(), without_mention);
+/// ```
+/// [`ContentSafeOptions`]: struct.ContentSafeOptions.html
+/// [`Cache`]: ../cache/struct.Cache.html
+#[cfg(feature = "cache")]
+pub fn content_safe(s: &str, options: &ContentSafeOptions) -> String {
+    let cache = &CACHE;
+
+    _content_safe(&cache, s, options)
+}
+
+
+#[cfg(feature = "cache")]
+fn _content_safe(cache: &RwLock<Cache>, s: &str, options: &ContentSafeOptions) -> String {
+    let mut s = s.to_string();
+
+    if options.clean_role {
+        clean_roles(&cache, &mut s);
+    }
+
+    if options.clean_channel {
+        clean_channels(&cache, &mut s);
+    }
+
+    if options.clean_user {
+        clean_users(&cache, &mut s, options.show_discriminator, options.guild_reference);
+    }
+
+    if options.clean_here {
+        s = s.replace("@here", "@\u{200B}here");
+    }
+
+    if options.clean_everyone {
+        s = s.replace("@everyone", "@\u{200B}everyone");
+    }
+
+    s
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -577,4 +910,197 @@ mod test {
         assert!(!is_nsfw("général"));
         assert!(is_nsfw("nsfw-général"));
     }
+
+    #[test]
+    fn test_content_safe() {
+        use model::{
+            user::User,
+            Permissions,
+            prelude::*,
+        };
+        use chrono::DateTime;
+        use std::{
+            collections::HashMap,
+            sync::Arc,
+        };
+
+        let user = User {
+            id: UserId(100000000000000000),
+            avatar: None,
+            bot: false,
+            discriminator: 0000,
+            name: "Crab".to_string(),
+        };
+
+        let mut guild = Guild {
+            afk_channel_id: None,
+            afk_timeout: 0,
+            application_id: None,
+            channels: HashMap::new(),
+            default_message_notifications: DefaultMessageNotificationLevel::All,
+            emojis: HashMap::new(),
+            explicit_content_filter: ExplicitContentFilter::None,
+            features: Vec::new(),
+            icon: None,
+            id: GuildId(381880193251409931),
+            joined_at: DateTime::parse_from_str(
+                "1983 Apr 13 12:09:14.274 +0000",
+                "%Y %b %d %H:%M:%S%.3f %z").unwrap(),
+            large: false,
+            member_count: 1,
+            members: HashMap::new(),
+            mfa_level: MfaLevel::None,
+            name: "serenity".to_string(),
+            owner_id: UserId(114941315417899012),
+            presences: HashMap::new(),
+            region: "Ferris Island".to_string(),
+            roles: HashMap::new(),
+            splash: None,
+            system_channel_id: None,
+            verification_level: VerificationLevel::None,
+            voice_states: HashMap::new(),
+        };
+
+        let member = Member {
+            deaf: false,
+            guild_id: guild.id,
+            joined_at: None,
+            mute: false,
+            nick: Some("Ferris".to_string()),
+            roles: Vec::new(),
+            user: Arc::new(RwLock::new(user.clone())),
+        };
+
+        let role = Role {
+            id: RoleId(333333333333333333),
+            colour: Colour::ORANGE,
+            hoist: true,
+            managed: false,
+            mentionable: true,
+            name: "ferris-club-member".to_string(),
+            permissions: Permissions::all(),
+            position: 0,
+        };
+
+        let channel = GuildChannel {
+            id: ChannelId(111880193700067777),
+            bitrate: None,
+            category_id: None,
+            guild_id: guild.id,
+            kind: ChannelType::Text,
+            last_message_id: None,
+            last_pin_timestamp: None,
+            name: "general".to_string(),
+            permission_overwrites: Vec::new(),
+            position: 0,
+            topic: None,
+            user_limit: None,
+            nsfw: false,
+        };
+
+        let cache = RwLock::new(Cache::default());
+
+        {
+            let mut cache = cache.try_write().unwrap();
+            guild.members.insert(user.id, member.clone());
+            guild.roles.insert(role.id, role.clone());
+            cache.users.insert(user.id, Arc::new(RwLock::new(user.clone())));
+            cache.guilds.insert(guild.id, Arc::new(RwLock::new(guild.clone())));
+            cache.channels.insert(channel.id, Arc::new(RwLock::new(channel.clone())));
+        }
+
+        let with_user_metions = "<@!100000000000000000> <@!000000000000000000> <@123> <@!123> \
+        <@!123123123123123123123> <@123> <@123123123123123123> <@!invalid> \
+        <@invalid> <@日本語 한국어$§)[/__#\\(/&2032$§#> \
+        <@!i)/==(<<>z/9080)> <@!1231invalid> <@invalid123> \
+        <@123invalid> <@> <@ ";
+
+        let without_user_mentions = "@Crab#0000 @invalid-user @invalid-user @invalid-user \
+        @invalid-user @invalid-user @invalid-user <@!invalid> \
+        <@invalid> <@日本語 한국어$§)[/__#\\(/&2032$§#> \
+        <@!i)/==(<<>z/9080)> <@!1231invalid> <@invalid123> \
+        <@123invalid> <@> <@ ";
+
+        // User mentions
+        let options = ContentSafeOptions::default();
+        assert_eq!(without_user_mentions, _content_safe(&cache, with_user_metions, &options));
+
+        let options = ContentSafeOptions::default();
+        assert_eq!(format!("@{}#{:04}", user.name, user.discriminator),
+            _content_safe(&cache, "<@!100000000000000000>", &options));
+
+        let options = ContentSafeOptions::default();
+        assert_eq!(format!("@{}#{:04}", user.name, user.discriminator),
+            _content_safe(&cache, "<@100000000000000000>", &options));
+
+        let options = options.show_discriminator(false);
+        assert_eq!(format!("@{}", user.name),
+            _content_safe(&cache, "<@!100000000000000000>", &options));
+
+        let options = options.show_discriminator(false);
+        assert_eq!(format!("@{}", user.name),
+            _content_safe(&cache, "<@100000000000000000>", &options));
+
+        let options = options.display_as_member_from(guild.id);
+        assert_eq!(format!("@{}", member.nick.unwrap()),
+            _content_safe(&cache, "<@!100000000000000000>", &options));
+
+        let options = options.clean_user(false);
+        assert_eq!(with_user_metions,
+            _content_safe(&cache, with_user_metions, &options));
+
+        // Channel mentions
+        let with_channel_mentions = "<#> <#deleted-channel> #deleted-channel <#0> \
+        #unsafe-club <#111880193700067777> <#ferrisferrisferris> \
+        <#000000000000000000>";
+
+        let without_channel_mentions = "<#> <#deleted-channel> #deleted-channel \
+        #deleted-channel #unsafe-club #general <#ferrisferrisferris> \
+        #deleted-channel";
+
+        assert_eq!(without_channel_mentions,
+            _content_safe(&cache, with_channel_mentions, &options));
+
+        let options = options.clean_channel(false);
+        assert_eq!(with_channel_mentions,
+            _content_safe(&cache, with_channel_mentions, &options));
+
+        // Role mentions
+        let with_role_mentions = "<@&> @deleted-role <@&9829> \
+        <@&333333333333333333> <@&000000000000000000>";
+
+        let without_role_mentions = "<@&> @deleted-role @deleted-role \
+        @ferris-club-member @deleted-role";
+
+        assert_eq!(without_role_mentions,
+            _content_safe(&cache, with_role_mentions, &options));
+
+        let options = options.clean_role(false);
+        assert_eq!(with_role_mentions,
+            _content_safe(&cache, with_role_mentions, &options));
+
+        // Everyone mentions
+        let with_everyone_mention = "@everyone";
+
+        let without_everyone_mention = "@\u{200B}everyone";
+
+        assert_eq!(without_everyone_mention,
+            _content_safe(&cache, with_everyone_mention, &options));
+
+        let options = options.clean_everyone(false);
+        assert_eq!(with_everyone_mention,
+            _content_safe(&cache, with_everyone_mention, &options));
+
+        // Here mentions
+        let with_here_mention = "@here";
+
+        let without_here_mention = "@\u{200B}here";
+
+        assert_eq!(without_here_mention,
+            _content_safe(&cache, with_here_mention, &options));
+
+        let options = options.clean_here(false);
+        assert_eq!(with_here_mention,
+            _content_safe(&cache, with_here_mention, &options));
+    }
 }