feat: allow any spl paid token flag (#175)

Yolley · web-flow · commit 7bf4fa919c44 · 2025-08-28T15:28:55.000-04:00
* feat: allow any spl paid token flag

* move spl paid token checks into a function, rebase, use mock rpc client

* switch to enum in allowed_spl_paid_tokens

* remove unused option from demo config

* rename params in `with`, add test for atas init when all payment tokens are allowed

* use mock client, make sure that signer exists in test initialize atas

* rename token config to make it more generic
diff --git a/crates/lib/src/admin/token_util.rs b/crates/lib/src/admin/token_util.rs
@@ -22,7 +22,7 @@ use {crate::cache::CacheUtil, crate::state::get_config};
 
 #[cfg(test)]
 use {
-    crate::tests::config_mock::mock_state::get_config,
+    crate::config::SplTokenConfig, crate::tests::config_mock::mock_state::get_config,
     crate::tests::redis_cache_mock::MockCacheUtil as CacheUtil,
 };
 
@@ -289,7 +289,9 @@ mod tests {
     async fn test_find_missing_atas_no_spl_tokens() {
         let _m = ConfigMockBuilder::new()
             .with_validation(
-                ValidationConfigBuilder::new().with_allowed_spl_paid_tokens(vec![]).build(),
+                ValidationConfigBuilder::new()
+                    .with_allowed_spl_paid_tokens(SplTokenConfig::Allowlist(vec![]))
+                    .build(),
             )
             .build_and_setup();
 
@@ -308,9 +310,9 @@ mod tests {
         let _m = ConfigMockBuilder::new()
             .with_validation(
                 ValidationConfigBuilder::new()
-                    .with_allowed_spl_paid_tokens(
+                    .with_allowed_spl_paid_tokens(SplTokenConfig::Allowlist(
                         allowed_spl_tokens.iter().map(|p| p.to_string()).collect(),
-                    )
+                    ))
                     .build(),
             )
             .build_and_setup();
@@ -397,4 +399,19 @@ mod tests {
             }
         }
     }
+
+    #[tokio::test]
+    async fn test_initialize_atas_when_all_tokens_are_allowed() {
+        let _m = ConfigMockBuilder::new()
+            .with_allowed_spl_paid_tokens(SplTokenConfig::All)
+            .build_and_setup();
+
+        let _ = setup_or_get_test_signer();
+
+        let rpc_client = RpcMockBuilder::new().build();
+
+        let result = initialize_atas(&rpc_client, None, None, None, None).await;
+
+        assert!(result.is_ok(), "Expected atas init to succeed");
+    }
 }
diff --git a/crates/lib/src/config.rs b/crates/lib/src/config.rs
@@ -58,13 +58,61 @@ impl Default for FeePayerBalanceMetricsConfig {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub enum SplTokenConfig {
+    All,
+    #[serde(untagged)]
+    Allowlist(Vec<String>),
+}
+
+impl Default for SplTokenConfig {
+    fn default() -> Self {
+        SplTokenConfig::Allowlist(vec![])
+    }
+}
+
+impl<'a> IntoIterator for &'a SplTokenConfig {
+    type Item = &'a String;
+    type IntoIter = std::slice::Iter<'a, String>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        match self {
+            SplTokenConfig::All => [].iter(),
+            SplTokenConfig::Allowlist(tokens) => tokens.iter(),
+        }
+    }
+}
+
+impl SplTokenConfig {
+    pub fn has_token(&self, token: &str) -> bool {
+        match self {
+            SplTokenConfig::All => true,
+            SplTokenConfig::Allowlist(tokens) => tokens.iter().any(|s| s == token),
+        }
+    }
+
+    pub fn has_tokens(&self) -> bool {
+        match self {
+            SplTokenConfig::All => true,
+            SplTokenConfig::Allowlist(tokens) => !tokens.is_empty(),
+        }
+    }
+
+    pub fn as_slice(&self) -> &[String] {
+        match self {
+            SplTokenConfig::All => &[],
+            SplTokenConfig::Allowlist(v) => v.as_slice(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
 pub struct ValidationConfig {
     pub max_allowed_lamports: u64,
     pub max_signatures: u64,
     pub allowed_programs: Vec<String>,
     pub allowed_tokens: Vec<String>,
-    pub allowed_spl_paid_tokens: Vec<String>,
+    pub allowed_spl_paid_tokens: SplTokenConfig,
     pub disallowed_accounts: Vec<String>,
     pub price_source: PriceSource,
     #[serde(default)] // Default for backward compatibility
@@ -81,7 +129,7 @@ impl ValidationConfig {
     }
 
     pub fn supports_token(&self, token: &str) -> bool {
-        self.allowed_spl_paid_tokens.iter().any(|s| s == token)
+        self.allowed_spl_paid_tokens.has_token(token)
     }
 }
 
@@ -378,7 +426,7 @@ mod tests {
         let config = ConfigBuilder::new()
             .with_programs(vec!["program1", "program2"])
             .with_tokens(vec!["token1", "token2"])
-            .with_spl_paid_tokens(vec!["token3"])
+            .with_spl_paid_tokens(SplTokenConfig::Allowlist(vec!["token3".to_string()]))
             .with_disallowed_accounts(vec!["account1"])
             .build_config()
             .unwrap();
@@ -387,7 +435,10 @@ mod tests {
         assert_eq!(config.validation.max_signatures, 10);
         assert_eq!(config.validation.allowed_programs, vec!["program1", "program2"]);
         assert_eq!(config.validation.allowed_tokens, vec!["token1", "token2"]);
-        assert_eq!(config.validation.allowed_spl_paid_tokens, vec!["token3"]);
+        assert_eq!(
+            config.validation.allowed_spl_paid_tokens,
+            SplTokenConfig::Allowlist(vec!["token3".to_string()])
+        );
         assert_eq!(config.validation.disallowed_accounts, vec!["account1"]);
         assert_eq!(config.validation.price_source, PriceSource::Jupiter);
         assert_eq!(config.kora.rate_limit, 100);
@@ -400,7 +451,7 @@ mod tests {
         let config = ConfigBuilder::new()
             .with_programs(vec!["program1", "program2"])
             .with_tokens(vec!["token1", "token2"])
-            .with_spl_paid_tokens(vec!["token3"])
+            .with_spl_paid_tokens(SplTokenConfig::Allowlist(vec!["token3".to_string()]))
             .with_disallowed_accounts(vec!["account1"])
             .with_enabled_methods(&[
                 ("liveness", true),
@@ -441,6 +492,14 @@ mod tests {
         assert!(result.is_err());
     }
 
+    #[test]
+    fn test_parse_spl_payment_config() {
+        let config =
+            ConfigBuilder::new().with_spl_paid_tokens(SplTokenConfig::All).build_config().unwrap();
+
+        assert_eq!(config.validation.allowed_spl_paid_tokens, SplTokenConfig::All);
+    }
+
     #[test]
     fn test_parse_margin_price_config() {
         let config = ConfigBuilder::new().with_margin_price(0.1).build_config().unwrap();
diff --git a/crates/lib/src/oracle/jupiter.rs b/crates/lib/src/oracle/jupiter.rs
@@ -246,4 +246,43 @@ mod tests {
         assert_eq!(price.price, 1.0);
         assert_eq!(price.source, PriceSource::Jupiter);
     }
+
+    #[tokio::test]
+    async fn test_jupiter_price_fetch_when_no_price_data() {
+        // No API key
+        {
+            let mut api_key_guard = GLOBAL_JUPITER_API_KEY.write();
+            *api_key_guard = None;
+        }
+
+        let mock_response = r#"{
+            "So11111111111111111111111111111111111111112": {
+                "usdPrice": 100.0,
+                "blockId": 12345,
+                "decimals": 9,
+                "priceChange24h": 2.5
+            }
+        }"#;
+        let mut server = Server::new_async().await;
+        let _m = server
+            .mock("GET", "/price/v3")
+            .match_query(Matcher::Any)
+            .with_status(200)
+            .with_header("content-type", "application/json")
+            .with_body(mock_response)
+            .create();
+
+        let client = Client::new();
+        // Test without API key - should use lite API
+        let mut oracle = JupiterPriceOracle::new();
+        oracle.lite_api_url = format!("{}/price/v3", server.url());
+
+        let result = oracle.get_price(&client, "JUPyiwrYJFskUPiHa7hkeR8VUtAeFoSYbKedZNsDvCN").await;
+
+        assert!(result.is_err());
+        assert_eq!(
+            result.err(),
+            Some(KoraError::RpcError("No price data from Jupiter".to_string()))
+        )
+    }
 }
diff --git a/crates/lib/src/rpc_server/method/get_config.rs b/crates/lib/src/rpc_server/method/get_config.rs
@@ -80,9 +80,9 @@ mod tests {
             response.validation_config.allowed_tokens[0],
             "4zMMC9srt5Ri5X14GAgXhaHii3GnPAEERYPJgZJDncDU"
         ); // USDC devnet
-        assert_eq!(response.validation_config.allowed_spl_paid_tokens.len(), 1);
+        assert_eq!(response.validation_config.allowed_spl_paid_tokens.as_slice().len(), 1);
         assert_eq!(
-            response.validation_config.allowed_spl_paid_tokens[0],
+            response.validation_config.allowed_spl_paid_tokens.as_slice()[0],
             "4zMMC9srt5Ri5X14GAgXhaHii3GnPAEERYPJgZJDncDU"
         ); // USDC devnet
         assert_eq!(response.validation_config.disallowed_accounts.len(), 0);
diff --git a/crates/lib/src/tests/config_mock.rs b/crates/lib/src/tests/config_mock.rs
@@ -1,7 +1,8 @@
 use crate::{
     config::{
         AuthConfig, CacheConfig, Config, EnabledMethods, FeePayerBalanceMetricsConfig,
-        FeePayerPolicy, KoraConfig, MetricsConfig, Token2022Config, ValidationConfig,
+        FeePayerPolicy, KoraConfig, MetricsConfig, SplTokenConfig, Token2022Config,
+        ValidationConfig,
     },
     fee::price::PriceConfig,
     oracle::PriceSource,
@@ -85,9 +86,9 @@ impl ConfigMockBuilder {
                     allowed_tokens: vec![
                         "4zMMC9srt5Ri5X14GAgXhaHii3GnPAEERYPJgZJDncDU".parse().unwrap(), // USDC devnet
                     ],
-                    allowed_spl_paid_tokens: vec![
+                    allowed_spl_paid_tokens: SplTokenConfig::Allowlist(vec![
                         "4zMMC9srt5Ri5X14GAgXhaHii3GnPAEERYPJgZJDncDU".parse().unwrap(), // USDC devnet
-                    ],
+                    ]),
                     disallowed_accounts: vec![],
                     price_source: PriceSource::Mock,
                     fee_payer_policy: FeePayerPolicy::default(),
@@ -160,6 +161,11 @@ impl ConfigMockBuilder {
         self
     }
 
+    pub fn with_allowed_spl_paid_tokens(mut self, spl_payment_config: SplTokenConfig) -> Self {
+        self.config.validation.allowed_spl_paid_tokens = spl_payment_config;
+        self
+    }
+
     pub fn with_payment_address(mut self, payment_address: Option<String>) -> Self {
         self.config.kora.payment_address = payment_address;
         self
@@ -232,7 +238,7 @@ impl ValidationConfigBuilder {
                 max_signatures: 10,
                 allowed_programs: vec![],
                 allowed_tokens: vec![],
-                allowed_spl_paid_tokens: vec![],
+                allowed_spl_paid_tokens: SplTokenConfig::Allowlist(vec![]),
                 disallowed_accounts: vec![],
                 price_source: PriceSource::Mock,
                 fee_payer_policy: FeePayerPolicy::default(),
@@ -261,8 +267,8 @@ impl ValidationConfigBuilder {
         self
     }
 
-    pub fn with_allowed_spl_paid_tokens(mut self, tokens: Vec<String>) -> Self {
-        self.config.allowed_spl_paid_tokens = tokens;
+    pub fn with_allowed_spl_paid_tokens(mut self, spl_payment_config: SplTokenConfig) -> Self {
+        self.config.allowed_spl_paid_tokens = spl_payment_config;
         self
     }
 
diff --git a/crates/lib/src/tests/toml_mock.rs b/crates/lib/src/tests/toml_mock.rs
@@ -1,7 +1,10 @@
 use std::fs;
 use tempfile::NamedTempFile;
 
-use crate::{config::Config, error::KoraError};
+use crate::{
+    config::{Config, SplTokenConfig},
+    error::KoraError,
+};
 
 /// TOML-specific configuration builder for testing TOML parsing and serialization
 ///
@@ -24,7 +27,7 @@ struct ValidationSection {
     max_signatures: u64,
     allowed_programs: Vec<String>,
     allowed_tokens: Vec<String>,
-    allowed_spl_paid_tokens: Vec<String>,
+    allowed_spl_paid_tokens: SplTokenConfig,
     disallowed_accounts: Vec<String>,
     price_source: String,
     price_config: Option<String>,
@@ -44,7 +47,7 @@ impl Default for ValidationSection {
             max_signatures: 10,
             allowed_programs: vec!["program1".to_string()],
             allowed_tokens: vec!["token1".to_string()],
-            allowed_spl_paid_tokens: vec!["token2".to_string()],
+            allowed_spl_paid_tokens: SplTokenConfig::Allowlist(vec!["token2".to_string()]),
             disallowed_accounts: vec![],
             price_source: "Jupiter".to_string(),
             price_config: None,
@@ -74,8 +77,8 @@ impl ConfigBuilder {
         self
     }
 
-    pub fn with_spl_paid_tokens(mut self, tokens: Vec<&str>) -> Self {
-        self.validation.allowed_spl_paid_tokens = tokens.iter().map(|s| s.to_string()).collect();
+    pub fn with_spl_paid_tokens(mut self, spl_payment_config: SplTokenConfig) -> Self {
+        self.validation.allowed_spl_paid_tokens = spl_payment_config;
         self
     }
 
@@ -170,13 +173,13 @@ impl ConfigBuilder {
             .collect::<Vec<_>>()
             .join(", ");
 
-        let spl_tokens_list = self
-            .validation
-            .allowed_spl_paid_tokens
-            .iter()
-            .map(|t| format!("\"{t}\""))
-            .collect::<Vec<_>>()
-            .join(", ");
+        let spl_tokens_config = match self.validation.allowed_spl_paid_tokens {
+            SplTokenConfig::Allowlist(ref tokens) => format!(
+                "[{}]",
+                tokens.iter().map(|t| format!("\"{t}\"")).collect::<Vec<_>>().join(", ")
+            ),
+            SplTokenConfig::All => format!("\"{}\"", "All"),
+        };
 
         let disallowed_list = if self.validation.disallowed_accounts.is_empty() {
             "[]".to_string()
@@ -198,14 +201,14 @@ impl ConfigBuilder {
             max_signatures = {}\n\
             allowed_programs = [{}]\n\
             allowed_tokens = [{}]\n\
-            allowed_spl_paid_tokens = [{}]\n\
+            allowed_spl_paid_tokens = {}\n\
             disallowed_accounts = {}\n\
             price_source = \"{}\"\n\n",
             self.validation.max_allowed_lamports,
             self.validation.max_signatures,
             programs_list,
             tokens_list,
-            spl_tokens_list,
+            spl_tokens_config,
             disallowed_list,
             self.validation.price_source
         );
diff --git a/crates/lib/src/token/token.rs b/crates/lib/src/token/token.rs
@@ -326,7 +326,7 @@ impl TokenUtil {
                     (token_state.mint(), actual_amount)
                 };
 
-                if !config.validation.allowed_spl_paid_tokens.contains(&mint_address.to_string()) {
+                if !config.validation.supports_token(&mint_address.to_string()) {
                     return Ok(false);
                 }
 
diff --git a/crates/lib/src/transaction/versioned_transaction.rs b/crates/lib/src/transaction/versioned_transaction.rs
@@ -377,6 +377,7 @@ impl LookupTableUtil {
 #[cfg(test)]
 mod tests {
     use crate::{
+        config::SplTokenConfig,
         tests::{
             common::RpcMockBuilder, config_mock::mock_state::setup_config_mock,
             toml_mock::ConfigBuilder,
@@ -403,7 +404,7 @@ mod tests {
         ConfigBuilder::new()
             .with_programs(vec![])
             .with_tokens(vec![])
-            .with_spl_paid_tokens(vec![])
+            .with_spl_paid_tokens(SplTokenConfig::Allowlist(vec![]))
             .with_free_price()
             .with_cache_config(None, false, 60, 30) // Disable cache for tests
             .build_config()
diff --git a/crates/lib/src/validator/config_validator.rs b/crates/lib/src/validator/config_validator.rs
diff --git a/kora.toml b/kora.toml

Original file line number	Diff line number	Diff line change
`@@ -326,7 +326,7 @@ impl TokenUtil {`
`326`	`326`	`(token_state.mint(), actual_amount)`
`327`	`327`	`};`
`328`	`328`
`329`		`- if !config.validation.allowed_spl_paid_tokens.contains(&mint_address.to_string()) {`
	`329`	`+ if !config.validation.supports_token(&mint_address.to_string()) {`
`330`	`330`	`return Ok(false);`
`331`	`331`	`}`
`332`	`332`