jwk #2: ensure jwk txns are expected in consensus

aptos-move/aptos-release-builder/src/components/feature_flags.rs

@@ -99,6 +99,7 @@
     ZkIdSignature,
     ZkIdZkLessSignature,
     RemoveDetailedError,
+    JWKConsensus,
 }
 
 fn generate_features_blob(writer: &CodeWriter, data: &[u64]) {
@@ -256,6 +257,7 @@
             FeatureFlag::ZkIdSignature => AptosFeatureFlag::ZK_ID_SIGNATURES,
             FeatureFlag::ZkIdZkLessSignature => AptosFeatureFlag::ZK_ID_ZKLESS_SIGNATURE,
             FeatureFlag::RemoveDetailedError => AptosFeatureFlag::REMOVE_DETAILED_ERROR_FROM_HASH,
+            FeatureFlag::JWKConsensus => AptosFeatureFlag::JWK_CONSENSUS,
         }
     }
 }
@@ -336,6 +338,7 @@
             AptosFeatureFlag::ZK_ID_SIGNATURES => FeatureFlag::ZkIdSignature,
             AptosFeatureFlag::ZK_ID_ZKLESS_SIGNATURE => FeatureFlag::ZkIdZkLessSignature,
             AptosFeatureFlag::REMOVE_DETAILED_ERROR_FROM_HASH => FeatureFlag::RemoveDetailedError,
+            AptosFeatureFlag::JWK_CONSENSUS => FeatureFlag::JWKConsensus,
         }
     }
 }

aptos-move/aptos-vm/Cargo.toml

@@ -15,6 +15,7 @@ rust-version = { workspace = true }
 [dependencies]
 anyhow = { workspace = true }
 aptos-aggregator = { workspace = true }
+aptos-bitvec = { workspace = true }
 aptos-block-executor = { workspace = true }
 aptos-block-partitioner = { workspace = true }
 aptos-crypto = { workspace = true }

aptos-move/aptos-vm/src/system_module_names.rs

@@ -39,6 +39,15 @@
 
 pub const FINISH_WITH_DKG_RESULT: &IdentStr = ident_str!("finish_with_dkg_result");
 
+pub static JWKS_MODULE: Lazy<ModuleId> = Lazy::new(|| {
+    ModuleId::new(
+        account_config::CORE_CODE_ADDRESS,
+        ident_str!("jwks").to_owned(),
+    )
+});
+
+pub const UPSERT_INTO_OBSERVED_JWKS: &IdentStr = ident_str!("upsert_into_observed_jwks");
+
 pub static MULTISIG_ACCOUNT_MODULE: Lazy<ModuleId> = Lazy::new(|| {
     ModuleId::new(
         account_config::CORE_CODE_ADDRESS,

aptos-move/aptos-vm/src/validator_txns/jwk.rs

@@ -0,0 +1,161 @@
+// Copyright © Aptos Foundation
+
+use crate::{
+    aptos_vm::get_or_vm_startup_failure,
+    errors::expect_only_successful_execution,
+    move_vm_ext::{AptosMoveResolver, SessionId},
+    system_module_names::{JWKS_MODULE, UPSERT_INTO_OBSERVED_JWKS},
+    validator_txns::jwk::{
+        ExecutionFailure::{Expected, Unexpected},
+        ExpectedFailure::{
+            IncorrectVersion, MissingResourceObservedJWKs, MissingResourceValidatorSet,
+            MultiSigVerificationFailed, NotEnoughVotingPower,
+        },
+    },
+    AptosVM,
+};
+use aptos_bitvec::BitVec;
+use aptos_types::{
+    aggregate_signature::AggregateSignature,
+    fee_statement::FeeStatement,
+    jwks,
+    jwks::{Issuer, ObservedJWKs, ProviderJWKs, QuorumCertifiedUpdate},
+    move_utils::as_move_value::AsMoveValue,
+    on_chain_config::{OnChainConfig, ValidatorSet},
+    transaction::{ExecutionStatus, TransactionStatus},
+    validator_verifier::ValidatorVerifier,
+};
+use aptos_vm_logging::log_schema::AdapterLogSchema;
+use aptos_vm_types::output::VMOutput;
+use move_core_types::{
+    account_address::AccountAddress,
+    value::{serialize_values, MoveValue},
+    vm_status::{AbortLocation, StatusCode, VMStatus},
+};
+use move_vm_types::gas::UnmeteredGasMeter;
+use std::collections::HashMap;
+
+enum ExpectedFailure {
+    // Move equivalent: `errors::invalid_argument(*)`
+    IncorrectVersion = 0x010103,
+    MultiSigVerificationFailed = 0x010104,
+    NotEnoughVotingPower = 0x010105,
+
+    // Move equivalent: `errors::invalid_state(*)`
+    MissingResourceValidatorSet = 0x30101,
+    MissingResourceObservedJWKs = 0x30102,
+}
+
+enum ExecutionFailure {
+    Expected(ExpectedFailure),
+    Unexpected(VMStatus),
+}
+
+impl AptosVM {
+    pub(crate) fn process_jwk_update(
+        &self,
+        resolver: &impl AptosMoveResolver,
+        log_context: &AdapterLogSchema,
+        session_id: SessionId,
+        update: jwks::QuorumCertifiedUpdate,
+    ) -> Result<(VMStatus, VMOutput), VMStatus> {
+        match self.process_jwk_update_inner(resolver, log_context, session_id, update) {
+            Ok((vm_status, vm_output)) => Ok((vm_status, vm_output)),
+            Err(Expected(failure)) => {
+                // Pretend we are inside Move, and expected failures are like Move aborts.
+                Ok((
+                    VMStatus::MoveAbort(AbortLocation::Script, failure as u64),
+                    VMOutput::empty_with_status(TransactionStatus::Discard(StatusCode::ABORTED)),
+                ))
+            },
+            Err(Unexpected(vm_status)) => Err(vm_status),
+        }
+    }
+
+    fn process_jwk_update_inner(
+        &self,
+        resolver: &impl AptosMoveResolver,
+        log_context: &AdapterLogSchema,
+        session_id: SessionId,
+        update: jwks::QuorumCertifiedUpdate,
+    ) -> Result<(VMStatus, VMOutput), ExecutionFailure> {
+        // Load resources.
+        let validator_set = ValidatorSet::fetch_config(resolver)
+            .ok_or_else(|| Expected(MissingResourceValidatorSet))?;
+        let observed_jwks = ObservedJWKs::fetch_config(resolver)
+            .ok_or_else(|| Expected(MissingResourceObservedJWKs))?;
+
+        let mut jwks_by_issuer: HashMap<Issuer, ProviderJWKs> =
+            observed_jwks.into_providers_jwks().into();
+        let issuer = update.update.issuer.clone();
+        let on_chain = jwks_by_issuer
+            .entry(issuer.clone())
+            .or_insert_with(|| ProviderJWKs::new(issuer));
+        let verifier = ValidatorVerifier::from(&validator_set);
+
+        let QuorumCertifiedUpdate {
+            authors,
+            update: observed,
+            multi_sig,
+        } = update;
+
+        // Check version.
+        if on_chain.version + 1 != observed.version {
+            return Err(Expected(IncorrectVersion));
+        }
+
+        let signer_bit_vec = BitVec::from(
+            verifier
+                .get_ordered_account_addresses()
+                .into_iter()
+                .map(|addr| authors.contains(&addr))
+                .collect::<Vec<_>>(),
+        );
+
+        // Verify multi-sig.
+        verifier
+            .verify_multi_signatures(
+                &observed,
+                &AggregateSignature::new(signer_bit_vec, Some(multi_sig)),
+            )
+            .map_err(|_| Expected(MultiSigVerificationFailed))?;
+
+        // Check voting power.
+        verifier
+            .check_voting_power(authors.iter(), true)
+            .map_err(|_| Expected(NotEnoughVotingPower))?;
+
+        // All verification passed. Apply the `observed`.
+        let mut gas_meter = UnmeteredGasMeter;
+        let mut session = self.new_session(resolver, session_id);
+        let args = vec![
+            MoveValue::Signer(AccountAddress::ONE),
+            vec![observed].as_move_value(),
+        ];
+
+        session
+            .execute_function_bypass_visibility(
+                &JWKS_MODULE,
+                UPSERT_INTO_OBSERVED_JWKS,
+                vec![],
+                serialize_values(&args),
+                &mut gas_meter,
+            )
+            .map_err(|e| {
+                expect_only_successful_execution(e, UPSERT_INTO_OBSERVED_JWKS.as_str(), log_context)
+            })
+            .map_err(|r| Unexpected(r.unwrap_err()))?;
+
+        let output = crate::aptos_vm::get_transaction_output(
+            session,
+            FeeStatement::zero(),
+            ExecutionStatus::Success,
+            &get_or_vm_startup_failure(&self.storage_gas_params, log_context)
+                .map_err(Unexpected)?
+                .change_set_configs,
+        )
+        .map_err(Unexpected)?;
+
+        Ok((VMStatus::Executed, output))
+    }
+}

aptos-move/aptos-vm/src/validator_txns/mod.rs

@@ -21,6 +21,9 @@
             ValidatorTransaction::DKGResult(dkg_node) => {
                 self.process_dkg_result(resolver, log_context, session_id, dkg_node)
             },
+            ValidatorTransaction::ObservedJWKUpdate(jwk_update) => {
+                self.process_jwk_update(resolver, log_context, session_id, jwk_update)
+            },
             ValidatorTransaction::DummyTopic1(dummy) | ValidatorTransaction::DummyTopic2(dummy) => {
                 self.process_dummy_validator_txn(resolver, log_context, session_id, dummy)
             },
@@ -30,3 +33,4 @@
 
 mod dkg;
 mod dummy;
+mod jwk;

aptos-move/framework/aptos-framework/doc/jwks.md

@@ -620,6 +620,51 @@ This is what applications should consume.
 
 
 
+<a id="0x1_jwks_ENATIVE_INCORRECT_VERSION"></a>
+
+
+
+<pre><code><b>const</b> <a href="jwks.md#0x1_jwks_ENATIVE_INCORRECT_VERSION">ENATIVE_INCORRECT_VERSION</a>: u64 = 259;
+</code></pre>
+
+
+
+<a id="0x1_jwks_ENATIVE_MISSING_RESOURCE_OBSERVED_JWKS"></a>
+
+
+
+<pre><code><b>const</b> <a href="jwks.md#0x1_jwks_ENATIVE_MISSING_RESOURCE_OBSERVED_JWKS">ENATIVE_MISSING_RESOURCE_OBSERVED_JWKS</a>: u64 = 258;
+</code></pre>
+
+
+
+<a id="0x1_jwks_ENATIVE_MISSING_RESOURCE_VALIDATOR_SET"></a>
+
+
+
+<pre><code><b>const</b> <a href="jwks.md#0x1_jwks_ENATIVE_MISSING_RESOURCE_VALIDATOR_SET">ENATIVE_MISSING_RESOURCE_VALIDATOR_SET</a>: u64 = 257;
+</code></pre>
+
+
+
+<a id="0x1_jwks_ENATIVE_MULTISIG_VERIFICATION_FAILED"></a>
+
+
+
+<pre><code><b>const</b> <a href="jwks.md#0x1_jwks_ENATIVE_MULTISIG_VERIFICATION_FAILED">ENATIVE_MULTISIG_VERIFICATION_FAILED</a>: u64 = 260;
+</code></pre>
+
+
+
+<a id="0x1_jwks_ENATIVE_NOT_ENOUGH_VOTING_POWER"></a>
+
+
+
+<pre><code><b>const</b> <a href="jwks.md#0x1_jwks_ENATIVE_NOT_ENOUGH_VOTING_POWER">ENATIVE_NOT_ENOUGH_VOTING_POWER</a>: u64 = 261;
+</code></pre>
+
+
+
 <a id="0x1_jwks_EUNEXPECTED_EPOCH"></a>
 
 

aptos-move/framework/aptos-framework/sources/jwks.move

@@ -29,6 +29,12 @@ module aptos_framework::jwks {
     const EISSUER_NOT_FOUND: u64 = 5;
     const EJWK_ID_NOT_FOUND: u64 = 6;
 
+    const ENATIVE_MISSING_RESOURCE_VALIDATOR_SET: u64 = 0x0101;
+    const ENATIVE_MISSING_RESOURCE_OBSERVED_JWKS: u64 = 0x0102;
+    const ENATIVE_INCORRECT_VERSION: u64 = 0x0103;
+    const ENATIVE_MULTISIG_VERIFICATION_FAILED: u64 = 0x0104;
+    const ENATIVE_NOT_ENOUGH_VOTING_POWER: u64 = 0x0105;
+
     /// An OIDC provider.
     struct OIDCProvider has drop, store {
         /// The utf-8 encoded issuer string. E.g., b"https://www.facebook.com".

aptos-move/framework/move-stdlib/doc/features.md

@@ -88,6 +88,8 @@ return true.
 -  [Function `zkid_feature_enabled`](#0x1_features_zkid_feature_enabled)
 -  [Function `get_zkid_zkless_feature`](#0x1_features_get_zkid_zkless_feature)
 -  [Function `zkid_zkless_feature_enabled`](#0x1_features_zkid_zkless_feature_enabled)
+-  [Function `get_jwk_consensus_feature`](#0x1_features_get_jwk_consensus_feature)
+-  [Function `jwk_consensus_enabled`](#0x1_features_jwk_consensus_enabled)
 -  [Function `change_feature_flags`](#0x1_features_change_feature_flags)
 -  [Function `is_enabled`](#0x1_features_is_enabled)
 -  [Function `set`](#0x1_features_set)
@@ -369,6 +371,18 @@ Lifetime: transient
 
 
 
+<a id="0x1_features_JWK_CONSENSUS"></a>
+
+The JWK consensus feature.
+
+Lifetime: permanent
+
+
+<pre><code><b>const</b> <a href="features.md#0x1_features_JWK_CONSENSUS">JWK_CONSENSUS</a>: u64 = 49;
+</code></pre>
+
+
+
 <a id="0x1_features_LIMIT_MAX_IDENTIFIER_LENGTH"></a>
 
 
@@ -1918,6 +1932,52 @@ Lifetime: transient
 
 
 
+</details>
+
+<a id="0x1_features_get_jwk_consensus_feature"></a>
+
+## Function `get_jwk_consensus_feature`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="features.md#0x1_features_get_jwk_consensus_feature">get_jwk_consensus_feature</a>(): u64
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="features.md#0x1_features_get_jwk_consensus_feature">get_jwk_consensus_feature</a>(): u64 { <a href="features.md#0x1_features_JWK_CONSENSUS">JWK_CONSENSUS</a> }
+</code></pre>
+
+
+
+</details>
+
+<a id="0x1_features_jwk_consensus_enabled"></a>
+
+## Function `jwk_consensus_enabled`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="features.md#0x1_features_jwk_consensus_enabled">jwk_consensus_enabled</a>(): bool
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="features.md#0x1_features_jwk_consensus_enabled">jwk_consensus_enabled</a>(): bool <b>acquires</b> <a href="features.md#0x1_features_Features">Features</a> {
+    <a href="features.md#0x1_features_is_enabled">is_enabled</a>(<a href="features.md#0x1_features_JWK_CONSENSUS">JWK_CONSENSUS</a>)
+}
+</code></pre>
+
+
+
 </details>
 
 <a id="0x1_features_change_feature_flags"></a>