Tests for pathologcial jwk cases (MystenLabs#13880)

Tests for handling duplicated / conflicting JWKs.

These cases are unlikely to happen in practice, as oauth providers will
almost certainly remain well-behaved. Nevertheless, we want to continue
behaving correctly even with pathological inputs.

crates/sui-benchmark/Cargo.toml

@@ -40,6 +40,7 @@ sui-swarm-config.workspace = true
 telemetry-subscribers.workspace = true
 roaring.workspace = true
 regex.workspace = true
+fastcrypto-zkp.workspace = true
 
 move-core-types.workspace = true
 mysten-metrics.workspace = true

crates/sui-benchmark/tests/simtest.rs

@@ -3,6 +3,7 @@
 
 #[cfg(msim)]
 mod test {
+    use fastcrypto_zkp::bn254::zk_login::{JwkId, JWK};
     use rand::{distributions::uniform::SampleRange, thread_rng, Rng};
     use std::collections::HashSet;
     use std::path::PathBuf;

crates/sui-config/src/node.rs

@@ -152,6 +152,9 @@ pub struct NodeConfig {
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub transaction_kv_store_write_config: Option<TransactionKeyValueStoreWriteConfig>,
+
+    #[serde(default = "default_jwk_fetch_interval_seconds")]
+    pub jwk_fetch_interval_seconds: u64,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize, Default)]
@@ -160,6 +163,10 @@ pub struct TransactionKeyValueStoreReadConfig {
     pub base_url: String,
 }
 
+fn default_jwk_fetch_interval_seconds() -> u64 {
+    3600
+}
+
 fn default_transaction_kv_store_config() -> TransactionKeyValueStoreReadConfig {
     TransactionKeyValueStoreReadConfig {
         base_url: "https://transactions.sui.io/".to_string(),

crates/sui-core/src/authority.rs

@@ -75,6 +75,7 @@ use sui_storage::indexes::{CoinInfo, ObjectIndexChanges};
 use sui_storage::key_value_store::{TransactionKeyValueStore, TransactionKeyValueStoreTrait};
 use sui_storage::key_value_store_metrics::KeyValueStoreMetrics;
 use sui_storage::IndexStore;
+use sui_types::authenticator_state::get_authenticator_state;
 use sui_types::committee::{EpochId, ProtocolVersion};
 use sui_types::crypto::{default_hash, AuthoritySignInfo, Signer};
 use sui_types::digests::ChainIdentifier;
@@ -1069,6 +1070,28 @@ impl AuthorityState {
             }
             debug_assert!(execution_error_opt.is_none());
             epoch_store.update_authenticator_state(auth_state);
+
+            // double check that the signature verifier always matches the authenticator state
+            if cfg!(debug_assertions) {
+                let authenticator_state = get_authenticator_state(&self.database)
+                    .expect("Read cannot fail")
+                    .expect("Authenticator state must exist");
+
+                let mut sys_jwks: Vec<_> = authenticator_state
+                    .active_jwks
+                    .into_iter()
+                    .map(|jwk| (jwk.jwk_id, jwk.jwk))
+                    .collect();
+                let mut active_jwks: Vec<_> = epoch_store
+                    .signature_verifier
+                    .get_jwks()
+                    .into_iter()
+                    .collect();
+                sys_jwks.sort();
+                active_jwks.sort();
+
+                assert_eq!(sys_jwks, active_jwks);
+            }
         }
 
         Ok((effects, execution_error_opt))

crates/sui-core/src/authority/authority_per_epoch_store.rs

@@ -29,7 +29,7 @@ use sui_types::transaction::{
     AuthenticatorStateUpdate, CertifiedTransaction, SenderSignedData, SharedInputObject,
     TransactionDataAPI, VerifiedCertificate, VerifiedSignedTransaction,
 };
-use tracing::{debug, error, info, trace, warn};
+use tracing::{debug, error, error_span, info, trace, warn};
 use typed_store::rocks::{
     default_db_options, DBBatch, DBMap, DBOptions, MetricConf, TypedStoreError,
 };
@@ -431,6 +431,10 @@ impl AuthorityPerEpochStore {
     ) -> Arc<Self> {
         let current_time = Instant::now();
         let epoch_id = committee.epoch;
+
+        let span = error_span!("AuthorityPerEpochStore::new", ?epoch_id);
+        let _guard = span.enter();
+
         let tables = AuthorityEpochTables::open(epoch_id, parent_path, db_options.clone());
         let end_of_publish =
             StakeAggregator::from_iter(committee.clone(), tables.end_of_publish.unbounded_iter());

crates/sui-core/src/signature_verifier.rs

@@ -306,19 +306,30 @@ impl SignatureVerifier {
         });
     }
 
-    /// Insert a JWK into the verifier state based on provider. Returns true if the kid of the JWK has not already
-    /// been inserted.
+    /// Insert a JWK into the verifier state. Pre-existing entries for a given JwkId will not be
+    /// overwritten.
     pub(crate) fn insert_jwk(&self, jwk_id: &JwkId, jwk: &JWK) {
-        debug!("insert JWK with kid: {:?}", jwk_id);
         let mut jwks = self.jwks.write();
-        jwks.insert(jwk_id.clone(), jwk.clone());
+        match jwks.entry(jwk_id.clone()) {
+            im::hashmap::Entry::Occupied(_) => {
+                debug!("JWK with kid {:?} already exists", jwk_id);
+            }
+            im::hashmap::Entry::Vacant(entry) => {
+                debug!("inserting JWK with kid: {:?}", jwk_id);
+                entry.insert(jwk.clone());
+            }
+        }
     }
 
     pub fn has_jwk(&self, jwk_id: &JwkId, jwk: &JWK) -> bool {
         let jwks = self.jwks.read();
         jwks.get(jwk_id) == Some(jwk)
     }
 
+    pub fn get_jwks(&self) -> ImHashMap<JwkId, JWK> {
+        self.jwks.read().clone()
+    }
+
     pub fn verify_tx(&self, signed_tx: &SenderSignedData) -> SuiResult {
         self.signed_data_cache
             .is_verified(signed_tx.full_message_digest(), || {

crates/sui-e2e-tests/Cargo.toml

@@ -34,6 +34,7 @@ move-binary-format.workspace = true
 move-package.workspace = true
 telemetry-subscribers.workspace = true
 fastcrypto.workspace = true
+fastcrypto-zkp.workspace = true
 move-core-types.workspace = true
 
 sui-core.workspace = true

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

@@ -66,7 +66,7 @@ async fn test_zklogin_provider_not_supported() {
 
 #[sim_test]
 async fn zklogin_end_to_end_test() {
-    run_zklogin_end_to_end_test(TestClusterBuilder::new().build().await).await;
+    run_zklogin_end_to_end_test(TestClusterBuilder::new().with_default_jwks().build().await).await;
 }
 
 #[sim_test]
@@ -75,6 +75,7 @@ async fn zklogin_end_to_end_test_with_auth_state_creation() {
     let test_cluster = TestClusterBuilder::new()
         .with_protocol_version(23.into())
         .with_epoch_duration_ms(10000)
+        .with_default_jwks()
         .build()
         .await;
 
@@ -177,3 +178,65 @@ async fn test_create_authenticator_state_object() {
         });
     }
 }
+
+// This test is intended to look for forks caused by conflicting / repeated JWK votes from
+// validators.
+#[cfg(msim)]
+#[sim_test]
+async fn test_conflicting_jwks() {
+    use futures::StreamExt;
+    use std::collections::HashSet;
+    use std::sync::{Arc, Mutex};
+    use sui_json_rpc_types::SuiTransactionBlockEffectsAPI;
+    use sui_json_rpc_types::TransactionFilter;
+    use sui_types::base_types::ObjectID;
+    use sui_types::transaction::{TransactionDataAPI, TransactionKind};
+
+    let test_cluster = TestClusterBuilder::new()
+        .with_epoch_duration_ms(45000)
+        .with_jwk_fetch_interval(Duration::from_secs(5))
+        .build()
+        .await;
+
+    let jwks = Arc::new(Mutex::new(Vec::new()));
+    let jwks_clone = jwks.clone();
+
+    test_cluster.fullnode_handle.sui_node.with(|node| {
+        let mut txns = node.state().subscription_handler.subscribe_transactions(
+            TransactionFilter::ChangedObject(ObjectID::from_hex_literal("0x7").unwrap()),
+        );
+        let state = node.state();
+
+        tokio::spawn(async move {
+            while let Some(tx) = txns.next().await {
+                let digest = *tx.transaction_digest();
+                let tx = state
+                    .database
+                    .get_transaction_block(&digest)
+                    .unwrap()
+                    .unwrap();
+                match &tx.data().intent_message().value.kind() {
+                    TransactionKind::EndOfEpochTransaction(_) => (),
+                    TransactionKind::AuthenticatorStateUpdate(update) => {
+                        let jwks = &mut *jwks_clone.lock().unwrap();
+                        for jwk in &update.new_active_jwks {
+                            jwks.push(jwk.clone());
+                        }
+                    }
+                    _ => panic!("{:?}", tx),
+                }
+            }
+        });
+    });
+
+    for _ in 0..5 {
+        test_cluster.wait_for_epoch(None).await;
+    }
+
+    let mut seen_jwks = HashSet::new();
+
+    // ensure no jwk is repeated.
+    for jwk in jwks.lock().unwrap().iter() {
+        assert!(seen_jwks.insert((jwk.jwk_id.clone(), jwk.jwk.clone(), jwk.epoch)));
+    }
+}

crates/sui-framework/packages/sui-framework/sources/authenticator_state.move

@@ -80,14 +80,20 @@ module sui::authenticator_state {
         }
     }
 
-    fun jwk_equal(a: &ActiveJwk, b: &ActiveJwk): bool {
+    fun active_jwk_equal(a: &ActiveJwk, b: &ActiveJwk): bool {
         // note: epoch is ignored
-        (&a.jwk.kty == &b.jwk.kty) &&
-           (&a.jwk.e == &b.jwk.e) &&
-           (&a.jwk.n == &b.jwk.n) &&
-           (&a.jwk.alg == &b.jwk.alg) &&
-           (&a.jwk_id.iss == &b.jwk_id.iss) &&
-           (&a.jwk_id.kid == &b.jwk_id.kid)
+        jwk_equal(&a.jwk, &b.jwk) && jwk_id_equal(&a.jwk_id, &b.jwk_id)
+    }
+
+    fun jwk_equal(a: &JWK, b: &JWK): bool {
+        (&a.kty == &b.kty) &&
+           (&a.e == &b.e) &&
+           (&a.n == &b.n) &&
+           (&a.alg == &b.alg)
+    }
+
+    fun jwk_id_equal(a: &JwkId, b: &JwkId): bool {
+        (&a.iss == &b.iss) && (&a.kid == &b.kid)
     }
 
     // Compare the underlying byte arrays lexicographically. Since the strings may be utf8 this
@@ -213,6 +219,7 @@ module sui::authenticator_state {
         assert!(tx_context::sender(ctx) == @0x0, ENotSystemAddress);
 
         check_sorted(&new_active_jwks);
+        let new_active_jwks = deduplicate(new_active_jwks);
 
         let inner = load_inner_mut(self);
 
@@ -227,12 +234,19 @@ module sui::authenticator_state {
             let new_jwk = vector::borrow(&new_active_jwks, j);
 
             // when they are equal, push only one, but use the max epoch of the two
-            if (jwk_equal(old_jwk, new_jwk)) {
+            if (active_jwk_equal(old_jwk, new_jwk)) {
                 let jwk = *old_jwk;
                 jwk.epoch = math::max(old_jwk.epoch, new_jwk.epoch);
                 vector::push_back(&mut res, jwk);
                 i = i + 1;
                 j = j + 1;
+            } else if (jwk_id_equal(&old_jwk.jwk_id, &new_jwk.jwk_id)) {
+                // if only jwk_id is equal, then the key has changed. Providers should not send
+                // JWKs like this, but if they do, we must ignore the new JWK to avoid having a
+                // liveness / forking issues
+                vector::push_back(&mut res, *old_jwk);
+                i = i + 1;
+                j = j + 1;
             } else if (jwk_lt(old_jwk, new_jwk)) {
                 vector::push_back(&mut res, *old_jwk);
                 i = i + 1;
@@ -254,6 +268,27 @@ module sui::authenticator_state {
         inner.active_jwks = res;
     }
 
+    fun deduplicate(jwks: vector<ActiveJwk>): vector<ActiveJwk> {
+        let res = vector[];
+        let i = 0;
+        let prev: Option<JwkId> = option::none();
+        while (i < vector::length(&jwks)) {
+            let jwk = vector::borrow(&jwks, i);
+            if (option::is_none(&prev)) {
+                option::fill(&mut prev, jwk.jwk_id);
+            } else if (jwk_id_equal(option::borrow(&prev), &jwk.jwk_id)) {
+                // skip duplicate jwks in input
+                i = i + 1;
+                continue;
+            } else {
+                *option::borrow_mut(&mut prev) = jwk.jwk_id;
+            };
+            vector::push_back(&mut res, *jwk);
+            i = i + 1;
+        };
+        res
+    }
+
     #[allow(unused_function)]
     // Called directly by rust when constructing the ChangeEpoch transaction.
     fun expire_jwks(

crates/sui-framework/packages/sui-framework/tests/authenticator_state_tests.move

@@ -16,6 +16,7 @@ module sui::authenticator_state_tests {
         update_authenticator_state_for_testing,
         get_active_jwks_for_testing,
         expire_jwks_for_testing,
+        ActiveJwk,
     };
     use sui::tx_context;
 
@@ -65,6 +66,45 @@ module sui::authenticator_state_tests {
         test_scenario::end(scenario_val);
     }
 
+    #[test]
+    fun authenticator_state_tests_deduplication() {
+        let scenario_val = test_scenario::begin(@0x0);
+        let scenario = &mut scenario_val;
+
+        authenticator_state::create_for_testing(test_scenario::ctx(scenario));
+        test_scenario::next_tx(scenario, @0x0);
+
+        let auth_state = test_scenario::take_shared<AuthenticatorState>(scenario);
+
+        let jwk1 = create_active_jwk(utf8(b"https://accounts.google.com"), utf8(b"kid2"), utf8(b"k1"), 0);
+        update_authenticator_state_for_testing(&mut auth_state, vector[jwk1], test_scenario::ctx(scenario));
+
+        let recorded_jwks = get_active_jwks_for_testing(&auth_state, test_scenario::ctx(scenario));
+        assert!(vector::length(&recorded_jwks) == 1, 0);
+        assert!(vector::borrow(&recorded_jwks, 0) == &jwk1, 0);
+
+        let jwk2 = create_active_jwk(utf8(b"https://www.facebook.com"), utf8(b"kid1"), utf8(b"k2"), 0);
+        let jwk3 = create_active_jwk(utf8(b"https://accounts.google.com"), utf8(b"kid2"), utf8(b"k3"), 0);
+        update_authenticator_state_for_testing(&mut auth_state, vector[jwk2, jwk3], test_scenario::ctx(scenario));
+
+        let recorded_jwks = get_active_jwks_for_testing(&auth_state, test_scenario::ctx(scenario));
+        assert!(vector::length(&recorded_jwks) == 2, 0);
+        // jwk2 sorts before 1, and 3 is dropped because its a duplicated
+        assert!(vector::borrow(&recorded_jwks, 0) == &jwk2, 0);
+        assert!(vector::borrow(&recorded_jwks, 1) == &jwk1, 0);
+
+        let jwk4 = create_active_jwk(utf8(b"https://accounts.google.com"), utf8(b"kid4"), utf8(b"k4"), 0);
+        update_authenticator_state_for_testing(&mut auth_state, vector[jwk4], test_scenario::ctx(scenario));
+        let recorded_jwks = get_active_jwks_for_testing(&auth_state, test_scenario::ctx(scenario));
+        assert!(vector::length(&recorded_jwks) == 3, 0);
+        assert!(vector::borrow(&recorded_jwks, 0) == &jwk2, 0);
+        assert!(vector::borrow(&recorded_jwks, 1) == &jwk1, 0);
+        assert!(vector::borrow(&recorded_jwks, 2) == &jwk4, 0);
+
+        test_scenario::return_shared(auth_state);
+        test_scenario::end(scenario_val);
+    }
+
     #[test]
     fun authenticator_state_tests_expiry_edge_cases() {
         let scenario_val = test_scenario::begin(@0x0);