Support TestNConnIP for deterministic testing; add simtest

crates/sui-core/src/traffic_controller/mod.rs

@@ -25,9 +25,7 @@ use mysten_metrics::spawn_monitored_task;
 use rand::Rng;
 use std::fmt::Debug;
 use std::time::{Duration, Instant, SystemTime};
-use sui_types::traffic_control::{
-    FreqThresholdConfig, PolicyConfig, PolicyType, RemoteFirewallConfig, Weight,
-};
+use sui_types::traffic_control::{PolicyConfig, PolicyType, RemoteFirewallConfig, Weight};
 use tokio::sync::broadcast;
 use tokio::sync::mpsc;
 use tokio::sync::mpsc::error::TrySendError;
@@ -122,29 +120,35 @@ impl TrafficController {
     ) -> SuiResult<()> {
         let mut config = self.config.lock();
         if let Some(error_threshold) = error_threshold {
-            if let PolicyType::FreqThreshold(threshold_config) = &mut config.error_policy_type {
-                threshold_config.client_threshold = error_threshold;
-            } else {
-                let error_policy_type = PolicyType::FreqThreshold(FreqThresholdConfig {
-                    client_threshold: error_threshold,
-                    window_size_secs: 5,
-                    update_interval_secs: 1,
-                    ..Default::default()
-                });
-                config.error_policy_type = error_policy_type;
+            match &mut config.error_policy_type {
+                PolicyType::FreqThreshold(threshold_config) => {
+                    threshold_config.client_threshold = error_threshold;
+                }
+                PolicyType::TestNConnIP(n) => {
+                    *n = error_threshold;
+                }
+                _ => {
+                    return Err(SuiError::InvalidAdminRequest(
+                        "Unsupported prior error policy type during traffic control reconfiguration"
+                            .to_string(),
+                    ));
+                }
             }
         }
         if let Some(spam_threshold) = spam_threshold {
-            if let PolicyType::FreqThreshold(threshold_config) = &mut config.spam_policy_type {
-                threshold_config.client_threshold = spam_threshold;
-            } else {
-                let spam_policy_type = PolicyType::FreqThreshold(FreqThresholdConfig {
-                    client_threshold: spam_threshold,
-                    window_size_secs: 5,
-                    update_interval_secs: 1,
-                    ..Default::default()
-                });
-                config.spam_policy_type = spam_policy_type;
+            match &mut config.spam_policy_type {
+                PolicyType::FreqThreshold(threshold_config) => {
+                    threshold_config.client_threshold = spam_threshold;
+                }
+                PolicyType::TestNConnIP(n) => {
+                    *n = spam_threshold;
+                }
+                _ => {
+                    return Err(SuiError::InvalidAdminRequest(
+                        "Unsupported prior spam policy type during traffic control reconfiguration"
+                            .to_string(),
+                    ));
+                }
             }
         }
         if let Some(dry_run) = dry_run {

crates/sui-e2e-tests/tests/traffic_control_tests.rs

@@ -254,6 +254,70 @@ async fn test_validator_traffic_control_error_blocked() -> Result<(), anyhow::Er
     panic!("Expected error policy to trigger within {n} requests");
 }
 
+#[tokio::test]
+async fn test_validator_traffic_control_error_blocked_with_policy_reconfig(
+) -> Result<(), anyhow::Error> {
+    let n = 5;
+    let policy_config = PolicyConfig {
+        connection_blocklist_ttl_sec: 1,
+        // Test that any N requests will cause an IP to be added to the blocklist.
+        error_policy_type: PolicyType::TestNConnIP(n - 1),
+        dry_run: true,
+        ..Default::default()
+    };
+    let network_config = ConfigBuilder::new_with_temp_dir()
+        .committee_size(NonZeroUsize::new(4).unwrap())
+        .with_policy_config(Some(policy_config))
+        .build();
+    let committee = network_config.committee_with_network();
+    let test_cluster = TestClusterBuilder::new()
+        .set_network_config(network_config)
+        .build()
+        .await;
+    let local_clients = make_network_authority_clients_with_network_config(
+        &committee,
+        &default_mysten_network_config(),
+    );
+    let (_, auth_client) = local_clients.first_key_value().unwrap();
+
+    let mut txns = batch_make_transfer_transactions(&test_cluster.wallet, n as usize).await;
+    let mut tx = txns.swap_remove(0);
+    let signatures = tx.tx_signatures_mut_for_testing();
+    signatures.pop();
+    signatures.push(GenericSignature::Signature(
+        sui_types::crypto::Signature::Ed25519SuiSignature(Ed25519SuiSignature::default()),
+    ));
+
+    // Before reconfiguring the policy, we should not block any requests due to dry run mode,
+    // even after far exceeding the threshold. However the blocklist should be updated.
+    for _ in 0..(2 * n) {
+        let response = auth_client.handle_transaction(tx.clone(), None).await;
+        if let Err(err) = response {
+            assert!(
+                !err.to_string().contains("Too many requests"),
+                "Expected no blocked requests due to dry run mode"
+            );
+        }
+    }
+    // Reconfigure traffic control to disable dry run mode
+    for node in test_cluster.all_node_handles() {
+        node.with(|node| {
+            node.state()
+                .reconfigure_traffic_control(None, None, Some(false))
+                .unwrap();
+        });
+    }
+    // If Node and TrafficController has not crashed, blocklist and policy freq state should still
+    // be intact. A single additional erroneous request from the client should trigger enforcement.
+    let response = auth_client.handle_transaction(tx.clone(), None).await;
+    if let Err(err) = response {
+        if err.to_string().contains("Too many requests") {
+            return Ok(());
+        }
+    }
+    panic!("Expected error policy to trigger on next requests after reconfiguration");
+}
+
 #[tokio::test]
 async fn test_fullnode_traffic_control_spam_blocked() -> Result<(), anyhow::Error> {
     let txn_count = 15;

crates/sui-node/src/admin.rs

@@ -68,6 +68,10 @@ use tracing::info;
 // Inject a full signature from another node, bypassing validity checks.
 //
 //  $ curl 'http://127.0.0.1:1337/randomness-inject-full-sig?round=123&sigs=base64encodedsig'
+//
+// Reconfigure traffic control policy
+//
+//  $ curl 'http://127.0.0.1:1337/traffic-control?error_threshold=100&spam_threshold=100&dry_run=true'
 
 const LOGGING_ROUTE: &str = "/logging";
 const TRACING_ROUTE: &str = "/enable-tracing";