From 3b70891b0a62667dd152529d2444b0414db4c278 Mon Sep 17 00:00:00 2001
From: projectmoon <projectmoon@agnos.is>
Date: Fri, 16 Oct 2020 22:02:25 +0000
Subject: [PATCH] Reject dice pool expressions over 100 elements.

---
 src/cofd/dice.rs | 34 ++++++++++++++++++++++++++++++++++
 src/error.rs     |  2 +-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/cofd/dice.rs b/src/cofd/dice.rs
index db645e7..ee9ff0d 100644
--- a/src/cofd/dice.rs
+++ b/src/cofd/dice.rs
@@ -11,6 +11,9 @@ use thiserror::Error;
 pub enum DiceRollingError {
     #[error("variable not found: {0}")]
     VariableNotFound(String),
+
+    #[error("dice pool expression too large")]
+    ExpressionTooLarge,
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -370,6 +373,10 @@ fn roll_dice<R: DieRoller>(
     pool: &DicePoolWithContext,
     roller: &mut R,
 ) -> Result<RolledDicePool, BotError> {
+    if pool.0.amounts.len() > 100 {
+        return Err(DiceRollingError::ExpressionTooLarge.into());
+    }
+
     let num_dice = calculate_dice_amount(&pool)?;
     let rolls: Vec<i32> = (0..num_dice)
         .flat_map(|_| roll_die(roller, &pool.0))
@@ -531,6 +538,33 @@ mod tests {
         assert_eq!(5, roll.num_dice);
     }
 
+    #[test]
+    fn rejects_large_expression_test() {
+        let db = Database::new(&sled::open(tempdir().unwrap()).unwrap());
+        let ctx = Context::new(&db, "roomid", "username", "message");
+
+        let mut amounts = vec![];
+
+        for _ in 0..500 {
+            amounts.push(Amount {
+                operator: Operator::Plus,
+                element: Element::Number(1),
+            });
+        }
+
+        let pool = DicePool::new(amounts, DicePoolModifiers::default());
+        let pool_with_ctx = DicePoolWithContext(&pool, &ctx);
+
+        let mut roller = SequentialDieRoller::new(vec![1, 2, 3, 4, 5]);
+        let result = roll_dice(&pool_with_ctx, &mut roller);
+        assert!(matches!(
+            result,
+            Err(BotError::DiceRollingError(
+                DiceRollingError::ExpressionTooLarge
+            ))
+        ));
+    }
+
     #[test]
     fn can_resolve_variables_test() {
         let db = Database::new(&sled::open(tempdir().unwrap()).unwrap());
diff --git a/src/error.rs b/src/error.rs
index f16292f..c0dd269 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -57,7 +57,7 @@ pub enum BotError {
     DiceParsingError(#[from] crate::cofd::parser::DiceParsingError),
 
     #[error("dice pool roll error: {0}")]
-    DicePoolError(#[from] DiceRollingError),
+    DiceRollingError(#[from] DiceRollingError),
 
     #[error("variable parsing error: {0}")]
     VariableParsingError(#[from] crate::variables::VariableParsingError),