feat: (PRO-262) Implement usage limit feature with Redis support

crates/lib/Cargo.toml

@@ -85,4 +85,3 @@ tempfile = "3.2"
 tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
 mockito = "1.2.0"
 serial_test = "3.2.0"
-redis-test = "0.12.0"

crates/lib/src/config.rs

@@ -10,6 +10,7 @@ use crate::{
         DEFAULT_CACHE_ACCOUNT_TTL, DEFAULT_CACHE_DEFAULT_TTL,
         DEFAULT_FEE_PAYER_BALANCE_METRICS_EXPIRY_SECONDS, DEFAULT_MAX_TIMESTAMP_AGE,
         DEFAULT_METRICS_ENDPOINT, DEFAULT_METRICS_PORT, DEFAULT_METRICS_SCRAPE_INTERVAL,
+        DEFAULT_USAGE_LIMIT_DEFAULT_MAX_TRANSACTIONS, DEFAULT_USAGE_LIMIT_FALLBACK_IF_UNAVAILABLE,
     },
     error::KoraError,
     fee::price::{PriceConfig, PriceModel},
@@ -349,6 +350,8 @@ pub struct KoraConfig {
     pub payment_address: Option<String>,
     #[serde(default)]
     pub cache: CacheConfig,
+    #[serde(default)]
+    pub usage_limit: UsageLimitConfig,
 }
 
 impl Default for KoraConfig {
@@ -359,6 +362,30 @@ impl Default for KoraConfig {
             auth: AuthConfig::default(),
             payment_address: None,
             cache: CacheConfig::default(),
+            usage_limit: UsageLimitConfig::default(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
+pub struct UsageLimitConfig {
+    /// Enable per-wallet usage limiting
+    pub enabled: bool,
+    /// Cache URL for shared usage limiting across multiple Kora instances
+    pub cache_url: Option<String>,
+    /// Default maximum transactions per wallet (0 = unlimited)
+    pub default_max_transactions: u64,
+    /// Fallback behavior when cache is unavailable
+    pub fallback_if_unavailable: bool,
+}
+
+impl Default for UsageLimitConfig {
+    fn default() -> Self {
+        Self {
+            enabled: false,
+            cache_url: None,
+            default_max_transactions: DEFAULT_USAGE_LIMIT_DEFAULT_MAX_TRANSACTIONS,
+            fallback_if_unavailable: DEFAULT_USAGE_LIMIT_FALLBACK_IF_UNAVAILABLE,
         }
     }
 }
@@ -671,4 +698,44 @@ mod tests {
         assert_eq!(config.kora.cache.default_ttl, 300);
         assert_eq!(config.kora.cache.account_ttl, 60);
     }
+
+    #[test]
+    fn test_usage_limit_config_parsing() {
+        let config = ConfigBuilder::new()
+            .with_usage_limit_config(true, Some("redis://localhost:6379"), 10, false)
+            .build_config()
+            .unwrap();
+
+        assert!(config.kora.usage_limit.enabled);
+        assert_eq!(config.kora.usage_limit.cache_url, Some("redis://localhost:6379".to_string()));
+        assert_eq!(config.kora.usage_limit.default_max_transactions, 10);
+        assert!(!config.kora.usage_limit.fallback_if_unavailable);
+    }
+
+    #[test]
+    fn test_usage_limit_config_default() {
+        let config = ConfigBuilder::new().build_config().unwrap();
+
+        assert!(!config.kora.usage_limit.enabled);
+        assert_eq!(config.kora.usage_limit.cache_url, None);
+        assert_eq!(
+            config.kora.usage_limit.default_max_transactions,
+            DEFAULT_USAGE_LIMIT_DEFAULT_MAX_TRANSACTIONS
+        );
+        assert_eq!(
+            config.kora.usage_limit.fallback_if_unavailable,
+            DEFAULT_USAGE_LIMIT_FALLBACK_IF_UNAVAILABLE
+        );
+    }
+
+    #[test]
+    fn test_usage_limit_config_unlimited() {
+        let config = ConfigBuilder::new()
+            .with_usage_limit_config(true, None, 0, true)
+            .build_config()
+            .unwrap();
+
+        assert!(config.kora.usage_limit.enabled);
+        assert_eq!(config.kora.usage_limit.default_max_transactions, 0); // 0 = unlimited
+    }
 }

crates/lib/src/constant.rs

@@ -25,6 +25,9 @@ pub const DEFAULT_CACHE_DEFAULT_TTL: u64 = 300; // 5 minutes
 pub const DEFAULT_CACHE_ACCOUNT_TTL: u64 = 60; // 1 minute for account data
 pub const DEFAULT_FEE_PAYER_BALANCE_METRICS_EXPIRY_SECONDS: u64 = 30; // 30 seconds
 
+pub const DEFAULT_USAGE_LIMIT_DEFAULT_MAX_TRANSACTIONS: u64 = 0; // 0 = unlimited
+pub const DEFAULT_USAGE_LIMIT_FALLBACK_IF_UNAVAILABLE: bool = false;
+
 // Account Indexes within instructions
 // Instruction indexes for the instructions that we support to parse from the transaction
 pub mod instruction_indexes {

crates/lib/src/error.rs

@@ -58,6 +58,9 @@ pub enum KoraError {
 
     #[error("Rate limit exceeded")]
     RateLimitExceeded,
+
+    #[error("Usage limit exceeded: {0}")]
+    UsageLimitExceeded(String),
 }
 
 impl From<ClientError> for KoraError {

crates/lib/src/mod.rs

@@ -17,6 +17,7 @@ pub mod signer;
 pub mod state;
 pub mod token;
 pub mod transaction;
+pub mod usage_limit;
 pub mod validator;
 pub use cache::CacheUtil;
 pub use config::Config;

crates/lib/src/rpc_server/method/estimate_transaction_fee.rs

@@ -5,13 +5,19 @@ use crate::{
     error::KoraError,
     fee::fee::FeeConfigUtil,
     rpc_server::middleware_utils::default_sig_verify,
-    state::{get_config, get_request_signer_with_signer_key},
+    state::get_request_signer_with_signer_key,
     transaction::{TransactionUtil, VersionedTransactionResolved},
 };
 
 use serde::{Deserialize, Serialize};
 use solana_client::nonblocking::rpc_client::RpcClient;
 
+#[cfg(not(test))]
+use crate::state::get_config;
+
+#[cfg(test)]
+use crate::tests::config_mock::mock_state::get_config;
+
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
 pub struct EstimateTransactionFeeRequest {
     pub transaction: String, // Base64 encoded serialized transaction

crates/lib/src/rpc_server/method/sign_and_send_transaction.rs

@@ -1,4 +1,4 @@
-use crate::rpc_server::middleware_utils::default_sig_verify;
+use crate::{rpc_server::middleware_utils::default_sig_verify, usage_limit::UsageTracker};
 use serde::{Deserialize, Serialize};
 use solana_client::nonblocking::rpc_client::RpcClient;
 use std::sync::Arc;
@@ -34,6 +34,10 @@ pub async fn sign_and_send_transaction(
     request: SignAndSendTransactionRequest,
 ) -> Result<SignAndSendTransactionResponse, KoraError> {
     let transaction = TransactionUtil::decode_b64_transaction(&request.transaction)?;
+
+    // Check usage limit for transaction sender
+    UsageTracker::check_transaction_usage_limit(&transaction).await?;
+
     let signer = get_request_signer_with_signer_key(request.signer_key.as_deref())?;
 
     let mut resolved_transaction = VersionedTransactionResolved::from_transaction(
@@ -57,7 +61,7 @@ pub async fn sign_and_send_transaction(
 mod tests {
     use super::*;
     use crate::tests::{
-        common::{setup_or_get_test_signer, RpcMockBuilder},
+        common::{setup_or_get_test_signer, setup_or_get_test_usage_limiter, RpcMockBuilder},
         config_mock::ConfigMockBuilder,
         transaction_mock::create_mock_encoded_transaction,
     };
@@ -67,6 +71,8 @@ mod tests {
         let _m = ConfigMockBuilder::new().build_and_setup();
         let _ = setup_or_get_test_signer();
 
+        let _ = setup_or_get_test_usage_limiter().await;
+
         let rpc_client = Arc::new(RpcMockBuilder::new().build());
 
         let request = SignAndSendTransactionRequest {
@@ -85,6 +91,8 @@ mod tests {
         let _m = ConfigMockBuilder::new().build_and_setup();
         let _ = setup_or_get_test_signer();
 
+        let _ = setup_or_get_test_usage_limiter().await;
+
         let rpc_client = Arc::new(RpcMockBuilder::new().build());
 
         let request = SignAndSendTransactionRequest {

crates/lib/src/rpc_server/method/sign_transaction.rs

@@ -2,6 +2,7 @@ use crate::{
     rpc_server::middleware_utils::default_sig_verify,
     state::get_request_signer_with_signer_key,
     transaction::{TransactionUtil, VersionedTransactionOps, VersionedTransactionResolved},
+    usage_limit::UsageTracker,
     KoraError,
 };
 use serde::{Deserialize, Serialize};
@@ -33,6 +34,10 @@ pub async fn sign_transaction(
     request: SignTransactionRequest,
 ) -> Result<SignTransactionResponse, KoraError> {
     let transaction = TransactionUtil::decode_b64_transaction(&request.transaction)?;
+
+    // Check usage limit for transaction sender
+    UsageTracker::check_transaction_usage_limit(&transaction).await?;
+
     let signer = get_request_signer_with_signer_key(request.signer_key.as_deref())?;
 
     let mut resolved_transaction = VersionedTransactionResolved::from_transaction(
@@ -58,7 +63,7 @@ pub async fn sign_transaction(
 mod tests {
     use super::*;
     use crate::tests::{
-        common::{setup_or_get_test_signer, RpcMockBuilder},
+        common::{setup_or_get_test_signer, setup_or_get_test_usage_limiter, RpcMockBuilder},
         config_mock::ConfigMockBuilder,
         transaction_mock::create_mock_encoded_transaction,
     };
@@ -68,6 +73,8 @@ mod tests {
         let _m = ConfigMockBuilder::new().build_and_setup();
         let _ = setup_or_get_test_signer();
 
+        let _ = setup_or_get_test_usage_limiter().await;
+
         let rpc_client = Arc::new(RpcMockBuilder::new().build());
 
         let request = SignTransactionRequest {
@@ -86,6 +93,8 @@ mod tests {
         let _m = ConfigMockBuilder::new().build_and_setup();
         let _ = setup_or_get_test_signer();
 
+        let _ = setup_or_get_test_usage_limiter().await;
+
         let rpc_client = Arc::new(RpcMockBuilder::new().build());
 
         let request = SignTransactionRequest {

crates/lib/src/rpc_server/method/sign_transaction_if_paid.rs

@@ -2,6 +2,7 @@ use crate::{
     rpc_server::middleware_utils::default_sig_verify,
     state::get_request_signer_with_signer_key,
     transaction::{TransactionUtil, VersionedTransactionOps, VersionedTransactionResolved},
+    usage_limit::UsageTracker,
     KoraError,
 };
 use serde::{Deserialize, Serialize};
@@ -33,6 +34,10 @@ pub async fn sign_transaction_if_paid(
     request: SignTransactionIfPaidRequest,
 ) -> Result<SignTransactionIfPaidResponse, KoraError> {
     let transaction_requested = TransactionUtil::decode_b64_transaction(&request.transaction)?;
+
+    // Check usage limit for transaction sender
+    UsageTracker::check_transaction_usage_limit(&transaction_requested).await?;
+
     let signer = get_request_signer_with_signer_key(request.signer_key.as_deref())?;
 
     let mut resolved_transaction = VersionedTransactionResolved::from_transaction(
@@ -58,7 +63,7 @@ pub async fn sign_transaction_if_paid(
 mod tests {
     use super::*;
     use crate::tests::{
-        common::{setup_or_get_test_signer, RpcMockBuilder},
+        common::{setup_or_get_test_signer, setup_or_get_test_usage_limiter, RpcMockBuilder},
         config_mock::ConfigMockBuilder,
         transaction_mock::create_mock_encoded_transaction,
     };
@@ -68,6 +73,8 @@ mod tests {
         let _m = ConfigMockBuilder::new().build_and_setup();
         let _ = setup_or_get_test_signer();
 
+        let _ = setup_or_get_test_usage_limiter().await;
+
         let rpc_client = Arc::new(RpcMockBuilder::new().build());
 
         let request = SignTransactionIfPaidRequest {
@@ -86,6 +93,8 @@ mod tests {
         let _m = ConfigMockBuilder::new().build_and_setup();
         let _ = setup_or_get_test_signer();
 
+        let _ = setup_or_get_test_usage_limiter().await;
+
         let rpc_client = Arc::new(RpcMockBuilder::new().build());
 
         let request = SignTransactionIfPaidRequest {

crates/lib/src/rpc_server/server.rs

@@ -5,6 +5,7 @@ use crate::{
         auth::{ApiKeyAuthLayer, HmacAuthLayer},
         rpc::KoraRpc,
     },
+    usage_limit::UsageTracker,
 };
 
 #[cfg(not(test))]
@@ -37,6 +38,12 @@ pub async fn run_rpc_server(rpc: KoraRpc, port: u16) -> Result<ServerHandles, an
     let addr = SocketAddr::from(([0, 0, 0, 0], port));
     log::info!("RPC server started on {addr}, port {port}");
 
+    // Initialize usage limiter
+    if let Err(e) = UsageTracker::init_usage_limiter().await {
+        log::error!("Failed to initialize usage limiter: {e}");
+        return Err(anyhow::anyhow!("Usage limiter initialization failed: {e}"));
+    }
+
     // Build middleware stack with tracing and CORS
     let cors = CorsLayer::new()
         .allow_origin(tower_http::cors::Any)

crates/lib/src/tests/common/mod.rs

@@ -8,7 +8,8 @@ use crate::{
     signer::{KoraSigner, SignerPool, SignerWithMetadata, SolanaMemorySigner},
     state::{get_config, update_config, update_signer_pool},
     tests::{account_mock, config_mock::ConfigMockBuilder, rpc_mock},
-    Config,
+    usage_limit::UsageTracker,
+    Config, KoraError,
 };
 use solana_sdk::{pubkey::Pubkey, signature::Keypair};
 
@@ -61,3 +62,19 @@ pub fn setup_or_get_test_config() -> Config {
         }
     }
 }
+
+/// Initialize or update the global usage limiter (test only)
+///
+/// This function ignores "already initialized" errors for test flexibility.
+/// Usage limiter initialization is optional and will not fail tests if unavailable.
+pub async fn setup_or_get_test_usage_limiter() -> Result<(), KoraError> {
+    match UsageTracker::init_usage_limiter().await {
+        Ok(()) => Ok(()),
+        Err(KoraError::InternalServerError(ref msg)) if msg.contains("already initialized") => {
+            // In tests, ignore the already initialized error
+            // The limiter is already set up from a previous test
+            Ok(())
+        }
+        Err(e) => Err(e),
+    }
+}