feat(ckbtc): Support KytMode in new KYT canister (#2280)

Support KytMode in the initialization and upgrade parameters to the new KYT canister. This simplifies work involved in testing with this canister, because the previous KYT canister had this feature and existing ckbtc minter tests are making use of it. --------- Co-authored-by: gregorydemay <[email protected]>
dfinity · Oct 30, 2024 · 844faf1 · 844faf1
1 parent dac2f36
commit 844faf1
Show file tree

Hide file tree

Showing 11 changed files with 367 additions and 123 deletions.
diff --git a/rs/bitcoin/kyt/btc_kyt_canister.did b/rs/bitcoin/kyt/btc_kyt_canister.did
@@ -49,10 +49,13 @@ type CheckTransactionIrrecoverableError = variant {
 };
 
 type InitArg = record {
-    btc_network : BtcNetwork
+    btc_network : BtcNetwork;
+    kyt_mode: KytMode;
 };
 
-type UpgradeArg = record {};
+type UpgradeArg = record {
+    kyt_mode: opt KytMode;
+};
 
 type KytArg = variant {
     InitArg : InitArg;
@@ -64,6 +67,12 @@ type BtcNetwork = variant {
     testnet
 };
 
+type KytMode = variant {
+    AcceptAll;
+    RejectAll;
+    Normal;
+};
+
 service : (KytArg) -> {
     // Check input addresses of a transaction matching the given transaction id.
     // See `CheckTransactionResponse` for more details on the return result.

diff --git a/rs/bitcoin/kyt/src/dashboard.rs b/rs/bitcoin/kyt/src/dashboard.rs
@@ -1,7 +1,7 @@
-use crate::{state, BtcNetwork};
+use crate::state;
 use askama::Template;
 use ic_btc_interface::Txid;
-use state::{FetchTxStatus, Timestamp};
+use state::{Config, FetchTxStatus, Timestamp};
 use std::fmt;
 
 #[cfg(test)]
@@ -25,7 +25,7 @@ mod filters {
 #[derive(Template)]
 #[template(path = "dashboard.html", whitespace = "suppress")]
 pub struct DashboardTemplate {
-    btc_network: BtcNetwork,
+    config: Config,
     outcall_capacity: u32,
     cached_entries: usize,
     tx_table_page_size: usize,
@@ -80,7 +80,7 @@ const DEFAULT_TX_TABLE_PAGE_SIZE: usize = 500;
 pub fn dashboard(page_index: usize) -> DashboardTemplate {
     let tx_table_page_size = DEFAULT_TX_TABLE_PAGE_SIZE;
     DashboardTemplate {
-        btc_network: state::get_config().btc_network,
+        config: state::get_config(),
         outcall_capacity: state::OUTCALL_CAPACITY.with(|capacity| *capacity.borrow()),
         cached_entries: state::FETCH_TX_CACHE.with(|cache| cache.borrow().iter().count()),
         tx_table_page_size,

diff --git a/rs/bitcoin/kyt/src/dashboard/tests.rs b/rs/bitcoin/kyt/src/dashboard/tests.rs
@@ -1,7 +1,7 @@
 use crate::dashboard::tests::assertions::DashboardAssert;
 use crate::dashboard::{filters, DashboardTemplate, Fetched, Status, DEFAULT_TX_TABLE_PAGE_SIZE};
-use crate::state::{Timestamp, TransactionKytData};
-use crate::{blocklist::BTC_ADDRESS_BLOCKLIST, dashboard, state, BtcNetwork};
+use crate::state::{Config, Timestamp, TransactionKytData};
+use crate::{blocklist::BTC_ADDRESS_BLOCKLIST, dashboard, state, BtcNetwork, KytMode};
 use bitcoin::Address;
 use bitcoin::{absolute::LockTime, transaction::Version, Transaction};
 use ic_btc_interface::Txid;
@@ -16,13 +16,16 @@ fn mock_txid(v: usize) -> Txid {
 
 #[test]
 fn should_display_metadata() {
-    let btc_network = BtcNetwork::Mainnet;
+    let config = Config {
+        btc_network: BtcNetwork::Mainnet,
+        kyt_mode: KytMode::Normal,
+    };
     let outcall_capacity = 50;
     let cached_entries = 0;
     let oldest_entry_time = 0;
     let latest_entry_time = 1_000_000_000_000;
     let dashboard = DashboardTemplate {
-        btc_network,
+        config: config.clone(),
         outcall_capacity,
         cached_entries,
         tx_table_page_size: 10,
@@ -33,7 +36,8 @@ fn should_display_metadata() {
     };
 
     DashboardAssert::assert_that(dashboard)
-        .has_btc_network_in_title(btc_network)
+        .has_btc_network_in_title(config.btc_network)
+        .has_kyt_mode(config.kyt_mode)
         .has_outcall_capacity(outcall_capacity)
         .has_cached_entries(cached_entries)
         .has_oldest_entry_time(oldest_entry_time)
@@ -73,7 +77,10 @@ fn should_display_statuses() {
     });
 
     let dashboard = DashboardTemplate {
-        btc_network: BtcNetwork::Mainnet,
+        config: Config {
+            btc_network: BtcNetwork::Mainnet,
+            kyt_mode: KytMode::Normal,
+        },
         outcall_capacity: 50,
         cached_entries: 6,
         tx_table_page_size: 10,
@@ -115,6 +122,7 @@ fn test_pagination() {
 
     state::set_config(state::Config {
         btc_network: BtcNetwork::Mainnet,
+        kyt_mode: KytMode::Normal,
     });
     let mock_transaction = TransactionKytData::from_transaction(
         &BtcNetwork::Mainnet,
@@ -193,6 +201,14 @@ mod assertions {
             )
         }
 
+        pub fn has_kyt_mode(&self, kyt_mode: KytMode) -> &Self {
+            self.has_string_value(
+                "#kyt-mode > td > code",
+                &format!("{}", kyt_mode),
+                "wrong KYT mode",
+            )
+        }
+
         pub fn has_outcall_capacity(&self, outcall_capacity: u32) -> &Self {
             self.has_string_value(
                 "#outcall-capacity > td > code",

diff --git a/rs/bitcoin/kyt/src/fetch.rs b/rs/bitcoin/kyt/src/fetch.rs
@@ -6,7 +6,7 @@ use crate::types::{
     CheckTransactionIrrecoverableError, CheckTransactionResponse, CheckTransactionRetriable,
     CheckTransactionStatus,
 };
-use crate::{blocklist_contains, providers, state, BtcNetwork};
+use crate::{blocklist_contains, providers, state, Config};
 use bitcoin::Transaction;
 use futures::future::try_join_all;
 use ic_btc_interface::Txid;
@@ -59,7 +59,7 @@ pub enum TryFetchResult<F> {
 pub trait FetchEnv {
     type FetchGuard;
     fn new_fetch_guard(&self, txid: Txid) -> Result<Self::FetchGuard, FetchGuardError>;
-    fn btc_network(&self) -> BtcNetwork;
+    fn config(&self) -> Config;
 
     async fn http_get_tx(
         &self,
@@ -81,13 +81,13 @@ pub trait FetchEnv {
     ) -> TryFetchResult<impl futures::Future<Output = Result<FetchResult, Infallible>>> {
         let (provider, max_response_bytes) = match state::get_fetch_status(txid) {
             None => (
-                providers::next_provider(self.btc_network()),
+                providers::next_provider(self.config().btc_network),
                 INITIAL_MAX_RESPONSE_BYTES,
             ),
             Some(FetchTxStatus::PendingRetry {
                 max_response_bytes, ..
             }) => (
-                providers::next_provider(self.btc_network()),
+                providers::next_provider(self.config().btc_network),
                 max_response_bytes,
             ),
             Some(FetchTxStatus::PendingOutcall { .. }) => return TryFetchResult::Pending,

diff --git a/rs/bitcoin/kyt/src/fetch/tests.rs b/rs/bitcoin/kyt/src/fetch/tests.rs
@@ -1,6 +1,9 @@
 use super::*;
 use crate::{
-    blocklist, providers::Provider, types::BtcNetwork, CheckTransactionIrrecoverableError,
+    blocklist,
+    providers::Provider,
+    types::{BtcNetwork, KytMode},
+    CheckTransactionIrrecoverableError,
 };
 use bitcoin::{
     absolute::LockTime, address::Address, hashes::Hash, transaction::Version, Amount, OutPoint,
@@ -33,8 +36,11 @@ impl FetchEnv for MockEnv {
         }
     }
 
-    fn btc_network(&self) -> BtcNetwork {
-        BtcNetwork::Mainnet
+    fn config(&self) -> Config {
+        Config {
+            btc_network: BtcNetwork::Mainnet,
+            kyt_mode: KytMode::Normal,
+        }
     }
 
     async fn http_get_tx(
@@ -166,7 +172,7 @@ fn mock_transaction_with_inputs(input_txids: Vec<(Txid, u32)>) -> Transaction {
 async fn test_mock_env() {
     // Test cycle mock functions
     let env = MockEnv::new(CHECK_TRANSACTION_CYCLES_REQUIRED);
-    let provider = providers::next_provider(env.btc_network());
+    let provider = providers::next_provider(env.config().btc_network);
     assert_eq!(
         env.cycles_accept(CHECK_TRANSACTION_CYCLES_SERVICE_FEE),
         CHECK_TRANSACTION_CYCLES_SERVICE_FEE
@@ -200,7 +206,7 @@ fn test_try_fetch_tx() {
     let txid_1 = mock_txid(1);
     let txid_2 = mock_txid(2);
     let from_tx = |tx: &bitcoin::Transaction| {
-        TransactionKytData::from_transaction(&env.btc_network(), tx.clone()).unwrap()
+        TransactionKytData::from_transaction(&env.config().btc_network, tx.clone()).unwrap()
     };
 
     // case Fetched
@@ -245,12 +251,12 @@ fn test_try_fetch_tx() {
 #[tokio::test]
 async fn test_fetch_tx() {
     let env = MockEnv::new(CHECK_TRANSACTION_CYCLES_REQUIRED);
-    let provider = providers::next_provider(env.btc_network());
+    let provider = providers::next_provider(env.config().btc_network);
     let txid_0 = mock_txid(0);
     let txid_1 = mock_txid(1);
     let txid_2 = mock_txid(2);
     let from_tx = |tx: &bitcoin::Transaction| {
-        TransactionKytData::from_transaction(&env.btc_network(), tx.clone()).unwrap()
+        TransactionKytData::from_transaction(&env.config().btc_network, tx.clone()).unwrap()
     };
 
     // case Fetched
@@ -318,7 +324,7 @@ async fn test_check_fetched() {
     let tx_0 = mock_transaction_with_inputs(vec![(txid_1, 0), (txid_2, 1)]);
     let tx_1 = mock_transaction_with_outputs(1);
     let tx_2 = mock_transaction_with_outputs(2);
-    let network = env.btc_network();
+    let network = env.config().btc_network;
     let from_tx = |tx: &bitcoin::Transaction| {
         TransactionKytData::from_transaction(&network, tx.clone()).unwrap()
     };
@@ -513,7 +519,7 @@ async fn test_check_fetched() {
 
     // case HttpGetTxError can be retried.
     let remaining_cycles = env.cycles_available();
-    let provider = providers::next_provider(env.btc_network());
+    let provider = providers::next_provider(env.config().btc_network);
     state::set_fetch_status(
         txid_2,
         FetchTxStatus::Error(FetchTxStatusError {

diff --git a/rs/bitcoin/kyt/src/lib.rs b/rs/bitcoin/kyt/src/lib.rs
@@ -54,8 +54,8 @@ impl FetchEnv for KytCanisterEnv {
         state::FetchGuard::new(txid)
     }
 
-    fn btc_network(&self) -> BtcNetwork {
-        get_config().btc_network
+    fn config(&self) -> Config {
+        get_config()
     }
 
     async fn http_get_tx(
@@ -126,17 +126,27 @@ impl FetchEnv for KytCanisterEnv {
 ///    in order to compute their corresponding addresses.
 pub async fn check_transaction_inputs(txid: Txid) -> CheckTransactionResponse {
     let env = &KytCanisterEnv;
-    match env.try_fetch_tx(txid) {
-        TryFetchResult::Pending => CheckTransactionRetriable::Pending.into(),
-        TryFetchResult::HighLoad => CheckTransactionRetriable::HighLoad.into(),
-        TryFetchResult::NotEnoughCycles => CheckTransactionStatus::NotEnoughCycles.into(),
-        TryFetchResult::Fetched(fetched) => env.check_fetched(txid, &fetched).await,
-        TryFetchResult::ToFetch(do_fetch) => {
-            match do_fetch.await {
-                Ok(FetchResult::Fetched(fetched)) => env.check_fetched(txid, &fetched).await,
-                Ok(FetchResult::Error(err)) => (txid, err).into(),
-                Ok(FetchResult::RetryWithBiggerBuffer) => CheckTransactionRetriable::Pending.into(),
-                Err(_) => unreachable!(), // should never happen
+    match env.config().kyt_mode {
+        KytMode::AcceptAll => CheckTransactionResponse::Passed,
+        KytMode::RejectAll => CheckTransactionResponse::Failed(Vec::new()),
+        KytMode::Normal => {
+            match env.try_fetch_tx(txid) {
+                TryFetchResult::Pending => CheckTransactionRetriable::Pending.into(),
+                TryFetchResult::HighLoad => CheckTransactionRetriable::HighLoad.into(),
+                TryFetchResult::NotEnoughCycles => CheckTransactionStatus::NotEnoughCycles.into(),
+                TryFetchResult::Fetched(fetched) => env.check_fetched(txid, &fetched).await,
+                TryFetchResult::ToFetch(do_fetch) => {
+                    match do_fetch.await {
+                        Ok(FetchResult::Fetched(fetched)) => {
+                            env.check_fetched(txid, &fetched).await
+                        }
+                        Ok(FetchResult::Error(err)) => (txid, err).into(),
+                        Ok(FetchResult::RetryWithBiggerBuffer) => {
+                            CheckTransactionRetriable::Pending.into()
+                        }
+                        Err(_) => unreachable!(), // should never happen
+                    }
+                }
             }
         }
     }

diff --git a/rs/bitcoin/kyt/src/main.rs b/rs/bitcoin/kyt/src/main.rs
@@ -4,7 +4,7 @@ use ic_btc_kyt::{
     blocklist_contains, check_transaction_inputs, dashboard, get_config, set_config,
     CheckAddressArgs, CheckAddressResponse, CheckTransactionArgs,
     CheckTransactionIrrecoverableError, CheckTransactionResponse, CheckTransactionStatus, Config,
-    KytArg, CHECK_TRANSACTION_CYCLES_REQUIRED, CHECK_TRANSACTION_CYCLES_SERVICE_FEE,
+    KytArg, KytMode, CHECK_TRANSACTION_CYCLES_REQUIRED, CHECK_TRANSACTION_CYCLES_SERVICE_FEE,
 };
 use ic_canisters_http_types as http;
 use ic_cdk::api::management_canister::http_request::{HttpResponse, TransformArgs};
@@ -15,18 +15,24 @@ use std::str::FromStr;
 /// `Failed` otherwise.
 /// May throw error (trap) if the given address is malformed or not a mainnet address.
 fn check_address(args: CheckAddressArgs) -> CheckAddressResponse {
-    let btc_network = get_config().btc_network;
+    let config = get_config();
+    let btc_network = config.btc_network;
     let address = Address::from_str(args.address.trim())
         .unwrap_or_else(|err| ic_cdk::trap(&format!("Invalid bitcoin address: {}", err)))
         .require_network(btc_network.into())
         .unwrap_or_else(|err| {
             ic_cdk::trap(&format!("Not a bitcoin {} address: {}", btc_network, err))
         });
 
-    if blocklist_contains(&address) {
-        CheckAddressResponse::Failed
-    } else {
-        CheckAddressResponse::Passed
+    match config.kyt_mode {
+        KytMode::AcceptAll => CheckAddressResponse::Passed,
+        KytMode::RejectAll => CheckAddressResponse::Failed,
+        KytMode::Normal => {
+            if blocklist_contains(&address) {
+                return CheckAddressResponse::Failed;
+            }
+            CheckAddressResponse::Passed
+        }
     }
 }
 
@@ -82,6 +88,7 @@ fn init(arg: KytArg) {
     match arg {
         KytArg::InitArg(init_arg) => set_config(Config {
             btc_network: init_arg.btc_network,
+            kyt_mode: init_arg.kyt_mode,
         }),
         KytArg::UpgradeArg(_) => {
             ic_cdk::trap("cannot init canister state without init args");
@@ -92,7 +99,15 @@ fn init(arg: KytArg) {
 #[ic_cdk::post_upgrade]
 fn post_upgrade(arg: KytArg) {
     match arg {
-        KytArg::UpgradeArg(_) => (),
+        KytArg::UpgradeArg(arg) => {
+            if let Some(kyt_mode) = arg.and_then(|arg| arg.kyt_mode) {
+                let config = Config {
+                    kyt_mode,
+                    ..get_config()
+                };
+                set_config(config);
+            }
+        }
         KytArg::InitArg(_) => ic_cdk::trap("cannot upgrade canister state without upgrade args"),
     }
 }

diff --git a/rs/bitcoin/kyt/src/state.rs b/rs/bitcoin/kyt/src/state.rs
@@ -1,4 +1,7 @@
-use crate::{providers::Provider, types::BtcNetwork};
+use crate::{
+    providers::Provider,
+    types::{BtcNetwork, KytMode},
+};
 use bitcoin::{Address, Transaction};
 use ic_btc_interface::Txid;
 use ic_cdk::api::call::RejectionCode;
@@ -258,6 +261,7 @@ impl Drop for FetchGuard {
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct Config {
     pub btc_network: BtcNetwork,
+    pub kyt_mode: KytMode,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]