Tests for pathologcial jwk cases (#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.

Author: mystenmark

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);