From 8e25224f11d67289d5a28a6ba434945bfd7fd42e Mon Sep 17 00:00:00 2001
From: Supanat Potiwarakorn <supanat.ptk@gmail.com>
Date: Mon, 7 Oct 2024 12:07:14 +0700
Subject: [PATCH] impose max created asset group

---
 contracts/transmuter/src/error.rs             |  3 ++
 .../src/transmuter_pool/asset_group.rs        | 46 +++++++++++++++++--
 .../transmuter/src/transmuter_pool/mod.rs     |  4 ++
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/contracts/transmuter/src/error.rs b/contracts/transmuter/src/error.rs
index c01beb1..56c3325 100644
--- a/contracts/transmuter/src/error.rs
+++ b/contracts/transmuter/src/error.rs
@@ -57,6 +57,9 @@ pub enum ContractError {
         actual: Uint64,
     },
 
+    #[error("Asset group count must be within {max} inclusive, but got: {actual}")]
+    AssetGroupCountOutOfRange { max: Uint64, actual: Uint64 },
+
     #[error("Insufficient pool asset: required: {required}, available: {available}")]
     InsufficientPoolAsset { required: Coin, available: Coin },
 
diff --git a/contracts/transmuter/src/transmuter_pool/asset_group.rs b/contracts/transmuter/src/transmuter_pool/asset_group.rs
index 88e3f8a..a2959bb 100644
--- a/contracts/transmuter/src/transmuter_pool/asset_group.rs
+++ b/contracts/transmuter/src/transmuter_pool/asset_group.rs
@@ -1,9 +1,9 @@
 use std::collections::BTreeMap;
 
 use cosmwasm_schema::cw_serde;
-use cosmwasm_std::{ensure, Decimal};
+use cosmwasm_std::{ensure, Decimal, Uint64};
 
-use crate::{corruptable::Corruptable, ContractError};
+use crate::{corruptable::Corruptable, transmuter_pool::MAX_ASSET_GROUPS, ContractError};
 
 use super::TransmuterPool;
 
@@ -109,10 +109,16 @@ impl TransmuterPool {
             );
         }
 
-        // TODO: limit sizes of asset groups
-
         self.asset_groups.insert(label, AssetGroup::new(denoms));
 
+        ensure!(
+            Uint64::from(self.asset_groups.len() as u64) <= MAX_ASSET_GROUPS,
+            ContractError::AssetGroupCountOutOfRange {
+                max: MAX_ASSET_GROUPS,
+                actual: Uint64::new(self.asset_groups.len() as u64)
+            }
+        );
+
         Ok(self)
     }
 
@@ -236,4 +242,36 @@ mod tests {
             &Decimal::raw(333333333333333333)
         );
     }
+
+    #[test]
+    fn test_create_asset_group_within_range() {
+        let mut pool = TransmuterPool::new(vec![
+            Asset::new(Uint128::new(100), "denom1", Uint128::new(1)).unwrap(),
+            Asset::new(Uint128::new(200), "denom2", Uint128::new(1)).unwrap(),
+            Asset::new(Uint128::new(300), "denom3", Uint128::new(1)).unwrap(),
+        ])
+        .unwrap();
+
+        // Test creating groups up to the maximum allowed
+        for i in 1..=MAX_ASSET_GROUPS.u64() {
+            let group_name = format!("group{}", i);
+            let result = pool.create_asset_group(group_name.clone(), vec!["denom1".to_string()]);
+            assert!(result.is_ok(), "Failed to create group {}", i);
+        }
+
+        // Attempt to create one more group, which should fail
+        let result = pool.create_asset_group("extra_group".to_string(), vec!["denom1".to_string()]);
+        assert!(
+            result.is_err(),
+            "Should not be able to create group beyond the maximum"
+        );
+        assert!(
+            matches!(
+                result.unwrap_err(),
+                ContractError::AssetGroupCountOutOfRange { max, actual }
+                if max == MAX_ASSET_GROUPS && actual == MAX_ASSET_GROUPS + Uint64::one()
+            ),
+            "Unexpected error when exceeding max asset groups"
+        );
+    }
 }
diff --git a/contracts/transmuter/src/transmuter_pool/mod.rs b/contracts/transmuter/src/transmuter_pool/mod.rs
index b503db8..7a2d110 100644
--- a/contracts/transmuter/src/transmuter_pool/mod.rs
+++ b/contracts/transmuter/src/transmuter_pool/mod.rs
@@ -25,6 +25,10 @@ const MIN_POOL_ASSET_DENOMS: Uint64 = Uint64::new(1u64);
 /// prevent the contract from running out of gas when iterating
 const MAX_POOL_ASSET_DENOMS: Uint64 = Uint64::new(20u64);
 
+/// Maximum number of asset groups allowed in a pool.
+/// This limit helps prevent excessive gas consumption when iterating over groups.
+const MAX_ASSET_GROUPS: Uint64 = Uint64::new(10u64);
+
 #[cw_serde]
 pub struct TransmuterPool {
     pub pool_assets: Vec<Asset>,