From 68e2f246a915239a9613fbf2025e1069916ad151 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 01:48:38 -0800 Subject: [PATCH 01/22] jwk types update --- types/src/jwks/{jwk.rs => jwk/mod.rs} | 105 +++++++++------- types/src/jwks/jwk/tests.rs | 79 ++++++++++++ types/src/jwks/mod.rs | 128 +++++++++++++++++-- types/src/jwks/{rsa.rs => rsa/mod.rs} | 133 +++++--------------- types/src/jwks/rsa/tests.rs | 98 +++++++++++++++ types/src/jwks/unsupported.rs | 89 ------------- types/src/jwks/unsupported/mod.rs | 73 +++++++++++ types/src/jwks/unsupported/tests.rs | 45 +++++++ types/src/on_chain_config/aptos_features.rs | 12 +- 9 files changed, 516 insertions(+), 246 deletions(-) rename types/src/jwks/{jwk.rs => jwk/mod.rs} (54%) create mode 100644 types/src/jwks/jwk/tests.rs rename types/src/jwks/{rsa.rs => rsa/mod.rs} (50%) create mode 100644 types/src/jwks/rsa/tests.rs delete mode 100644 types/src/jwks/unsupported.rs create mode 100644 types/src/jwks/unsupported/mod.rs create mode 100644 types/src/jwks/unsupported/tests.rs diff --git a/types/src/jwks/jwk.rs b/types/src/jwks/jwk/mod.rs similarity index 54% rename from types/src/jwks/jwk.rs rename to types/src/jwks/jwk/mod.rs index 44e710a547be1..524bdf597e463 100644 --- a/types/src/jwks/jwk.rs +++ b/types/src/jwks/jwk/mod.rs @@ -3,26 +3,82 @@ use crate::{ jwks::{rsa::RSA_JWK, unsupported::UnsupportedJWK}, move_any::{Any as MoveAny, AsMoveAny}, + move_utils::as_move_value::AsMoveValue, }; use anyhow::anyhow; +use aptos_crypto_derive::{BCSCryptoHash, CryptoHasher}; +use move_core_types::value::{MoveStruct, MoveValue}; use serde::{Deserialize, Serialize}; +use std::{ + cmp::Ordering, + fmt::{Debug, Formatter}, +}; /// Reflection of Move type `0x1::jwks::JWK`. /// When you load an on-chain config that contains some JWK(s), the JWK will be of this type. /// When you call a Move function from rust that takes some JWKs as input, pass in JWKs of this type. /// Otherwise, it is recommended to convert this to the rust enum `JWK` below for better rust experience. -#[derive(Debug, PartialEq, Serialize, Deserialize)] +/// See its doc in Move for more details. +#[derive(Clone, Eq, PartialEq, Serialize, Deserialize, CryptoHasher, BCSCryptoHash)] pub struct JWKMoveStruct { pub variant: MoveAny, } +impl Debug for JWKMoveStruct { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let jwk = JWK::try_from(self); + f.debug_struct("JWKMoveStruct") + .field("variant", &jwk) + .finish() + } +} + +impl AsMoveValue for JWKMoveStruct { + fn as_move_value(&self) -> MoveValue { + MoveValue::Struct(MoveStruct::Runtime(vec![self.variant.as_move_value()])) + } +} + /// The JWK type that can be converted from/to `JWKMoveStruct` but easier to use in rust. -#[derive(Debug, PartialEq)] +#[derive(Debug, PartialEq, Eq)] pub enum JWK { RSA(RSA_JWK), Unsupported(UnsupportedJWK), } +impl JWK { + pub fn id(&self) -> Vec { + match self { + JWK::RSA(rsa) => rsa.id(), + JWK::Unsupported(unsupported) => unsupported.id(), + } + } +} + +impl PartialOrd for JWK { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for JWK { + fn cmp(&self, other: &Self) -> Ordering { + self.id().cmp(&other.id()) + } +} + +impl From for JWK { + fn from(value: serde_json::Value) -> Self { + match RSA_JWK::try_from(&value) { + Ok(rsa) => Self::RSA(rsa), + Err(_) => { + let unsupported = UnsupportedJWK::from(value); + Self::Unsupported(unsupported) + }, + } + } +} + impl From for JWKMoveStruct { fn from(jwk: JWK) -> Self { let variant = match jwk { @@ -55,46 +111,5 @@ impl TryFrom<&JWKMoveStruct> for JWK { } } -#[test] -fn convert_jwk_move_struct_to_jwk() { - let unsupported_jwk = UnsupportedJWK::new_for_testing("id1", "payload1"); - let jwk_move_struct = JWKMoveStruct { - variant: unsupported_jwk.as_move_any(), - }; - assert_eq!( - JWK::Unsupported(unsupported_jwk), - JWK::try_from(&jwk_move_struct).unwrap() - ); - - let rsa_jwk = RSA_JWK::new_for_testing("kid1", "kty1", "alg1", "e1", "n1"); - let jwk_move_struct = JWKMoveStruct { - variant: rsa_jwk.as_move_any(), - }; - assert_eq!(JWK::RSA(rsa_jwk), JWK::try_from(&jwk_move_struct).unwrap()); - - let unknown_jwk_variant = MoveAny { - type_name: "type1".to_string(), - data: vec![], - }; - assert!(JWK::try_from(&JWKMoveStruct { - variant: unknown_jwk_variant - }) - .is_err()); -} - -#[test] -fn convert_jwk_to_jwk_move_struct() { - let unsupported_jwk = UnsupportedJWK::new_for_testing("id1", "payload1"); - let jwk = JWK::Unsupported(unsupported_jwk.clone()); - let jwk_move_struct = JWKMoveStruct { - variant: unsupported_jwk.as_move_any(), - }; - assert_eq!(jwk_move_struct, JWKMoveStruct::from(jwk)); - - let rsa_jwk = RSA_JWK::new_for_testing("kid1", "kty1", "alg1", "e1", "n1"); - let jwk = JWK::RSA(rsa_jwk.clone()); - let jwk_move_struct = JWKMoveStruct { - variant: rsa_jwk.as_move_any(), - }; - assert_eq!(jwk_move_struct, JWKMoveStruct::from(jwk)); -} +#[cfg(test)] +mod tests; diff --git a/types/src/jwks/jwk/tests.rs b/types/src/jwks/jwk/tests.rs new file mode 100644 index 0000000000000..07ede2771ad49 --- /dev/null +++ b/types/src/jwks/jwk/tests.rs @@ -0,0 +1,79 @@ +// Copyright © Aptos Foundation + +use crate::{ + jwks::{ + jwk::{JWKMoveStruct, JWK}, + rsa::RSA_JWK, + unsupported::UnsupportedJWK, + }, + move_any::{Any as MoveAny, AsMoveAny}, +}; +use aptos_crypto::HashValue; +use std::str::FromStr; + +#[test] +fn convert_jwk_move_struct_to_jwk() { + let unsupported_jwk = UnsupportedJWK::new_for_testing("id1", "payload1"); + let jwk_move_struct = JWKMoveStruct { + variant: unsupported_jwk.as_move_any(), + }; + assert_eq!( + JWK::Unsupported(unsupported_jwk), + JWK::try_from(&jwk_move_struct).unwrap() + ); + + let rsa_jwk = RSA_JWK::new_from_strs("kid1", "kty1", "alg1", "e1", "n1"); + let jwk_move_struct = JWKMoveStruct { + variant: rsa_jwk.as_move_any(), + }; + assert_eq!(JWK::RSA(rsa_jwk), JWK::try_from(&jwk_move_struct).unwrap()); + + let unknown_jwk_variant = MoveAny { + type_name: "type1".to_string(), + data: vec![], + }; + assert!(JWK::try_from(&JWKMoveStruct { + variant: unknown_jwk_variant + }) + .is_err()); +} + +#[test] +fn convert_jwk_to_jwk_move_struct() { + let unsupported_jwk = UnsupportedJWK::new_for_testing("id1", "payload1"); + let jwk = JWK::Unsupported(unsupported_jwk.clone()); + let jwk_move_struct = JWKMoveStruct { + variant: unsupported_jwk.as_move_any(), + }; + assert_eq!(jwk_move_struct, JWKMoveStruct::from(jwk)); + + let rsa_jwk = RSA_JWK::new_from_strs("kid1", "kty1", "alg1", "e1", "n1"); + let jwk = JWK::RSA(rsa_jwk.clone()); + let jwk_move_struct = JWKMoveStruct { + variant: rsa_jwk.as_move_any(), + }; + assert_eq!(jwk_move_struct, JWKMoveStruct::from(jwk)); +} + +#[test] +fn convert_json_value_to_jwk() { + let json_str = + r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + let actual = JWK::from(json); + let expected = JWK::RSA(RSA_JWK::new_from_strs( + "kid1", "RSA", "RS256", "AQAB", "13131", + )); + assert_eq!(expected, actual); + + let compact_json_str = r#"{"alg":13131}"#; + let json = serde_json::Value::from_str(compact_json_str).unwrap(); + let actual = JWK::from(json); + let expected_payload = compact_json_str.as_bytes().to_vec(); + let expected_id = HashValue::sha3_256_of(expected_payload.as_slice()).to_vec(); + let expected = JWK::Unsupported(UnsupportedJWK { + id: expected_id, + payload: expected_payload, + }); + assert_eq!(expected, actual); +} diff --git a/types/src/jwks/mod.rs b/types/src/jwks/mod.rs index cdf5505c52ec4..307502b398fe4 100644 --- a/types/src/jwks/mod.rs +++ b/types/src/jwks/mod.rs @@ -1,10 +1,23 @@ // Copyright © Aptos Foundation use self::jwk::JWK; -use anyhow::{bail, Context, Ok, Result}; +use crate::{move_utils::as_move_value::AsMoveValue, on_chain_config::OnChainConfig}; +use anyhow::{bail, Context}; +use aptos_crypto::bls12381; +use aptos_crypto_derive::{BCSCryptoHash, CryptoHasher}; use jwk::JWKMoveStruct; -use move_core_types::{ident_str, identifier::IdentStr, move_resource::MoveStructType}; +use move_core_types::{ + account_address::AccountAddress, + ident_str, + identifier::IdentStr, + move_resource::MoveStructType, + value::{MoveStruct, MoveValue}, +}; use serde::{Deserialize, Serialize}; +use std::{ + collections::{BTreeSet, HashMap}, + fmt::{Debug, Formatter}, +}; pub mod jwk; pub mod rsa; @@ -17,30 +30,66 @@ pub fn issuer_from_str(s: &str) -> Issuer { } /// Move type `0x1::jwks::OIDCProvider` in rust. +/// See its doc in Move for more details. +#[derive(Default, Serialize, Deserialize)] pub struct OIDCProvider { pub name: Issuer, pub config_url: Vec, } /// Move type `0x1::jwks::SupportedOIDCProviders` in rust. +/// See its doc in Move for more details. +#[derive(Default, Serialize, Deserialize)] pub struct SupportedOIDCProviders { pub providers: Vec, } +impl SupportedOIDCProviders { + pub fn into_provider_vec(self) -> Vec { + self.providers + } +} + +impl OnChainConfig for SupportedOIDCProviders { + const MODULE_IDENTIFIER: &'static str = "jwks"; + const TYPE_IDENTIFIER: &'static str = "SupportedOIDCProviders"; +} + /// Move type `0x1::jwks::ProviderJWKs` in rust. -#[derive(Debug, PartialEq, Serialize, Deserialize)] +/// See its doc in Move for more details. +#[derive(Clone, Default, Eq, PartialEq, Serialize, Deserialize, CryptoHasher, BCSCryptoHash)] pub struct ProviderJWKs { pub issuer: Issuer, pub version: u64, pub jwks: Vec, } +impl ProviderJWKs { + pub fn new(issuer: Issuer) -> Self { + Self { + issuer, + version: 0, + jwks: vec![], + } + } +} + +impl Debug for ProviderJWKs { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ProviderJWKs") + .field("issuer", &String::from_utf8(self.issuer.clone())) + .field("version", &self.version) + .field("jwks", &self.jwks) + .finish() + } +} + impl ProviderJWKs { pub fn jwks(&self) -> &Vec { &self.jwks } - pub fn get_jwk(&self, id: &str) -> Result<&JWKMoveStruct> { + pub fn get_jwk(&self, id: &str) -> anyhow::Result<&JWKMoveStruct> { for jwk_move in self.jwks() { let jwk = JWK::try_from(jwk_move)?; match jwk { @@ -60,18 +109,51 @@ impl ProviderJWKs { } } +impl AsMoveValue for ProviderJWKs { + fn as_move_value(&self) -> MoveValue { + MoveValue::Struct(MoveStruct::Runtime(vec![ + self.issuer.as_move_value(), + self.version.as_move_value(), + self.jwks.as_move_value(), + ])) + } +} /// Move type `0x1::jwks::JWKs` in rust. -#[derive(Debug, Serialize, Deserialize, PartialEq)] +/// See its doc in Move for more details. +#[derive(Debug, Default, Serialize, Deserialize, PartialEq)] pub struct AllProvidersJWKs { pub entries: Vec, } +impl From for HashMap { + fn from(value: AllProvidersJWKs) -> Self { + let AllProvidersJWKs { entries } = value; + entries + .into_iter() + .map(|entry| (entry.issuer.clone(), entry)) + .collect() + } +} + /// Move type `0x1::jwks::ObservedJWKs` in rust. -#[derive(Debug, Serialize, Deserialize, PartialEq)] +/// See its doc in Move for more details. +#[derive(Debug, Default, Serialize, Deserialize, PartialEq)] pub struct ObservedJWKs { pub jwks: AllProvidersJWKs, } +impl ObservedJWKs { + pub fn into_providers_jwks(self) -> AllProvidersJWKs { + let Self { jwks } = self; + jwks + } +} + +impl OnChainConfig for ObservedJWKs { + const MODULE_IDENTIFIER: &'static str = "jwks"; + const TYPE_IDENTIFIER: &'static str = "ObservedJWKs"; +} + /// Reflection of Move type `0x1::jwks::ObservedJWKs`. #[derive(Debug, Serialize, Deserialize, PartialEq)] pub struct PatchedJWKs { @@ -86,7 +168,7 @@ impl PatchedJWKs { .find(|&provider_jwk_set| provider_jwk_set.issuer.eq(&issuer_from_str(iss))) } - pub fn get_jwk(&self, iss: &str, kid: &str) -> Result<&JWKMoveStruct> { + pub fn get_jwk(&self, iss: &str, kid: &str) -> anyhow::Result<&JWKMoveStruct> { let provider_jwk_set = self .get_provider_jwks(iss) .context("JWK not found for issuer")?; @@ -99,3 +181,35 @@ impl MoveStructType for PatchedJWKs { const MODULE_NAME: &'static IdentStr = ident_str!("jwks"); const STRUCT_NAME: &'static IdentStr = ident_str!("PatchedJWKs"); } + +/// A JWK update in format of `ProviderJWKs` and a multi-signature of it as a quorum certificate. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, CryptoHasher, BCSCryptoHash)] +pub struct QuorumCertifiedUpdate { + pub authors: BTreeSet, + pub update: ProviderJWKs, + pub multi_sig: bls12381::Signature, +} + +impl QuorumCertifiedUpdate { + #[cfg(any(test, feature = "fuzzing"))] + pub fn dummy() -> Self { + Self { + authors: Default::default(), + update: Default::default(), + multi_sig: bls12381::Signature::dummy_signature(), + } + } +} + +/// Move event type `0x1::jwks::ObservedJWKsUpdated` in rust. +/// See its doc in Move for more details. +#[derive(Serialize, Deserialize)] +pub struct ObservedJWKsUpdated { + pub epoch: u64, + pub jwks: AllProvidersJWKs, +} + +impl MoveStructType for ObservedJWKsUpdated { + const MODULE_NAME: &'static IdentStr = ident_str!("jwks"); + const STRUCT_NAME: &'static IdentStr = ident_str!("ObservedJWKsUpdated"); +} diff --git a/types/src/jwks/rsa.rs b/types/src/jwks/rsa/mod.rs similarity index 50% rename from types/src/jwks/rsa.rs rename to types/src/jwks/rsa/mod.rs index 42337b13325b2..d3fd836f5c8b6 100644 --- a/types/src/jwks/rsa.rs +++ b/types/src/jwks/rsa/mod.rs @@ -1,7 +1,5 @@ // Copyright © Aptos Foundation -#[cfg(test)] -use crate::move_any::Any as MoveAny; use crate::{move_any::AsMoveAny, move_utils::as_move_value::AsMoveValue, zkid::Claims}; use anyhow::{anyhow, bail, ensure, Result}; use aptos_crypto::poseidon_bn254; @@ -9,14 +7,13 @@ use base64::URL_SAFE_NO_PAD; use jsonwebtoken::{Algorithm, DecodingKey, TokenData, Validation}; use move_core_types::value::{MoveStruct, MoveValue}; use serde::{Deserialize, Serialize}; -#[cfg(test)] -use std::str::FromStr; pub const RSA_MODULUS_BYTES: usize = 256; /// Move type `0x1::jwks::RSA_JWK` in rust. +/// See its doc in Move for more details. #[allow(non_camel_case_types)] -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct RSA_JWK { pub kid: String, pub kty: String, @@ -26,8 +23,19 @@ pub struct RSA_JWK { } impl RSA_JWK { - #[cfg(test)] - pub fn new_for_testing(kid: &str, kty: &str, alg: &str, e: &str, n: &str) -> Self { + /// Make an `RSA_JWK` from `kty="RSA", alg="RS256", e="AQAB"` (a popular setting) + /// and caller-specified `kid` and `n`. + pub fn new_256_aqab(kid: &str, n: &str) -> Self { + Self { + kid: kid.to_string(), + kty: "RSA".to_string(), + alg: "RS256".to_string(), + e: "AQAB".to_string(), + n: n.to_string(), + } + } + + pub fn new_from_strs(kid: &str, kty: &str, alg: &str, e: &str, n: &str) -> Self { Self { kid: kid.to_string(), kty: kty.to_string(), @@ -37,6 +45,18 @@ impl RSA_JWK { } } + pub fn verify_signature(&self, jwt_token: &str) -> Result> { + let mut validation = Validation::new(Algorithm::RS256); + validation.validate_exp = false; + let key = &DecodingKey::from_rsa_components(&self.n, &self.e)?; + let claims = jsonwebtoken::decode::(jwt_token, key, &validation)?; + Ok(claims) + } + + pub fn id(&self) -> Vec { + self.kid.as_bytes().to_vec() + } + // TODO(zkid): Move this to aptos-crypto so other services can use this pub fn to_poseidon_scalar(&self) -> Result { let mut modulus = base64::decode_config(&self.n, URL_SAFE_NO_PAD)?; @@ -55,14 +75,6 @@ impl RSA_JWK { scalars.push(ark_bn254::Fr::from(RSA_MODULUS_BYTES as i32)); poseidon_bn254::hash_scalars(scalars) } - - pub fn verify_signature(&self, jwt_token: &str) -> Result> { - let mut validation = Validation::new(Algorithm::RS256); - validation.validate_exp = false; - let key = &DecodingKey::from_rsa_components(&self.n, &self.e)?; - let claims = jsonwebtoken::decode::(jwt_token, key, &validation)?; - Ok(claims) - } } impl AsMoveAny for RSA_JWK { @@ -129,92 +141,5 @@ impl AsMoveValue for RSA_JWK { } } -#[test] -fn test_rsa_jwk_from_json() { - // Valid JWK JSON should be accepted. - let json_str = - r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - let actual = RSA_JWK::try_from(&json); - let expected = RSA_JWK::new_for_testing("kid1", "RSA", "RS256", "AQAB", "13131"); - assert_eq!(expected, actual.unwrap()); - - // JWK JSON without `kid` should be rejected. - let json_str = r#"{"alg": "RS256", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON with wrong `kid` type should be rejected. - let json_str = - r#"{"alg": "RS256", "kid": {}, "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON without `alg` should be rejected. - let json_str = r#"{"kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON with wrong `alg` type should be rejected. - let json_str = - r#"{"alg": 0, "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON without `kty` should be rejected. - let json_str = r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON with wrong `kty` value should be rejected. - let json_str = - r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSB", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON without `e` should be rejected. - let json_str = r#"{"alg": "RS256", "kid": "kid1", "use": "sig", "kty": "RSA", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON with wrong `e` type should be rejected. - let json_str = - r#"{"alg": "RS256", "kid": "kid1", "e": 65537, "use": "sig", "kty": "RSA", "n": "13131"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON without `n` should be rejected. - let json_str = r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA"}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); - - // JWK JSON with wrong `n` type should be rejected. - let json_str = - r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": false}"#; - let json = serde_json::Value::from_str(json_str).unwrap(); - assert!(RSA_JWK::try_from(&json).is_err()); -} - -#[test] -fn test_rsa_jwk_as_move_value() { - let rsa_jwk = RSA_JWK::new_for_testing("kid1", "RSA", "RS256", "AQAB", "13131"); - let move_value = rsa_jwk.as_move_value(); - assert_eq!( - vec![ - 4, 107, 105, 100, 49, 3, 82, 83, 65, 5, 82, 83, 50, 53, 54, 4, 65, 81, 65, 66, 5, 49, - 51, 49, 51, 49 - ], - move_value.simple_serialize().unwrap() - ); -} - -#[test] -fn test_rsa_jwk_as_move_any() { - let rsa_jwk = RSA_JWK::new_for_testing("kid1", "RSA", "RS256", "AQAB", "1313131313131"); - let actual = rsa_jwk.as_move_any(); - let expected = MoveAny { - type_name: "0x1::jwks::RSA_JWK".to_string(), - data: bcs::to_bytes(&rsa_jwk).unwrap(), - }; - assert_eq!(expected, actual); -} +#[cfg(test)] +mod tests; diff --git a/types/src/jwks/rsa/tests.rs b/types/src/jwks/rsa/tests.rs new file mode 100644 index 0000000000000..38b4958110e1a --- /dev/null +++ b/types/src/jwks/rsa/tests.rs @@ -0,0 +1,98 @@ +// Copyright © Aptos Foundation + +use crate::{ + jwks::rsa::RSA_JWK, + move_any::{Any as MoveAny, AsMoveAny}, + move_utils::as_move_value::AsMoveValue, +}; +use std::str::FromStr; + +#[test] +fn convert_json_to_rsa_jwk() { + // Valid JWK JSON should be accepted. + let json_str = + r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + let actual = RSA_JWK::try_from(&json); + let expected = RSA_JWK::new_from_strs("kid1", "RSA", "RS256", "AQAB", "13131"); + assert_eq!(expected, actual.unwrap()); + + // JWK JSON without `kid` should be rejected. + let json_str = r#"{"alg": "RS256", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON with wrong `kid` type should be rejected. + let json_str = + r#"{"alg": "RS256", "kid": {}, "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON without `alg` should be rejected. + let json_str = r#"{"kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON with wrong `alg` type should be rejected. + let json_str = + r#"{"alg": 0, "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON without `kty` should be rejected. + let json_str = r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON with wrong `kty` value should be rejected. + let json_str = + r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSB", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON without `e` should be rejected. + let json_str = r#"{"alg": "RS256", "kid": "kid1", "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON with wrong `e` type should be rejected. + let json_str = + r#"{"alg": "RS256", "kid": "kid1", "e": 65537, "use": "sig", "kty": "RSA", "n": "13131"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON without `n` should be rejected. + let json_str = r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA"}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); + + // JWK JSON with wrong `n` type should be rejected. + let json_str = + r#"{"alg": "RS256", "kid": "kid1", "e": "AQAB", "use": "sig", "kty": "RSA", "n": false}"#; + let json = serde_json::Value::from_str(json_str).unwrap(); + assert!(RSA_JWK::try_from(&json).is_err()); +} + +#[test] +fn rsa_jwk_as_move_value() { + let rsa_jwk = RSA_JWK::new_from_strs("kid1", "RSA", "RS256", "AQAB", "13131"); + let move_value = rsa_jwk.as_move_value(); + assert_eq!( + vec![ + 4, 107, 105, 100, 49, 3, 82, 83, 65, 5, 82, 83, 50, 53, 54, 4, 65, 81, 65, 66, 5, 49, + 51, 49, 51, 49 + ], + move_value.simple_serialize().unwrap() + ); +} + +#[test] +fn rsa_jwk_as_move_any() { + let rsa_jwk = RSA_JWK::new_from_strs("kid1", "RSA", "RS256", "AQAB", "1313131313131"); + let actual = rsa_jwk.as_move_any(); + let expected = MoveAny { + type_name: "0x1::jwks::RSA_JWK".to_string(), + data: bcs::to_bytes(&rsa_jwk).unwrap(), + }; + assert_eq!(expected, actual); +} diff --git a/types/src/jwks/unsupported.rs b/types/src/jwks/unsupported.rs deleted file mode 100644 index 989d79e913217..0000000000000 --- a/types/src/jwks/unsupported.rs +++ /dev/null @@ -1,89 +0,0 @@ -// Copyright © Aptos Foundation - -#[cfg(test)] -use crate::move_any::Any as MoveAny; -use crate::{move_any::AsMoveAny, move_utils::as_move_value::AsMoveValue}; -use aptos_crypto::HashValue; -use move_core_types::value::{MoveStruct, MoveValue}; -use serde::{Deserialize, Serialize}; -#[cfg(test)] -use std::str::FromStr; - -/// Move type `0x1::jwks::UnsupportedJWK` in rust. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct UnsupportedJWK { - pub id: Vec, - pub payload: Vec, -} - -impl UnsupportedJWK { - #[cfg(any(test, feature = "fuzzing"))] - pub fn new_for_testing(id: &str, payload: &str) -> Self { - Self { - id: id.as_bytes().to_vec(), - payload: payload.as_bytes().to_vec(), - } - } -} - -impl TryFrom<&serde_json::Value> for UnsupportedJWK { - type Error = anyhow::Error; - - fn try_from(json_value: &serde_json::Value) -> Result { - let payload = json_value.to_string().into_bytes(); //TODO: canonical to_string. - let ret = Self { - id: HashValue::sha3_256_of(payload.as_slice()).to_vec(), - payload, - }; - Ok(ret) - } -} - -impl AsMoveValue for UnsupportedJWK { - fn as_move_value(&self) -> MoveValue { - MoveValue::Struct(MoveStruct::Runtime(vec![ - self.id.as_move_value(), - self.payload.as_move_value(), - ])) - } -} - -impl AsMoveAny for UnsupportedJWK { - const MOVE_TYPE_NAME: &'static str = "0x1::jwks::UnsupportedJWK"; -} - -#[test] -fn test_unsupported_jwk_from_json() { - // Some unknown JWK format - let compact_json_str = "{\"key0\":\"val0\",\"key1\":999}"; - let expected_payload = compact_json_str.as_bytes().to_vec(); - let expected_id = HashValue::sha3_256_of(expected_payload.as_slice()).to_vec(); - let json = serde_json::Value::from_str(compact_json_str).unwrap(); - let actual = UnsupportedJWK::try_from(&json).unwrap(); - let expected = UnsupportedJWK { - id: expected_id, - payload: expected_payload, - }; - assert_eq!(expected, actual); -} - -#[test] -fn test_unsupported_jwk_as_move_value() { - let unsupported_jwk = UnsupportedJWK::new_for_testing("AAA", "BBBB"); - let move_value = unsupported_jwk.as_move_value(); - assert_eq!( - vec![3, 65, 65, 65, 4, 66, 66, 66, 66], - move_value.simple_serialize().unwrap() - ); -} - -#[test] -fn test_unsupported_jwk_as_move_any() { - let unsupported_jwk = UnsupportedJWK::new_for_testing("AAA", "BBBB"); - let actual = unsupported_jwk.as_move_any(); - let expected = MoveAny { - type_name: "0x1::jwks::UnsupportedJWK".to_string(), - data: bcs::to_bytes(&unsupported_jwk).unwrap(), - }; - assert_eq!(expected, actual); -} diff --git a/types/src/jwks/unsupported/mod.rs b/types/src/jwks/unsupported/mod.rs new file mode 100644 index 0000000000000..68cafe2cd867f --- /dev/null +++ b/types/src/jwks/unsupported/mod.rs @@ -0,0 +1,73 @@ +// Copyright © Aptos Foundation + +use crate::{move_any::AsMoveAny, move_utils::as_move_value::AsMoveValue}; +use aptos_crypto::HashValue; +use move_core_types::value::{MoveStruct, MoveValue}; +use serde::{Deserialize, Serialize}; +use std::fmt::{Debug, Formatter}; + +/// Move type `0x1::jwks::UnsupportedJWK` in rust. +/// See its doc in Move for more details. +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct UnsupportedJWK { + pub id: Vec, + pub payload: Vec, +} + +impl Debug for UnsupportedJWK { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UnsupportedJWK") + .field("id", &hex::encode(self.id.as_slice())) + .field("payload", &String::from_utf8(self.payload.clone())) + .finish() + } +} + +impl UnsupportedJWK { + #[cfg(any(test, feature = "fuzzing"))] + pub fn new_for_testing(id: &str, payload: &str) -> Self { + Self { + id: id.as_bytes().to_vec(), + payload: payload.as_bytes().to_vec(), + } + } + + #[cfg(any(test, feature = "fuzzing"))] + pub fn new_with_payload(payload: &str) -> Self { + let id = HashValue::sha3_256_of(payload.as_bytes()).to_vec(); + Self { + id, + payload: payload.as_bytes().to_vec(), + } + } + + pub fn id(&self) -> Vec { + self.id.clone() + } +} + +impl From for UnsupportedJWK { + fn from(json_value: serde_json::Value) -> Self { + let payload = json_value.to_string().into_bytes(); //TODO: canonical to_string. + Self { + id: HashValue::sha3_256_of(payload.as_slice()).to_vec(), + payload, + } + } +} + +impl AsMoveValue for UnsupportedJWK { + fn as_move_value(&self) -> MoveValue { + MoveValue::Struct(MoveStruct::Runtime(vec![ + self.id.as_move_value(), + self.payload.as_move_value(), + ])) + } +} + +impl AsMoveAny for UnsupportedJWK { + const MOVE_TYPE_NAME: &'static str = "0x1::jwks::UnsupportedJWK"; +} + +#[cfg(test)] +mod tests; diff --git a/types/src/jwks/unsupported/tests.rs b/types/src/jwks/unsupported/tests.rs new file mode 100644 index 0000000000000..28810b028c44e --- /dev/null +++ b/types/src/jwks/unsupported/tests.rs @@ -0,0 +1,45 @@ +// Copyright © Aptos Foundation + +use crate::{ + jwks::unsupported::UnsupportedJWK, + move_any::{Any as MoveAny, AsMoveAny}, + move_utils::as_move_value::AsMoveValue, +}; +use aptos_crypto::HashValue; +use std::str::FromStr; + +#[test] +fn convert_json_to_unsupported_jwk() { + // Some unknown JWK format + let compact_json_str = "{\"key0\":\"val0\",\"key1\":999}"; + let expected_payload = compact_json_str.as_bytes().to_vec(); + let expected_id = HashValue::sha3_256_of(expected_payload.as_slice()).to_vec(); + let json = serde_json::Value::from_str(compact_json_str).unwrap(); + let actual = UnsupportedJWK::from(json); + let expected = UnsupportedJWK { + id: expected_id, + payload: expected_payload, + }; + assert_eq!(expected, actual); +} + +#[test] +fn unsupported_jwk_as_move_value() { + let unsupported_jwk = UnsupportedJWK::new_for_testing("AAA", "BBBB"); + let move_value = unsupported_jwk.as_move_value(); + assert_eq!( + vec![3, 65, 65, 65, 4, 66, 66, 66, 66], + move_value.simple_serialize().unwrap() + ); +} + +#[test] +fn unsupported_jwk_as_move_any() { + let unsupported_jwk = UnsupportedJWK::new_for_testing("AAA", "BBBB"); + let actual = unsupported_jwk.as_move_any(); + let expected = MoveAny { + type_name: "0x1::jwks::UnsupportedJWK".to_string(), + data: bcs::to_bytes(&unsupported_jwk).unwrap(), + }; + assert_eq!(expected, actual); +} diff --git a/types/src/on_chain_config/aptos_features.rs b/types/src/on_chain_config/aptos_features.rs index a72027b926c04..657126ba35e85 100644 --- a/types/src/on_chain_config/aptos_features.rs +++ b/types/src/on_chain_config/aptos_features.rs @@ -56,6 +56,7 @@ pub enum FeatureFlag { ZK_ID_SIGNATURES = 46, ZK_ID_ZKLESS_SIGNATURE = 47, REMOVE_DETAILED_ERROR_FROM_HASH = 48, + JWK_CONSENSUS = 49, } /// Representation of features on chain as a bitset. @@ -90,16 +91,25 @@ impl OnChainConfig for Features { } impl Features { - pub fn enable(&mut self, flag: FeatureFlag) { + fn resize_for_flag(&mut self, flag: FeatureFlag) -> (usize, u8) { let byte_index = (flag as u64 / 8) as usize; let bit_mask = 1 << (flag as u64 % 8); while self.features.len() <= byte_index { self.features.push(0); } + (byte_index, bit_mask) + } + pub fn enable(&mut self, flag: FeatureFlag) { + let (byte_index, bit_mask) = self.resize_for_flag(flag); self.features[byte_index] |= bit_mask; } + pub fn disable(&mut self, flag: FeatureFlag) { + let (byte_index, bit_mask) = self.resize_for_flag(flag); + self.features[byte_index] &= !bit_mask; + } + pub fn is_enabled(&self, flag: FeatureFlag) -> bool { let val = flag as u64; let byte_index = (val / 8) as usize; From 777450d9ac69f31d732eb442394a595cd481e8ec Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 01:50:38 -0800 Subject: [PATCH 02/22] update --- types/src/move_any.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/types/src/move_any.rs b/types/src/move_any.rs index fb160a3c70ab0..ba425f8ec8586 100644 --- a/types/src/move_any.rs +++ b/types/src/move_any.rs @@ -1,6 +1,8 @@ // Copyright © Aptos Foundation +use crate::move_utils::as_move_value::AsMoveValue; use anyhow::bail; +use move_core_types::value::{MoveStruct, MoveValue}; use serde::{de::DeserializeOwned, Deserialize, Serialize}; /// Rust representation of the Move Any type @@ -30,6 +32,15 @@ impl Any { } } +impl AsMoveValue for Any { + fn as_move_value(&self) -> MoveValue { + MoveValue::Struct(MoveStruct::Runtime(vec![ + self.type_name.as_move_value(), + self.data.as_move_value(), + ])) + } +} + pub trait AsMoveAny: Serialize { const MOVE_TYPE_NAME: &'static str; From ef680a95ead81199236bd052a3373d79a7d872c3 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 01:54:54 -0800 Subject: [PATCH 03/22] update --- .../aptos-release-builder/src/components/feature_flags.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aptos-move/aptos-release-builder/src/components/feature_flags.rs b/aptos-move/aptos-release-builder/src/components/feature_flags.rs index 6f001fc75e15d..0e972c6e93ee4 100644 --- a/aptos-move/aptos-release-builder/src/components/feature_flags.rs +++ b/aptos-move/aptos-release-builder/src/components/feature_flags.rs @@ -99,6 +99,7 @@ pub enum FeatureFlag { ZkIdSignature, ZkIdZkLessSignature, RemoveDetailedError, + JWKConsensus, } fn generate_features_blob(writer: &CodeWriter, data: &[u64]) { @@ -256,6 +257,7 @@ impl From for AptosFeatureFlag { 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 @@ impl From for FeatureFlag { 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, } } } From d0957db3afd2ad68a6a004fc140cc5361fad1853 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 02:04:23 -0800 Subject: [PATCH 04/22] jwk txn and execution --- aptos-move/aptos-vm/Cargo.toml | 1 + .../aptos-vm/src/system_module_names.rs | 9 + aptos-move/aptos-vm/src/validator_txns/jwk.rs | 161 ++++++++++++++++++ aptos-move/aptos-vm/src/validator_txns/mod.rs | 4 + .../framework/aptos-framework/doc/jwks.md | 45 +++++ .../aptos-framework/sources/jwks.move | 6 + .../framework/move-stdlib/doc/features.md | 60 +++++++ .../move-stdlib/sources/configs/features.move | 10 ++ aptos-move/vm-genesis/src/lib.rs | 1 + types/src/validator_txn.rs | 15 +- 10 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 aptos-move/aptos-vm/src/validator_txns/jwk.rs diff --git a/aptos-move/aptos-vm/Cargo.toml b/aptos-move/aptos-vm/Cargo.toml index 812a312148e4a..17362472fd3c5 100644 --- a/aptos-move/aptos-vm/Cargo.toml +++ b/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 } diff --git a/aptos-move/aptos-vm/src/system_module_names.rs b/aptos-move/aptos-vm/src/system_module_names.rs index 2a85cde828572..e6945a6114c20 100644 --- a/aptos-move/aptos-vm/src/system_module_names.rs +++ b/aptos-move/aptos-vm/src/system_module_names.rs @@ -39,6 +39,15 @@ pub static RECONFIGURATION_WITH_DKG_MODULE: Lazy = Lazy::new(|| { pub const FINISH_WITH_DKG_RESULT: &IdentStr = ident_str!("finish_with_dkg_result"); +pub static JWKS_MODULE: Lazy = 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 = Lazy::new(|| { ModuleId::new( account_config::CORE_CODE_ADDRESS, diff --git a/aptos-move/aptos-vm/src/validator_txns/jwk.rs b/aptos-move/aptos-vm/src/validator_txns/jwk.rs new file mode 100644 index 0000000000000..cdd9088d4e2dd --- /dev/null +++ b/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 = + 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::>(), + ); + + // 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)) + } +} diff --git a/aptos-move/aptos-vm/src/validator_txns/mod.rs b/aptos-move/aptos-vm/src/validator_txns/mod.rs index 13de1fda5f9e2..eea258ffcf08f 100644 --- a/aptos-move/aptos-vm/src/validator_txns/mod.rs +++ b/aptos-move/aptos-vm/src/validator_txns/mod.rs @@ -21,6 +21,9 @@ impl AptosVM { 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 @@ impl AptosVM { mod dkg; mod dummy; +mod jwk; diff --git a/aptos-move/framework/aptos-framework/doc/jwks.md b/aptos-move/framework/aptos-framework/doc/jwks.md index 778cf010deb63..2166154629728 100644 --- a/aptos-move/framework/aptos-framework/doc/jwks.md +++ b/aptos-move/framework/aptos-framework/doc/jwks.md @@ -620,6 +620,51 @@ This is what applications should consume. + + + + +
const ENATIVE_INCORRECT_VERSION: u64 = 259;
+
+ + + + + + + +
const ENATIVE_MISSING_RESOURCE_OBSERVED_JWKS: u64 = 258;
+
+ + + + + + + +
const ENATIVE_MISSING_RESOURCE_VALIDATOR_SET: u64 = 257;
+
+ + + + + + + +
const ENATIVE_MULTISIG_VERIFICATION_FAILED: u64 = 260;
+
+ + + + + + + +
const ENATIVE_NOT_ENOUGH_VOTING_POWER: u64 = 261;
+
+ + + diff --git a/aptos-move/framework/aptos-framework/sources/jwks.move b/aptos-move/framework/aptos-framework/sources/jwks.move index f405906207288..b87a71c85f4b5 100644 --- a/aptos-move/framework/aptos-framework/sources/jwks.move +++ b/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". diff --git a/aptos-move/framework/move-stdlib/doc/features.md b/aptos-move/framework/move-stdlib/doc/features.md index a56f3f691e13c..87f2ec8117eda 100644 --- a/aptos-move/framework/move-stdlib/doc/features.md +++ b/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 + + +The JWK consensus feature. + +Lifetime: permanent + + +
const JWK_CONSENSUS: u64 = 49;
+
+ + + @@ -1918,6 +1932,52 @@ Lifetime: transient + + + + +## Function `get_jwk_consensus_feature` + + + +
public fun get_jwk_consensus_feature(): u64
+
+ + + +
+Implementation + + +
public fun get_jwk_consensus_feature(): u64 { JWK_CONSENSUS }
+
+ + + +
+ + + +## Function `jwk_consensus_enabled` + + + +
public fun jwk_consensus_enabled(): bool
+
+ + + +
+Implementation + + +
public fun jwk_consensus_enabled(): bool acquires Features {
+    is_enabled(JWK_CONSENSUS)
+}
+
+ + +
diff --git a/aptos-move/framework/move-stdlib/sources/configs/features.move b/aptos-move/framework/move-stdlib/sources/configs/features.move index 8db36f003ee42..53189fc2408ed 100644 --- a/aptos-move/framework/move-stdlib/sources/configs/features.move +++ b/aptos-move/framework/move-stdlib/sources/configs/features.move @@ -355,6 +355,16 @@ module std::features { is_enabled(ZK_ID_ZKLESS_SIGNATURE) } + /// The JWK consensus feature. + /// + /// Lifetime: permanent + const JWK_CONSENSUS: u64 = 49; + + public fun get_jwk_consensus_feature(): u64 { JWK_CONSENSUS } + + public fun jwk_consensus_enabled(): bool acquires Features { + is_enabled(JWK_CONSENSUS) + } // ============================================================================================ // Feature Flag Implementation diff --git a/aptos-move/vm-genesis/src/lib.rs b/aptos-move/vm-genesis/src/lib.rs index 17788f6d64cae..de1792bacbd59 100644 --- a/aptos-move/vm-genesis/src/lib.rs +++ b/aptos-move/vm-genesis/src/lib.rs @@ -446,6 +446,7 @@ pub fn default_features() -> Vec { // FeatureFlag::RECONFIGURE_WITH_DKG, //TODO: re-enable once randomness is ready. FeatureFlag::ZK_ID_SIGNATURES, FeatureFlag::ZK_ID_ZKLESS_SIGNATURE, + FeatureFlag::JWK_CONSENSUS, ] } diff --git a/types/src/validator_txn.rs b/types/src/validator_txn.rs index 7960ae0da7867..a037930b7df25 100644 --- a/types/src/validator_txn.rs +++ b/types/src/validator_txn.rs @@ -11,6 +11,7 @@ pub enum ValidatorTransaction { DummyTopic1(DummyValidatorTransaction), DKGResult(DKGTranscript), DummyTopic2(DummyValidatorTransaction), + ObservedJWKUpdate(jwks::QuorumCertifiedUpdate), } #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, CryptoHasher, BCSCryptoHash)] @@ -40,14 +41,24 @@ impl ValidatorTransaction { pub fn size_in_bytes(&self) -> usize { bcs::serialized_size(self).unwrap() } + + pub fn topic(&self) -> Topic { + match self { + ValidatorTransaction::DummyTopic1(_) => Topic::DUMMY1, + ValidatorTransaction::DKGResult(_) => Topic::DKG, + ValidatorTransaction::DummyTopic2(_) => Topic::DUMMY2, + ValidatorTransaction::ObservedJWKUpdate(update) => { + Topic::JWK_CONSENSUS(update.update.issuer.clone()) + }, + } + } } -#[derive(Clone, Eq, Hash, PartialEq)] +#[derive(Clone, Debug, Eq, Hash, PartialEq)] #[allow(non_camel_case_types)] pub enum Topic { DKG, JWK_CONSENSUS(jwks::Issuer), DUMMY1, - #[cfg(any(test, feature = "fuzzing"))] DUMMY2, } From 3a467336891fcbea33ff2042d582b2d428f293d6 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 02:21:15 -0800 Subject: [PATCH 05/22] consensus ensure jwk txns are expected --- consensus/safety-rules/Cargo.toml | 1 + consensus/safety-rules/src/lib.rs | 2 +- .../safety-rules/src/safety_rules_manager.rs | 14 +++ consensus/src/dag/bootstrap.rs | 7 +- consensus/src/dag/rb_handler.rs | 42 ++++++--- consensus/src/dag/tests/rb_handler_tests.rs | 10 +- consensus/src/epoch_manager.rs | 16 +++- consensus/src/round_manager.rs | 16 +++- consensus/src/round_manager_fuzzing.rs | 3 +- consensus/src/round_manager_test.rs | 93 ++++++++++++++++++- consensus/src/util/mod.rs | 17 ++++ 11 files changed, 200 insertions(+), 21 deletions(-) diff --git a/consensus/safety-rules/Cargo.toml b/consensus/safety-rules/Cargo.toml index 0b85516acacf8..cb81702a69fd7 100644 --- a/consensus/safety-rules/Cargo.toml +++ b/consensus/safety-rules/Cargo.toml @@ -13,6 +13,7 @@ repository = { workspace = true } rust-version = { workspace = true } [dependencies] +anyhow = { workspace = true } aptos-config = { workspace = true } aptos-consensus-types = { workspace = true } aptos-crypto = { workspace = true } diff --git a/consensus/safety-rules/src/lib.rs b/consensus/safety-rules/src/lib.rs index 0cd3f245fd411..4b66b4c489a82 100644 --- a/consensus/safety-rules/src/lib.rs +++ b/consensus/safety-rules/src/lib.rs @@ -14,7 +14,7 @@ mod process; mod remote_service; mod safety_rules; mod safety_rules_2chain; -mod safety_rules_manager; +pub mod safety_rules_manager; mod serializer; mod t_safety_rules; mod thread; diff --git a/consensus/safety-rules/src/safety_rules_manager.rs b/consensus/safety-rules/src/safety_rules_manager.rs index 4537cdfb2d550..c73664be47344 100644 --- a/consensus/safety-rules/src/safety_rules_manager.rs +++ b/consensus/safety-rules/src/safety_rules_manager.rs @@ -11,7 +11,10 @@ use crate::{ thread::ThreadService, SafetyRules, TSafetyRules, }; +use anyhow::anyhow; use aptos_config::config::{InitialSafetyRulesConfig, SafetyRulesConfig, SafetyRulesService}; +use aptos_crypto::bls12381::PrivateKey; +use aptos_global_constants::CONSENSUS_KEY; use aptos_infallible::RwLock; use aptos_secure_storage::{KVStorage, Storage}; use std::{net::SocketAddr, sync::Arc}; @@ -73,6 +76,17 @@ pub fn storage(config: &SafetyRulesConfig) -> PersistentSafetyStorage { } } +pub fn load_consensus_key_from_secure_storage( + config: &SafetyRulesConfig, +) -> anyhow::Result { + let storage: Storage = (&config.backend).into(); + let storage = Box::new(storage); + let response = storage.get::(CONSENSUS_KEY).map_err(|e| { + anyhow!("load_consensus_key_from_secure_storage failed with storage read error: {e}") + })?; + Ok(response.value) +} + enum SafetyRulesWrapper { Local(Arc>), Process(ProcessService), diff --git a/consensus/src/dag/bootstrap.rs b/consensus/src/dag/bootstrap.rs index a84175f215f9a..4e98dd70c801c 100644 --- a/consensus/src/dag/bootstrap.rs +++ b/consensus/src/dag/bootstrap.rs @@ -46,7 +46,7 @@ use aptos_logger::{debug, info}; use aptos_reliable_broadcast::{RBNetworkSender, ReliableBroadcast}; use aptos_types::{ epoch_state::EpochState, - on_chain_config::{DagConsensusConfigV1, ValidatorTxnConfig}, + on_chain_config::{DagConsensusConfigV1, Features, ValidatorTxnConfig}, validator_signer::ValidatorSigner, }; use async_trait::async_trait; @@ -330,6 +330,7 @@ pub struct DagBootstrapper { quorum_store_enabled: bool, vtxn_config: ValidatorTxnConfig, executor: BoundedExecutor, + features: Features, } impl DagBootstrapper { @@ -352,6 +353,7 @@ impl DagBootstrapper { quorum_store_enabled: bool, vtxn_config: ValidatorTxnConfig, executor: BoundedExecutor, + features: Features, ) -> Self { Self { self_peer, @@ -371,6 +373,7 @@ impl DagBootstrapper { quorum_store_enabled, vtxn_config, executor, + features, } } @@ -557,6 +560,7 @@ impl DagBootstrapper { fetch_requester, self.config.node_payload_config.clone(), self.vtxn_config.clone(), + self.features.clone(), ); let fetch_handler = FetchRequestHandler::new(dag_store.clone(), self.epoch_state.clone()); @@ -665,6 +669,7 @@ pub(super) fn bootstrap_dag_for_test( false, ValidatorTxnConfig::default_enabled(), BoundedExecutor::new(2, Handle::current()), + Features::default(), ); let (_base_state, handler, fetch_service) = bootstraper.full_bootstrap(); diff --git a/consensus/src/dag/rb_handler.rs b/consensus/src/dag/rb_handler.rs index 25e2d4a35496b..46cfeb76ed3fb 100644 --- a/consensus/src/dag/rb_handler.rs +++ b/consensus/src/dag/rb_handler.rs @@ -1,18 +1,21 @@ // Copyright © Aptos Foundation // SPDX-License-Identifier: Apache-2.0 -use crate::dag::{ - dag_fetcher::TFetchRequester, - dag_network::RpcHandler, - dag_store::Dag, - errors::NodeBroadcastHandleError, - observability::{ - logging::{LogEvent, LogSchema}, - tracing::{observe_node, NodeStage}, +use crate::{ + dag::{ + dag_fetcher::TFetchRequester, + dag_network::RpcHandler, + dag_store::Dag, + errors::NodeBroadcastHandleError, + observability::{ + logging::{LogEvent, LogSchema}, + tracing::{observe_node, NodeStage}, + }, + storage::DAGStorage, + types::{Node, NodeCertificate, Vote}, + NodeId, }, - storage::DAGStorage, - types::{Node, NodeCertificate, Vote}, - NodeId, + util::is_vtxn_expected, }; use anyhow::{bail, ensure}; use aptos_config::config::DagPayloadConfig; @@ -20,8 +23,10 @@ use aptos_consensus_types::common::{Author, Round}; use aptos_infallible::RwLock; use aptos_logger::{debug, error}; use aptos_types::{ - epoch_state::EpochState, on_chain_config::ValidatorTxnConfig, - validator_signer::ValidatorSigner, validator_txn::ValidatorTransaction, + epoch_state::EpochState, + on_chain_config::{Features, ValidatorTxnConfig}, + validator_signer::ValidatorSigner, + validator_txn::ValidatorTransaction, }; use async_trait::async_trait; use std::{collections::BTreeMap, mem, sync::Arc}; @@ -35,6 +40,7 @@ pub(crate) struct NodeBroadcastHandler { fetch_requester: Arc, payload_config: DagPayloadConfig, vtxn_config: ValidatorTxnConfig, + features: Features, } impl NodeBroadcastHandler { @@ -46,6 +52,7 @@ impl NodeBroadcastHandler { fetch_requester: Arc, payload_config: DagPayloadConfig, vtxn_config: ValidatorTxnConfig, + features: Features, ) -> Self { let epoch = epoch_state.epoch; let votes_by_round_peer = read_votes_from_storage(&storage, epoch); @@ -59,6 +66,7 @@ impl NodeBroadcastHandler { fetch_requester, payload_config, vtxn_config, + features, } } @@ -87,6 +95,14 @@ impl NodeBroadcastHandler { fn validate(&self, node: Node) -> anyhow::Result { let num_vtxns = node.validator_txns().len() as u64; ensure!(num_vtxns <= self.vtxn_config.per_block_limit_txn_count()); + for vtxn in node.validator_txns() { + ensure!( + is_vtxn_expected(&self.features, vtxn), + "unexpected validator transaction: {:?}", + vtxn.topic() + ); + } + let vtxn_total_bytes = node .validator_txns() .iter() diff --git a/consensus/src/dag/tests/rb_handler_tests.rs b/consensus/src/dag/tests/rb_handler_tests.rs index d19b3e29fa4de..47e47d5ab421f 100644 --- a/consensus/src/dag/tests/rb_handler_tests.rs +++ b/consensus/src/dag/tests/rb_handler_tests.rs @@ -17,8 +17,10 @@ use crate::dag::{ use aptos_config::config::DagPayloadConfig; use aptos_infallible::RwLock; use aptos_types::{ - aggregate_signature::PartialSignatures, epoch_state::EpochState, - on_chain_config::ValidatorTxnConfig, validator_verifier::random_validator_verifier, + aggregate_signature::PartialSignatures, + epoch_state::EpochState, + on_chain_config::{Features, ValidatorTxnConfig}, + validator_verifier::random_validator_verifier, }; use claims::{assert_ok, assert_ok_eq}; use futures::executor::block_on; @@ -68,6 +70,7 @@ async fn test_node_broadcast_receiver_succeed() { Arc::new(MockFetchRequester {}), DagPayloadConfig::default(), ValidatorTxnConfig::default_disabled(), + Features::default(), ); let expected_result = Vote::new( @@ -114,6 +117,7 @@ async fn test_node_broadcast_receiver_failure() { Arc::new(MockFetchRequester {}), DagPayloadConfig::default(), ValidatorTxnConfig::default_disabled(), + Features::default(), ) }) .collect(); @@ -197,6 +201,7 @@ async fn test_node_broadcast_receiver_storage() { Arc::new(MockFetchRequester {}), DagPayloadConfig::default(), ValidatorTxnConfig::default_disabled(), + Features::default(), ); let sig = rb_receiver.process(node).await.expect("must succeed"); @@ -213,6 +218,7 @@ async fn test_node_broadcast_receiver_storage() { Arc::new(MockFetchRequester {}), DagPayloadConfig::default(), ValidatorTxnConfig::default_disabled(), + Features::default(), ); assert_ok!(rb_receiver.gc_before_round(2)); assert_eq!(storage.get_votes().unwrap().len(), 0); diff --git a/consensus/src/epoch_manager.rs b/consensus/src/epoch_manager.rs index 29ee8b844e830..6ffd58555c4f4 100644 --- a/consensus/src/epoch_manager.rs +++ b/consensus/src/epoch_manager.rs @@ -823,6 +823,7 @@ impl EpochManager

{ network_sender: NetworkSender, payload_client: Arc, payload_manager: Arc, + features: Features, ) { let epoch = epoch_state.epoch; info!( @@ -930,6 +931,7 @@ impl EpochManager

{ onchain_consensus_config, buffered_proposal_tx, self.config.clone(), + features, ); round_manager.init(last_vote).await; @@ -975,7 +977,7 @@ impl EpochManager

{ let onchain_consensus_config: anyhow::Result = payload.get(); let onchain_execution_config: anyhow::Result = payload.get(); - let features = payload.get::().ok().unwrap_or_default(); + let features = payload.get::(); if let Err(error) = &onchain_consensus_config { error!("Failed to read on-chain consensus config {}", error); @@ -985,11 +987,17 @@ impl EpochManager

{ error!("Failed to read on-chain execution config {}", error); } + if let Err(error) = &features { + error!("Failed to read on-chain features {}", error); + } + self.epoch_state = Some(epoch_state.clone()); let consensus_config = onchain_consensus_config.unwrap_or_default(); let execution_config = onchain_execution_config .unwrap_or_else(|_| OnChainExecutionConfig::default_if_missing()); + let features = features.unwrap_or_default(); + let (network_sender, payload_client, payload_manager) = self .initialize_shared_component( &epoch_state, @@ -1006,6 +1014,7 @@ impl EpochManager

{ network_sender, payload_client, payload_manager, + features, ) .await } else { @@ -1015,6 +1024,7 @@ impl EpochManager

{ network_sender, payload_client, payload_manager, + features, ) .await } @@ -1061,6 +1071,7 @@ impl EpochManager

{ network_sender: NetworkSender, payload_client: Arc, payload_manager: Arc, + features: Features, ) { match self.storage.start() { LivenessStorageData::FullRecoveryData(initial_data) => { @@ -1072,6 +1083,7 @@ impl EpochManager

{ network_sender, payload_client, payload_manager, + features, ) .await }, @@ -1095,6 +1107,7 @@ impl EpochManager

{ network_sender: NetworkSender, payload_client: Arc, payload_manager: Arc, + features: Features, ) { let epoch = epoch_state.epoch; @@ -1147,6 +1160,7 @@ impl EpochManager

{ onchain_consensus_config.quorum_store_enabled(), onchain_consensus_config.effective_validator_txn_config(), self.bounded_executor.clone(), + features, ); let (dag_rpc_tx, dag_rpc_rx) = aptos_channel::new(QueueStyle::FIFO, 10, None); diff --git a/consensus/src/round_manager.rs b/consensus/src/round_manager.rs index 6c83aad53ad8e..b0880f377e0ab 100644 --- a/consensus/src/round_manager.rs +++ b/consensus/src/round_manager.rs @@ -23,6 +23,7 @@ use crate::{ pending_votes::VoteReceptionResult, persistent_liveness_storage::PersistentLivenessStorage, quorum_store::types::BatchMsg, + util::is_vtxn_expected, }; use anyhow::{bail, ensure, Context}; use aptos_channels::aptos_channel; @@ -47,7 +48,7 @@ use aptos_safety_rules::ConsensusState; use aptos_safety_rules::TSafetyRules; use aptos_types::{ epoch_state::EpochState, - on_chain_config::{OnChainConsensusConfig, ValidatorTxnConfig}, + on_chain_config::{Features, OnChainConsensusConfig, ValidatorTxnConfig}, validator_verifier::ValidatorVerifier, PeerId, }; @@ -187,6 +188,7 @@ pub struct RoundManager { vtxn_config: ValidatorTxnConfig, buffered_proposal_tx: aptos_channel::Sender, local_config: ConsensusConfig, + features: Features, } impl RoundManager { @@ -202,6 +204,7 @@ impl RoundManager { onchain_config: OnChainConsensusConfig, buffered_proposal_tx: aptos_channel::Sender, local_config: ConsensusConfig, + features: Features, ) -> Self { // when decoupled execution is false, // the counter is still static. @@ -226,6 +229,7 @@ impl RoundManager { vtxn_config, buffered_proposal_tx, local_config, + features, } } @@ -652,6 +656,16 @@ impl RoundManager { bail!("ProposalExt unexpected while the feature is disabled."); } + if let Some(vtxns) = proposal.validator_txns() { + for vtxn in vtxns { + ensure!( + is_vtxn_expected(&self.features, vtxn), + "unexpected validator txn: {:?}", + vtxn.topic() + ); + } + } + let (num_validator_txns, validator_txns_total_bytes): (usize, usize) = proposal.validator_txns().map_or((0, 0), |txns| { txns.iter().fold((0, 0), |(count_acc, size_acc), txn| { diff --git a/consensus/src/round_manager_fuzzing.rs b/consensus/src/round_manager_fuzzing.rs index 2c720b4439acf..0d8507ec5238c 100644 --- a/consensus/src/round_manager_fuzzing.rs +++ b/consensus/src/round_manager_fuzzing.rs @@ -38,7 +38,7 @@ use aptos_types::{ epoch_change::EpochChangeProof, epoch_state::EpochState, ledger_info::{LedgerInfo, LedgerInfoWithSignatures}, - on_chain_config::{OnChainConsensusConfig, ValidatorSet, ValidatorTxnConfig}, + on_chain_config::{Features, OnChainConsensusConfig, ValidatorSet, ValidatorTxnConfig}, validator_info::ValidatorInfo, validator_signer::ValidatorSigner, validator_verifier::ValidatorVerifier, @@ -209,6 +209,7 @@ fn create_node_for_fuzzing() -> RoundManager { OnChainConsensusConfig::default(), round_manager_tx, ConsensusConfig::default(), + Features::default(), ) } diff --git a/consensus/src/round_manager_test.rs b/consensus/src/round_manager_test.rs index 8efe90991bedf..cbca9989ca93c 100644 --- a/consensus/src/round_manager_test.rs +++ b/consensus/src/round_manager_test.rs @@ -62,9 +62,11 @@ use aptos_safety_rules::{PersistentSafetyStorage, SafetyRulesManager}; use aptos_secure_storage::Storage; use aptos_types::{ epoch_state::EpochState, + jwks::QuorumCertifiedUpdate, ledger_info::LedgerInfo, on_chain_config::{ - ConsensusAlgorithmConfig, ConsensusConfigV1, OnChainConsensusConfig, ValidatorTxnConfig, + ConsensusAlgorithmConfig, ConsensusConfigV1, FeatureFlag, Features, OnChainConsensusConfig, + ValidatorTxnConfig, }, transaction::SignedTransaction, validator_signer::ValidatorSigner, @@ -110,6 +112,7 @@ pub struct NodeSetup { id: usize, onchain_consensus_config: OnChainConsensusConfig, local_consensus_config: ConsensusConfig, + features: Features, } impl NodeSetup { @@ -138,6 +141,7 @@ impl NodeSetup { proposer_indices: Option>, onchain_consensus_config: Option, local_consensus_config: Option, + features: Option, ) -> Vec { let onchain_consensus_config = onchain_consensus_config.unwrap_or_default(); let local_consensus_config = local_consensus_config.unwrap_or_default(); @@ -190,6 +194,7 @@ impl NodeSetup { id, onchain_consensus_config.clone(), local_consensus_config.clone(), + features.clone().unwrap_or_default(), )); } nodes @@ -206,6 +211,7 @@ impl NodeSetup { id: usize, onchain_consensus_config: OnChainConsensusConfig, local_consensus_config: ConsensusConfig, + features: Features, ) -> Self { let _entered_runtime = executor.enter(); let epoch_state = Arc::new(EpochState { @@ -296,6 +302,7 @@ impl NodeSetup { onchain_consensus_config.clone(), round_manager_tx, local_consensus_config.clone(), + features.clone(), ); block_on(round_manager.init(last_vote_sent)); Self { @@ -313,6 +320,7 @@ impl NodeSetup { id, onchain_consensus_config, local_consensus_config, + features, } } @@ -332,6 +340,7 @@ impl NodeSetup { self.id, self.onchain_consensus_config.clone(), self.local_consensus_config.clone(), + self.features, ) } @@ -629,6 +638,7 @@ fn new_round_on_quorum_cert() { None, None, None, + None, ); let node = &mut nodes[0]; let genesis = node.block_store.ordered_root(); @@ -673,6 +683,7 @@ fn vote_on_successful_proposal() { None, None, None, + None, ); let node = &mut nodes[0]; @@ -717,6 +728,7 @@ fn delay_proposal_processing_in_sync_only() { None, None, None, + None, ); let node = &mut nodes[0]; @@ -785,6 +797,7 @@ fn no_vote_on_old_proposal() { None, None, None, + None, ); let node = &mut nodes[0]; let genesis_qc = certificate_for_genesis(); @@ -841,6 +854,7 @@ fn no_vote_on_mismatch_round() { None, None, None, + None, ) .pop() .unwrap(); @@ -897,6 +911,7 @@ fn sync_info_carried_on_timeout_vote() { None, None, None, + None, ); let mut node = nodes.pop().unwrap(); @@ -959,6 +974,7 @@ fn no_vote_on_invalid_proposer() { None, None, None, + None, ); let incorrect_proposer = nodes.pop().unwrap(); let mut node = nodes.pop().unwrap(); @@ -1017,6 +1033,7 @@ fn new_round_on_timeout_certificate() { None, None, None, + None, ) .pop() .unwrap(); @@ -1083,6 +1100,7 @@ fn reject_invalid_failed_authors() { None, None, None, + None, ) .pop() .unwrap(); @@ -1177,6 +1195,7 @@ fn response_on_block_retrieval() { None, None, None, + None, ) .pop() .unwrap(); @@ -1288,6 +1307,7 @@ fn recover_on_restart() { None, None, None, + None, ) .pop() .unwrap(); @@ -1363,6 +1383,7 @@ fn nil_vote_on_timeout() { None, None, None, + None, ); let node = &mut nodes[0]; let genesis = node.block_store.ordered_root(); @@ -1404,6 +1425,7 @@ fn vote_resent_on_timeout() { None, None, None, + None, ); let node = &mut nodes[0]; timed_block_on(&runtime, async { @@ -1443,6 +1465,7 @@ fn sync_on_partial_newer_sync_info() { None, None, None, + None, ); let mut node = nodes.pop().unwrap(); runtime.spawn(playground.start()); @@ -1502,6 +1525,7 @@ fn safety_rules_crash() { None, None, None, + None, ); let mut node = nodes.pop().unwrap(); runtime.spawn(playground.start()); @@ -1565,6 +1589,7 @@ fn echo_timeout() { None, None, None, + None, ); runtime.spawn(playground.start()); timed_block_on(&runtime, async { @@ -1618,6 +1643,7 @@ fn no_next_test() { None, None, None, + None, ); runtime.spawn(playground.start()); @@ -1653,6 +1679,7 @@ fn commit_pipeline_test() { Some(proposers.clone()), None, None, + None, ); runtime.spawn(playground.start()); let behind_node = 6; @@ -1692,6 +1719,7 @@ fn block_retrieval_test() { Some(vec![0, 1]), None, None, + None, ); runtime.spawn(playground.start()); @@ -1760,6 +1788,7 @@ pub fn forking_retrieval_test() { ]), None, None, + None, ); runtime.spawn(playground.start()); @@ -2003,6 +2032,7 @@ fn no_vote_on_proposal_ext_when_feature_disabled() { None, None, None, + None, ); let node = &mut nodes[0]; let genesis_qc = certificate_for_genesis(); @@ -2046,6 +2076,66 @@ fn no_vote_on_proposal_ext_when_feature_disabled() { }); } +#[test] +fn no_vote_on_proposal_with_unexpected_vtxns() { + let vtxns = vec![ValidatorTransaction::ObservedJWKUpdate( + QuorumCertifiedUpdate::dummy(), + )]; + let mut features = Features::default(); + features.disable(FeatureFlag::JWK_CONSENSUS); + assert_process_proposal_result(Some(features.clone()), vtxns.clone(), false); + + features.enable(FeatureFlag::JWK_CONSENSUS); + assert_process_proposal_result(Some(features), vtxns, true); +} + +/// Setup a node with default configs and an optional `Features` override. +/// Create a block, fill it with the given vtxns, and process it with the `RoundManager` from the setup. +/// Assert the processing result. +fn assert_process_proposal_result( + features: Option, + vtxns: Vec, + expected_result: bool, +) { + let runtime = consensus_runtime(); + let mut playground = NetworkPlayground::new(runtime.handle().clone()); + let mut nodes = NodeSetup::create_nodes( + &mut playground, + runtime.handle().clone(), + 1, + None, + Some(OnChainConsensusConfig::default_for_genesis()), + None, + features, + ); + + let node = &mut nodes[0]; + let genesis_qc = certificate_for_genesis(); + let block = Block::new_proposal_ext( + vtxns, + Payload::empty(false), + 1, + 1, + genesis_qc.clone(), + &node.signer, + Vec::new(), + ) + .unwrap(); + + timed_block_on(&runtime, async { + // clear the message queue + node.next_proposal().await; + + assert_eq!( + expected_result, + node.round_manager + .process_proposal(block.clone()) + .await + .is_ok() + ); + }); +} + #[test] /// If receiving txn num/block size limit is exceeded, ProposalExt should be rejected. fn no_vote_on_proposal_ext_when_receiving_limit_exceeded() { @@ -2077,6 +2167,7 @@ fn no_vote_on_proposal_ext_when_receiving_limit_exceeded() { vtxn: vtxn_config, }), Some(local_config), + None, ); let node = &mut nodes[0]; let genesis_qc = certificate_for_genesis(); diff --git a/consensus/src/util/mod.rs b/consensus/src/util/mod.rs index 9761f85b6b5fa..7e570eb11f190 100644 --- a/consensus/src/util/mod.rs +++ b/consensus/src/util/mod.rs @@ -2,7 +2,24 @@ // Parts of the project are originally copyright © Meta Platforms, Inc. // SPDX-License-Identifier: Apache-2.0 +use aptos_types::{ + on_chain_config::{FeatureFlag, Features}, + validator_txn::ValidatorTransaction, +}; + pub mod db_tool; #[cfg(any(test, feature = "fuzzing"))] pub mod mock_time_service; pub mod time_service; + +pub fn is_vtxn_expected(features: &Features, vtxn: &ValidatorTransaction) -> bool { + match vtxn { + ValidatorTransaction::DummyTopic1(_) | ValidatorTransaction::DummyTopic2(_) => true, + ValidatorTransaction::DKGResult(_) => { + features.is_enabled(FeatureFlag::RECONFIGURE_WITH_DKG) + }, + ValidatorTransaction::ObservedJWKUpdate(_) => { + features.is_enabled(FeatureFlag::JWK_CONSENSUS) + }, + } +} From 6452075cc44f683530edc011a22d3adf6debfb99 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 10:40:45 +0000 Subject: [PATCH 06/22] update --- Cargo.lock | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index a38195b3653b5..7159b39abc9d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3516,6 +3516,7 @@ dependencies = [ name = "aptos-safety-rules" version = "0.1.0" dependencies = [ + "anyhow", "aptos-config", "aptos-consensus-types", "aptos-crypto", @@ -4254,6 +4255,7 @@ version = "0.1.0" dependencies = [ "anyhow", "aptos-aggregator", + "aptos-bitvec", "aptos-block-executor", "aptos-block-partitioner", "aptos-crypto", From 873ab981d8367ed77bafe6c07f125193893222da Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 02:50:53 -0800 Subject: [PATCH 07/22] jwk consensus network type defs --- crates/aptos-jwk-consensus/src/lib.rs | 28 +-- crates/aptos-jwk-consensus/src/network.rs | 184 ++++++++++++++++++ .../src/network_interface.rs | 34 +++- crates/aptos-jwk-consensus/src/types.rs | 53 +++++ 4 files changed, 272 insertions(+), 27 deletions(-) create mode 100644 crates/aptos-jwk-consensus/src/network.rs create mode 100644 crates/aptos-jwk-consensus/src/types.rs diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 1b873aeeba873..372b20b630875 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -35,30 +35,6 @@ pub fn start_jwk_consensus_runtime( runtime } -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct JWKConsensusMsg {} - -#[derive(Clone)] -pub struct JWKNetworkClient { - network_client: NetworkClient, -} - -impl> JWKNetworkClient { - pub fn new(network_client: NetworkClient) -> Self { - Self { network_client } - } - - pub async fn send_rpc( - &self, - peer: PeerId, - message: JWKConsensusMsg, - rpc_timeout: Duration, - ) -> Result { - let peer_network_id = PeerNetworkId::new(NetworkId::Validator, peer); - self.network_client - .send_to_peer_rpc(message, rpc_timeout, peer_network_id) - .await - } -} - +pub mod network; pub mod network_interface; +pub mod types; diff --git a/crates/aptos-jwk-consensus/src/network.rs b/crates/aptos-jwk-consensus/src/network.rs new file mode 100644 index 0000000000000..ffea47f1f2114 --- /dev/null +++ b/crates/aptos-jwk-consensus/src/network.rs @@ -0,0 +1,184 @@ +// Copyright © Aptos Foundation + +use crate::{ + network_interface::{JWKConsensusNetworkClient, RPC}, + types::JWKConsensusMsg, +}; +use anyhow::bail; +use aptos_channels::{aptos_channel, message_queues::QueueStyle}; +use aptos_config::network_id::NetworkId; +use aptos_consensus_types::common::Author; +#[cfg(test)] +use aptos_infallible::RwLock; +use aptos_logger::warn; +use aptos_network::{ + application::interface::{NetworkClient, NetworkServiceEvents}, + protocols::network::{Event, RpcError}, + ProtocolId, +}; +use aptos_reliable_broadcast::RBNetworkSender; +use aptos_types::account_address::AccountAddress; +use bytes::Bytes; +use futures::Stream; +use futures_channel::oneshot; +use futures_util::{ + stream::{select, select_all, StreamExt}, + SinkExt, +}; +#[cfg(test)] +use std::sync::Arc; +use std::time::Duration; +use tokio::time::timeout; + +pub struct IncomingRpcRequest { + pub msg: JWKConsensusMsg, + pub sender: AccountAddress, + pub response_sender: Box, +} + +pub struct NetworkSender { + author: AccountAddress, + jwk_network_client: JWKConsensusNetworkClient>, + self_sender: aptos_channels::Sender>, +} + +impl NetworkSender { + pub fn new( + author: AccountAddress, + jwk_network_client: JWKConsensusNetworkClient>, + self_sender: aptos_channels::Sender>, + ) -> Self { + Self { + author, + jwk_network_client, + self_sender, + } + } +} + +#[async_trait::async_trait] +impl RBNetworkSender for NetworkSender { + async fn send_rb_rpc( + &self, + receiver: Author, + msg: JWKConsensusMsg, + time_limit: Duration, + ) -> anyhow::Result { + if receiver == self.author { + let (tx, rx) = oneshot::channel(); + let self_msg = Event::RpcRequest(receiver, msg, RPC[0], tx); + self.self_sender.clone().send(self_msg).await?; + if let Ok(Ok(Ok(bytes))) = timeout(time_limit, rx).await { + Ok(RPC[0].from_bytes(&bytes)?) + } else { + bail!("self rpc failed"); + } + } else { + let result = self + .jwk_network_client + .send_rpc(receiver, msg, time_limit) + .await?; + Ok(result) + } + } +} + +pub trait RpcResponseSender: Send + Sync { + fn send(&mut self, response: anyhow::Result); +} + +pub struct RealRpcResponseSender { + pub inner: Option>>, + pub protocol: ProtocolId, +} + +impl RpcResponseSender for RealRpcResponseSender { + fn send(&mut self, response: anyhow::Result) { + let rpc_response = response + .and_then(|msg| self.protocol.to_bytes(&msg).map(Bytes::from)) + .map_err(RpcError::ApplicationError); + if let Some(tx) = self.inner.take() { + let _ = tx.send(rpc_response); + } + } +} + +#[cfg(test)] +pub struct DummyRpcResponseSender { + pub rpc_response_collector: Arc>>>, +} + +#[cfg(test)] +impl DummyRpcResponseSender { + pub fn new(rpc_response_collector: Arc>>>) -> Self { + Self { + rpc_response_collector, + } + } +} + +#[cfg(test)] +impl RpcResponseSender for DummyRpcResponseSender { + fn send(&mut self, response: anyhow::Result) { + self.rpc_response_collector.write().push(response); + } +} + +pub struct NetworkReceivers { + pub rpc_rx: aptos_channel::Receiver, +} + +pub struct NetworkTask { + all_events: Box> + Send + Unpin>, + rpc_tx: aptos_channel::Sender, +} + +impl NetworkTask { + /// Establishes the initial connections with the peers and returns the receivers. + pub fn new( + network_service_events: NetworkServiceEvents, + self_receiver: aptos_channels::Receiver>, + ) -> (NetworkTask, NetworkReceivers) { + let (rpc_tx, rpc_rx) = aptos_channel::new(QueueStyle::FIFO, 10, None); + + let network_and_events = network_service_events.into_network_and_events(); + if (network_and_events.values().len() != 1) + || !network_and_events.contains_key(&NetworkId::Validator) + { + panic!("The network has not been setup correctly for JWK consensus!"); + } + + // Collect all the network events into a single stream + let network_events: Vec<_> = network_and_events.into_values().collect(); + let network_events = select_all(network_events).fuse(); + let all_events = Box::new(select(network_events, self_receiver)); + + (NetworkTask { rpc_tx, all_events }, NetworkReceivers { + rpc_rx, + }) + } + + pub async fn start(mut self) { + while let Some(message) = self.all_events.next().await { + match message { + Event::RpcRequest(peer_id, msg, protocol, response_sender) => { + let req = IncomingRpcRequest { + msg, + sender: peer_id, + response_sender: Box::new(RealRpcResponseSender { + inner: Some(response_sender), + protocol, + }), + }; + + if let Err(e) = self.rpc_tx.push(peer_id, (peer_id, req)) { + warn!(error = ?e, "aptos channel closed"); + }; + }, + _ => { + // Ignore + }, + } + } + } +} diff --git a/crates/aptos-jwk-consensus/src/network_interface.rs b/crates/aptos-jwk-consensus/src/network_interface.rs index 3e29a1bec328c..5fe9ecf56f8df 100644 --- a/crates/aptos-jwk-consensus/src/network_interface.rs +++ b/crates/aptos-jwk-consensus/src/network_interface.rs @@ -1,6 +1,13 @@ // Copyright © Aptos Foundation -use aptos_network::ProtocolId; +use crate::types::JWKConsensusMsg; +use aptos_config::network_id::{NetworkId, PeerNetworkId}; +use aptos_network::{ + application::{error::Error, interface::NetworkClientInterface}, + ProtocolId, +}; +use move_core_types::account_address::AccountAddress as PeerId; +use std::time::Duration; /// Supported protocols in preferred order (from highest priority to lowest). pub const DIRECT_SEND: &[ProtocolId] = &[ @@ -15,3 +22,28 @@ pub const RPC: &[ProtocolId] = &[ ProtocolId::JWKConsensusRpcBcs, ProtocolId::JWKConsensusRpcJson, ]; + +#[derive(Clone)] +pub struct JWKConsensusNetworkClient { + network_client: NetworkClient, +} + +impl> + JWKConsensusNetworkClient +{ + pub fn new(network_client: NetworkClient) -> Self { + Self { network_client } + } + + pub async fn send_rpc( + &self, + peer: PeerId, + message: JWKConsensusMsg, + rpc_timeout: Duration, + ) -> Result { + let peer_network_id = PeerNetworkId::new(NetworkId::Validator, peer); + self.network_client + .send_to_peer_rpc(message, rpc_timeout, peer_network_id) + .await + } +} diff --git a/crates/aptos-jwk-consensus/src/types.rs b/crates/aptos-jwk-consensus/src/types.rs new file mode 100644 index 0000000000000..31223ec4c746f --- /dev/null +++ b/crates/aptos-jwk-consensus/src/types.rs @@ -0,0 +1,53 @@ +// Copyright © Aptos Foundation + +use aptos_crypto::bls12381::Signature; +use aptos_enum_conversion_derive::EnumConversion; +use aptos_reliable_broadcast::RBMessage; +use aptos_types::{ + account_address::AccountAddress, + jwks::{Issuer, ProviderJWKs}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, EnumConversion, Deserialize, Serialize, PartialEq)] +pub enum JWKConsensusMsg { + ObservationRequest(ObservedUpdateRequest), + ObservationResponse(ObservedUpdateResponse), +} + +impl JWKConsensusMsg { + pub fn name(&self) -> &str { + match self { + JWKConsensusMsg::ObservationRequest(_) => "ObservationRequest", + JWKConsensusMsg::ObservationResponse(_) => "ObservationResponse", + } + } + + pub fn epoch(&self) -> u64 { + match self { + JWKConsensusMsg::ObservationRequest(request) => request.epoch, + JWKConsensusMsg::ObservationResponse(response) => response.epoch, + } + } +} + +impl RBMessage for JWKConsensusMsg {} + +#[derive(Clone, Debug, Deserialize, Serialize, Eq, PartialEq)] +pub struct ObservedUpdate { + pub author: AccountAddress, + pub observed: ProviderJWKs, + pub signature: Signature, +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] +pub struct ObservedUpdateRequest { + pub epoch: u64, + pub issuer: Issuer, +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] +pub struct ObservedUpdateResponse { + pub epoch: u64, + pub update: ObservedUpdate, +} From 54c5c9af37aec84a34b02c1aeba4958c5eaa50b1 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 02:56:41 -0800 Subject: [PATCH 08/22] update cargo.toml --- crates/aptos-jwk-consensus/Cargo.toml | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/crates/aptos-jwk-consensus/Cargo.toml b/crates/aptos-jwk-consensus/Cargo.toml index da4d6a204ff20..f77612bbcddfa 100644 --- a/crates/aptos-jwk-consensus/Cargo.toml +++ b/crates/aptos-jwk-consensus/Cargo.toml @@ -13,15 +13,40 @@ repository = { workspace = true } rust-version = { workspace = true } [dependencies] +anyhow = { workspace = true } +aptos-bitvec = { workspace = true } +aptos-bounded-executor = { workspace = true } +aptos-channels = { workspace = true } aptos-config = { workspace = true } +aptos-consensus-types = { workspace = true } +aptos-crypto = { workspace = true } +aptos-enum-conversion-derive = { workspace = true } aptos-event-notifications = { workspace = true } +aptos-global-constants = { workspace = true } +aptos-infallible = { workspace = true } +aptos-logger = { workspace = true } +aptos-metrics-core = { workspace = true } aptos-network = { workspace = true } +aptos-reliable-broadcast = { workspace = true } aptos-runtimes = { workspace = true } +aptos-time-service = { workspace = true } aptos-types = { workspace = true } aptos-validator-transaction-pool = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +futures = { workspace = true } +futures-channel = { workspace = true } futures-util = { workspace = true } +move-core-types = { workspace = true } +once_cell = { workspace = true } +reqwest = { workspace = true } serde = { workspace = true } +serde_json = { workspace = true } tokio = { workspace = true } +tokio-retry = { workspace = true } +[dev-dependencies] +aptos-types = { workspace = true, features = ["fuzzing"] } +aptos-validator-transaction-pool = { workspace = true, features = ["fuzzing"] } [features] smoke-test = [] From f7ef04315e6e03427392f8f74902f5607ee8a1aa Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 10:59:25 +0000 Subject: [PATCH 09/22] update --- Cargo.lock | 22 ++++++++++++++++++++ crates/validator-transaction-pool/Cargo.toml | 3 +++ crates/validator-transaction-pool/src/lib.rs | 12 +++++++++++ 3 files changed, 37 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 7159b39abc9d4..c6abe05be8c93 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2496,15 +2496,37 @@ dependencies = [ name = "aptos-jwk-consensus" version = "0.1.0" dependencies = [ + "anyhow", + "aptos-bitvec", + "aptos-bounded-executor", + "aptos-channels", "aptos-config", + "aptos-consensus-types", + "aptos-crypto", + "aptos-enum-conversion-derive", "aptos-event-notifications", + "aptos-global-constants", + "aptos-infallible", + "aptos-logger", + "aptos-metrics-core", "aptos-network", + "aptos-reliable-broadcast", "aptos-runtimes", + "aptos-time-service", "aptos-types", "aptos-validator-transaction-pool", + "async-trait", + "bytes", + "futures", + "futures-channel", "futures-util", + "move-core-types", + "once_cell", + "reqwest", "serde", + "serde_json", "tokio", + "tokio-retry", ] [[package]] diff --git a/crates/validator-transaction-pool/Cargo.toml b/crates/validator-transaction-pool/Cargo.toml index c3eac3de8d54c..c0604b9d34f4b 100644 --- a/crates/validator-transaction-pool/Cargo.toml +++ b/crates/validator-transaction-pool/Cargo.toml @@ -23,3 +23,6 @@ tokio = { workspace = true } [dev-dependencies] aptos-types = { workspace = true, features = ["fuzzing"] } + +[features] +fuzzing = [] diff --git a/crates/validator-transaction-pool/src/lib.rs b/crates/validator-transaction-pool/src/lib.rs index dbe9d3065a17b..5ed1cd10dd56c 100644 --- a/crates/validator-transaction-pool/src/lib.rs +++ b/crates/validator-transaction-pool/src/lib.rs @@ -22,6 +22,10 @@ impl TransactionFilter { } impl TransactionFilter { + pub fn empty() -> Self { + Self::PendingTxnHashSet(HashSet::new()) + } + pub fn should_exclude(&self, txn: &ValidatorTransaction) -> bool { match self { TransactionFilter::PendingTxnHashSet(set) => set.contains(&txn.hash()), @@ -87,6 +91,14 @@ impl VTxnPoolState { .lock() .pull(deadline, max_items, max_bytes, filter) } + + #[cfg(any(test, feature = "fuzzing"))] + pub fn dummy_txn_guard(&self) -> TxnGuard { + TxnGuard { + pool: self.inner.clone(), + seq_num: u64::MAX, + } + } } struct PoolItem { From 20ccbc665421745a1b4600c7fe50be2a9ae1c06a Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 03:01:58 -0800 Subject: [PATCH 10/22] update --- crates/aptos-jwk-consensus/src/lib.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 372b20b630875..5f9cf0adf4252 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -1,19 +1,15 @@ // Copyright © Aptos Foundation -use aptos_config::network_id::{NetworkId, PeerNetworkId}; use aptos_event_notifications::{ DbBackedOnChainConfig, EventNotificationListener, ReconfigNotificationListener, }; use aptos_network::application::{ - error::Error, - interface::{NetworkClient, NetworkClientInterface, NetworkServiceEvents}, + interface::{NetworkClient, NetworkServiceEvents}, }; -use aptos_types::PeerId; use aptos_validator_transaction_pool::VTxnPoolState; use futures_util::StreamExt; -use serde::{Deserialize, Serialize}; -use std::time::Duration; use tokio::runtime::Runtime; +use crate::types::JWKConsensusMsg; #[allow(clippy::let_and_return)] pub fn start_jwk_consensus_runtime( From 27b7d804e2ff7d5a7baf393b9ad977e170c0fc89 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:20:00 -0800 Subject: [PATCH 11/22] update --- aptos-node/src/network.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aptos-node/src/network.rs b/aptos-node/src/network.rs index 4167e88ffe09f..30db63d7e121a 100644 --- a/aptos-node/src/network.rs +++ b/aptos-node/src/network.rs @@ -10,7 +10,6 @@ use aptos_config::{ use aptos_consensus::network_interface::ConsensusMsg; use aptos_dkg_runtime::DKGMessage; use aptos_event_notifications::EventSubscriptionService; -use aptos_jwk_consensus::JWKConsensusMsg; use aptos_logger::debug; use aptos_mempool::network::MempoolSyncMsg; use aptos_network::{ @@ -33,6 +32,7 @@ use aptos_types::chain_id::ChainId; use serde::{Deserialize, Serialize}; use std::{collections::HashMap, sync::Arc}; use tokio::runtime::Runtime; +use aptos_jwk_consensus::types::JWKConsensusMsg; /// A simple struct that holds both the network client /// and receiving interfaces for an application. From 99cc957249c07cf5fda9854e2d0ab35999c2bdda Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 15:21:39 +0000 Subject: [PATCH 12/22] lint --- aptos-node/src/network.rs | 2 +- crates/aptos-jwk-consensus/src/lib.rs | 6 ++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/aptos-node/src/network.rs b/aptos-node/src/network.rs index 30db63d7e121a..9c95ec7a30d0c 100644 --- a/aptos-node/src/network.rs +++ b/aptos-node/src/network.rs @@ -10,6 +10,7 @@ use aptos_config::{ use aptos_consensus::network_interface::ConsensusMsg; use aptos_dkg_runtime::DKGMessage; use aptos_event_notifications::EventSubscriptionService; +use aptos_jwk_consensus::types::JWKConsensusMsg; use aptos_logger::debug; use aptos_mempool::network::MempoolSyncMsg; use aptos_network::{ @@ -32,7 +33,6 @@ use aptos_types::chain_id::ChainId; use serde::{Deserialize, Serialize}; use std::{collections::HashMap, sync::Arc}; use tokio::runtime::Runtime; -use aptos_jwk_consensus::types::JWKConsensusMsg; /// A simple struct that holds both the network client /// and receiving interfaces for an application. diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 5f9cf0adf4252..737b1a1f8a1ca 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -1,15 +1,13 @@ // Copyright © Aptos Foundation +use crate::types::JWKConsensusMsg; use aptos_event_notifications::{ DbBackedOnChainConfig, EventNotificationListener, ReconfigNotificationListener, }; -use aptos_network::application::{ - interface::{NetworkClient, NetworkServiceEvents}, -}; +use aptos_network::application::interface::{NetworkClient, NetworkServiceEvents}; use aptos_validator_transaction_pool::VTxnPoolState; use futures_util::StreamExt; use tokio::runtime::Runtime; -use crate::types::JWKConsensusMsg; #[allow(clippy::let_and_return)] pub fn start_jwk_consensus_runtime( From b019753fa80dfab6483eb66b87c7a0b90ca8910d Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:30:45 -0800 Subject: [PATCH 13/22] jwk update quorum certification --- .../src/certified_update_producer.rs | 65 ++++++++ crates/aptos-jwk-consensus/src/lib.rs | 2 + .../src/observation_aggregation/mod.rs | 108 +++++++++++++ .../src/observation_aggregation/tests.rs | 143 ++++++++++++++++++ 4 files changed, 318 insertions(+) create mode 100644 crates/aptos-jwk-consensus/src/certified_update_producer.rs create mode 100644 crates/aptos-jwk-consensus/src/observation_aggregation/mod.rs create mode 100644 crates/aptos-jwk-consensus/src/observation_aggregation/tests.rs diff --git a/crates/aptos-jwk-consensus/src/certified_update_producer.rs b/crates/aptos-jwk-consensus/src/certified_update_producer.rs new file mode 100644 index 0000000000000..d60c241573f4e --- /dev/null +++ b/crates/aptos-jwk-consensus/src/certified_update_producer.rs @@ -0,0 +1,65 @@ +// Copyright © Aptos Foundation + +use crate::{ + observation_aggregation::ObservationAggregationState, + types::{JWKConsensusMsg, ObservedUpdateRequest}, +}; +use aptos_channels::aptos_channel; +use aptos_reliable_broadcast::ReliableBroadcast; +use aptos_types::{ + epoch_state::EpochState, + jwks::{ProviderJWKs, QuorumCertifiedUpdate}, +}; +use futures_util::future::{AbortHandle, Abortable}; +use std::sync::Arc; +use tokio_retry::strategy::ExponentialBackoff; + +/// A sub-process of the whole JWK consensus process. +/// Once invoked by `JWKConsensusManager` to `start_produce`, +/// it starts producing a `QuorumCertifiedUpdate` and returns an abort handle. +/// Once an `QuorumCertifiedUpdate` is available, it is sent back via a channel given earlier. +pub trait CertifiedUpdateProducer: Send + Sync { + fn start_produce( + &self, + epoch_state: Arc, + payload: ProviderJWKs, + qc_update_tx: Option>, + ) -> AbortHandle; +} + +pub struct RealCertifiedUpdateProducer { + reliable_broadcast: Arc>, +} + +impl RealCertifiedUpdateProducer { + pub fn new(reliable_broadcast: ReliableBroadcast) -> Self { + Self { + reliable_broadcast: Arc::new(reliable_broadcast), + } + } +} + +impl CertifiedUpdateProducer for RealCertifiedUpdateProducer { + fn start_produce( + &self, + epoch_state: Arc, + payload: ProviderJWKs, + qc_update_tx: Option>, + ) -> AbortHandle { + let rb = self.reliable_broadcast.clone(); + let req = ObservedUpdateRequest { + epoch: epoch_state.epoch, + issuer: payload.issuer.clone(), + }; + let agg_state = Arc::new(ObservationAggregationState::new(epoch_state, payload)); + let task = async move { + let qc_update = rb.broadcast(req, agg_state).await; + if let Some(tx) = qc_update_tx { + let _ = tx.push((), qc_update); + } + }; + let (abort_handle, abort_registration) = AbortHandle::new_pair(); + tokio::spawn(Abortable::new(task, abort_registration)); + abort_handle + } +} diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 737b1a1f8a1ca..d51816ec21376 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -29,6 +29,8 @@ pub fn start_jwk_consensus_runtime( runtime } +pub mod certified_update_producer; pub mod network; pub mod network_interface; +pub mod observation_aggregation; pub mod types; diff --git a/crates/aptos-jwk-consensus/src/observation_aggregation/mod.rs b/crates/aptos-jwk-consensus/src/observation_aggregation/mod.rs new file mode 100644 index 0000000000000..ae850996fa40e --- /dev/null +++ b/crates/aptos-jwk-consensus/src/observation_aggregation/mod.rs @@ -0,0 +1,108 @@ +// Copyright © Aptos Foundation + +use crate::types::{ + JWKConsensusMsg, ObservedUpdate, ObservedUpdateRequest, ObservedUpdateResponse, +}; +use anyhow::ensure; +use aptos_consensus_types::common::Author; +use aptos_crypto::bls12381; +use aptos_infallible::Mutex; +use aptos_reliable_broadcast::BroadcastStatus; +use aptos_types::{ + epoch_state::EpochState, + jwks::{ProviderJWKs, QuorumCertifiedUpdate}, +}; +use move_core_types::account_address::AccountAddress; +use std::{collections::HashSet, sync::Arc}; + +/// The aggregation state of reliable broadcast where a validator broadcast JWK observation requests +/// and produce quorum-certified JWK updates. +pub struct ObservationAggregationState { + epoch_state: Arc, + local_view: ProviderJWKs, + inner_state: Mutex, +} + +#[derive(Default)] +struct InnerState { + pub contributors: HashSet, + pub multi_sig: Option, +} + +impl ObservationAggregationState { + pub fn new(epoch_state: Arc, local_view: ProviderJWKs) -> Self { + Self { + epoch_state, + local_view, + inner_state: Mutex::new(InnerState::default()), + } + } +} + +impl BroadcastStatus for Arc { + type Aggregated = QuorumCertifiedUpdate; + type Message = ObservedUpdateRequest; + type Response = ObservedUpdateResponse; + + fn add( + &self, + sender: Author, + response: Self::Response, + ) -> anyhow::Result> { + let ObservedUpdateResponse { epoch, update } = response; + let ObservedUpdate { + author, + observed: peer_view, + signature, + } = update; + ensure!( + epoch == self.epoch_state.epoch, + "adding peer observation failed with invalid epoch", + ); + ensure!( + author == sender, + "adding peer observation failed with mismatched author", + ); + + let mut aggregator = self.inner_state.lock(); + if aggregator.contributors.contains(&sender) { + return Ok(None); + } + + ensure!( + self.local_view == peer_view, + "adding peer observation failed with mismatched view" + ); + + // Verify the quorum-cert. + self.epoch_state + .verifier + .verify(sender, &peer_view, &signature)?; + + // All checks passed. Aggregating. + aggregator.contributors.insert(sender); + let new_multi_sig = if let Some(existing) = aggregator.multi_sig.take() { + bls12381::Signature::aggregate(vec![existing, signature])? + } else { + signature + }; + + let maybe_qc_update = self + .epoch_state + .verifier + .check_voting_power(aggregator.contributors.iter(), true) + .ok() + .map(|_| QuorumCertifiedUpdate { + authors: aggregator.contributors.clone().into_iter().collect(), + update: peer_view, + multi_sig: new_multi_sig.clone(), + }); + + aggregator.multi_sig = Some(new_multi_sig); + + Ok(maybe_qc_update) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/aptos-jwk-consensus/src/observation_aggregation/tests.rs b/crates/aptos-jwk-consensus/src/observation_aggregation/tests.rs new file mode 100644 index 0000000000000..60c4ca16db450 --- /dev/null +++ b/crates/aptos-jwk-consensus/src/observation_aggregation/tests.rs @@ -0,0 +1,143 @@ +// Copyright © Aptos Foundation + +use crate::{ + observation_aggregation::ObservationAggregationState, + types::{ObservedUpdate, ObservedUpdateResponse}, +}; +use aptos_bitvec::BitVec; +use aptos_crypto::{bls12381, SigningKey, Uniform}; +use aptos_reliable_broadcast::BroadcastStatus; +use aptos_types::{ + aggregate_signature::AggregateSignature, + epoch_state::EpochState, + jwks::{ + jwk::{JWKMoveStruct, JWK}, + unsupported::UnsupportedJWK, + ProviderJWKs, QuorumCertifiedUpdate, + }, + validator_verifier::{ValidatorConsensusInfo, ValidatorVerifier}, +}; +use move_core_types::account_address::AccountAddress; +use std::sync::Arc; + +#[test] +fn test_observation_aggregation_state() { + let num_validators = 5; + let epoch = 999; + let addrs: Vec = (0..num_validators) + .map(|_| AccountAddress::random()) + .collect(); + let private_keys: Vec = (0..num_validators) + .map(|_| bls12381::PrivateKey::generate_for_testing()) + .collect(); + let public_keys: Vec = (0..num_validators) + .map(|i| bls12381::PublicKey::from(&private_keys[i])) + .collect(); + let voting_powers = [1, 1, 1, 6, 6]; // total voting power: 15, default threshold: 11 + let validator_infos: Vec = (0..num_validators) + .map(|i| ValidatorConsensusInfo::new(addrs[i], public_keys[i].clone(), voting_powers[i])) + .collect(); + let verifier = ValidatorVerifier::new(validator_infos); + let epoch_state = Arc::new(EpochState { epoch, verifier }); + let view_0 = ProviderJWKs { + issuer: b"https::/alice.com".to_vec(), + version: 123, + jwks: vec![JWKMoveStruct::from(JWK::Unsupported( + UnsupportedJWK::new_for_testing("id1", "payload1"), + ))], + }; + let view_1 = ProviderJWKs { + issuer: b"https::/alice.com".to_vec(), + version: 123, + jwks: vec![JWKMoveStruct::from(JWK::Unsupported( + UnsupportedJWK::new_for_testing("id2", "payload2"), + ))], + }; + let ob_agg_state = Arc::new(ObservationAggregationState::new( + epoch_state.clone(), + view_0.clone(), + )); + + // `ObservedUpdate` with incorrect epoch should be rejected. + let result = ob_agg_state.add(addrs[0], ObservedUpdateResponse { + epoch: 998, + update: ObservedUpdate { + author: addrs[0], + observed: view_0.clone(), + signature: private_keys[0].sign(&view_0).unwrap(), + }, + }); + assert!(result.is_err()); + + // `ObservedUpdate` authored by X but sent by Y should be rejected. + let result = ob_agg_state.add(addrs[1], ObservedUpdateResponse { + epoch: 999, + update: ObservedUpdate { + author: addrs[0], + observed: view_0.clone(), + signature: private_keys[0].sign(&view_0).unwrap(), + }, + }); + assert!(result.is_err()); + + // `ObservedUpdate` that cannot be verified should be rejected. + let result = ob_agg_state.add(addrs[2], ObservedUpdateResponse { + epoch: 999, + update: ObservedUpdate { + author: addrs[2], + observed: view_0.clone(), + signature: private_keys[2].sign(&view_1).unwrap(), + }, + }); + assert!(result.is_err()); + + // Good `ObservedUpdate` should be accepted. + let result = ob_agg_state.add(addrs[3], ObservedUpdateResponse { + epoch: 999, + update: ObservedUpdate { + author: addrs[3], + observed: view_0.clone(), + signature: private_keys[3].sign(&view_0).unwrap(), + }, + }); + assert!(matches!(result, Ok(None))); + + // `ObservedUpdate` from contributed author should be ignored. + let result = ob_agg_state.add(addrs[3], ObservedUpdateResponse { + epoch: 999, + update: ObservedUpdate { + author: addrs[3], + observed: view_0.clone(), + signature: private_keys[3].sign(&view_0).unwrap(), + }, + }); + assert!(matches!(result, Ok(None))); + + // Quorum-certified update should be returned if after adding an `ObservedUpdate`, the threshold is exceeded. + let result = ob_agg_state.add(addrs[4], ObservedUpdateResponse { + epoch: 999, + update: ObservedUpdate { + author: addrs[4], + observed: view_0.clone(), + signature: private_keys[4].sign(&view_0).unwrap(), + }, + }); + let QuorumCertifiedUpdate { + authors, + update: observed, + multi_sig, + } = result.unwrap().unwrap(); + assert_eq!(view_0, observed); + let bits: Vec = epoch_state + .verifier + .get_ordered_account_addresses() + .into_iter() + .map(|addr| authors.contains(&addr)) + .collect(); + let bit_vec = BitVec::from(bits); + let multi_sig = AggregateSignature::new(bit_vec, Some(multi_sig)); + assert!(epoch_state + .verifier + .verify_multi_signatures(&observed, &multi_sig) + .is_ok()); +} From 1c7a9b9265274dc8cbfa95bb3af190b474a6797b Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:36:17 -0800 Subject: [PATCH 14/22] jwk observation --- .../aptos-jwk-consensus/src/jwk_observer.rs | 136 ++++++++++++++++++ crates/aptos-jwk-consensus/src/lib.rs | 1 + 2 files changed, 137 insertions(+) create mode 100644 crates/aptos-jwk-consensus/src/jwk_observer.rs diff --git a/crates/aptos-jwk-consensus/src/jwk_observer.rs b/crates/aptos-jwk-consensus/src/jwk_observer.rs new file mode 100644 index 0000000000000..798869de3bf8e --- /dev/null +++ b/crates/aptos-jwk-consensus/src/jwk_observer.rs @@ -0,0 +1,136 @@ +// Copyright © Aptos Foundation + +use anyhow::Result; +use aptos_channels::aptos_channel; +use aptos_logger::{debug, info}; +use aptos_types::jwks::{jwk::JWK, Issuer}; +use futures::{FutureExt, StreamExt}; +use move_core_types::account_address::AccountAddress; +use serde::{Deserialize, Serialize}; +use std::time::Duration; +use tokio::{sync::oneshot, task::JoinHandle, time::MissedTickBehavior}; + +#[derive(Serialize, Deserialize)] +struct OpenIDConfiguration { + issuer: String, + jwks_uri: String, +} + +#[derive(Serialize, Deserialize)] +struct JWKsResponse { + keys: Vec, +} + +/// Given an Open ID configuration URL, fetch its JWKs. +pub async fn fetch_jwks(my_addr: AccountAddress, config_url: Vec) -> Result> { + if cfg!(feature = "smoke-test") { + use reqwest::header; + let maybe_url = String::from_utf8(config_url); + let jwk_url = maybe_url?; + let client = reqwest::Client::new(); + let JWKsResponse { keys } = client + .get(jwk_url.as_str()) + .header(header::COOKIE, my_addr.to_hex()) + .send() + .await? + .json() + .await?; + let jwks = keys.into_iter().map(JWK::from).collect(); + Ok(jwks) + } else { + let maybe_url = String::from_utf8(config_url); + let config_url = maybe_url?; + let client = reqwest::Client::new(); + let OpenIDConfiguration { jwks_uri, .. } = + client.get(config_url.as_str()).send().await?.json().await?; + let JWKsResponse { keys } = client.get(jwks_uri.as_str()).send().await?.json().await?; + let jwks = keys.into_iter().map(JWK::from).collect(); + Ok(jwks) + } +} + +/// A process thread that periodically fetch JWKs of a provider and push it back to JWKManager. +pub struct JWKObserver { + close_tx: oneshot::Sender<()>, + join_handle: JoinHandle<()>, +} + +impl JWKObserver { + pub fn spawn( + my_addr: AccountAddress, + issuer: Issuer, + config_url: Vec, + fetch_interval: Duration, + observation_tx: aptos_channel::Sender<(), (Issuer, Vec)>, + ) -> Self { + let (close_tx, close_rx) = oneshot::channel(); + let join_handle = tokio::spawn(Self::thread_main( + fetch_interval, + my_addr, + issuer.clone(), + config_url.clone(), + observation_tx, + close_rx, + )); + info!( + "[JWK] observer spawned, issuer={:?}, config_url={:?}", + String::from_utf8(issuer), + String::from_utf8(config_url) + ); + Self { + close_tx, + join_handle, + } + } + + async fn thread_main( + fetch_interval: Duration, + my_addr: AccountAddress, + issuer: Issuer, + open_id_config_url: Vec, + observation_tx: aptos_channel::Sender<(), (Issuer, Vec)>, + close_rx: oneshot::Receiver<()>, + ) { + let mut interval = tokio::time::interval(fetch_interval); + interval.set_missed_tick_behavior(MissedTickBehavior::Delay); + let mut close_rx = close_rx.into_stream(); + loop { + tokio::select! { + _ = interval.tick().fuse() => { + let result = fetch_jwks(my_addr, open_id_config_url.clone()).await; + debug!("observe_result={:?}", result); + if let Ok(mut jwks) = result { + jwks.sort(); + let _ = observation_tx.push((), (issuer.clone(), jwks)); + } + }, + _ = close_rx.select_next_some() => { + break; + } + } + } + } + + pub async fn shutdown(self) { + let Self { + close_tx, + join_handle, + } = self; + let _ = close_tx.send(()); + let _ = join_handle.await; + } +} + +#[ignore] +#[tokio::test] +async fn test_fetch_real_jwks() { + let jwks = fetch_jwks( + AccountAddress::ZERO, + "https://www.facebook.com/.well-known/openid-configuration/" + .as_bytes() + .to_vec(), + ) + .await + .unwrap(); + println!("{:?}", jwks); +} diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index d51816ec21376..25f214dcb76bb 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -30,6 +30,7 @@ pub fn start_jwk_consensus_runtime( } pub mod certified_update_producer; +pub mod jwk_observer; pub mod network; pub mod network_interface; pub mod observation_aggregation; From bc901bc4e1dd2fcb1fa4661a84261d27c245ad7f Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:41:00 -0800 Subject: [PATCH 15/22] the main jwk consensus state machine --- .../src/jwk_manager/mod.rs | 408 +++++++++++++++ .../src/jwk_manager/tests.rs | 475 ++++++++++++++++++ crates/aptos-jwk-consensus/src/lib.rs | 1 + 3 files changed, 884 insertions(+) create mode 100644 crates/aptos-jwk-consensus/src/jwk_manager/mod.rs create mode 100644 crates/aptos-jwk-consensus/src/jwk_manager/tests.rs diff --git a/crates/aptos-jwk-consensus/src/jwk_manager/mod.rs b/crates/aptos-jwk-consensus/src/jwk_manager/mod.rs new file mode 100644 index 0000000000000..b0722a1070a50 --- /dev/null +++ b/crates/aptos-jwk-consensus/src/jwk_manager/mod.rs @@ -0,0 +1,408 @@ +// Copyright © Aptos Foundation + +use crate::{ + certified_update_producer::CertifiedUpdateProducer, + jwk_observer::JWKObserver, + network::IncomingRpcRequest, + types::{JWKConsensusMsg, ObservedUpdate, ObservedUpdateResponse}, +}; +use anyhow::{anyhow, bail, Result}; +use aptos_channels::{aptos_channel, message_queues::QueueStyle}; +use aptos_crypto::{bls12381::PrivateKey, SigningKey}; +use aptos_logger::{debug, error, info}; +use aptos_types::{ + account_address::AccountAddress, + epoch_state::EpochState, + jwks::{ + jwk::JWKMoveStruct, AllProvidersJWKs, Issuer, ObservedJWKs, ObservedJWKsUpdated, + ProviderJWKs, QuorumCertifiedUpdate, SupportedOIDCProviders, + }, + validator_txn::{Topic, ValidatorTransaction}, +}; +use aptos_validator_transaction_pool::{TxnGuard, VTxnPoolState}; +use futures_channel::oneshot; +use futures_util::{ + future::{join_all, AbortHandle}, + FutureExt, StreamExt, +}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Duration, +}; + +/// `JWKManager` executes per-issuer JWK consensus sessions +/// and updates validator txn pool with quorum-certified JWK updates. +pub struct JWKManager { + /// Some useful metadata. + my_addr: AccountAddress, + epoch_state: Arc, + + /// Used to sign JWK observations before sharing them with peers. + consensus_key: Arc, + + /// The sub-process that collects JWK updates from peers and aggregate them into a quorum-certified JWK update. + certified_update_producer: Arc, + + /// When a quorum-certified JWK update is available, use this to put it into the validator transaction pool. + vtxn_pool: VTxnPoolState, + + /// The JWK consensus states of all the issuers. + states_by_issuer: HashMap, + + /// Whether a CLOSE command has been received. + stopped: bool, + + qc_update_tx: Option>, + jwk_observers: Vec, +} + +impl JWKManager { + pub fn new( + consensus_key: Arc, + my_addr: AccountAddress, + epoch_state: Arc, + certified_update_producer: Arc, + vtxn_pool: VTxnPoolState, + ) -> Self { + Self { + consensus_key, + my_addr, + epoch_state, + certified_update_producer, + vtxn_pool, + states_by_issuer: HashMap::default(), + stopped: false, + qc_update_tx: None, + jwk_observers: vec![], + } + } + + pub async fn run( + mut self, + oidc_providers: Option, + observed_jwks: Option, + mut jwk_updated_rx: aptos_channel::Receiver<(), ObservedJWKsUpdated>, + mut rpc_req_rx: aptos_channel::Receiver<(), (AccountAddress, IncomingRpcRequest)>, + close_rx: oneshot::Receiver>, + ) { + self.reset_with_on_chain_state(observed_jwks.unwrap_or_default().into_providers_jwks()) + .unwrap(); + let (qc_update_tx, mut qc_update_rx) = aptos_channel::new(QueueStyle::FIFO, 100, None); + self.qc_update_tx = Some(qc_update_tx); + + let (local_observation_tx, mut local_observation_rx) = + aptos_channel::new(QueueStyle::KLAST, 100, None); + + self.jwk_observers = oidc_providers + .unwrap_or_default() + .into_provider_vec() + .into_iter() + .map(|provider| { + JWKObserver::spawn( + self.my_addr, + provider.name.clone(), + provider.config_url.clone(), + Duration::from_secs(10), + local_observation_tx.clone(), + ) + }) + .collect(); + + let mut close_rx = close_rx.into_stream(); + + while !self.stopped { + let handle_result = tokio::select! { + jwk_updated = jwk_updated_rx.select_next_some() => { + let ObservedJWKsUpdated { jwks, .. } = jwk_updated; + self.reset_with_on_chain_state(jwks) + }, + (_sender, msg) = rpc_req_rx.select_next_some() => { + self.process_peer_request(msg) + }, + qc_update = qc_update_rx.select_next_some() => { + self.process_quorum_certified_update(qc_update) + }, + (issuer, jwks) = local_observation_rx.select_next_some() => { + let jwks = jwks.into_iter().map(JWKMoveStruct::from).collect(); + self.process_new_observation(issuer, jwks) + }, + ack_tx = close_rx.select_next_some() => { + self.tear_down(ack_tx.ok()).await + } + }; + + if let Err(e) = handle_result { + error!("[JWK] handling_err={}", e); + } + } + } + + async fn tear_down(&mut self, ack_tx: Option>) -> Result<()> { + self.stopped = true; + let futures = std::mem::take(&mut self.jwk_observers) + .into_iter() + .map(JWKObserver::shutdown) + .collect::>(); + join_all(futures).await; + if let Some(tx) = ack_tx { + let _ = tx.send(()); + } + Ok(()) + } + + /// Triggered by an observation thread periodically. + pub fn process_new_observation( + &mut self, + issuer: Issuer, + jwks: Vec, + ) -> Result<()> { + let state = self.states_by_issuer.entry(issuer.clone()).or_default(); + state.observed = Some(jwks.clone()); + if state.observed.as_ref() != state.on_chain.as_ref().map(ProviderJWKs::jwks) { + let observed = ProviderJWKs { + issuer: issuer.clone(), + version: state.on_chain_version() + 1, + jwks, + }; + let signature = self + .consensus_key + .sign(&observed) + .map_err(|e| anyhow!("crypto material error occurred duing signing: {}", e))?; + let abort_handle = self.certified_update_producer.start_produce( + self.epoch_state.clone(), + observed.clone(), + self.qc_update_tx.clone(), + ); + state.consensus_state = ConsensusState::InProgress { + my_proposal: ObservedUpdate { + author: self.my_addr, + observed: observed.clone(), + signature, + }, + abort_handle_wrapper: QuorumCertProcessGuard::new(abort_handle), + }; + info!("[JWK] update observed, update={:?}", observed); + } + + Ok(()) + } + + /// Invoked on start, or on on-chain JWK updated event. + pub fn reset_with_on_chain_state(&mut self, on_chain_state: AllProvidersJWKs) -> Result<()> { + debug!( + "[JWK] reset_with_on_chain_state: BEGIN: on_chain_state={:?}", + on_chain_state + ); + let onchain_issuer_set: HashSet = on_chain_state + .entries + .iter() + .map(|entry| entry.issuer.clone()) + .collect(); + self.states_by_issuer + .retain(|issuer, _| onchain_issuer_set.contains(issuer)); + for provider_jwks in on_chain_state.entries { + let x = self + .states_by_issuer + .get(&provider_jwks.issuer) + .and_then(|s| s.on_chain.as_ref()); + if x != Some(&provider_jwks) { + self.states_by_issuer.insert( + provider_jwks.issuer.clone(), + PerProviderState::new(provider_jwks), + ); + } + } + Ok(()) + } + + pub fn process_peer_request(&mut self, rpc_req: IncomingRpcRequest) -> Result<()> { + let IncomingRpcRequest { + msg, + mut response_sender, + sender, + } = rpc_req; + debug!( + "[JWK] process_peer_request: sender={}, is_self={}", + sender, + sender == self.my_addr + ); + match msg { + JWKConsensusMsg::ObservationRequest(request) => { + let state = self.states_by_issuer.entry(request.issuer).or_default(); + let response: Result = match &state.consensus_state { + ConsensusState::NotStarted => Err(anyhow!("observed update unavailable")), + ConsensusState::InProgress { my_proposal, .. } + | ConsensusState::Finished { my_proposal, .. } => Ok( + JWKConsensusMsg::ObservationResponse(ObservedUpdateResponse { + epoch: self.epoch_state.epoch, + update: my_proposal.clone(), + }), + ), + }; + response_sender.send(response); + Ok(()) + }, + _ => { + bail!("unexpected rpc: {}", msg.name()); + }, + } + } + + /// Triggered once the `certified_update_producer` produced a quorum-certified update. + pub fn process_quorum_certified_update(&mut self, update: QuorumCertifiedUpdate) -> Result<()> { + let issuer = update.update.issuer.clone(); + let state = self.states_by_issuer.entry(issuer.clone()).or_default(); + match &state.consensus_state { + ConsensusState::InProgress { my_proposal, .. } => { + //TODO: counters + let txn = ValidatorTransaction::ObservedJWKUpdate(update.clone()); + let vtxn_guard = + self.vtxn_pool + .put(Topic::JWK_CONSENSUS(issuer), Arc::new(txn), None); + state.consensus_state = ConsensusState::Finished { + vtxn_guard, + my_proposal: my_proposal.clone(), + quorum_certified: update.clone(), + }; + info!("[JWK] qc update obtained, update={:?}", update); + Ok(()) + }, + _ => Err(anyhow!( + "qc update not expected for issuer {:?} in state {}", + update.update.issuer, + state.consensus_state.name() + )), + } + } +} + +/// An instance of this resource is created when `JWKManager` starts the QC update building process for an issuer. +/// Then `JWKManager` needs to hold it. Once this resource is dropped, the corresponding QC update process will be cancelled. +#[derive(Clone, Debug)] +pub struct QuorumCertProcessGuard { + handle: Option, +} + +impl QuorumCertProcessGuard { + pub fn new(handle: AbortHandle) -> Self { + Self { + handle: Some(handle), + } + } + + #[cfg(test)] + pub fn dummy() -> Self { + Self { handle: None } + } +} + +impl Drop for QuorumCertProcessGuard { + fn drop(&mut self) { + let QuorumCertProcessGuard { handle } = self; + if let Some(handle) = handle { + handle.abort(); + } + } +} + +#[derive(Debug, Clone)] +pub enum ConsensusState { + NotStarted, + InProgress { + my_proposal: ObservedUpdate, + abort_handle_wrapper: QuorumCertProcessGuard, + }, + Finished { + vtxn_guard: TxnGuard, + my_proposal: ObservedUpdate, + quorum_certified: QuorumCertifiedUpdate, + }, +} + +impl PartialEq for ConsensusState { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (ConsensusState::NotStarted, ConsensusState::NotStarted) => true, + ( + ConsensusState::InProgress { + my_proposal: update_0, + .. + }, + ConsensusState::InProgress { + my_proposal: update_1, + .. + }, + ) if update_0 == update_1 => true, + ( + ConsensusState::Finished { + my_proposal: update_0, + .. + }, + ConsensusState::Finished { + my_proposal: update_1, + .. + }, + ) if update_0 == update_1 => true, + _ => false, + } + } +} + +impl Eq for ConsensusState {} + +impl ConsensusState { + pub fn name(&self) -> &str { + match self { + ConsensusState::NotStarted => "NotStarted", + ConsensusState::InProgress { .. } => "InProgress", + ConsensusState::Finished { .. } => "Finished", + } + } + + pub fn my_proposal_cloned(&self) -> ObservedUpdate { + match self { + ConsensusState::InProgress { my_proposal, .. } + | ConsensusState::Finished { my_proposal, .. } => my_proposal.clone(), + _ => panic!("my_proposal unavailable"), + } + } +} + +impl Default for ConsensusState { + fn default() -> Self { + Self::NotStarted + } +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct PerProviderState { + pub on_chain: Option, + pub observed: Option>, + pub consensus_state: ConsensusState, +} + +impl PerProviderState { + pub fn new(provider_jwks: ProviderJWKs) -> Self { + Self { + on_chain: Some(provider_jwks), + observed: None, + consensus_state: ConsensusState::NotStarted, + } + } + + pub fn on_chain_version(&self) -> u64 { + self.on_chain + .as_ref() + .map_or(0, |provider_jwks| provider_jwks.version) + } + + pub fn reset_with_onchain_state(&mut self, onchain_state: ProviderJWKs) { + if self.on_chain.as_ref() != Some(&onchain_state) { + *self = Self::new(onchain_state) + } + } +} + +#[cfg(test)] +pub mod tests; diff --git a/crates/aptos-jwk-consensus/src/jwk_manager/tests.rs b/crates/aptos-jwk-consensus/src/jwk_manager/tests.rs new file mode 100644 index 0000000000000..eb1c7428515ee --- /dev/null +++ b/crates/aptos-jwk-consensus/src/jwk_manager/tests.rs @@ -0,0 +1,475 @@ +// Copyright © Aptos Foundation + +use crate::{ + certified_update_producer::CertifiedUpdateProducer, + jwk_manager::{ConsensusState, JWKManager, PerProviderState, QuorumCertProcessGuard}, + network::{DummyRpcResponseSender, IncomingRpcRequest}, + types::{JWKConsensusMsg, ObservedUpdate, ObservedUpdateRequest, ObservedUpdateResponse}, +}; +use aptos_channels::aptos_channel; +use aptos_crypto::{ + bls12381::{PrivateKey, PublicKey, Signature}, + hash::CryptoHash, + SigningKey, Uniform, +}; +use aptos_infallible::{Mutex, RwLock}; +use aptos_types::{ + account_address::AccountAddress, + epoch_state::EpochState, + jwks::{ + issuer_from_str, jwk::JWK, unsupported::UnsupportedJWK, AllProvidersJWKs, Issuer, + ProviderJWKs, QuorumCertifiedUpdate, + }, + validator_txn::ValidatorTransaction, + validator_verifier::{ValidatorConsensusInfo, ValidatorVerifier}, +}; +use aptos_validator_transaction_pool::{TransactionFilter, VTxnPoolState}; +use futures_util::future::AbortHandle; +use std::{ + collections::{BTreeSet, HashMap, HashSet}, + sync::Arc, + time::{Duration, Instant}, +}; + +#[tokio::test] +async fn test_jwk_manager_state_transition() { + // Setting up an epoch of 4 validators, and simulate the JWKManager in validator 0. + let private_keys: Vec> = (0..4) + .map(|_| Arc::new(PrivateKey::generate_for_testing())) + .collect(); + let public_keys: Vec = private_keys + .iter() + .map(|sk| PublicKey::from(sk.as_ref())) + .collect(); + let addrs: Vec = (0..4).map(|_| AccountAddress::random()).collect(); + let voting_powers: Vec = vec![1, 1, 1, 1]; + let validator_consensus_infos: Vec = (0..4) + .map(|i| ValidatorConsensusInfo::new(addrs[i], public_keys[i].clone(), voting_powers[i])) + .collect(); + let epoch_state = EpochState { + epoch: 999, + verifier: ValidatorVerifier::new(validator_consensus_infos.clone()), + }; + + let certified_update_producer = DummyCertifiedUpdateProducer::default(); + let vtxn_pool = VTxnPoolState::default(); + let mut jwk_manager = JWKManager::new( + private_keys[0].clone(), + addrs[0], + Arc::new(epoch_state), + Arc::new(certified_update_producer), + vtxn_pool.clone(), + ); + + // In this example, Alice and Bob are 2 existing issuers; Carl was added in the last epoch so no JWKs of Carl is on chain. + let issuer_alice = issuer_from_str("https://alice.info"); + let issuer_bob = issuer_from_str("https://bob.io"); + let issuer_carl = issuer_from_str("https://carl.dev"); + let alice_jwks = vec![ + JWK::Unsupported(UnsupportedJWK::new_for_testing( + "alice_jwk_id_0", + "jwk_payload_0", + )) + .into(), + JWK::Unsupported(UnsupportedJWK::new_for_testing( + "alice_jwk_id_1", + "jwk_payload_1", + )) + .into(), + ]; + let bob_jwks = vec![ + JWK::Unsupported(UnsupportedJWK::new_for_testing( + "bob_jwk_id_0", + "jwk_payload_2", + )) + .into(), + JWK::Unsupported(UnsupportedJWK::new_for_testing( + "bob_jwk_id_1", + "jwk_payload_3", + )) + .into(), + ]; + let on_chain_state_alice_v111 = ProviderJWKs { + issuer: issuer_alice.clone(), + version: 111, + jwks: alice_jwks.clone(), + }; + + let on_chain_state_bob_v222 = ProviderJWKs { + issuer: issuer_bob.clone(), + version: 222, + jwks: bob_jwks.clone(), + }; + + let initial_on_chain_state = AllProvidersJWKs { + entries: vec![ + on_chain_state_alice_v111.clone(), + on_chain_state_bob_v222.clone(), + ], + }; + + // On start, JWKManager is always initialized with the on-chain state. + assert!(jwk_manager + .reset_with_on_chain_state(initial_on_chain_state) + .is_ok()); + let mut expected_states = HashMap::from([ + ( + issuer_alice.clone(), + PerProviderState::new(on_chain_state_alice_v111.clone()), + ), + ( + issuer_bob.clone(), + PerProviderState::new(on_chain_state_bob_v222.clone()), + ), + ]); + assert_eq!(expected_states, jwk_manager.states_by_issuer); + + let rpc_response_collector = Arc::new(RwLock::new(vec![])); + + // When JWK consensus is `NotStarted` for issuer Bob, JWKConsensusManager should: + // reply an error to any observation request and keep the state unchanged. + let bob_ob_req = new_rpc_observation_request( + 999, + issuer_bob.clone(), + addrs[3], + rpc_response_collector.clone(), + ); + assert!(jwk_manager.process_peer_request(bob_ob_req).is_ok()); + let last_invocations = std::mem::take(&mut *rpc_response_collector.write()); + assert!(last_invocations.len() == 1 && last_invocations[0].is_err()); + assert_eq!(expected_states, jwk_manager.states_by_issuer); + + // When JWK consensus is `NotStarted` for issuer Carl, JWKConsensusManager should: + // reply an error to any observation request and keep the state unchanged; + // also create an entry in the state table on the fly. + let carl_ob_req = new_rpc_observation_request( + 999, + issuer_carl.clone(), + addrs[3], + rpc_response_collector.clone(), + ); + assert!(jwk_manager.process_peer_request(carl_ob_req).is_ok()); + let last_invocations = std::mem::take(&mut *rpc_response_collector.write()); + assert!(last_invocations.len() == 1 && last_invocations[0].is_err()); + expected_states.insert(issuer_carl.clone(), PerProviderState::default()); + assert_eq!(expected_states, jwk_manager.states_by_issuer); + + // When JWK consensus is `NotStarted` for issuer Bob, JWKConsensusManager should: + // do nothing to an observation equal to on-chain state (except storing it, which may be unnecessary). + assert!(jwk_manager + .process_new_observation(issuer_bob.clone(), bob_jwks.clone()) + .is_ok()); + expected_states.get_mut(&issuer_bob).unwrap().observed = Some(bob_jwks.clone()); + assert_eq!(expected_states, jwk_manager.states_by_issuer); + + // When JWK consensus is `NotStarted` for issuer Alice, JWKConsensusManager should: + // initiate a JWK consensus session if an update was observed. + let alice_jwks_new = vec![JWK::Unsupported(UnsupportedJWK::new_for_testing( + "alice_jwk_id_1", + "jwk_payload_1", + )) + .into()]; + assert!(jwk_manager + .process_new_observation(issuer_alice.clone(), alice_jwks_new.clone()) + .is_ok()); + { + let expected_alice_state = expected_states.get_mut(&issuer_alice).unwrap(); + expected_alice_state.observed = Some(alice_jwks_new.clone()); + let observed = ProviderJWKs { + issuer: issuer_alice.clone(), + version: 112, // on-chain baseline is at version 111. + jwks: alice_jwks_new.clone(), + }; + let signature = private_keys[0].sign(&observed).unwrap(); + expected_alice_state.consensus_state = ConsensusState::InProgress { + my_proposal: ObservedUpdate { + author: addrs[0], + observed, + signature, + }, + abort_handle_wrapper: QuorumCertProcessGuard::dummy(), + }; + } + assert_eq!(expected_states, jwk_manager.states_by_issuer); + + // If we also found a JWK update for issuer Carl, a separate JWK consensus session should be started. + let carl_jwks_new = vec![JWK::Unsupported(UnsupportedJWK::new_for_testing( + "carl_jwk_id_0", + "jwk_payload_4", + )) + .into()]; + assert!(jwk_manager + .process_new_observation(issuer_carl.clone(), carl_jwks_new.clone()) + .is_ok()); + { + let expected_carl_state = expected_states.get_mut(&issuer_carl).unwrap(); + expected_carl_state.observed = Some(carl_jwks_new.clone()); + let observed = ProviderJWKs { + issuer: issuer_carl.clone(), + version: 1, + jwks: carl_jwks_new.clone(), + }; + let signature = private_keys[0].sign(&observed).unwrap(); + expected_carl_state.consensus_state = ConsensusState::InProgress { + my_proposal: ObservedUpdate { + author: addrs[0], + observed, + signature, + }, + abort_handle_wrapper: QuorumCertProcessGuard::dummy(), + }; + } + assert_eq!(expected_states, jwk_manager.states_by_issuer); + + // Now that there are in-progress consensus sessions for Alice/Carl, + // if receiving an observation request for issuer Alice/Carl, JWKConsensusManager should reply with their signed observation. + let alice_ob_req = new_rpc_observation_request( + 999, + issuer_alice.clone(), + addrs[3], + rpc_response_collector.clone(), + ); + let carl_ob_req = new_rpc_observation_request( + 999, + issuer_carl.clone(), + addrs[3], + rpc_response_collector.clone(), + ); + assert!(jwk_manager.process_peer_request(alice_ob_req).is_ok()); + assert!(jwk_manager.process_peer_request(carl_ob_req).is_ok()); + assert_eq!(expected_states, jwk_manager.states_by_issuer); + let last_invocations: Vec = + std::mem::take(&mut *rpc_response_collector.write()) + .into_iter() + .map(|maybe_msg| maybe_msg.unwrap()) + .collect(); + let expected_responses = vec![ + JWKConsensusMsg::ObservationResponse(ObservedUpdateResponse { + epoch: 999, + update: expected_states + .get(&issuer_alice) + .unwrap() + .consensus_state + .my_proposal_cloned(), + }), + JWKConsensusMsg::ObservationResponse(ObservedUpdateResponse { + epoch: 999, + update: expected_states + .get(&issuer_carl) + .unwrap() + .consensus_state + .my_proposal_cloned(), + }), + ]; + assert_eq!(expected_responses, last_invocations); + + // If Alice rotates again while the consensus session for Alice is in progress, the existing session should be discarded and a new session should start. + let alice_jwks_new_2 = vec![ + JWK::Unsupported(UnsupportedJWK::new_for_testing( + "alice_jwk_id_1", + "jwk_payload_1", + )) + .into(), + JWK::Unsupported(UnsupportedJWK::new_for_testing( + "alice_jwk_id_3", + "jwk_payload_5", + )) + .into(), + ]; + assert!(jwk_manager + .process_new_observation(issuer_alice.clone(), alice_jwks_new_2.clone()) + .is_ok()); + { + let expected_alice_state = expected_states.get_mut(&issuer_alice).unwrap(); + expected_alice_state.observed = Some(alice_jwks_new_2.clone()); + let observed = ProviderJWKs { + issuer: issuer_alice.clone(), + version: 112, + jwks: alice_jwks_new_2.clone(), + }; + let signature = private_keys[0].sign(&observed).unwrap(); + expected_alice_state.consensus_state = ConsensusState::InProgress { + my_proposal: ObservedUpdate { + author: addrs[0], + observed, + signature, + }, + abort_handle_wrapper: QuorumCertProcessGuard::dummy(), + }; + } + assert_eq!(expected_states, jwk_manager.states_by_issuer); + + // For issuer Carl, in state `InProgress`, when receiving a quorum-certified update from the the aggregator: + // the state should be switched to `Finished`; + // Carl's update should be available in validator txn pool. + let qc_jwks_for_carl = expected_states + .get(&issuer_carl) + .unwrap() + .consensus_state + .my_proposal_cloned() + .observed; + let multi_sig = Signature::aggregate( + private_keys + .iter() + .map(|sk| sk.sign(&qc_jwks_for_carl).unwrap()) + .collect::>(), + ) + .unwrap(); + let qc_update_for_carl = QuorumCertifiedUpdate { + authors: BTreeSet::from_iter(addrs.clone()), + update: qc_jwks_for_carl, + multi_sig: multi_sig.clone(), + }; + assert!(jwk_manager + .process_quorum_certified_update(qc_update_for_carl.clone()) + .is_ok()); + { + let expected_carl_state = expected_states.get_mut(&issuer_carl).unwrap(); + expected_carl_state.consensus_state = ConsensusState::Finished { + vtxn_guard: vtxn_pool.dummy_txn_guard(), + my_proposal: expected_carl_state.consensus_state.my_proposal_cloned(), + quorum_certified: qc_update_for_carl.clone(), + }; + } + assert_eq!(expected_states, jwk_manager.states_by_issuer); + let expected_vtxns = vec![ValidatorTransaction::ObservedJWKUpdate( + qc_update_for_carl.clone(), + )]; + let actual_vtxns = vtxn_pool.pull( + Instant::now() + Duration::from_secs(3600), + 999, + 2048, + TransactionFilter::empty(), + ); + assert_eq!(expected_vtxns, actual_vtxns); + + // For issuer Carl, in state 'Finished`, JWKConsensusManager should still reply to observation requests with its own proposal. + let carl_ob_req = new_rpc_observation_request( + 999, + issuer_carl.clone(), + addrs[3], + rpc_response_collector.clone(), + ); + assert!(jwk_manager.process_peer_request(carl_ob_req).is_ok()); + assert_eq!(expected_states, jwk_manager.states_by_issuer); + let expected_responses = vec![JWKConsensusMsg::ObservationResponse( + ObservedUpdateResponse { + epoch: 999, + update: expected_states + .get(&issuer_carl) + .unwrap() + .consensus_state + .my_proposal_cloned(), + }, + )]; + let actual_responses: Vec = + std::mem::take(&mut *rpc_response_collector.write()) + .into_iter() + .map(|maybe_msg| maybe_msg.unwrap()) + .collect(); + assert_eq!(expected_responses, actual_responses); + + // If the consensus session for Alice is also done, JWKConsensusManager should: + // update the state for Alice to `Finished`; + // update the validator txn in the pool to also include the update for Alice. + let qc_jwks_for_alice = expected_states + .get(&issuer_alice) + .unwrap() + .consensus_state + .my_proposal_cloned() + .observed; + let multi_sig = Signature::aggregate( + private_keys + .iter() + .take(3) + .map(|sk| sk.sign(&qc_jwks_for_alice).unwrap()) + .collect::>(), + ) + .unwrap(); + let qc_update_for_alice = QuorumCertifiedUpdate { + authors: BTreeSet::from_iter(addrs[0..3].to_vec()), + update: qc_jwks_for_alice, + multi_sig: multi_sig.clone(), + }; + assert!(jwk_manager + .process_quorum_certified_update(qc_update_for_alice.clone()) + .is_ok()); + { + let expected_alice_state = expected_states.get_mut(&issuer_alice).unwrap(); + expected_alice_state.consensus_state = ConsensusState::Finished { + vtxn_guard: vtxn_pool.dummy_txn_guard(), + my_proposal: expected_alice_state.consensus_state.my_proposal_cloned(), + quorum_certified: qc_update_for_alice.clone(), + }; + } + assert_eq!(expected_states, jwk_manager.states_by_issuer); + let expected_vtxn_hashes = vec![ + ValidatorTransaction::ObservedJWKUpdate(qc_update_for_alice), + ValidatorTransaction::ObservedJWKUpdate(qc_update_for_carl), + ] + .iter() + .map(CryptoHash::hash) + .collect::>(); + + let actual_vtxn_hashes = vtxn_pool + .pull( + Instant::now() + Duration::from_secs(3600), + 999, + 2048, + TransactionFilter::empty(), + ) + .iter() + .map(CryptoHash::hash) + .collect::>(); + assert_eq!(expected_vtxn_hashes, actual_vtxn_hashes); + + // At any time, JWKConsensusManager should fully follow on-chain update notification and re-initialize. + let second_on_chain_state = AllProvidersJWKs { + entries: vec![on_chain_state_alice_v111.clone()], + }; + + assert!(jwk_manager + .reset_with_on_chain_state(second_on_chain_state) + .is_ok()); + expected_states.remove(&issuer_bob); + expected_states.remove(&issuer_carl); + assert_eq!(expected_states, jwk_manager.states_by_issuer); +} + +fn new_rpc_observation_request( + epoch: u64, + issuer: Issuer, + sender: AccountAddress, + response_collector: Arc>>>, +) -> IncomingRpcRequest { + IncomingRpcRequest { + msg: JWKConsensusMsg::ObservationRequest(ObservedUpdateRequest { epoch, issuer }), + sender, + response_sender: Box::new(DummyRpcResponseSender::new(response_collector)), + } +} + +pub struct DummyCertifiedUpdateProducer { + pub invocations: Mutex, ProviderJWKs)>>, +} + +impl Default for DummyCertifiedUpdateProducer { + fn default() -> Self { + Self { + invocations: Mutex::new(vec![]), + } + } +} + +impl CertifiedUpdateProducer for DummyCertifiedUpdateProducer { + fn start_produce( + &self, + epoch_state: Arc, + payload: ProviderJWKs, + _agg_node_tx: Option>, + ) -> AbortHandle { + self.invocations.lock().push((epoch_state, payload)); + let (abort_handle, _) = AbortHandle::new_pair(); + abort_handle + } +} diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 25f214dcb76bb..60c5371f1f931 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -30,6 +30,7 @@ pub fn start_jwk_consensus_runtime( } pub mod certified_update_producer; +pub mod jwk_manager; pub mod jwk_observer; pub mod network; pub mod network_interface; From 166ef1519bd929f7272338188b349501eda8bb92 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:43:36 -0800 Subject: [PATCH 16/22] jwk consensus epoch manager --- .../aptos-jwk-consensus/src/epoch_manager.rs | 225 ++++++++++++++++++ crates/aptos-jwk-consensus/src/lib.rs | 1 + 2 files changed, 226 insertions(+) create mode 100644 crates/aptos-jwk-consensus/src/epoch_manager.rs diff --git a/crates/aptos-jwk-consensus/src/epoch_manager.rs b/crates/aptos-jwk-consensus/src/epoch_manager.rs new file mode 100644 index 0000000000000..8b5497f6fcf44 --- /dev/null +++ b/crates/aptos-jwk-consensus/src/epoch_manager.rs @@ -0,0 +1,225 @@ +// Copyright © Aptos Foundation + +use crate::{ + certified_update_producer::RealCertifiedUpdateProducer, + jwk_manager::JWKManager, + network::{IncomingRpcRequest, NetworkReceivers, NetworkSender}, + network_interface::JWKConsensusNetworkClient, + types::JWKConsensusMsg, +}; +use anyhow::Result; +use aptos_bounded_executor::BoundedExecutor; +use aptos_channels::{aptos_channel, message_queues::QueueStyle}; +use aptos_consensus_types::common::Author; +use aptos_crypto::bls12381::PrivateKey; +use aptos_event_notifications::{ + EventNotification, EventNotificationListener, ReconfigNotification, + ReconfigNotificationListener, +}; +use aptos_logger::{error, info}; +use aptos_network::{application::interface::NetworkClient, protocols::network::Event}; +use aptos_reliable_broadcast::ReliableBroadcast; +use aptos_types::{ + account_address::AccountAddress, + epoch_state::EpochState, + jwks::{ObservedJWKs, ObservedJWKsUpdated, SupportedOIDCProviders}, + on_chain_config::{ + FeatureFlag, Features, OnChainConfigPayload, OnChainConfigProvider, ValidatorSet, + }, +}; +use aptos_validator_transaction_pool::VTxnPoolState; +use futures::StreamExt; +use futures_channel::oneshot; +use std::{sync::Arc, time::Duration}; +use tokio_retry::strategy::ExponentialBackoff; + +pub struct EpochManager { + // some useful metadata + my_addr: AccountAddress, + epoch_state: Option>, + + // credential + consensus_key: Arc, + + // events we subscribe + reconfig_events: ReconfigNotificationListener

, + jwk_updated_events: EventNotificationListener, + + // message channels to JWK manager + jwk_updated_event_txs: Option>, + jwk_rpc_msg_tx: Option>, + jwk_manager_close_tx: Option>>, + + // network utils + self_sender: aptos_channels::Sender>, + network_sender: JWKConsensusNetworkClient>, + + // vtxn pool handle + vtxn_pool: VTxnPoolState, +} + +impl EpochManager

{ + pub fn new( + my_addr: AccountAddress, + consensus_key: PrivateKey, + reconfig_events: ReconfigNotificationListener

, + jwk_updated_events: EventNotificationListener, + self_sender: aptos_channels::Sender>, + network_sender: JWKConsensusNetworkClient>, + vtxn_pool: VTxnPoolState, + ) -> Self { + Self { + my_addr, + consensus_key: Arc::new(consensus_key), + epoch_state: None, + reconfig_events, + jwk_updated_events, + self_sender, + network_sender, + vtxn_pool, + jwk_updated_event_txs: None, + jwk_rpc_msg_tx: None, + jwk_manager_close_tx: None, + } + } + + /// On a new RPC request, forward to JWK consensus manager, if it is alive. + fn process_rpc_request( + &mut self, + peer_id: Author, + rpc_request: IncomingRpcRequest, + ) -> Result<()> { + if Some(rpc_request.msg.epoch()) == self.epoch_state.as_ref().map(|s| s.epoch) { + if let Some(tx) = &self.jwk_rpc_msg_tx { + let _ = tx.push((), (peer_id, rpc_request)); + } + } + Ok(()) + } + + /// On a on-chain JWK updated events, forward to JWK consensus manager if it is alive. + fn process_onchain_event(&mut self, notification: EventNotification) -> Result<()> { + let EventNotification { + subscribed_events, .. + } = notification; + for event in subscribed_events { + if let Ok(jwk_event) = ObservedJWKsUpdated::try_from(&event) { + if let Some(tx) = self.jwk_updated_event_txs.as_ref() { + let _ = tx.push((), jwk_event); + } + } + } + Ok(()) + } + + pub async fn start(mut self, mut network_receivers: NetworkReceivers) { + self.await_reconfig_notification().await; + loop { + let handle_result = tokio::select! { + reconfig_notification = self.reconfig_events.select_next_some() => { + self.on_new_epoch(reconfig_notification).await + }, + event = self.jwk_updated_events.select_next_some() => { + self.process_onchain_event(event) + }, + (peer, rpc_request) = network_receivers.rpc_rx.select_next_some() => { + self.process_rpc_request(peer, rpc_request) + } + }; + + if let Err(e) = handle_result { + error!("{}", e); + } + } + } + + async fn await_reconfig_notification(&mut self) { + let reconfig_notification = self + .reconfig_events + .next() + .await + .expect("Reconfig sender dropped, unable to start new epoch"); + self.start_new_epoch(reconfig_notification.on_chain_configs) + .await; + } + + async fn start_new_epoch(&mut self, payload: OnChainConfigPayload

) { + let validator_set: ValidatorSet = payload + .get() + .expect("failed to get ValidatorSet from payload"); + + let epoch_state = Arc::new(EpochState { + epoch: payload.epoch(), + verifier: (&validator_set).into(), + }); + self.epoch_state = Some(epoch_state.clone()); + info!("[JWK] start_new_epoch: new_epoch={}", epoch_state.epoch); + + let features = payload.get::().unwrap_or_default(); + + if features.is_enabled(FeatureFlag::JWK_CONSENSUS) { + let onchain_oidc_provider_set = payload.get::().ok(); + let onchain_observed_jwks = payload.get::().ok(); + info!("[JWK] JWK manager init, epoch={}", epoch_state.epoch); + let network_sender = NetworkSender::new( + self.my_addr, + self.network_sender.clone(), + self.self_sender.clone(), + ); + let rb = ReliableBroadcast::new( + epoch_state.verifier.get_ordered_account_addresses(), + Arc::new(network_sender), + ExponentialBackoff::from_millis(5), + aptos_time_service::TimeService::real(), + Duration::from_millis(1000), + BoundedExecutor::new(8, tokio::runtime::Handle::current()), + ); + let qc_update_producer = RealCertifiedUpdateProducer::new(rb); + + let jwk_consensus_manager = JWKManager::new( + self.consensus_key.clone(), + self.my_addr, + epoch_state.clone(), + Arc::new(qc_update_producer), + self.vtxn_pool.clone(), + ); + + let (jwk_event_tx, jwk_event_rx) = aptos_channel::new(QueueStyle::KLAST, 1, None); + self.jwk_updated_event_txs = Some(jwk_event_tx); + let (jwk_rpc_msg_tx, jwk_rpc_msg_rx) = aptos_channel::new(QueueStyle::FIFO, 100, None); + + let (jwk_manager_close_tx, jwk_manager_close_rx) = oneshot::channel(); + self.jwk_rpc_msg_tx = Some(jwk_rpc_msg_tx); + self.jwk_manager_close_tx = Some(jwk_manager_close_tx); + + tokio::spawn(jwk_consensus_manager.run( + onchain_oidc_provider_set, + onchain_observed_jwks, + jwk_event_rx, + jwk_rpc_msg_rx, + jwk_manager_close_rx, + )); + info!( + "jwk consensus manager spawned for epoch {}", + epoch_state.epoch + ); + } + } + + async fn on_new_epoch(&mut self, reconfig_notification: ReconfigNotification

) -> Result<()> { + self.shutdown_current_processor().await; + self.start_new_epoch(reconfig_notification.on_chain_configs) + .await; + Ok(()) + } + + async fn shutdown_current_processor(&mut self) { + if let Some(tx) = self.jwk_manager_close_tx.take() { + let (ack_tx, ack_rx) = oneshot::channel(); + let _ = tx.send(ack_tx); + let _ = ack_rx.await; + } + + self.jwk_updated_event_txs = None; + } +} diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 60c5371f1f931..9e7cc232321d1 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -30,6 +30,7 @@ pub fn start_jwk_consensus_runtime( } pub mod certified_update_producer; +pub mod epoch_manager; pub mod jwk_manager; pub mod jwk_observer; pub mod network; From f5295926829febf399161ee4d33765b4c60fc513 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:46:01 -0800 Subject: [PATCH 17/22] update --- types/src/contract_event.rs | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/types/src/contract_event.rs b/types/src/contract_event.rs index c3f1d9addc472..10fbdaf41aab4 100644 --- a/types/src/contract_event.rs +++ b/types/src/contract_event.rs @@ -6,6 +6,7 @@ use crate::{ account_config::{DepositEvent, NewBlockEvent, NewEpochEvent, WithdrawEvent}, dkg::DKGStartEvent, event::EventKey, + jwks::ObservedJWKsUpdated, on_chain_config::new_epoch_event_key, transaction::Version, }; @@ -360,6 +361,24 @@ impl TryFrom<&ContractEvent> for DepositEvent { } } +impl TryFrom<&ContractEvent> for ObservedJWKsUpdated { + type Error = Error; + + fn try_from(event: &ContractEvent) -> Result { + match event { + ContractEvent::V1(_) => { + bail!("conversion to `ObservedJWKsUpdated` failed with wrong event version") + }, + ContractEvent::V2(v2) => { + if v2.type_tag != TypeTag::Struct(Box::new(Self::struct_tag())) { + bail!("conversion to `ObservedJWKsUpdated` failed with wrong type tag"); + } + bcs::from_bytes(&v2.event_data).map_err(Into::into) + }, + } + } +} + impl std::fmt::Debug for ContractEvent { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { From 57bc3a182d8c5a6a7289ae51af8666857df344c4 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:53:30 -0800 Subject: [PATCH 18/22] jwk consensus wired into node --- Cargo.lock | 1 + aptos-node/Cargo.toml | 1 + aptos-node/src/lib.rs | 62 ++++++++++++++++------------ aptos-node/src/state_sync.rs | 4 +- config/src/config/identity_config.rs | 17 +------- execution/executor-types/src/lib.rs | 2 +- 6 files changed, 41 insertions(+), 46 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c6abe05be8c93..baa5ee74953cb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3081,6 +3081,7 @@ dependencies = [ "aptos-peer-monitoring-service-server", "aptos-peer-monitoring-service-types", "aptos-runtimes", + "aptos-safety-rules", "aptos-secure-storage", "aptos-state-sync-driver", "aptos-storage-interface", diff --git a/aptos-node/Cargo.toml b/aptos-node/Cargo.toml index 2b5ca01938778..9e87b7bc9cb7d 100644 --- a/aptos-node/Cargo.toml +++ b/aptos-node/Cargo.toml @@ -51,6 +51,7 @@ aptos-peer-monitoring-service-client = { workspace = true } aptos-peer-monitoring-service-server = { workspace = true } aptos-peer-monitoring-service-types = { workspace = true } aptos-runtimes = { workspace = true } +aptos-safety-rules = { workspace = true } aptos-secure-storage = { workspace = true } aptos-state-sync-driver = { workspace = true } aptos-storage-interface = { workspace = true } diff --git a/aptos-node/src/lib.rs b/aptos-node/src/lib.rs index a98369f493f4a..5a30624975797 100644 --- a/aptos-node/src/lib.rs +++ b/aptos-node/src/lib.rs @@ -25,6 +25,7 @@ use aptos_dkg_runtime::start_dkg_runtime; use aptos_framework::ReleaseBundle; use aptos_jwk_consensus::start_jwk_consensus_runtime; use aptos_logger::{prelude::*, telemetry_log_writer::TelemetryLog, Level, LoggerFilterUpdater}; +use aptos_safety_rules::safety_rules_manager::load_consensus_key_from_secure_storage; use aptos_state_sync_driver::driver_factory::StateSyncRuntimes; use aptos_types::chain_id::ChainId; use aptos_validator_transaction_pool::VTxnPoolState; @@ -657,17 +658,16 @@ pub fn setup_environment_and_start_node( mempool_client_receiver, peers_and_metadata, ); - let vtxn_pool = VTxnPoolState::default(); - let maybe_dkg_dealer_sk = node_config - .consensus - .safety_rules - .initial_safety_rules_config - .identity_blob() - .ok() - .and_then(|blob| blob.try_into_dkg_dealer_private_key()); + // Ensure consensus key in secure DB. + aptos_safety_rules::safety_rules_manager::storage(&node_config.consensus.safety_rules); + + let vtxn_pool = VTxnPoolState::default(); + let maybe_dkg_dealer_sk = + load_consensus_key_from_secure_storage(&node_config.consensus.safety_rules); + debug!("maybe_dkg_dealer_sk={:?}", maybe_dkg_dealer_sk); let dkg_runtime = match (dkg_network_interfaces, maybe_dkg_dealer_sk) { - (Some(interfaces), Some(dkg_dealer_sk)) => { + (Some(interfaces), Ok(dkg_dealer_sk)) => { let ApplicationNetworkInterfaces { network_client, network_service_events, @@ -689,24 +689,32 @@ pub fn setup_environment_and_start_node( _ => None, }; - let jwk_consensus_runtime = if let Some(obj) = jwk_consensus_network_interfaces { - let ApplicationNetworkInterfaces { - network_client, - network_service_events, - } = obj; - let (reconfig_events, onchain_jwk_updated_events) = jwk_consensus_subscriptions.expect( - "JWK consensus needs to listen to NewEpochEvents and OnChainJWKMapUpdated events.", - ); - let jwk_consensus_runtime = start_jwk_consensus_runtime( - network_client, - network_service_events, - vtxn_pool.clone(), - reconfig_events, - onchain_jwk_updated_events, - ); - Some(jwk_consensus_runtime) - } else { - None + let maybe_jwk_consensus_key = + load_consensus_key_from_secure_storage(&node_config.consensus.safety_rules); + debug!("maybe_jwk_consensus_key={:?}", maybe_jwk_consensus_key); + + let jwk_consensus_runtime = match (jwk_consensus_network_interfaces, maybe_jwk_consensus_key) { + (Some(interfaces), Ok(consensus_key)) => { + let ApplicationNetworkInterfaces { + network_client, + network_service_events, + } = interfaces; + let (reconfig_events, onchain_jwk_updated_events) = jwk_consensus_subscriptions.expect( + "JWK consensus needs to listen to NewEpochEvents and OnChainJWKMapUpdated events.", + ); + let my_addr = node_config.validator_network.as_ref().unwrap().peer_id(); + let jwk_consensus_runtime = start_jwk_consensus_runtime( + my_addr, + consensus_key, + network_client, + network_service_events, + reconfig_events, + onchain_jwk_updated_events, + vtxn_pool.clone(), + ); + Some(jwk_consensus_runtime) + }, + _ => None, }; // Create the consensus runtime (this blocks on state sync first) diff --git a/aptos-node/src/state_sync.rs b/aptos-node/src/state_sync.rs index 83cc9483e57e5..4ae8f6492f948 100644 --- a/aptos-node/src/state_sync.rs +++ b/aptos-node/src/state_sync.rs @@ -53,7 +53,7 @@ pub fn create_event_subscription_service( Option<( ReconfigNotificationListener, EventNotificationListener, - )>, // (reconfig_events, jwk_map_updated_events) for JWK consensus + )>, // (reconfig_events, jwk_updated_events) for JWK consensus ) { // Create the event subscription service let mut event_subscription_service = @@ -92,7 +92,7 @@ pub fn create_event_subscription_service( .subscribe_to_reconfigurations() .expect("JWK consensus must subscribe to reconfigurations"); let jwk_updated_events = event_subscription_service - .subscribe_to_events(vec![], vec!["0x1::jwks::OnChainJWKMapUpdated".to_string()]) + .subscribe_to_events(vec![], vec!["0x1::jwks::ObservedJWKsUpdated".to_string()]) .expect("JWK consensus must subscribe to DKG events"); Some((reconfig_events, jwk_updated_events)) } else { diff --git a/config/src/config/identity_config.rs b/config/src/config/identity_config.rs index 7026a35105505..f113fdb2c82ec 100644 --- a/config/src/config/identity_config.rs +++ b/config/src/config/identity_config.rs @@ -3,10 +3,7 @@ use crate::{config::SecureBackend, keys::ConfigKey}; use aptos_crypto::{bls12381, ed25519::Ed25519PrivateKey, x25519}; -use aptos_types::{ - account_address::{AccountAddress, AccountAddress as PeerId}, - dkg::{DKGTrait, DefaultDKG}, -}; +use aptos_types::account_address::{AccountAddress, AccountAddress as PeerId}; use serde::{Deserialize, Serialize}; use std::{ fs, @@ -40,18 +37,6 @@ impl IdentityBlob { let mut file = File::open(path)?; Ok(file.write_all(serde_yaml::to_string(self)?.as_bytes())?) } - - pub fn try_into_dkg_dealer_private_key( - self, - ) -> Option<::DealerPrivateKey> { - self.consensus_private_key - } - - pub fn try_into_dkg_new_validator_decrypt_key( - self, - ) -> Option<::NewValidatorDecryptKey> { - self.consensus_private_key - } } #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] diff --git a/execution/executor-types/src/lib.rs b/execution/executor-types/src/lib.rs index 0d588e86b4871..579fbca31280c 100644 --- a/execution/executor-types/src/lib.rs +++ b/execution/executor-types/src/lib.rs @@ -530,6 +530,6 @@ pub fn should_forward_to_subscription_service(event: &ContractEvent) -> bool { event.type_tag().to_string().as_str(), "0x1::reconfiguration::NewEpochEvent" | "0x1::dkg::DKGStartEvent" - | "0x1::jwks::OnChainJWKMapUpdated" + | "0x1::jwks::ObservedJWKsUpdated" ) } From 0f4de8d24050f5e3f826c372f39a7793ed4832a2 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:56:48 -0800 Subject: [PATCH 19/22] update --- crates/aptos-jwk-consensus/src/lib.rs | 43 +++++++++++++++++---------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 9e7cc232321d1..42ed3d47f9d72 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -1,31 +1,42 @@ // Copyright © Aptos Foundation -use crate::types::JWKConsensusMsg; +use crate::{ + epoch_manager::EpochManager, network::NetworkTask, network_interface::JWKConsensusNetworkClient, +}; +use aptos_crypto::bls12381::PrivateKey; use aptos_event_notifications::{ DbBackedOnChainConfig, EventNotificationListener, ReconfigNotificationListener, }; use aptos_network::application::interface::{NetworkClient, NetworkServiceEvents}; +use aptos_types::account_address::AccountAddress; use aptos_validator_transaction_pool::VTxnPoolState; -use futures_util::StreamExt; use tokio::runtime::Runtime; +use types::JWKConsensusMsg; -#[allow(clippy::let_and_return)] pub fn start_jwk_consensus_runtime( - _network_client: NetworkClient, - _network_service_events: NetworkServiceEvents, - _vtxn_pool: VTxnPoolState, - mut reconfig_events: ReconfigNotificationListener, - mut onchain_jwk_updated_events: EventNotificationListener, + my_addr: AccountAddress, + consensus_key: PrivateKey, + network_client: NetworkClient, + network_service_events: NetworkServiceEvents, + reconfig_events: ReconfigNotificationListener, + jwk_updated_events: EventNotificationListener, + vtxn_pool_writer: VTxnPoolState, ) -> Runtime { let runtime = aptos_runtimes::spawn_named_runtime("jwk".into(), Some(4)); - runtime.spawn(async move { - loop { - tokio::select! { - _ = reconfig_events.select_next_some() => {}, - _ = onchain_jwk_updated_events.select_next_some() => {}, - } - } - }); + let (self_sender, self_receiver) = aptos_channels::new(1_024, &counters::PENDING_SELF_MESSAGES); + let jwk_consensus_network_client = JWKConsensusNetworkClient::new(network_client); + let epoch_manager = EpochManager::new( + my_addr, + consensus_key, + reconfig_events, + jwk_updated_events, + self_sender, + jwk_consensus_network_client, + vtxn_pool_writer, + ); + let (network_task, network_receiver) = NetworkTask::new(network_service_events, self_receiver); + runtime.spawn(network_task.start()); + runtime.spawn(epoch_manager.start(network_receiver)); runtime } From f5a3471f72dc6889480b0dd21a5fa89842438e7b Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 07:58:28 -0800 Subject: [PATCH 20/22] update --- crates/aptos-jwk-consensus/src/counters.rs | 13 +++++++++++++ crates/aptos-jwk-consensus/src/lib.rs | 1 + 2 files changed, 14 insertions(+) create mode 100644 crates/aptos-jwk-consensus/src/counters.rs diff --git a/crates/aptos-jwk-consensus/src/counters.rs b/crates/aptos-jwk-consensus/src/counters.rs new file mode 100644 index 0000000000000..18ec2bf99a862 --- /dev/null +++ b/crates/aptos-jwk-consensus/src/counters.rs @@ -0,0 +1,13 @@ +// Copyright © Aptos Foundation + +use aptos_metrics_core::{register_int_gauge, IntGauge}; +use once_cell::sync::Lazy; + +/// Count of the pending messages sent to itself in the channel +pub static PENDING_SELF_MESSAGES: Lazy = Lazy::new(|| { + register_int_gauge!( + "aptos_jwk_consensus_pending_self_messages", + "Count of the pending JWK consensus messages sent to itself in the channel" + ) + .unwrap() +}); diff --git a/crates/aptos-jwk-consensus/src/lib.rs b/crates/aptos-jwk-consensus/src/lib.rs index 42ed3d47f9d72..6ed4127d58b47 100644 --- a/crates/aptos-jwk-consensus/src/lib.rs +++ b/crates/aptos-jwk-consensus/src/lib.rs @@ -41,6 +41,7 @@ pub fn start_jwk_consensus_runtime( } pub mod certified_update_producer; +pub mod counters; pub mod epoch_manager; pub mod jwk_manager; pub mod jwk_observer; From ca1bc9b2de0e1933902a3076e6374cbb957143e2 Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 08:02:14 -0800 Subject: [PATCH 21/22] jwk consensus smoke tests --- testsuite/forge/src/backend/local/cargo.rs | 4 +- testsuite/smoke-test/Cargo.toml | 1 + .../smoke-test/src/jwks/dummy_provider/mod.rs | 109 +++++++++++++ .../jwks/dummy_provider/request_handler.rs | 118 ++++++++++++++ .../src/jwks/jwk_consensus_basic.rs | 147 ++++++++++++++++++ .../src/jwks/jwk_consensus_per_issuer.rs | 94 +++++++++++ .../jwk_consensus_provider_change_mind.rs | 104 +++++++++++++ .../smoke-test/src/{jwks.rs => jwks/mod.rs} | 52 ++++++- 8 files changed, 624 insertions(+), 5 deletions(-) create mode 100644 testsuite/smoke-test/src/jwks/dummy_provider/mod.rs create mode 100644 testsuite/smoke-test/src/jwks/dummy_provider/request_handler.rs create mode 100644 testsuite/smoke-test/src/jwks/jwk_consensus_basic.rs create mode 100644 testsuite/smoke-test/src/jwks/jwk_consensus_per_issuer.rs create mode 100644 testsuite/smoke-test/src/jwks/jwk_consensus_provider_change_mind.rs rename testsuite/smoke-test/src/{jwks.rs => jwks/mod.rs} (61%) diff --git a/testsuite/forge/src/backend/local/cargo.rs b/testsuite/forge/src/backend/local/cargo.rs index cdb97b0a53dbc..93af06ac2d8c7 100644 --- a/testsuite/forge/src/backend/local/cargo.rs +++ b/testsuite/forge/src/backend/local/cargo.rs @@ -171,9 +171,9 @@ pub fn git_merge_base>(rev: R) -> Result { pub fn cargo_build_common_args() -> Vec<&'static str> { let mut args = if build_aptos_node_without_indexer() { - vec!["build", "--features=failpoints"] + vec!["build", "--features=failpoints,smoke-test"] } else { - vec!["build", "--features=failpoints,indexer"] + vec!["build", "--features=failpoints,indexer,smoke-test"] }; if build_consensus_only_node() { args.push("--features=consensus-only-perf-test"); diff --git a/testsuite/smoke-test/Cargo.toml b/testsuite/smoke-test/Cargo.toml index 3f9cb5338d7e0..acbb13ca84256 100644 --- a/testsuite/smoke-test/Cargo.toml +++ b/testsuite/smoke-test/Cargo.toml @@ -50,6 +50,7 @@ diesel = { workspace = true, features = [ "serde_json", ] } hex = { workspace = true } +hyper = { workspace = true } move-core-types = { workspace = true } proptest = { workspace = true } reqwest = { workspace = true } diff --git a/testsuite/smoke-test/src/jwks/dummy_provider/mod.rs b/testsuite/smoke-test/src/jwks/dummy_provider/mod.rs new file mode 100644 index 0000000000000..f82a6b9bee95c --- /dev/null +++ b/testsuite/smoke-test/src/jwks/dummy_provider/mod.rs @@ -0,0 +1,109 @@ +// Copyright © Aptos Foundation + +use aptos_infallible::RwLock; +use hyper::{ + service::{make_service_fn, service_fn}, + Body, Request, Response, Server, +}; +use request_handler::RequestHandler; +use std::{convert::Infallible, mem, net::SocketAddr, sync::Arc}; +use tokio::{ + sync::{ + oneshot, + oneshot::{Receiver, Sender}, + }, + task::JoinHandle, +}; + +pub(crate) mod request_handler; + +/// A dummy OIDC provider. +pub struct DummyProvider { + close_tx: Sender<()>, + open_id_config_url: String, + handler_holder: Arc>>>, + server_join_handle: JoinHandle<()>, +} + +impl DummyProvider { + pub(crate) async fn spawn() -> Self { + let addr = SocketAddr::from(([127, 0, 0, 1], 0)); + let handler_holder = Arc::new(RwLock::new(None)); + let (port_tx, port_rx) = oneshot::channel::(); + let (close_tx, close_rx) = oneshot::channel::<()>(); + let server_join_handle = tokio::spawn(Self::run_server( + addr, + handler_holder.clone(), + port_tx, + close_rx, + )); + let actual_port = port_rx.await.unwrap(); + let open_id_config_url = format!("http://127.0.0.1:{}", actual_port); + Self { + close_tx, + open_id_config_url, + handler_holder, + server_join_handle, + } + } + + pub fn open_id_config_url(&self) -> String { + self.open_id_config_url.clone() + } + + pub fn update_request_handler( + &self, + handler: Option>, + ) -> Option> { + mem::replace(&mut *self.handler_holder.write(), handler) + } + + pub async fn shutdown(self) { + let DummyProvider { + close_tx, + server_join_handle, + .. + } = self; + close_tx.send(()).unwrap(); + server_join_handle.await.unwrap(); + } +} + +// Private functions. +impl DummyProvider { + async fn run_server( + addr: SocketAddr, + handler_holder: Arc>>>, + port_tx: Sender, + close_rx: Receiver<()>, + ) { + let make_svc = make_service_fn(move |_| { + let handler_holder_clone = handler_holder.clone(); + async move { + Ok::<_, Infallible>(service_fn(move |req| { + Self::handle_request(req, handler_holder_clone.clone()) + })) + } + }); + + let server = Server::bind(&addr).serve(make_svc); + let actual_addr = server.local_addr(); + port_tx.send(actual_addr.port()).unwrap(); + + // Graceful shutdown + let graceful = server.with_graceful_shutdown(async { + close_rx.await.unwrap(); + }); + + graceful.await.unwrap(); + } + + async fn handle_request( + request: Request, + handler_holder: Arc>>>, + ) -> Result, Infallible> { + let handler = handler_holder.write(); + let raw_response = handler.as_ref().unwrap().handle(request); + Ok(Response::new(Body::from(raw_response))) + } +} diff --git a/testsuite/smoke-test/src/jwks/dummy_provider/request_handler.rs b/testsuite/smoke-test/src/jwks/dummy_provider/request_handler.rs new file mode 100644 index 0000000000000..5310931ef1af7 --- /dev/null +++ b/testsuite/smoke-test/src/jwks/dummy_provider/request_handler.rs @@ -0,0 +1,118 @@ +// Copyright © Aptos Foundation + +use aptos_infallible::Mutex; +use hyper::{Body, Request}; +use move_core_types::account_address::AccountAddress; +use std::{collections::HashSet, str::FromStr}; + +/// A handler that handles JWK requests from a validator, +/// assuming the validator account address is written as the COOKIE. +pub trait RequestHandler: Send + Sync { + fn handle(&self, request: Request) -> Vec; +} + +pub struct StaticContentServer { + content: Vec, +} + +impl StaticContentServer { + pub fn new(content: Vec) -> Self { + Self { content } + } + + pub fn new_str(content: &str) -> Self { + Self::new(content.as_bytes().to_vec()) + } +} + +impl RequestHandler for StaticContentServer { + fn handle(&self, _origin: Request) -> Vec { + self.content.clone() + } +} + +fn origin_from_cookie(request: &Request) -> AccountAddress { + let cookie = request + .headers() + .get(hyper::header::COOKIE) + .unwrap() + .to_str() + .unwrap(); + AccountAddress::from_str(cookie).unwrap() +} + +/// The first `k` requesters will get content A forever, the rest will get content B forever. +pub struct EquivocatingServer { + content_a: Vec, + content_b: Vec, + k: usize, + requesters_observed: Mutex>, +} + +impl EquivocatingServer { + pub fn new(content_a: Vec, content_b: Vec, k: usize) -> Self { + Self { + content_a, + content_b, + k, + requesters_observed: Mutex::new(HashSet::new()), + } + } +} + +impl RequestHandler for EquivocatingServer { + fn handle(&self, request: Request) -> Vec { + let mut requesters_observed = self.requesters_observed.lock(); + let origin = origin_from_cookie(&request); + if requesters_observed.len() < self.k { + requesters_observed.insert(origin); + } + + if requesters_observed.contains(&origin) { + self.content_a.clone() + } else { + self.content_b.clone() + } + } +} + +/// This server first replies with `initial_thoughts`. +/// After enough audience receives it for at least once, it switches its reply to `second_thoughts`. +/// +/// This behavior simulates the situation where a provider performs a 2nd key rotation right after the 1st. +pub struct MindChangingServer { + initial_thoughts: Vec, + second_thoughts: Vec, + change_mind_threshold: usize, + requesters_observed: Mutex>, +} + +impl MindChangingServer { + pub fn new( + initial_thoughts: Vec, + second_thoughts: Vec, + change_mind_threshold: usize, + ) -> Self { + Self { + initial_thoughts, + second_thoughts, + change_mind_threshold, + requesters_observed: Mutex::new(HashSet::new()), + } + } +} + +impl RequestHandler for MindChangingServer { + fn handle(&self, request: Request) -> Vec { + let mut requesters_observed = self.requesters_observed.lock(); + let origin = origin_from_cookie(&request); + if requesters_observed.contains(&origin) + || requesters_observed.len() >= self.change_mind_threshold + { + self.second_thoughts.clone() + } else { + requesters_observed.insert(origin); + self.initial_thoughts.clone() + } + } +} diff --git a/testsuite/smoke-test/src/jwks/jwk_consensus_basic.rs b/testsuite/smoke-test/src/jwks/jwk_consensus_basic.rs new file mode 100644 index 0000000000000..0df449657d22b --- /dev/null +++ b/testsuite/smoke-test/src/jwks/jwk_consensus_basic.rs @@ -0,0 +1,147 @@ +// Copyright © Aptos Foundation + +use crate::{ + jwks::{ + dummy_provider::{ + request_handler::{EquivocatingServer, StaticContentServer}, + DummyProvider, + }, + get_patched_jwks, put_provider_on_chain, + }, + smoke_test_environment::SwarmBuilder, +}; +use aptos_forge::{NodeExt, Swarm, SwarmExt}; +use aptos_logger::{debug, info}; +use aptos_types::jwks::{ + jwk::JWK, rsa::RSA_JWK, unsupported::UnsupportedJWK, AllProvidersJWKs, OIDCProvider, + ProviderJWKs, +}; +use std::{sync::Arc, time::Duration}; +use tokio::time::sleep; + +/// The validators should agree on the JWK after provider set is changed/JWK is rotated. +#[tokio::test] +async fn jwk_consensus_basic() { + let epoch_duration_secs = 30; + + let (mut swarm, mut cli, _faucet) = SwarmBuilder::new_local(4) + .with_aptos() + .with_init_genesis_config(Arc::new(move |conf| { + conf.epoch_duration_secs = epoch_duration_secs; + })) + .build_with_cli(0) + .await; + let client = swarm.validators().next().unwrap().rest_client(); + let root_idx = cli.add_account_with_address_to_cli( + swarm.root_key(), + swarm.chain_info().root_account().address(), + ); + swarm + .wait_for_all_nodes_to_catchup_to_epoch(2, Duration::from_secs(epoch_duration_secs * 2)) + .await + .expect("Epoch 2 taking too long to arrive!"); + + info!("Initially the provider set is empty. So should be the JWK map."); + + sleep(Duration::from_secs(10)).await; + let patched_jwks = get_patched_jwks(&client).await; + debug!("patched_jwks={:?}", patched_jwks); + assert!(patched_jwks.jwks.entries.is_empty()); + + info!("Adding some providers."); + let (provider_alice, provider_bob) = + tokio::join!(DummyProvider::spawn(), DummyProvider::spawn()); + + provider_alice.update_request_handler(Some(Arc::new(StaticContentServer::new_str( + r#" +{ + "keys": [ + {"kid":"kid1", "kty":"RSA", "e":"AQAB", "n":"n1", "alg":"RS384", "use":"sig"}, + {"n":"n0", "kty":"RSA", "use":"sig", "alg":"RS256", "e":"AQAB", "kid":"kid0"} + ] +} +"#, + )))); + provider_bob.update_request_handler(Some(Arc::new(StaticContentServer::new( + r#"{"keys": ["BOB_JWK_V0"]}"#.as_bytes().to_vec(), + )))); + let providers = vec![ + OIDCProvider { + name: b"https://alice.io".to_vec(), + config_url: provider_alice.open_id_config_url().into_bytes(), + }, + OIDCProvider { + name: b"https://bob.dev".to_vec(), + config_url: provider_bob.open_id_config_url().into_bytes(), + }, + ]; + let txn_summary = put_provider_on_chain(cli, root_idx, providers).await; + debug!("txn_summary={:?}", txn_summary); + + info!("Waiting for an on-chain update. 10 sec should be enough."); + sleep(Duration::from_secs(10)).await; + let patched_jwks = get_patched_jwks(&client).await; + debug!("patched_jwks={:?}", patched_jwks); + assert_eq!( + AllProvidersJWKs { + entries: vec![ + ProviderJWKs { + issuer: b"https://alice.io".to_vec(), + version: 1, + jwks: vec![ + JWK::RSA(RSA_JWK::new_256_aqab("kid0", "n0")).into(), + JWK::RSA(RSA_JWK::new_from_strs("kid1", "RSA", "RS384", "AQAB", "n1")) + .into(), + ], + }, + ProviderJWKs { + issuer: b"https://bob.dev".to_vec(), + version: 1, + jwks: vec![JWK::Unsupported(UnsupportedJWK::new_with_payload( + "\"BOB_JWK_V0\"" + )) + .into()], + }, + ] + }, + patched_jwks.jwks + ); + + info!("Rotating Alice keys. Also making https://alice.io gently equivocate."); + provider_alice.update_request_handler(Some(Arc::new(EquivocatingServer::new( + r#"{"keys": ["ALICE_JWK_V1A"]}"#.as_bytes().to_vec(), + r#"{"keys": ["ALICE_JWK_V1B"]}"#.as_bytes().to_vec(), + 1, + )))); + + info!("Waiting for an on-chain update. 30 sec should be enough."); + sleep(Duration::from_secs(30)).await; + let patched_jwks = get_patched_jwks(&client).await; + debug!("patched_jwks={:?}", patched_jwks); + assert_eq!( + AllProvidersJWKs { + entries: vec![ + ProviderJWKs { + issuer: b"https://alice.io".to_vec(), + version: 2, + jwks: vec![JWK::Unsupported(UnsupportedJWK::new_with_payload( + "\"ALICE_JWK_V1B\"" + )) + .into()], + }, + ProviderJWKs { + issuer: b"https://bob.dev".to_vec(), + version: 1, + jwks: vec![JWK::Unsupported(UnsupportedJWK::new_with_payload( + "\"BOB_JWK_V0\"" + )) + .into()], + }, + ] + }, + patched_jwks.jwks + ); + + info!("Tear down."); + provider_alice.shutdown().await; +} diff --git a/testsuite/smoke-test/src/jwks/jwk_consensus_per_issuer.rs b/testsuite/smoke-test/src/jwks/jwk_consensus_per_issuer.rs new file mode 100644 index 0000000000000..e271c2bb9cd67 --- /dev/null +++ b/testsuite/smoke-test/src/jwks/jwk_consensus_per_issuer.rs @@ -0,0 +1,94 @@ +// Copyright © Aptos Foundation + +use crate::{ + jwks::{ + dummy_provider::{ + request_handler::{EquivocatingServer, StaticContentServer}, + DummyProvider, + }, + get_patched_jwks, put_provider_on_chain, + }, + smoke_test_environment::SwarmBuilder, +}; +use aptos_forge::{NodeExt, Swarm, SwarmExt}; +use aptos_logger::{debug, info}; +use aptos_types::jwks::{ + jwk::JWK, unsupported::UnsupportedJWK, AllProvidersJWKs, OIDCProvider, ProviderJWKs, +}; +use std::{sync::Arc, time::Duration}; +use tokio::time::sleep; + +/// The validators should do JWK consensus per issuer: +/// one problematic issuer should not block valid updates of other issuers. +#[tokio::test] +async fn jwk_consensus_per_issuer() { + let epoch_duration_secs = 30; + + let (mut swarm, mut cli, _faucet) = SwarmBuilder::new_local(4) + .with_aptos() + .with_init_genesis_config(Arc::new(move |conf| { + conf.epoch_duration_secs = epoch_duration_secs; + })) + .build_with_cli(0) + .await; + let client = swarm.validators().next().unwrap().rest_client(); + let root_idx = cli.add_account_with_address_to_cli( + swarm.root_key(), + swarm.chain_info().root_account().address(), + ); + swarm + .wait_for_all_nodes_to_catchup_to_epoch(2, Duration::from_secs(epoch_duration_secs * 2)) + .await + .expect("Epoch 2 taking too long to arrive!"); + + info!("Initially the provider set is empty. So should be the JWK map."); + + sleep(Duration::from_secs(10)).await; + let patched_jwks = get_patched_jwks(&client).await; + debug!("patched_jwks={:?}", patched_jwks); + assert!(patched_jwks.jwks.entries.is_empty()); + + info!("Adding some providers, one seriously equivocating, the other well behaving."); + let (provider_alice, provider_bob) = + tokio::join!(DummyProvider::spawn(), DummyProvider::spawn()); + provider_alice.update_request_handler(Some(Arc::new(EquivocatingServer::new( + r#"{"keys": ["ALICE_JWK_V1A"]}"#.as_bytes().to_vec(), + r#"{"keys": ["ALICE_JWK_V1B"]}"#.as_bytes().to_vec(), + 2, + )))); + provider_bob.update_request_handler(Some(Arc::new(StaticContentServer::new( + r#"{"keys": ["BOB_JWK_V0"]}"#.as_bytes().to_vec(), + )))); + let providers = vec![ + OIDCProvider { + name: b"https://alice.io".to_vec(), + config_url: provider_alice.open_id_config_url().into_bytes(), + }, + OIDCProvider { + name: b"https://bob.dev".to_vec(), + config_url: provider_bob.open_id_config_url().into_bytes(), + }, + ]; + let txn_summary = put_provider_on_chain(cli, root_idx, providers).await; + debug!("txn_summary={:?}", txn_summary); + + info!("Wait for 60 secs and there should only update for Bob, not Alice."); + sleep(Duration::from_secs(60)).await; + let patched_jwks = get_patched_jwks(&client).await; + debug!("patched_jwks={:?}", patched_jwks); + assert_eq!( + AllProvidersJWKs { + entries: vec![ProviderJWKs { + issuer: b"https://bob.dev".to_vec(), + version: 1, + jwks: vec![ + JWK::Unsupported(UnsupportedJWK::new_with_payload("\"BOB_JWK_V0\"")).into() + ], + }] + }, + patched_jwks.jwks + ); + + info!("Tear down."); + provider_alice.shutdown().await; +} diff --git a/testsuite/smoke-test/src/jwks/jwk_consensus_provider_change_mind.rs b/testsuite/smoke-test/src/jwks/jwk_consensus_provider_change_mind.rs new file mode 100644 index 0000000000000..7fad520c48ce1 --- /dev/null +++ b/testsuite/smoke-test/src/jwks/jwk_consensus_provider_change_mind.rs @@ -0,0 +1,104 @@ +// Copyright © Aptos Foundation + +use crate::{ + jwks::{ + dummy_provider::{ + request_handler::{MindChangingServer, StaticContentServer}, + DummyProvider, + }, + get_patched_jwks, put_provider_on_chain, + }, + smoke_test_environment::SwarmBuilder, +}; +use aptos_forge::{NodeExt, Swarm, SwarmExt}; +use aptos_logger::{debug, info}; +use aptos_types::jwks::{ + jwk::JWK, unsupported::UnsupportedJWK, AllProvidersJWKs, OIDCProvider, ProviderJWKs, +}; +use std::{sync::Arc, time::Duration}; +use tokio::time::sleep; + +/// The validators should be able to reach JWK consensus +/// even if a provider double-rotates its key in a very short period of time. +/// First rotation may have been observed by some validators. +#[tokio::test] +async fn jwk_consensus_provider_change_mind() { + // Big epoch duration to ensure epoch change does not help reset validators if they are stuck. + let epoch_duration_secs = 1800; + + let (mut swarm, mut cli, _faucet) = SwarmBuilder::new_local(4) + .with_aptos() + .with_init_genesis_config(Arc::new(move |conf| { + conf.epoch_duration_secs = epoch_duration_secs; + })) + .build_with_cli(0) + .await; + let client = swarm.validators().next().unwrap().rest_client(); + let root_idx = cli.add_account_with_address_to_cli( + swarm.root_key(), + swarm.chain_info().root_account().address(), + ); + swarm + .wait_for_all_nodes_to_catchup_to_epoch(2, Duration::from_secs(epoch_duration_secs * 2)) + .await + .expect("Epoch 2 taking too long to arrive!"); + + info!("Initially the provider set is empty. So should be the ObservedJWKs."); + + sleep(Duration::from_secs(10)).await; + let patched_jwks = get_patched_jwks(&client).await; + debug!("patched_jwks={:?}", patched_jwks); + assert!(patched_jwks.jwks.entries.is_empty()); + + info!("Adding some providers."); + let (provider_alice, provider_bob) = + tokio::join!(DummyProvider::spawn(), DummyProvider::spawn()); + provider_alice.update_request_handler(Some(Arc::new(StaticContentServer::new( + r#"{"keys": ["ALICE_JWK_V0"]}"#.as_bytes().to_vec(), + )))); + provider_bob.update_request_handler(Some(Arc::new(MindChangingServer::new( + r#"{"keys": ["BOB_JWK_V0"]}"#.as_bytes().to_vec(), + r#"{"keys": ["BOB_JWK_V0_1"]}"#.as_bytes().to_vec(), + 2, + )))); + let providers = vec![ + OIDCProvider { + name: b"https://alice.io".to_vec(), + config_url: provider_alice.open_id_config_url().into_bytes(), + }, + OIDCProvider { + name: b"https://bob.dev".to_vec(), + config_url: provider_bob.open_id_config_url().into_bytes(), + }, + ]; + let txn_summary = put_provider_on_chain(cli, root_idx, providers).await; + debug!("txn_summary={:?}", txn_summary); + + info!("Waiting for an on-chain update. 30 secs should be enough."); + sleep(Duration::from_secs(30)).await; + let patched_jwks = get_patched_jwks(&client).await; + debug!("patched_jwks={:?}", patched_jwks); + assert_eq!( + AllProvidersJWKs { + entries: vec![ + ProviderJWKs { + issuer: b"https://alice.io".to_vec(), + version: 1, + jwks: vec![JWK::Unsupported(UnsupportedJWK::new_with_payload( + "\"ALICE_JWK_V0\"" + )) + .into()], + }, + ProviderJWKs { + issuer: b"https://bob.dev".to_vec(), + version: 1, + jwks: vec![JWK::Unsupported(UnsupportedJWK::new_with_payload( + "\"BOB_JWK_V0_1\"" + )) + .into()], + }, + ] + }, + patched_jwks.jwks + ); +} diff --git a/testsuite/smoke-test/src/jwks.rs b/testsuite/smoke-test/src/jwks/mod.rs similarity index 61% rename from testsuite/smoke-test/src/jwks.rs rename to testsuite/smoke-test/src/jwks/mod.rs index 6509a35b1d32c..2740ef1fc7df9 100644 --- a/testsuite/smoke-test/src/jwks.rs +++ b/testsuite/smoke-test/src/jwks/mod.rs @@ -1,18 +1,64 @@ // Copyright © Aptos Foundation +mod dummy_provider; +mod jwk_consensus_basic; +mod jwk_consensus_per_issuer; +mod jwk_consensus_provider_change_mind; + use crate::smoke_test_environment::SwarmBuilder; +use aptos::{common::types::TransactionSummary, test::CliTestFramework}; use aptos_forge::{NodeExt, Swarm, SwarmExt}; use aptos_logger::{debug, info}; use aptos_rest_client::Client; use aptos_types::jwks::{ jwk::{JWKMoveStruct, JWK}, unsupported::UnsupportedJWK, - AllProvidersJWKs, PatchedJWKs, ProviderJWKs, + AllProvidersJWKs, OIDCProvider, PatchedJWKs, ProviderJWKs, }; use move_core_types::account_address::AccountAddress; use std::time::Duration; -async fn get_latest_jwkset(rest_client: &Client) -> PatchedJWKs { +pub async fn put_provider_on_chain( + cli: CliTestFramework, + account_idx: usize, + providers: Vec, +) -> TransactionSummary { + let implementation = providers + .into_iter() + .map(|provider| { + let OIDCProvider { name, config_url } = provider; + format!( + r#" + let issuer = b"{}"; + let config_url = b"{}"; + jwks::upsert_oidc_provider(&framework_signer, issuer, config_url); +"#, + String::from_utf8(name).unwrap(), + String::from_utf8(config_url).unwrap(), + ) + }) + .collect::>() + .join(""); + + let add_dummy_provider_script = format!( + r#" +script {{ + use aptos_framework::aptos_governance; + use aptos_framework::jwks; + fun main(core_resources: &signer) {{ + let framework_signer = aptos_governance::get_signer_testnet_only(core_resources, @0000000000000000000000000000000000000000000000000000000000000001); + {implementation} + aptos_governance::reconfigure(&framework_signer); + }} +}} +"#, + ); + cli.run_script(account_idx, &add_dummy_provider_script) + .await + .unwrap() +} + +async fn get_patched_jwks(rest_client: &Client) -> PatchedJWKs { let maybe_response = rest_client .get_account_resource_bcs::(AccountAddress::ONE, "0x1::jwks::PatchedJWKs") .await; @@ -58,7 +104,7 @@ script { debug!("txn_summary={:?}", txn_summary); info!("Use resource API to check the patch result."); - let patched_jwks = get_latest_jwkset(&client).await; + let patched_jwks = get_patched_jwks(&client).await; debug!("patched_jwks={:?}", patched_jwks); let expected_providers_jwks = AllProvidersJWKs { From 3df6061933cd4013a0a7cd6bd07c97ea0d13f0df Mon Sep 17 00:00:00 2001 From: "zhoujun.ma" Date: Thu, 1 Feb 2024 16:05:36 +0000 Subject: [PATCH 22/22] update --- Cargo.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/Cargo.lock b/Cargo.lock index baa5ee74953cb..4ebcc29d4d4b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14906,6 +14906,7 @@ dependencies = [ "diesel", "futures", "hex", + "hyper", "move-core-types", "num_cpus", "once_cell",