feat: check risk strategies

Author: Dodecahedr0x

config.example.toml

@@ -236,6 +236,14 @@ request-timeout = "5s"
 # Env: MBV_CHAINLINK__RISK__RISK_SCORE_THRESHOLD
 # risk-score-threshold = 5
 
+# Which post-delegation action signers to risk check.
+#   "all-signers":       check every signer of every action (conservative).
+#   "relevant-programs": only check when an action involves the SPL Token,
+#                        ephemeral SPL (eATA/ESPL), or Magic program.
+# Default: "relevant-programs"
+# Env: MBV_CHAINLINK__RISK__CHECK_STRATEGY
+check-strategy = "relevant-programs"
+
 # ==============================================================================
 # Leader Administration (Optional)
 # ==============================================================================

magicblock-aml/src/lib.rs

@@ -6,7 +6,7 @@ use std::{
 };
 
 use futures_util::future::{BoxFuture, FutureExt, Shared, try_join_all};
-use magicblock_config::config::RiskConfig;
+use magicblock_config::config::{AmlCheckStrategy, RiskConfig};
 use reqwest::Client;
 use rusqlite::{Connection, params};
 use serde_json::Value;
@@ -68,6 +68,7 @@ pub struct RiskService {
     api_key: String,
     cache_ttl: Duration,
     risk_score_threshold: u64,
+    check_strategy: AmlCheckStrategy,
 }
 
 impl RiskService {
@@ -120,9 +121,16 @@ impl RiskService {
             api_key,
             cache_ttl: config.cache_ttl,
             risk_score_threshold: config.risk_score_threshold,
+            check_strategy: config.check_strategy,
         }))
     }
 
+    /// The configured strategy for deciding which post-delegation action
+    /// signers to risk check.
+    pub fn check_strategy(&self) -> AmlCheckStrategy {
+        self.check_strategy
+    }
+
     pub async fn check_addresses(
         &self,
         addresses: Vec<String>,
@@ -450,6 +458,7 @@ mod tests {
             cache_ttl: Duration::from_secs(60),
             request_timeout: Duration::from_secs(2),
             risk_score_threshold: 7,
+            check_strategy: AmlCheckStrategy::AllSigners,
         }
     }
 

magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs

@@ -22,10 +22,10 @@ use engine::Engine;
 use keeper::MissingAccount;
 use lru::LruCache;
 use magicblock_aml::RiskService;
-use magicblock_config::config::AllowedProgram;
+use magicblock_config::config::{AllowedProgram, AmlCheckStrategy};
 use magicblock_core::token_programs::{
-    ASSOCIATED_TOKEN_PROGRAM_ID, EATA_PROGRAM_ID, TOKEN_PROGRAM_ID, is_ata,
-    normalize_native_token_account_for_local_clone,
+    ASSOCIATED_TOKEN_PROGRAM_ID, EATA_PROGRAM_ID, TOKEN_2022_PROGRAM_ID,
+    TOKEN_PROGRAM_ID, is_ata, normalize_native_token_account_for_local_clone,
 };
 use magicblock_metrics::metrics::{
     self, AccountFetchContext, AccountFetchReason, BankPrecheckOutcome,
@@ -333,6 +333,42 @@ fn log_companion_fetch_failure<E: std::fmt::Display + ?Sized>(
     );
 }
 
+/// Programs whose presence in a post-delegation action triggers a risk check
+/// under [`AmlCheckStrategy::RelevantPrograms`]: SPL Token (legacy and 2022),
+/// the ephemeral SPL / eATA program (ESPL), and the Magic program.
+const RISK_RELEVANT_PROGRAMS: [Pubkey; 4] = [
+    TOKEN_PROGRAM_ID,
+    TOKEN_2022_PROGRAM_ID,
+    EATA_PROGRAM_ID,
+    magicblock_magic_program_api::ID,
+];
+
+/// Decides whether the configured [`AmlCheckStrategy`] requires risk checking
+/// the signers of these post-delegation actions.
+fn delegation_actions_require_risk_check(
+    strategy: AmlCheckStrategy,
+    delegation_actions: &DelegationActions,
+) -> bool {
+    match strategy {
+        AmlCheckStrategy::AllSigners => true,
+        AmlCheckStrategy::RelevantPrograms => delegation_actions
+            .iter()
+            .any(instruction_involves_risk_relevant_program),
+    }
+}
+
+/// Returns true when a risk-relevant program is invoked by the instruction or
+/// referenced by any of its accounts (e.g. as the target of a CPI).
+fn instruction_involves_risk_relevant_program(
+    instruction: &solana_instruction::Instruction,
+) -> bool {
+    RISK_RELEVANT_PROGRAMS.contains(&instruction.program_id)
+        || instruction
+            .accounts
+            .iter()
+            .any(|meta| RISK_RELEVANT_PROGRAMS.contains(&meta.pubkey))
+}
+
 impl<T, U> FetchCloner<T, U>
 where
     T: ChainRpcClient,
@@ -1538,6 +1574,13 @@ where
             return Ok(());
         };
 
+        if !delegation_actions_require_risk_check(
+            risk_service.check_strategy(),
+            delegation_actions,
+        ) {
+            return Ok(());
+        }
+
         let mut signers = delegation_actions
             .iter()
             .flat_map(|instruction| {

magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs

@@ -82,3 +82,89 @@ fn clone_request_classification() {
         ChainlinkCloneIntent::ActionDependency
     );
 }
+
+mod aml_check_strategy {
+    use solana_instruction::AccountMeta;
+
+    use super::*;
+
+    fn action_for_program(program_id: Pubkey) -> Instruction {
+        Instruction::new_with_bytes(program_id, &[], vec![])
+    }
+
+    fn action_referencing(program_id: Pubkey, account: Pubkey) -> Instruction {
+        Instruction::new_with_bytes(
+            program_id,
+            &[],
+            vec![AccountMeta::new_readonly(account, false)],
+        )
+    }
+
+    #[test]
+    fn all_signers_strategy_always_requires_check() {
+        // An action that touches no risk-relevant program still gets checked.
+        let actions: DelegationActions =
+            vec![action_for_program(Pubkey::new_unique())].into();
+        assert!(delegation_actions_require_risk_check(
+            AmlCheckStrategy::AllSigners,
+            &actions,
+        ));
+    }
+
+    #[test]
+    fn relevant_programs_strategy_skips_unrelated_actions() {
+        let actions: DelegationActions = vec![
+            action_for_program(Pubkey::new_unique()),
+            action_for_program(Pubkey::new_unique()),
+        ]
+        .into();
+        assert!(!delegation_actions_require_risk_check(
+            AmlCheckStrategy::RelevantPrograms,
+            &actions,
+        ));
+    }
+
+    #[test]
+    fn relevant_programs_strategy_matches_each_relevant_program() {
+        for program in [
+            TOKEN_PROGRAM_ID,
+            TOKEN_2022_PROGRAM_ID,
+            EATA_PROGRAM_ID,
+            magicblock_magic_program_api::ID,
+        ] {
+            let actions: DelegationActions =
+                vec![action_for_program(program)].into();
+            assert!(
+                delegation_actions_require_risk_check(
+                    AmlCheckStrategy::RelevantPrograms,
+                    &actions,
+                ),
+                "program {program} invoked as program_id should require check",
+            );
+
+            // Referenced as a CPI target account, not the invoked program.
+            let actions: DelegationActions =
+                vec![action_referencing(Pubkey::new_unique(), program)].into();
+            assert!(
+                delegation_actions_require_risk_check(
+                    AmlCheckStrategy::RelevantPrograms,
+                    &actions,
+                ),
+                "program {program} referenced as account should require check",
+            );
+        }
+    }
+
+    #[test]
+    fn relevant_programs_strategy_matches_when_any_action_is_relevant() {
+        let actions: DelegationActions = vec![
+            action_for_program(Pubkey::new_unique()),
+            action_for_program(EATA_PROGRAM_ID),
+        ]
+        .into();
+        assert!(delegation_actions_require_risk_check(
+            AmlCheckStrategy::RelevantPrograms,
+            &actions,
+        ));
+    }
+}

magicblock-chainlink/tests/10_aml_undelegation.rs

@@ -16,7 +16,7 @@ use dlp_api::{
 };
 use magicblock_aml::RiskService;
 use magicblock_chainlink::testing::{context::TestContext, init_logger};
-use magicblock_config::config::RiskConfig;
+use magicblock_config::config::{AmlCheckStrategy, RiskConfig};
 use solana_account::{Account, AccountMode};
 use solana_pubkey::Pubkey;
 use tokio::task::JoinHandle;
@@ -101,6 +101,7 @@ fn risk_config(base_url: String) -> RiskConfig {
         cache_ttl: Duration::from_secs(60),
         request_timeout: Duration::from_secs(2),
         risk_score_threshold: 7,
+        check_strategy: AmlCheckStrategy::AllSigners,
     }
 }
 

magicblock-config/src/config/chain.rs

@@ -73,6 +73,23 @@ impl Default for ChainLinkConfig {
     }
 }
 
+/// Strategy for deciding which post-delegation action signers get AML/risk
+/// checked.
+#[derive(
+    Deserialize, Serialize, Debug, Clone, Copy, PartialEq, Eq, Default,
+)]
+#[serde(rename_all = "kebab-case")]
+pub enum AmlCheckStrategy {
+    /// Check every signer of every post-delegation action, regardless of which
+    /// programs the action invokes.
+    AllSigners,
+    /// Only check signers when a post-delegation action involves the SPL Token,
+    /// ephemeral SPL (eATA/ESPL), or Magic program. Actions touching none of
+    /// these programs skip the risk check entirely.
+    #[default]
+    RelevantPrograms,
+}
+
 /// Configuration for checking account risk with the Range API and a local sqlite cache.
 #[derive(Deserialize, Serialize, Debug, Clone)]
 #[serde(default, rename_all = "kebab-case", deny_unknown_fields)]
@@ -91,6 +108,8 @@ pub struct RiskConfig {
     pub request_timeout: Duration,
     /// Threshold on a scale of 0-10 for the risk score.
     pub risk_score_threshold: u64,
+    /// Which post-delegation action signers to risk check.
+    pub check_strategy: AmlCheckStrategy,
 }
 
 impl Default for RiskConfig {
@@ -104,6 +123,7 @@ impl Default for RiskConfig {
                 consts::DEFAULT_RISK_REQUEST_TIMEOUT_SEC,
             ),
             risk_score_threshold: consts::DEFAULT_RISK_SCORE_THRESHOLD,
+            check_strategy: AmlCheckStrategy::default(),
         }
     }
 }

magicblock-config/src/config/mod.rs

@@ -11,7 +11,8 @@ pub mod scheduler;
 
 pub use aperture::ApertureConfig;
 pub use chain::{
-    AdminConfig, AllowedProgram, ChainLinkConfig, CommittorConfig, RiskConfig,
+    AdminConfig, AllowedProgram, AmlCheckStrategy, ChainLinkConfig,
+    CommittorConfig, RiskConfig,
 };
 pub use engine::{EngineConfig, FollowerReplication, LeaderReplication};
 pub use grpc::GrpcConfig;