@layerzerolabs/protocol-stellar-v2 0.2.19 → 0.2.21

package/contracts/oapps/oft/src/extensions/pausable.rs

@@ -1,6 +1,10 @@
 use common_macros::{contract_error, contract_trait, only_auth, storage};
 use soroban_sdk::{assert_with_error, contractevent, Env};
-use utils::ownable::Ownable;
+use utils::auth::Auth;
+
+// =========================================================================
+// Storage
+// =========================================================================
 
 #[storage]
 pub enum OFTPausableStorage {
@@ -9,42 +13,87 @@ pub enum OFTPausableStorage {
     Paused,
 }
 
+// =========================================================================
+// Errors
+// =========================================================================
+
 #[contract_error]
 pub enum OFTPausableError {
-    Paused = 2300,
+    Paused = 3110,
     PauseStatusUnchanged,
 }
 
+// =========================================================================
+// Events
+// =========================================================================
+
 #[contractevent]
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct PausedSet {
     pub paused: bool,
 }
 
+// =========================================================================
+// Trait With Default Implementations
+// =========================================================================
+
 #[contract_trait]
-pub trait OFTPausable: OFTPausableInternal + Ownable + Sized {
+pub trait OFTPausable: OFTPausableInternal + Auth {
+    /// Sets the paused state of the OFT.
+    ///
+    /// When paused, the OFT will reject new send/receive/quote_send/quote_oft operations.
+    ///
+    /// # Arguments
+    /// * `paused` - `true` to pause, `false` to unpause
     #[only_auth]
-    fn set_paused(env: &Env, paused: bool) {
-        let current_paused = OFTPausableStorage::paused(env);
-        assert_with_error!(env, current_paused != paused, OFTPausableError::PauseStatusUnchanged);
-        OFTPausableStorage::set_paused(env, &paused);
-        PausedSet { paused }.publish(env);
+    fn set_paused(env: &soroban_sdk::Env, paused: bool) {
+        Self::__set_paused(env, paused);
     }
 
-    fn is_paused(env: &Env) -> bool {
-        OFTPausableStorage::paused(env)
+    /// Returns the paused state of the OFT.
+    fn is_paused(env: &soroban_sdk::Env) -> bool {
+        Self::__is_paused(env)
     }
 }
 
 /// Internal trait for pausable operations used by OFT hooks.
 /// Contains only truly internal methods that are called from OFTInternal implementations.
 pub trait OFTPausableInternal {
+    // =========================================================================
+    // OFT Hooks
+    // =========================================================================
+
     /// Asserts that the OFT is not paused, panics otherwise.
-    /// Used internally by `send`, `quote_send`, and `__lz_receive` to enforce pause state.
+    /// Used internally by `send`, `quote_send`, `quote_oft`, and `__lz_receive` to enforce pause state.
     ///
     /// # Errors
     /// * `Paused` - If the OFT is currently paused.
     fn __assert_not_paused(env: &Env) {
-        assert_with_error!(env, !OFTPausableStorage::paused(env), OFTPausableError::Paused);
+        assert_with_error!(env, !Self::__is_paused(env), OFTPausableError::Paused);
+    }
+
+    // =========================================================================
+    // Management Functions
+    // =========================================================================
+
+    /// Sets the paused state of the OFT.
+    ///
+    /// When paused, the OFT will reject new send/receive/quote_send/quote_oft operations.
+    ///
+    /// # Arguments
+    /// * `paused` - `true` to pause, `false` to unpause
+    fn __set_paused(env: &Env, paused: bool) {
+        assert_with_error!(env, Self::__is_paused(env) != paused, OFTPausableError::PauseStatusUnchanged);
+        OFTPausableStorage::set_paused(env, &paused);
+        PausedSet { paused }.publish(env);
+    }
+
+    // =========================================================================
+    // View Functions
+    // =========================================================================
+
+    /// Returns the paused state of the OFT.
+    fn __is_paused(env: &Env) -> bool {
+        OFTPausableStorage::paused(env)
     }
 }

package/contracts/oapps/oft/src/extensions/rate_limiter.rs

@@ -1,6 +1,11 @@
+use crate as oft;
 use common_macros::{contract_error, contract_trait, only_auth, storage};
-use soroban_sdk::{assert_with_error, contractevent, contracttype, panic_with_error, Env};
-use utils::ownable::Ownable;
+use soroban_sdk::{assert_with_error, contractevent, contracttype, Env};
+use utils::auth::Auth;
+
+// =========================================================================
+// Types
+// =========================================================================
 
 #[contracttype]
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -10,191 +15,250 @@ pub enum Direction {
     Outbound,
 }
 
+/// Configuration for rate limiting, used as input parameter.
 #[contracttype]
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct RateLimit {
-    limit: i128,
-    window_seconds: u64,
+pub struct RateLimitConfig {
+    pub limit: i128,
+    pub window_seconds: u64,
+}
+
+/// Internal storage struct for rate limit state.
+///
+/// The rate limiter uses a "leaky bucket" algorithm where:
+/// - `config.limit` defines the maximum tokens that can be "in flight" at any time
+/// - `config.window_seconds` defines how long it takes for the bucket to fully drain
+/// - Tokens decay linearly over time: `decay = elapsed_time * limit / window_seconds`
+/// - Current in-flight = `in_flight_on_last_update - decay` (clamped to 0)
+#[contracttype]
+#[derive(Clone, Debug, Eq, PartialEq)]
+struct RateLimitState {
+    /// The rate limit configuration (limit and window_seconds)
+    config: RateLimitConfig,
+    /// The in-flight amount at the time of `last_update` (before decay is applied)
     in_flight_on_last_update: i128,
+    /// Timestamp of the last update (used to calculate decay)
     last_update: u64,
 }
 
+// =========================================================================
+// Storage
+// =========================================================================
+
 #[storage]
-pub enum RateLimitStorage {
-    #[persistent(RateLimit)]
+enum RateLimitStorage {
+    #[persistent(RateLimitState)]
     RateLimit { direction: Direction, eid: u32 },
 }
 
+// =========================================================================
+// Errors
+// =========================================================================
+
 #[contract_error]
 pub enum RateLimitError {
-    ExceededRateLimit = 2400,
+    ExceededRateLimit = 3120,
+    InvalidAmount,
     InvalidTimestamp,
-    InvalidWindowSeconds,
-    InvalidLimit,
+    InvalidConfig,
     SameValue,
 }
 
-#[contractevent]
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub struct RateLimitSet {
-    pub direction: Direction,
-    pub eid: u32,
-    pub limit: i128,
-    pub window_seconds: u64,
-}
-
-#[contractevent]
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub struct RateLimitUpdated {
-    pub direction: Direction,
-    pub eid: u32,
-    pub limit: i128,
-    pub window_seconds: u64,
-}
+// =========================================================================
+// Events
+// =========================================================================
 
 #[contractevent]
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct RateLimitUnset {
+pub struct RateLimitSet {
     pub direction: Direction,
     pub eid: u32,
+    /// The rate limit configuration, or None if the rate limit is removed
+    pub config: Option<RateLimitConfig>,
 }
 
-/// Helper function to calculate the current in-flight amount with decay.
-/// Used by both public trait methods and internal implementations.
-fn calculate_in_flight(env: &Env, direction: &Direction, eid: u32) -> i128 {
-    if !RateLimitStorage::has_rate_limit(env, direction, eid) {
-        return 0;
-    }
-    let rate_limit = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-    let timestamp = env.ledger().timestamp();
-    assert_with_error!(env, timestamp >= rate_limit.last_update, RateLimitError::InvalidTimestamp);
-    let elapsed = timestamp - rate_limit.last_update;
-    let decay = (elapsed as i128) * rate_limit.limit / (rate_limit.window_seconds as i128);
-    if decay < rate_limit.in_flight_on_last_update {
-        rate_limit.in_flight_on_last_update - decay
-    } else {
-        0
-    }
-}
-
-/// Checkpoints the current in-flight amount by applying decay and updating timestamp.
-/// Used internally before modifying rate limit state.
-fn checkpoint_rate_limit_in_flight(env: &Env, direction: &Direction, eid: u32) {
-    let in_flight = calculate_in_flight(env, direction, eid);
-    let mut rate_limit = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-    rate_limit.in_flight_on_last_update = in_flight;
-    rate_limit.last_update = env.ledger().timestamp();
-    RateLimitStorage::set_rate_limit(env, direction, eid, &rate_limit);
-}
+// =========================================================================
+// Trait With Default Implementations
+// =========================================================================
 
 #[contract_trait]
-pub trait RateLimiter: RateLimiterInternal + Ownable + Sized {
-    #[only_auth]
-    fn set_rate_limit(env: &Env, direction: &Direction, eid: u32, limit: i128, window_seconds: u64) {
-        assert_with_error!(env, limit > 0, RateLimitError::InvalidLimit);
-        assert_with_error!(env, window_seconds > 0, RateLimitError::InvalidWindowSeconds);
-        if RateLimitStorage::has_rate_limit(env, direction, eid) {
-            let rate_limit_data = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-            assert_with_error!(
-                env,
-                limit != rate_limit_data.limit || window_seconds != rate_limit_data.window_seconds,
-                RateLimitError::SameValue
-            );
-            checkpoint_rate_limit_in_flight(env, direction, eid);
-            let mut rate_limit = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-            rate_limit.limit = limit;
-            rate_limit.window_seconds = window_seconds;
-            RateLimitStorage::set_rate_limit(env, direction, eid, &rate_limit);
-            RateLimitUpdated { direction: direction.clone(), eid, limit, window_seconds }.publish(env);
-        } else {
-            RateLimitStorage::set_rate_limit(
-                env,
-                direction,
-                eid,
-                &RateLimit {
-                    limit,
-                    window_seconds,
-                    in_flight_on_last_update: 0,
-                    last_update: env.ledger().timestamp(),
-                },
-            );
-            RateLimitSet { direction: direction.clone(), eid, limit, window_seconds }.publish(env);
-        }
-    }
+pub trait RateLimiter: RateLimiterInternal + Auth {
+    // =========================================================================
+    // Management Functions
+    // =========================================================================
 
+    /// Sets or removes a rate limit for a specific direction and endpoint.
+    ///
+    /// # Arguments
+    /// * `direction` - The direction (Inbound or Outbound)
+    /// * `eid` - The endpoint ID
+    /// * `config` - The rate limit configuration, or None to remove the rate limit
     #[only_auth]
-    fn unset_rate_limit(env: &Env, direction: &Direction, eid: u32) {
-        assert_with_error!(env, RateLimitStorage::has_rate_limit(env, direction, eid), RateLimitError::SameValue);
-        RateLimitStorage::remove_rate_limit(env, direction, eid);
-        RateLimitUnset { direction: direction.clone(), eid }.publish(env);
+    fn set_rate_limit(
+        env: &soroban_sdk::Env,
+        direction: &oft::rate_limiter::Direction,
+        eid: u32,
+        config: Option<oft::rate_limiter::RateLimitConfig>,
+    ) {
+        Self::__set_rate_limit(env, direction, eid, config);
     }
 
-    fn rate_limit_config(env: &Env, direction: &Direction, eid: u32) -> (i128, u64) {
-        if !RateLimitStorage::has_rate_limit(env, direction, eid) {
-            return (0, 0);
-        }
-        let rate_limit = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-        (rate_limit.limit, rate_limit.window_seconds)
+    // =========================================================================
+    // View Functions
+    // =========================================================================
+
+    /// Returns the rate limit configuration for a direction and endpoint.
+    /// Returns None if no rate limit is configured.
+    fn rate_limit_config(
+        env: &soroban_sdk::Env,
+        direction: &oft::rate_limiter::Direction,
+        eid: u32,
+    ) -> Option<oft::rate_limiter::RateLimitConfig> {
+        Self::__rate_limit_config(env, direction, eid)
     }
 
-    fn rate_limit_in_flight(env: &Env, direction: &Direction, eid: u32) -> i128 {
-        calculate_in_flight(env, direction, eid)
+    /// Returns the current in-flight amount for a direction and endpoint.
+    fn rate_limit_in_flight(env: &soroban_sdk::Env, direction: &oft::rate_limiter::Direction, eid: u32) -> i128 {
+        Self::__rate_limit_in_flight(env, direction, eid)
     }
 
-    fn rate_limit_capacity(env: &Env, direction: &Direction, eid: u32) -> i128 {
-        if !RateLimitStorage::has_rate_limit(env, direction, eid) {
-            return i128::MAX;
-        }
-        let rate_limit = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-        let in_flight = calculate_in_flight(env, direction, eid);
-        if rate_limit.limit > in_flight {
-            rate_limit.limit - in_flight
-        } else {
-            0
-        }
+    /// Returns the available capacity for a direction and endpoint.
+    /// Returns i128::MAX if no rate limit is configured.
+    fn rate_limit_capacity(env: &soroban_sdk::Env, direction: &oft::rate_limiter::Direction, eid: u32) -> i128 {
+        Self::__rate_limit_capacity(env, direction, eid)
     }
 }
 
 /// Internal trait for rate limiter operations used by OFT hooks.
 /// Contains only truly internal methods that are called from OFTInternal implementations.
 pub trait RateLimiterInternal {
+    // =========================================================================
+    // OFT Hooks
+    // =========================================================================
+
     /// Consumes the specified amount from the rate limit capacity.
     /// Used internally by `__debit` and `__credit` to enforce rate limits.
     ///
     /// # Errors
     /// * `ExceededRateLimit` - If the amount exceeds the available capacity.
     fn __consume_rate_limit_capacity(env: &Env, direction: &Direction, eid: u32, amount: i128) {
-        if !RateLimitStorage::has_rate_limit(env, direction, eid) {
+        assert_with_error!(env, amount >= 0, RateLimitError::InvalidAmount);
+
+        let Some(mut state) = RateLimitStorage::rate_limit(env, direction, eid) else {
             return;
-        }
-        checkpoint_rate_limit_in_flight(env, direction, eid);
-        let mut rate_limit = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-        // Check against remaining capacity (limit - current in_flight), not total limit
-        let capacity = if rate_limit.limit > rate_limit.in_flight_on_last_update {
-            rate_limit.limit - rate_limit.in_flight_on_last_update
-        } else {
-            0
         };
-        if amount > capacity {
-            panic_with_error!(env, RateLimitError::ExceededRateLimit);
-        }
-        rate_limit.in_flight_on_last_update += amount;
-        RateLimitStorage::set_rate_limit(env, direction, eid, &rate_limit);
+
+        // Apply decay and update timestamp
+        let in_flight = calculate_decayed_in_flight(env, &state);
+        state.in_flight_on_last_update = in_flight;
+        state.last_update = env.ledger().timestamp();
+
+        // Check capacity and consume
+        let capacity = (state.config.limit - in_flight).max(0);
+        assert_with_error!(env, amount <= capacity, RateLimitError::ExceededRateLimit);
+        state.in_flight_on_last_update += amount;
+
+        RateLimitStorage::set_rate_limit(env, direction, eid, &state);
     }
 
     /// Releases the specified amount back to the rate limit capacity.
     /// Used internally by `__credit` to release outbound capacity on inbound messages.
     fn __release_rate_limit_capacity(env: &Env, direction: &Direction, eid: u32, amount: i128) {
-        if !RateLimitStorage::has_rate_limit(env, direction, eid) {
+        assert_with_error!(env, amount >= 0, RateLimitError::InvalidAmount);
+
+        let Some(mut state) = RateLimitStorage::rate_limit(env, direction, eid) else {
             return;
+        };
+
+        // Apply decay and update timestamp
+        let in_flight = calculate_decayed_in_flight(env, &state);
+        state.in_flight_on_last_update = (in_flight - amount).max(0);
+        state.last_update = env.ledger().timestamp();
+
+        RateLimitStorage::set_rate_limit(env, direction, eid, &state);
+    }
+
+    // =========================================================================
+    // Management Functions
+    // =========================================================================
+
+    /// Sets or removes a rate limit for a specific direction and endpoint.
+    ///
+    /// # Arguments
+    /// * `direction` - The direction (Inbound or Outbound)
+    /// * `eid` - The endpoint ID
+    /// * `config` - The rate limit configuration, or None to remove the rate limit
+    fn __set_rate_limit(env: &Env, direction: &Direction, eid: u32, config: Option<RateLimitConfig>) {
+        let current_state = RateLimitStorage::rate_limit(env, direction, eid);
+        let current_config = current_state.as_ref().map(|s| s.config.clone());
+        assert_with_error!(env, current_config != config, RateLimitError::SameValue);
+
+        match config {
+            Some(cfg) => {
+                assert_with_error!(env, cfg.limit >= 0 && cfg.window_seconds > 0, RateLimitError::InvalidConfig);
+
+                let state = if let Some(mut existing) = current_state {
+                    // Update existing: checkpoint in-flight before changing config
+                    existing.in_flight_on_last_update = calculate_decayed_in_flight(env, &existing);
+                    existing.last_update = env.ledger().timestamp();
+                    existing.config = cfg.clone();
+                    existing
+                } else {
+                    // Create new
+                    RateLimitState {
+                        config: cfg.clone(),
+                        in_flight_on_last_update: 0,
+                        last_update: env.ledger().timestamp(),
+                    }
+                };
+
+                RateLimitStorage::set_rate_limit(env, direction, eid, &state);
+                RateLimitSet { direction: direction.clone(), eid, config: Some(cfg) }.publish(env);
+            }
+            None => {
+                RateLimitStorage::remove_rate_limit(env, direction, eid);
+                RateLimitSet { direction: direction.clone(), eid, config: None }.publish(env);
+            }
         }
-        checkpoint_rate_limit_in_flight(env, direction, eid);
-        let mut rate_limit = RateLimitStorage::rate_limit(env, direction, eid).unwrap();
-        if amount >= rate_limit.in_flight_on_last_update {
-            rate_limit.in_flight_on_last_update = 0;
-        } else {
-            rate_limit.in_flight_on_last_update -= amount;
-        }
-        RateLimitStorage::set_rate_limit(env, direction, eid, &rate_limit);
     }
+
+    // =========================================================================
+    // View Functions
+    // =========================================================================
+
+    /// Returns the rate limit configuration for a direction and endpoint.
+    /// Returns None if no rate limit is configured.
+    fn __rate_limit_config(env: &Env, direction: &Direction, eid: u32) -> Option<RateLimitConfig> {
+        RateLimitStorage::rate_limit(env, direction, eid).map(|s| s.config)
+    }
+
+    /// Returns the current in-flight amount for a direction and endpoint.
+    fn __rate_limit_in_flight(env: &Env, direction: &Direction, eid: u32) -> i128 {
+        RateLimitStorage::rate_limit(env, direction, eid).map(|s| calculate_decayed_in_flight(env, &s)).unwrap_or(0)
+    }
+
+    /// Returns the available capacity for a direction and endpoint.
+    /// Returns i128::MAX if no rate limit is configured.
+    fn __rate_limit_capacity(env: &Env, direction: &Direction, eid: u32) -> i128 {
+        RateLimitStorage::rate_limit(env, direction, eid)
+            .map(|s| (s.config.limit - calculate_decayed_in_flight(env, &s)).max(0))
+            .unwrap_or(i128::MAX)
+    }
+}
+
+// =========================================================================
+// Helper Functions
+// =========================================================================
+
+/// Calculates the current in-flight amount with decay applied.
+/// This is a pure function that doesn't access storage.
+fn calculate_decayed_in_flight(env: &Env, state: &RateLimitState) -> i128 {
+    let timestamp = env.ledger().timestamp();
+    assert_with_error!(env, timestamp >= state.last_update, RateLimitError::InvalidTimestamp);
+
+    let elapsed = timestamp - state.last_update;
+    let decay = (elapsed as i128) * state.config.limit / (state.config.window_seconds as i128);
+
+    // Ensure the decayed in-flight amount is not negative
+    (state.in_flight_on_last_update - decay).max(0)
 }