feat: check risk strategies

Author: Dodecahedr0x

config.example.toml

@@ -222,6 +222,14 @@ risk-server-url = "http://127.0.0.1:3001"
 # Env: MBV_CHAINLINK__RISK__REQUEST_TIMEOUT
 request-timeout = "5s"
 
+# 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 @@
 //! asks the server whether a set of addresses is risky.
 
 use futures_util::future::try_join_all;
-use magicblock_config::config::RiskConfig;
+use magicblock_config::config::{AmlCheckStrategy, RiskConfig};
 use reqwest::{Client, redirect};
 use serde::Deserialize;
 use thiserror::Error;
@@ -68,6 +68,7 @@ fn validate_base_url(base_url: &str) -> RiskResult<()> {
 pub struct RiskService {
     client: Client,
     base_url: String,
+    check_strategy: AmlCheckStrategy,
 }
 
 impl RiskService {
@@ -90,7 +91,17 @@ impl RiskService {
             .build()
             .map_err(RiskError::ClientBuild)?;
 
-        Ok(Some(Self { client, base_url }))
+        Ok(Some(Self {
+            client,
+            base_url,
+            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
     }
 
     /// Asks the risk server about each address concurrently, returning
@@ -234,6 +245,7 @@ mod tests {
             enabled: true,
             risk_server_url,
             request_timeout: Duration::from_secs(2),
+            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,
@@ -1570,6 +1606,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;
@@ -107,6 +107,7 @@ fn risk_config(risk_server_url: String) -> RiskConfig {
         enabled: true,
         risk_server_url,
         request_timeout: Duration::from_secs(2),
+        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 address risk against the risk server. The risk
 /// server owns the upstream provider credentials, caching, and threshold; the
 /// validator is a thin client.
@@ -86,6 +103,8 @@ pub struct RiskConfig {
     /// Request timeout for risk server calls.
     #[serde(with = "humantime")]
     pub request_timeout: Duration,
+    /// Which post-delegation action signers to risk check.
+    pub check_strategy: AmlCheckStrategy,
 }
 
 impl Default for RiskConfig {
@@ -96,6 +115,7 @@ impl Default for RiskConfig {
             request_timeout: Duration::from_secs(
                 consts::DEFAULT_RISK_REQUEST_TIMEOUT_SEC,
             ),
+            check_strategy: AmlCheckStrategy::default(),
         }
     }
 }

magicblock-config/src/config/mod.rs

@@ -10,7 +10,8 @@ pub mod program;
 
 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;