From 7512ca069460af2a7910b17a5da819587eff3559 Mon Sep 17 00:00:00 2001
From: projectmoon <projectmoon@agnos.is>
Date: Tue, 2 Feb 2021 21:41:16 +0000
Subject: [PATCH] Allow up to 50 commands per message.

If the amount of commands in a single message is greater, the bot will
now return an error. Includes slight refactoring of command execution
code to make use of streams for async iter-like mapping of the command
list.

Fixes #24.
---
 src/bot.rs      | 45 +++++++++++++++++++++++++++++----------------
 src/commands.rs |  2 +-
 src/error.rs    |  3 +++
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/src/bot.rs b/src/bot.rs
index 23b9d85..7e26a8c 100644
--- a/src/bot.rs
+++ b/src/bot.rs
@@ -6,9 +6,9 @@ use crate::error::BotError;
 use crate::matrix;
 use crate::state::DiceBotState;
 use dirs;
+use futures::stream::{self, StreamExt};
 use log::info;
 use matrix_sdk::{self, identifiers::RoomId, Client, ClientConfig, JoinedRoom, SyncSettings};
-//use matrix_sdk_common_macros::async_trait;
 use std::clone::Clone;
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
@@ -16,6 +16,10 @@ use url::Url;
 
 pub mod event_handlers;
 
+/// How many commands can be in one message. If the amount is higher
+/// than this, we reject execution.
+const MAX_COMMANDS_PER_MESSAGE: usize = 50;
+
 /// The DiceBot struct represents an active dice bot. The bot is not
 /// connected to Matrix until its run() function is called.
 pub struct DiceBot {
@@ -153,25 +157,34 @@ impl DiceBot {
     }
 
     async fn execute_commands(&self, room: &JoinedRoom, sender_username: &str, msg_body: &str) {
-        let room_name = room.display_name().await.ok().unwrap_or_default();
+        let room_name: &str = &room.display_name().await.ok().unwrap_or_default();
         let room_id = room.room_id().clone();
 
-        let mut results: Vec<(&str, CommandResult)> = Vec::with_capacity(msg_body.lines().count());
+        let commands: Vec<&str> = msg_body
+            .lines()
+            .filter(|line| line.starts_with("!"))
+            .collect();
 
-        let commands = msg_body.trim().lines().filter(|line| line.starts_with("!"));
+        //Up to 50 commands allowed, otherwise we send back an error.
+        let results: Vec<(&str, CommandResult)> = if commands.len() < MAX_COMMANDS_PER_MESSAGE {
+            stream::iter(commands)
+                .then(|command| async move {
+                    let ctx = Context {
+                        db: self.db.clone(),
+                        matrix_client: &self.client,
+                        room: RoomContext::new_with_name(&room, room_name),
+                        username: &sender_username,
+                        message_body: &command,
+                    };
 
-        for command in commands {
-            let ctx = Context {
-                db: self.db.clone(),
-                matrix_client: &self.client,
-                room: RoomContext::new_with_name(&room, &room_name),
-                username: &sender_username,
-                message_body: &command,
-            };
-
-            let cmd_result = execute_command(&ctx).await;
-            results.push((&command, cmd_result));
-        }
+                    let cmd_result = execute_command(&ctx).await;
+                    (command, cmd_result)
+                })
+                .collect()
+                .await
+        } else {
+            vec![("", Err(ExecutionError(BotError::MessageTooLarge)))]
+        };
 
         if results.len() >= 1 {
             if results.len() == 1 {
diff --git a/src/commands.rs b/src/commands.rs
index 71a275f..955ea83 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -46,7 +46,7 @@ impl Execution {
 /// provides formatting for successfully executed commands.
 #[derive(Error, Debug)]
 #[error("{0}")]
-pub struct ExecutionError(#[from] BotError);
+pub struct ExecutionError(#[from] pub BotError);
 
 impl From<crate::db::errors::DataError> for ExecutionError {
     fn from(error: crate::db::errors::DataError) -> Self {
diff --git a/src/error.rs b/src/error.rs
index 849866d..2f31943 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -69,6 +69,9 @@ pub enum BotError {
 
     #[error("database error")]
     DatabaseErrror(#[from] sled::Error),
+
+    #[error("too many commands or message was too large")]
+    MessageTooLarge,
 }
 
 #[derive(Error, Debug)]