From 844faf15ae2fe45cd676c67115353f460828c869 Mon Sep 17 00:00:00 2001 From: Paul Liu Date: Wed, 30 Oct 2024 09:49:10 +0800 Subject: [PATCH] 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 <112856886+gregorydemay@users.noreply.github.com> --- rs/bitcoin/kyt/btc_kyt_canister.did | 13 +- rs/bitcoin/kyt/src/dashboard.rs | 8 +- rs/bitcoin/kyt/src/dashboard/tests.rs | 28 ++- rs/bitcoin/kyt/src/fetch.rs | 8 +- rs/bitcoin/kyt/src/fetch/tests.rs | 24 +- rs/bitcoin/kyt/src/lib.rs | 36 ++- rs/bitcoin/kyt/src/main.rs | 29 ++- rs/bitcoin/kyt/src/state.rs | 6 +- rs/bitcoin/kyt/src/types.rs | 22 +- rs/bitcoin/kyt/templates/dashboard.html | 14 +- rs/bitcoin/kyt/tests/tests.rs | 302 ++++++++++++++++++------ 11 files changed, 367 insertions(+), 123 deletions(-) diff --git a/rs/bitcoin/kyt/btc_kyt_canister.did b/rs/bitcoin/kyt/btc_kyt_canister.did index a02f16b64f6..7552b3b084c 100644 --- 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 index cbeca4c2d9a..544413cf056 100644 --- 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 index d26fdd09150..d882db285e5 100644 --- 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 index 6a38fe97144..8d1a9da6d4c 100644 --- 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 { pub trait FetchEnv { type FetchGuard; fn new_fetch_guard(&self, txid: Txid) -> Result; - fn btc_network(&self) -> BtcNetwork; + fn config(&self) -> Config; async fn http_get_tx( &self, @@ -81,13 +81,13 @@ pub trait FetchEnv { ) -> TryFetchResult>> { 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 index dd4e8d0db47..4e491e7eba5 100644 --- 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 index dff59c88f85..40e13d8523d 100644 --- 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 index 913060caf61..35c0f528154 100644 --- 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,7 +15,8 @@ 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()) @@ -23,10 +24,15 @@ fn check_address(args: CheckAddressArgs) -> CheckAddressResponse { 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 index 905e701708e..7c19ee3ab91 100644 --- 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)] diff --git a/rs/bitcoin/kyt/src/types.rs b/rs/bitcoin/kyt/src/types.rs index 1f898c530ad..6d995cd1c9a 100644 --- a/rs/bitcoin/kyt/src/types.rs +++ b/rs/bitcoin/kyt/src/types.rs @@ -82,6 +82,7 @@ impl From for CheckTransactionResponse { #[derive(CandidType, Clone, Debug, Deserialize, Serialize)] pub struct InitArg { pub btc_network: BtcNetwork, + pub kyt_mode: KytMode, } #[derive(CandidType, Clone, Copy, Deserialize, Debug, Eq, PartialEq, Serialize, Hash)] @@ -92,6 +93,23 @@ pub enum BtcNetwork { Testnet, } +#[derive(CandidType, Clone, Copy, Deserialize, Debug, Eq, PartialEq, Serialize, Hash)] +pub enum KytMode { + AcceptAll, + RejectAll, + Normal, +} + +impl fmt::Display for KytMode { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Self::AcceptAll => write!(f, "AcceptAll"), + Self::RejectAll => write!(f, "RejectAll"), + Self::Normal => write!(f, "Normal"), + } + } +} + impl From for bitcoin::Network { fn from(btc_network: BtcNetwork) -> Self { match btc_network { @@ -111,7 +129,9 @@ impl fmt::Display for BtcNetwork { } #[derive(CandidType, Debug, Deserialize, Serialize)] -pub struct UpgradeArg {} +pub struct UpgradeArg { + pub kyt_mode: Option, +} #[derive(CandidType, Debug, Deserialize, Serialize)] pub enum KytArg { diff --git a/rs/bitcoin/kyt/templates/dashboard.html b/rs/bitcoin/kyt/templates/dashboard.html index 78269e2b3d2..338666b0535 100644 --- a/rs/bitcoin/kyt/templates/dashboard.html +++ b/rs/bitcoin/kyt/templates/dashboard.html @@ -7,10 +7,10 @@ {%- endmacro %} {% macro btc_tx_link(txid) -%} - {% match btc_network %} - {%- when BtcNetwork::Mainnet -%} + {% match config.btc_network %} + {%- when crate::BtcNetwork::Mainnet -%} {{txid}} - {%- when BtcNetwork::Testnet -%} + {%- when crate::BtcNetwork::Testnet -%} {{txid}} {% endmatch %} {%- endmacro %} @@ -32,7 +32,7 @@ - KYT Canister Dashboard for Bitcoin ({{ btc_network }}) + KYT Canister Dashboard for Bitcoin ({{ config.btc_network }})