From a90af894c463f32783b075bdf7e52e4edc5bd367 Mon Sep 17 00:00:00 2001
From: Attiss Ngo <92927591+AttissNgo@users.noreply.github.com>
Date: Mon, 13 Jan 2025 20:24:27 -0600
Subject: [PATCH] feat(interchain-token-service): add flow limit

- improve docstring describing flow direction
- remove unneeded variables in add_flow
---
 .../interchain-token-service/src/error.rs      |  1 +
 .../interchain-token-service/src/flow_limit.rs | 18 ++++++++++--------
 .../interchain-token-service/src/interface.rs  |  5 ++---
 .../tests/flow_limit.rs                        |  2 +-
 4 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/contracts/interchain-token-service/src/error.rs b/contracts/interchain-token-service/src/error.rs
index c7db7224..2d18f11a 100644
--- a/contracts/interchain-token-service/src/error.rs
+++ b/contracts/interchain-token-service/src/error.rs
@@ -24,4 +24,5 @@ pub enum ContractError {
     TokenAlreadyDeployed = 18,
     InvalidFlowLimit = 19,
     FlowLimitExceeded = 20,
+    FlowAmountOverflow = 21,
 }
diff --git a/contracts/interchain-token-service/src/flow_limit.rs b/contracts/interchain-token-service/src/flow_limit.rs
index 8f177f49..84d8e33b 100644
--- a/contracts/interchain-token-service/src/flow_limit.rs
+++ b/contracts/interchain-token-service/src/flow_limit.rs
@@ -52,7 +52,7 @@ impl FlowDirection {
     /// Checks that:
     /// - Flow amount doesn't exceed the flow limit
     /// - Adding flows won't cause overflow
-    /// - Total flow in one direction doesn't exceed flow in opposite direction plus limit
+    /// - Net flow (outgoing minus incoming flow) doesn't exceed the limit
     pub fn add_flow(
         &self,
         env: &Env,
@@ -63,16 +63,18 @@ impl FlowDirection {
             return Ok(());
         };
 
-        let flow_to_add = self.flow(env, token_id.clone());
-        let flow_to_compare = self.reverse_flow(env, token_id.clone());
-
         ensure!(flow_amount <= flow_limit, ContractError::FlowLimitExceeded);
-        let new_flow = flow_to_add
+
+        let new_flow = self
+            .flow(env, token_id.clone())
             .checked_add(flow_amount)
-            .ok_or(ContractError::FlowLimitExceeded)?;
-        let max_allowed = flow_to_compare
+            .ok_or(ContractError::FlowAmountOverflow)?;
+        let max_allowed = self
+            .reverse_flow(env, token_id.clone())
             .checked_add(flow_limit)
-            .ok_or(ContractError::FlowLimitExceeded)?;
+            .ok_or(ContractError::FlowAmountOverflow)?;
+
+        // Equivalent to flow_amount + flow - reverse_flow <= flow_limit
         ensure!(new_flow <= max_allowed, ContractError::FlowLimitExceeded);
 
         self.update_flow(env, token_id.clone(), new_flow);
diff --git a/contracts/interchain-token-service/src/interface.rs b/contracts/interchain-token-service/src/interface.rs
index 08443191..01e6c711 100644
--- a/contracts/interchain-token-service/src/interface.rs
+++ b/contracts/interchain-token-service/src/interface.rs
@@ -38,11 +38,11 @@ pub trait InterchainTokenServiceInterface: AxelarExecutableInterface {
     /// Returns `None` if no limit is set.
     fn flow_limit(env: &Env, token_id: BytesN<32>) -> Option<i128>;
 
-    /// Retrieves the amount that has flowed out of the contract during the current epoch
+    /// Retrieves the amount that has flowed out of the chain to other chains during the current epoch
     /// for the token associated with the specified token ID.
     fn flow_out_amount(env: &Env, token_id: BytesN<32>) -> i128;
 
-    /// Retrieves the amount that has flowed into the contract during the current epoch
+    /// Retrieves the amount that has flowed into the chain from other chains during the current epoch
     /// for the token associated with the specified token ID.
     fn flow_in_amount(env: &Env, token_id: BytesN<32>) -> i128;
 
@@ -53,7 +53,6 @@ pub trait InterchainTokenServiceInterface: AxelarExecutableInterface {
     /// Setting the limit to 0 effectively freezes the token by preventing any flow.
     ///
     /// # Arguments
-    /// - `env`: Reference to the contract environment.
     /// - `token_id`: Unique identifier of the token.
     /// - `flow_limit`: The new flow limit value. Must be positive if Some.
     ///
diff --git a/contracts/interchain-token-service/tests/flow_limit.rs b/contracts/interchain-token-service/tests/flow_limit.rs
index d2febd2c..1882f761 100644
--- a/contracts/interchain-token-service/tests/flow_limit.rs
+++ b/contracts/interchain-token-service/tests/flow_limit.rs
@@ -253,7 +253,7 @@ fn add_flow_out_fails_exceeds_flow_limit() {
 }
 
 #[test]
-#[should_panic(expected = "Error(Contract, #20)")] // ContractError::FlowLimitExceeded
+#[should_panic(expected = "Error(Contract, #21)")] // ContractError::FlowAmountOverflow
 fn add_flow_fails_on_flow_comparison_overflow() {
     let (env, client, gateway_client, _, signers) = setup_env();
     register_chains(&env, &client);