From 56ffb5c0c37dda4aae8545b90d093f0c1c96c6bd Mon Sep 17 00:00:00 2001
From: piobab <piobab0x@proton.me>
Date: Sat, 9 Sep 2023 12:29:02 +0200
Subject: [PATCH] Assert pool contains only two assets.

---
 contracts/oracle/osmosis/src/helpers.rs       | 49 +++++++++++--------
 .../oracle/osmosis/tests/tests/helpers/mod.rs |  5 ++
 .../tests/tests/test_set_price_source.rs      | 49 ++++++++++++++++++-
 3 files changed, 80 insertions(+), 23 deletions(-)

diff --git a/contracts/oracle/osmosis/src/helpers.rs b/contracts/oracle/osmosis/src/helpers.rs
index c0937bb4c..3d6d9a67e 100644
--- a/contracts/oracle/osmosis/src/helpers.rs
+++ b/contracts/oracle/osmosis/src/helpers.rs
@@ -15,22 +15,25 @@ pub fn assert_osmosis_pool_assets(
     denom: &str,
     base_denom: &str,
 ) -> ContractResult<()> {
+    assert_pool_has_two_assets(pool)?;
+    assert_pool_contains_assets(pool, denom, base_denom)?;
+
     match pool {
         Pool::Balancer(balancer_pool) => {
-            assert_osmosis_xyk_pool(balancer_pool)?;
+            assert_equal_asset_weights(balancer_pool)?;
         }
         Pool::StableSwap(_) => {}
     };
 
-    assert_osmosis_pool_contains_two_assets(pool, denom, base_denom)?;
-
     Ok(())
 }
 
 /// Assert the Osmosis pool indicated by `pool_id` is Balancer XYK type
 pub fn assert_osmosis_xyk_lp_pool(pool: &Pool) -> ContractResult<()> {
+    assert_pool_has_two_assets(pool)?;
+
     match pool {
-        Pool::Balancer(balancer_pool) => assert_osmosis_xyk_pool(balancer_pool)?,
+        Pool::Balancer(balancer_pool) => assert_equal_asset_weights(balancer_pool)?,
         Pool::StableSwap(stable_swap_pool) => {
             return Err(ContractError::InvalidPriceSource {
                 reason: format!("StableSwap pool not supported. Pool id {}", stable_swap_pool.id),
@@ -41,11 +44,25 @@ pub fn assert_osmosis_xyk_lp_pool(pool: &Pool) -> ContractResult<()> {
     Ok(())
 }
 
-fn assert_osmosis_pool_contains_two_assets(
-    pool: &Pool,
-    denom: &str,
-    base_denom: &str,
-) -> ContractResult<()> {
+/// Assert the Osmosis pool has exactly two assets
+fn assert_pool_has_two_assets(pool: &Pool) -> ContractResult<()> {
+    let pool_id = pool.get_pool_id();
+    let pool_denoms = pool.get_pool_denoms();
+    if pool_denoms.len() != 2 {
+        return Err(ContractError::InvalidPriceSource {
+            reason: format!(
+                "expecting pool {} to contain exactly two coins; found {}",
+                pool_id,
+                pool_denoms.len()
+            ),
+        });
+    }
+
+    Ok(())
+}
+
+/// Assert the Osmosis pool contains both `denom` and `base_denom`, and they are not the same
+fn assert_pool_contains_assets(pool: &Pool, denom: &str, base_denom: &str) -> ContractResult<()> {
     let pool_id = pool.get_pool_id();
     let pool_denoms = pool.get_pool_denoms();
 
@@ -70,18 +87,8 @@ fn assert_osmosis_pool_contains_two_assets(
     Ok(())
 }
 
-/// Assert the Osmosis pool indicated by `pool_id` is of XYK type
-pub fn assert_osmosis_xyk_pool(pool: &BalancerPool) -> ContractResult<()> {
-    if pool.pool_assets.len() != 2 {
-        return Err(ContractError::InvalidPriceSource {
-            reason: format!(
-                "expecting pool {} to contain exactly two coins; found {}",
-                pool.id,
-                pool.pool_assets.len()
-            ),
-        });
-    }
-
+/// Assert the Osmosis pool has assets with equal weights (for XYK pools)
+fn assert_equal_asset_weights(pool: &BalancerPool) -> ContractResult<()> {
     if pool.pool_assets[0].weight != pool.pool_assets[1].weight {
         return Err(ContractError::InvalidPriceSource {
             reason: format!("assets in pool {} do not have equal weights", pool.id),
diff --git a/contracts/oracle/osmosis/tests/tests/helpers/mod.rs b/contracts/oracle/osmosis/tests/tests/helpers/mod.rs
index 82bb4ea76..a022146e9 100644
--- a/contracts/oracle/osmosis/tests/tests/helpers/mod.rs
+++ b/contracts/oracle/osmosis/tests/tests/helpers/mod.rs
@@ -91,6 +91,11 @@ pub fn setup_test_with_pools() -> OwnedDeps<MockStorage, MockApi, MarsMockQuerie
     deps.querier
         .set_query_pool_response(5555, prepare_query_stable_swap_pool_response(5555, &assets));
 
+    // Set StableSwap pool with more than 3 assets
+    let assets = vec![coin(42069, "uatom"), coin(69420, "uosmo"), coin(69420, "uusdc")];
+    deps.querier
+        .set_query_pool_response(6666, prepare_query_stable_swap_pool_response(6666, &assets));
+
     deps
 }
 
diff --git a/contracts/oracle/osmosis/tests/tests/test_set_price_source.rs b/contracts/oracle/osmosis/tests/tests/test_set_price_source.rs
index e47d703ec..f0f38ee5a 100644
--- a/contracts/oracle/osmosis/tests/tests/test_set_price_source.rs
+++ b/contracts/oracle/osmosis/tests/tests/test_set_price_source.rs
@@ -190,6 +190,15 @@ fn setting_price_source_spot() {
         }
     );
 
+    // attempting to use a StableSwap pool that contains more than two assets; should fail
+    let err = set_price_source_spot("uatom", 6666).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::InvalidPriceSource {
+            reason: "expecting pool 6666 to contain exactly two coins; found 3".to_string()
+        }
+    );
+
     // properly set spot price source
     let res = set_price_source_spot("umars", 89).unwrap();
     assert_eq!(res.messages.len(), 0);
@@ -277,6 +286,15 @@ fn setting_price_source_arithmetic_twap_with_invalid_params() {
         }
     );
 
+    // attempting to use a StableSwap pool that contains more than two assets; should fail
+    let err = set_price_source_twap("uatom", 6666, 86400, None).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::InvalidPriceSource {
+            reason: "expecting pool 6666 to contain exactly two coins; found 3".to_string()
+        }
+    );
+
     // attempting to set window_size bigger than 172800 sec (48h)
     let err = set_price_source_twap("umars", 89, 172801, None).unwrap_err();
     assert_eq!(
@@ -449,6 +467,15 @@ fn setting_price_source_geometric_twap_with_invalid_params() {
         }
     );
 
+    // attempting to use a StableSwap pool that contains more than two assets; should fail
+    let err = set_price_source_twap("uatom", 6666, 86400, None).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::InvalidPriceSource {
+            reason: "expecting pool 6666 to contain exactly two coins; found 3".to_string()
+        }
+    );
+
     // attempting to set window_size bigger than 172800 sec (48h)
     let err = set_price_source_twap("umars", 89, 172801, None).unwrap_err();
     assert_eq!(
@@ -622,7 +649,7 @@ fn setting_price_source_staked_geometric_twap_with_invalid_params() {
     );
 
     // attempting to use not XYK pool
-    let err = set_price_source_twap("ustatom", "uatom", 4444, 86400, None).unwrap_err();
+    let err = set_price_source_twap("uion", "uosmo", 4444, 86400, None).unwrap_err();
     assert_eq!(
         err,
         ContractError::InvalidPriceSource {
@@ -630,6 +657,15 @@ fn setting_price_source_staked_geometric_twap_with_invalid_params() {
         }
     );
 
+    // attempting to use a StableSwap pool that contains more than two assets; should fail
+    let err = set_price_source_twap("uatom", "uusdc", 6666, 86400, None).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::InvalidPriceSource {
+            reason: "expecting pool 6666 to contain exactly two coins; found 3".to_string()
+        }
+    );
+
     // attempting to set window_size bigger than 172800 sec (48h)
     let err = set_price_source_twap("ustatom", "uatom", 803, 172801, None).unwrap_err();
     assert_eq!(
@@ -814,7 +850,7 @@ fn setting_price_source_lsd_with_invalid_params() {
     );
 
     // attempting to use not XYK pool
-    let err = set_price_source_twap("ustatom", "uatom", 4444, 86400, None).unwrap_err();
+    let err = set_price_source_twap("uion", "uosmo", 4444, 86400, None).unwrap_err();
     assert_eq!(
         err,
         ContractError::InvalidPriceSource {
@@ -822,6 +858,15 @@ fn setting_price_source_lsd_with_invalid_params() {
         }
     );
 
+    // attempting to use a StableSwap pool that contains more than two assets; should fail
+    let err = set_price_source_twap("uatom", "uusdc", 6666, 86400, None).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::InvalidPriceSource {
+            reason: "expecting pool 6666 to contain exactly two coins; found 3".to_string()
+        }
+    );
+
     // attempting to set window_size bigger than 172800 sec (48h)
     let err = set_price_source_twap("ustatom", "uatom", 803, 172801, None).unwrap_err();
     assert_eq!(