add test (in progress)

crates/sui-bridge/src/e2e_tests/basic.rs

@@ -21,6 +21,7 @@ use ethers::prelude::*;
 use ethers::types::Address as EthAddress;
 use move_core_types::ident_str;
 use std::collections::{HashMap, HashSet};
+use sui_types::traffic_control::{PolicyConfig, PolicyType, Weight};
 
 use std::path::Path;
 
@@ -327,6 +328,156 @@ async fn test_add_new_coins_on_sui_and_eth() {
     .unwrap();
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+async fn test_bridge_from_eth_to_sui_rate_limited() {
+    telemetry_subscribers::init_for_testing();
+
+    let eth_chain_id = BridgeChainId::EthCustom as u8;
+    let sui_chain_id = BridgeChainId::SuiCustom as u8;
+    let timer = std::time::Instant::now();
+    let txn_count = 5;
+    let policy_config = PolicyConfig {
+        connection_blocklist_ttl_sec: 1,
+        proxy_blocklist_ttl_sec: 5,
+        spam_policy_type: PolicyType::TestNConnIP(txn_count - 1),
+        spam_sample_rate: Weight::one(),
+        // This should never be invoked when set as an error policy
+        // as we are not sending requests that error
+        error_policy_type: PolicyType::TestPanicOnInvocation,
+        dry_run: true,
+        ..Default::default()
+    };
+    let mut bridge_test_cluster = BridgeTestClusterBuilder::new()
+        .with_eth_env(true)
+        .with_bridge_cluster(true)
+        .with_num_validators(3)
+        .with_policy_config(policy_config)
+        .build()
+        .await;
+    panic!("TESTING");
+    info!(
+        "[Timer] Bridge test cluster started in {:?}",
+        timer.elapsed()
+    );
+    let timer = std::time::Instant::now();
+    let (eth_signer, _) = bridge_test_cluster
+        .get_eth_signer_and_address()
+        .await
+        .unwrap();
+
+    let sui_address = bridge_test_cluster.sui_user_address();
+    let amount = 42;
+    let sui_amount = amount * 100_000_000;
+
+    for _ in 0..txn_count {
+        initiate_bridge_eth_to_sui(&bridge_test_cluster, amount, 0)
+            .await
+            .unwrap();
+        let events = bridge_test_cluster
+            .new_bridge_events(
+                HashSet::from_iter([
+                    TokenTransferApproved.get().unwrap().clone(),
+                    TokenTransferClaimed.get().unwrap().clone(),
+                ]),
+                true,
+            )
+            .await;
+        // There are exactly 1 approved and 1 claimed event
+        assert_eq!(events.len(), 2);
+
+        let eth_coin = bridge_test_cluster
+            .sui_client()
+            .coin_read_api()
+            .get_all_coins(sui_address, None, None)
+            .await
+            .unwrap()
+            .data
+            .iter()
+            .find(|c| c.coin_type.contains("ETH"))
+            .expect("Recipient should have received ETH coin now")
+            .clone();
+        assert_eq!(eth_coin.balance, sui_amount);
+        info!(
+            "[Timer] Eth to Sui bridge transfer finished in {:?}",
+            timer.elapsed()
+        );
+        let timer = std::time::Instant::now();
+
+        // Now let the recipient send the coin back to ETH
+        let eth_address_1 = EthAddress::random();
+        let nonce = 0;
+
+        let sui_to_eth_bridge_action = initiate_bridge_sui_to_eth(
+            &bridge_test_cluster,
+            eth_address_1,
+            eth_coin.object_ref(),
+            nonce,
+            sui_amount,
+        )
+        .await
+        .unwrap();
+        let events = bridge_test_cluster
+            .new_bridge_events(
+                HashSet::from_iter([
+                    SuiToEthTokenBridgeV1.get().unwrap().clone(),
+                    TokenTransferApproved.get().unwrap().clone(),
+                    TokenTransferClaimed.get().unwrap().clone(),
+                ]),
+                true,
+            )
+            .await;
+        // There are exactly 1 deposit and 1 approved event
+        assert_eq!(events.len(), 2);
+        info!(
+            "[Timer] Sui to Eth bridge transfer approved in {:?}",
+            timer.elapsed()
+        );
+        let timer = std::time::Instant::now();
+
+        // Test `get_parsed_token_transfer_message`
+        let parsed_msg = bridge_test_cluster
+            .bridge_client()
+            .get_parsed_token_transfer_message(sui_chain_id, nonce)
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(parsed_msg.source_chain as u8, sui_chain_id);
+        assert_eq!(parsed_msg.seq_num, nonce);
+        assert_eq!(
+            parsed_msg.parsed_payload.sender_address,
+            sui_address.to_vec()
+        );
+        assert_eq!(
+            &parsed_msg.parsed_payload.target_address,
+            eth_address_1.as_bytes()
+        );
+        assert_eq!(parsed_msg.parsed_payload.target_chain, eth_chain_id);
+        assert_eq!(parsed_msg.parsed_payload.token_type, TOKEN_ID_ETH);
+        assert_eq!(parsed_msg.parsed_payload.amount, sui_amount);
+
+        let message = eth_sui_bridge::Message::from(sui_to_eth_bridge_action);
+        let signatures =
+            get_signatures(bridge_test_cluster.bridge_client(), nonce, sui_chain_id).await;
+
+        let eth_sui_bridge = EthSuiBridge::new(
+            bridge_test_cluster.contracts().sui_bridge,
+            eth_signer.clone().into(),
+        );
+        let call = eth_sui_bridge.transfer_bridged_tokens_with_signatures(signatures, message);
+        let eth_claim_tx_receipt = send_eth_tx_and_get_tx_receipt(call).await;
+        assert_eq!(eth_claim_tx_receipt.status.unwrap().as_u64(), 1);
+        info!(
+            "[Timer] Sui to Eth bridge transfer claimed in {:?}",
+            timer.elapsed()
+        );
+        // Assert eth_address_1 has received ETH
+        assert_eq!(
+            eth_signer.get_balance(eth_address_1, None).await.unwrap(),
+            U256::from(amount) * U256::exp10(18)
+        );
+    }
+}
+
 pub(crate) async fn deposit_native_eth_to_sol_contract(
     signer: &EthSigner,
     contract_address: EthAddress,

crates/sui-bridge/src/e2e_tests/test_utils.rs

@@ -39,6 +39,7 @@ use sui_types::bridge::BridgeChainId;
 use sui_types::committee::TOTAL_VOTING_POWER;
 use sui_types::crypto::get_key_pair;
 use sui_types::digests::TransactionDigest;
+use sui_types::traffic_control::PolicyConfig;
 use sui_types::transaction::{ObjectArg, TransactionData};
 use sui_types::SUI_BRIDGE_OBJECT_ID;
 use tokio::join;
@@ -87,6 +88,7 @@ pub struct BridgeTestCluster {
     bridge_tx_cursor: Option<TransactionDigest>,
     eth_chain_id: BridgeChainId,
     sui_chain_id: BridgeChainId,
+    traffic_policy_config: Option<PolicyConfig>,
 }
 
 pub struct BridgeTestClusterBuilder {
@@ -96,6 +98,7 @@ pub struct BridgeTestClusterBuilder {
     approved_governance_actions: Option<Vec<Vec<BridgeAction>>>,
     eth_chain_id: BridgeChainId,
     sui_chain_id: BridgeChainId,
+    traffic_policy_config: Option<PolicyConfig>,
 }
 
 impl Default for BridgeTestClusterBuilder {
@@ -113,6 +116,7 @@ impl BridgeTestClusterBuilder {
             approved_governance_actions: None,
             eth_chain_id: BridgeChainId::EthCustom,
             sui_chain_id: BridgeChainId::SuiCustom,
+            traffic_policy_config: None,
         }
     }
 
@@ -150,6 +154,11 @@ impl BridgeTestClusterBuilder {
         self
     }
 
+    pub fn with_policy_config(mut self, policy_config: PolicyConfig) -> Self {
+        self.traffic_policy_config = Some(policy_config);
+        self
+    }
+
     pub async fn build(self) -> BridgeTestCluster {
         init_all_struct_tags();
         std::env::set_var("__TEST_ONLY_CONSENSUS_USE_LONG_MIN_ROUND_DELAY", "1");
@@ -175,8 +184,18 @@ impl BridgeTestClusterBuilder {
                 .clone()
                 .unwrap_or(vec![vec![]; self.num_validators]);
             bridge_node_handles = Some(
-                start_bridge_cluster(&test_cluster, &eth_environment, approved_governace_actions)
-                    .await,
+                start_bridge_cluster(
+                    &test_cluster,
+                    &eth_environment,
+                    approved_governace_actions,
+                    self.traffic_policy_config.clone(),
+                )
+                .await,
+            );
+        } else {
+            assert_eq!(
+                self.traffic_policy_config, None,
+                "Traffic Controller policy config is only used with bridge cluster (with_bridge_cluster=true)"
             );
         }
         let bridge_client = SuiBridgeClient::new(&test_cluster.fullnode_handle.rpc_url, metrics)
@@ -200,6 +219,7 @@ impl BridgeTestClusterBuilder {
             bridge_tx_cursor: None,
             sui_chain_id: self.sui_chain_id,
             eth_chain_id: self.eth_chain_id,
+            traffic_policy_config: self.traffic_policy_config.clone(),
         }
     }
 
@@ -355,12 +375,13 @@ impl BridgeTestCluster {
                 &self.test_cluster,
                 &self.eth_environment,
                 approved_governace_actions,
+                self.traffic_policy_config.clone(),
             )
             .await,
         );
     }
 
-    /// Returns new bridge transaction. It advanaces the stored tx digest cursor.
+    /// Returns new bridge transaction. It advances the stored tx digest cursor.
     /// When `assert_success` is true, it asserts all transactions are successful.
     pub async fn new_bridge_transactions(
         &mut self,
@@ -720,6 +741,7 @@ pub(crate) async fn start_bridge_cluster(
     test_cluster: &TestCluster,
     eth_environment: &EthBridgeEnvironment,
     approved_governance_actions: Vec<Vec<BridgeAction>>,
+    traffic_policy_config: Option<PolicyConfig>,
 ) -> Vec<JoinHandle<()>> {
     let bridge_authority_keys = test_cluster
         .bridge_authority_keys
@@ -780,7 +802,7 @@ pub(crate) async fn start_bridge_cluster(
             },
             metrics_key_pair: default_ed25519_key_pair(),
             metrics: None,
-            traffic_policy_config: None,
+            traffic_policy_config: traffic_policy_config.clone(),
         };
         // Spawn bridge node in memory
         handles.push(

crates/sui-bridge/src/server/mock_handler.rs

@@ -19,6 +19,7 @@ use crate::types::SignedBridgeAction;
 use arc_swap::ArcSwap;
 use async_trait::async_trait;
 use axum::Json;
+use sui_core::traffic_controller::TrafficController;
 use sui_types::digests::TransactionDigest;
 
 use super::handler::BridgeRequestHandlerTrait;
@@ -126,6 +127,10 @@ impl BridgeRequestHandlerTrait for BridgeRequestMockHandler {
         let signed_action = SignedBridgeAction::new_from_data_and_sig(action, sig);
         Ok(Json(signed_action))
     }
+
+    fn traffic_controller(&self) -> Option<Arc<TrafficController>> {
+        None
+    }
 }
 
 pub fn run_mock_server(

crates/sui-bridge/src/server/mod.rs

@@ -684,13 +684,11 @@ fn handle_traffic_resp(
     });
 }
 
-fn normalize(error_code: BridgeError) -> Weight {
-    match error_code {
-        // TODO: for now, we weight all error types equally.
-        // Later we may want to provide a weight distribution
-        // based on the error type.
-        _ => Weight::one(),
-    }
+fn normalize(_error_code: BridgeError) -> Weight {
+    // TODO: for now, we weight all error types equally.
+    // Later we may want to provide a weight distribution
+    // based on the error type.
+    Weight::one()
 }
 
 async fn handle_with_decoration<F, Fut>(
@@ -703,12 +701,8 @@ where
     Fut: std::future::Future<Output = Result<Json<SignedBridgeAction>, BridgeError>>,
 {
     if let Some(traffic_controller) = bridge_req_handler.traffic_controller() {
-        if let Err(blocked_response) = handle_traffic_req(traffic_controller.clone(), &client).await
-        {
-            return Err(blocked_response);
-        }
+        handle_traffic_req(traffic_controller.clone(), &client).await?
     }
-
     let response = fn_handler().await;
     if let Some(traffic_controller) = bridge_req_handler.traffic_controller() {
         handle_traffic_resp(traffic_controller.clone(), client, response.clone());

crates/sui-bridge/src/utils.rs

@@ -200,6 +200,7 @@ pub fn generate_bridge_node_config_and_write_to_file(
             push_interval_seconds: None, // use default value
             push_url: "metrics_proxy_url".to_string(),
         }),
+        traffic_policy_config: None,
     };
     if run_client {
         config.sui.bridge_client_key_path = Some(PathBuf::from("/path/to/your/bridge_client_key"));

crates/sui-types/src/traffic_control.rs

@@ -63,7 +63,7 @@ const TRAFFIC_SINK_TIMEOUT_SEC: u64 = 300;
 ///     [<known client IP>] <--- number of hops is 1
 ///     ["1.2.3.4", <known client IP>, "5.6.7.8", "9.10.11.12"] <--- number of hops is 3
 /// ```
-#[derive(Clone, Debug, Deserialize, Serialize, Default)]
+#[derive(Clone, Debug, Deserialize, Serialize, Default, PartialEq)]
 #[serde(rename_all = "kebab-case")]
 pub enum ClientIdSource {
     #[default]
@@ -135,7 +135,7 @@ fn default_drain_timeout() -> u64 {
 }
 
 #[serde_as]
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 #[serde(rename_all = "kebab-case")]
 pub struct FreqThresholdConfig {
     #[serde(default = "default_client_threshold")]
@@ -204,7 +204,7 @@ fn default_sketch_tolerance() -> f64 {
 
 // Serializable representation of policy types, used in config
 // in order to easily change in tests or to killswitch
-#[derive(Clone, Serialize, Deserialize, Debug, Default)]
+#[derive(Clone, Serialize, Deserialize, Debug, Default, PartialEq)]
 pub enum PolicyType {
     /// Does nothing
     #[default]
@@ -227,7 +227,7 @@ pub enum PolicyType {
 }
 
 #[serde_as]
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 #[serde(rename_all = "kebab-case")]
 pub struct PolicyConfig {
     #[serde(default = "default_client_id_source")]