diff --git a/keyless/pepper/common/src/lib.rs b/keyless/pepper/common/src/lib.rs index bf579e0ad396e..8f92714536222 100644 --- a/keyless/pepper/common/src/lib.rs +++ b/keyless/pepper/common/src/lib.rs @@ -223,7 +223,7 @@ impl PepperV0VufPubKey { } } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct PepperInput { pub iss: String, pub aud: String, diff --git a/keyless/pepper/service/src/accounts/account_recovery_db.rs b/keyless/pepper/service/src/accounts/account_recovery_db.rs index 9f3f28f823d82..b4007b7ca9936 100644 --- a/keyless/pepper/service/src/accounts/account_recovery_db.rs +++ b/keyless/pepper/service/src/accounts/account_recovery_db.rs @@ -1,12 +1,12 @@ // Copyright (c) Aptos Foundation // SPDX-License-Identifier: Apache-2.0 -use crate::error::PepperServiceError; +use crate::{error::PepperServiceError, metrics}; use aptos_infallible::duration_since_epoch; use aptos_keyless_pepper_common::{account_recovery_db::AccountRecoveryDbEntry, PepperInput}; -use aptos_logger::info; +use aptos_logger::{error, info}; use firestore::{path, paths, FirestoreDb, FirestoreDbOptions}; -use std::io::Write; +use std::{io::Write, time::Instant}; use tempfile::NamedTempFile; // Firestore DB transaction constants @@ -24,11 +24,19 @@ pub trait AccountRecoveryDBInterface { /// A GCP firestore that holds all account recovery data pub struct FirestoreAccountRecoveryDB { + /// Whether to disable async DB updates (e.g., for testing purposes) + disable_async_db_updates: bool, + + /// The Firestore DB instance firestore_db: FirestoreDb, } impl FirestoreAccountRecoveryDB { - pub async fn new(google_project_id: String, firestore_database_id: String) -> Self { + pub async fn new( + google_project_id: String, + firestore_database_id: String, + disable_async_db_updates: bool, + ) -> Self { // Create the FirestoreDB options let firestore_db_options = FirestoreDbOptions { google_project_id, @@ -39,7 +47,10 @@ impl FirestoreAccountRecoveryDB { // Create the FirestoreDB instance match FirestoreDb::with_options(firestore_db_options).await { - Ok(firestore_db) => Self { firestore_db }, + Ok(firestore_db) => Self { + disable_async_db_updates, + firestore_db, + }, Err(error) => { panic!( "Failed to create Firestore account recovery database! Error: {}", @@ -56,89 +67,36 @@ impl AccountRecoveryDBInterface for FirestoreAccountRecoveryDB { &self, pepper_input: &PepperInput, ) -> Result<(), PepperServiceError> { - // Create the database entry and get the document ID - let entry = AccountRecoveryDbEntry { - iss: pepper_input.iss.clone(), - aud: pepper_input.aud.clone(), - uid_key: pepper_input.uid_key.clone(), - uid_val: pepper_input.uid_val.clone(), - first_request_unix_ms_minus_1q: None, - last_request_unix_ms: None, - num_requests: None, + let firestore_db = self.firestore_db.clone(); + let pepper_input = pepper_input.clone(); + + // Create the DB update task + let db_update_task = async move { + // Start the db update timer + let db_update_timer = Instant::now(); + + // Update the firestore DB + let update_result = + update_firestore_with_pepper_input(firestore_db, pepper_input.clone()).await; + if let Err(error) = &update_result { + error!( + "Failed to update the firestore database for pepper input: {:?}! Error: {:?}", + pepper_input, error + ); + } + + // Update the write metrics + metrics::update_account_recovery_db_metrics(update_result.is_ok(), db_update_timer); + + update_result }; - let document_id = entry.document_id(); - - // Get the time now - let now_unix_ms = duration_since_epoch().as_millis() as i64; - - // To update the DB, we use the following strategy: - // 1. If the document doesn't exist, create the document for the user identifier `(iss, aud, uid_key, uid_val)`, - // but leave counter/time fields unspecified. - // 2. `num_requests += 1`, assuming the default value is 0. - // 3. `last_request_unix_ms = max(last_request_unix_ms, now)`, assuming the default value is 0. - // 4. `first_request_unix_ms = min(first_request_unix_ms, now)`, assuming the default value is +inf. - // - // This strategy is preferred because all the operations can be done on the server-side, - // which means the transaction should require only 1 RTT (this is better than using - // the read-compute-write pattern that requires 2 RTTs). - // - // Note: this strategy requires some modifications: - // - In firestore, the default value of a number field is 0, and we do not have a way to customize it for step 4. - // - The workaround here is to apply an offset so 0 becomes a legitimate default value. - // - So we work with `first_request_unix_ms_minus_1q` instead, which is defined as - // `first_request_unix_ms - 1_000_000_000_000_000`, where 1_000_000_000_000_000 milliseconds is roughly 31710 years. - - // Create the firestore DB transaction - let mut firestore_transaction = - self.firestore_db - .begin_transaction() - .await - .map_err(|error| { - PepperServiceError::InternalError(format!( - "Firestore DB begin_transaction() error: {}", - error - )) - })?; - self.firestore_db - .fluent() - .update() - .fields(paths!(AccountRecoveryDbEntry::{iss, aud, uid_key, uid_val})) - .in_col(FIRESTORE_DB_COLLECTION_ID) - .document_id(&document_id) - .object(&entry) // Step 1 - .transforms(|builder| { - builder.fields([ - builder - .field(path!(AccountRecoveryDbEntry::num_requests)) - .increment(1), // Step 2 - builder - .field(path!(AccountRecoveryDbEntry::last_request_unix_ms)) - .maximum(now_unix_ms), // Step 3 - builder - .field(path!( - AccountRecoveryDbEntry::first_request_unix_ms_minus_1q - )) - .minimum(now_unix_ms - 1_000_000_000_000_000), // Step 4 - ]) - }) - .add_to_transaction(&mut firestore_transaction) - .map_err(|error| { - PepperServiceError::InternalError(format!( - "Firestore DB add_to_transaction() error: {}", - error - )) - })?; - // Commit the DB transaction - match firestore_transaction.commit().await { - Ok(_) => Ok(()), - Err(error) => { - let error_message = format!( - "Failed to commit Firestore DB transaction for pepper input {:?}! Document ID {}, Error: {}", - pepper_input, document_id, error - ); - Err(PepperServiceError::InternalError(error_message)) - }, + // Run the DB update task + if self.disable_async_db_updates { + db_update_task.await // Run the task and wait for the result + } else { + tokio::spawn(db_update_task); // Run the task asynchronously + Ok(()) } } } @@ -200,3 +158,90 @@ impl Default for TestAccountRecoveryDB { Self::new() } } + +/// Updates the firestore DB with the given pepper input +async fn update_firestore_with_pepper_input( + firestore_db: FirestoreDb, + pepper_input: PepperInput, +) -> Result<(), PepperServiceError> { + // Create the database entry and get the document ID + let entry = AccountRecoveryDbEntry { + iss: pepper_input.iss.clone(), + aud: pepper_input.aud.clone(), + uid_key: pepper_input.uid_key.clone(), + uid_val: pepper_input.uid_val.clone(), + first_request_unix_ms_minus_1q: None, + last_request_unix_ms: None, + num_requests: None, + }; + let document_id = entry.document_id(); + + // Get the time now + let now_unix_ms = duration_since_epoch().as_millis() as i64; + + // To update the DB, we use the following strategy: + // 1. If the document doesn't exist, create the document for the user identifier `(iss, aud, uid_key, uid_val)`, + // but leave counter/time fields unspecified. + // 2. `num_requests += 1`, assuming the default value is 0. + // 3. `last_request_unix_ms = max(last_request_unix_ms, now)`, assuming the default value is 0. + // 4. `first_request_unix_ms = min(first_request_unix_ms, now)`, assuming the default value is +inf. + // + // This strategy is preferred because all the operations can be done on the server-side, + // which means the transaction should require only 1 RTT (this is better than using + // the read-compute-write pattern that requires 2 RTTs). + // + // Note: this strategy requires some modifications: + // - In firestore, the default value of a number field is 0, and we do not have a way to customize it for step 4. + // - The workaround here is to apply an offset so 0 becomes a legitimate default value. + // - So we work with `first_request_unix_ms_minus_1q` instead, which is defined as + // `first_request_unix_ms - 1_000_000_000_000_000`, where 1_000_000_000_000_000 milliseconds is roughly 31710 years. + + // Create the firestore DB transaction + let mut firestore_transaction = firestore_db.begin_transaction().await.map_err(|error| { + PepperServiceError::InternalError(format!( + "Firestore DB begin_transaction() error: {}", + error + )) + })?; + firestore_db + .fluent() + .update() + .fields(paths!(AccountRecoveryDbEntry::{iss, aud, uid_key, uid_val})) + .in_col(FIRESTORE_DB_COLLECTION_ID) + .document_id(&document_id) + .object(&entry) // Step 1 + .transforms(|builder| { + builder.fields([ + builder + .field(path!(AccountRecoveryDbEntry::num_requests)) + .increment(1), // Step 2 + builder + .field(path!(AccountRecoveryDbEntry::last_request_unix_ms)) + .maximum(now_unix_ms), // Step 3 + builder + .field(path!( + AccountRecoveryDbEntry::first_request_unix_ms_minus_1q + )) + .minimum(now_unix_ms - 1_000_000_000_000_000), // Step 4 + ]) + }) + .add_to_transaction(&mut firestore_transaction) + .map_err(|error| { + PepperServiceError::InternalError(format!( + "Firestore DB add_to_transaction() error: {}", + error + )) + })?; + + // Commit the DB transaction + match firestore_transaction.commit().await { + Ok(_) => Ok(()), + Err(error) => { + let error_message = format!( + "Failed to commit Firestore DB transaction for pepper input {:?}! Document ID {}, Error: {}", + pepper_input, document_id, error + ); + Err(PepperServiceError::InternalError(error_message)) + }, + } +} diff --git a/keyless/pepper/service/src/dedicated_handlers/handlers.rs b/keyless/pepper/service/src/dedicated_handlers/handlers.rs index b114a8d39fa90..f90fa3207b43f 100644 --- a/keyless/pepper/service/src/dedicated_handlers/handlers.rs +++ b/keyless/pepper/service/src/dedicated_handlers/handlers.rs @@ -40,7 +40,7 @@ pub trait HandlerTrait: Send + Sync { // TODO: is there a way we can remove the vuf_private_key param here? async fn handle_request( &self, - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, jwk_cache: JWKCache, cached_resources: CachedResources, request: TRequest, @@ -56,7 +56,7 @@ pub struct V0DelegatedFetchHandler; impl HandlerTrait for V0DelegatedFetchHandler { async fn handle_request( &self, - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, jwk_cache: JWKCache, cached_resources: CachedResources, request: PepperRequestV2, @@ -109,7 +109,7 @@ pub struct V0FetchHandler; impl HandlerTrait for V0FetchHandler { async fn handle_request( &self, - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, jwk_cache: JWKCache, cached_resources: CachedResources, request: PepperRequest, @@ -158,7 +158,7 @@ pub struct V0SignatureHandler; impl HandlerTrait for V0SignatureHandler { async fn handle_request( &self, - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, jwk_cache: JWKCache, cached_resources: CachedResources, request: PepperRequest, @@ -207,7 +207,7 @@ pub struct V0VerifyHandler; impl HandlerTrait for V0VerifyHandler { async fn handle_request( &self, - _: &ark_bls12_381::Fr, + _: Arc, jwk_cache: JWKCache, cached_resources: CachedResources, request: VerifyRequest, diff --git a/keyless/pepper/service/src/dedicated_handlers/pepper_request.rs b/keyless/pepper/service/src/dedicated_handlers/pepper_request.rs index 2d7e3320a9868..8b5a66e76fbd0 100644 --- a/keyless/pepper/service/src/dedicated_handlers/pepper_request.rs +++ b/keyless/pepper/service/src/dedicated_handlers/pepper_request.rs @@ -7,6 +7,7 @@ use crate::{ }, error::PepperServiceError, external_resources::{jwk_fetcher, jwk_fetcher::JWKCache, resource_fetcher::CachedResources}, + metrics, }; use aptos_keyless_pepper_common::{ jwt::Claims, @@ -27,8 +28,9 @@ use aptos_types::{ }; use jsonwebtoken::{Algorithm::RS256, DecodingKey, TokenData, Validation}; use std::{ + ops::Deref, sync::Arc, - time::{SystemTime, UNIX_EPOCH}, + time::{Instant, SystemTime, UNIX_EPOCH}, }; // The default derivation path (if none is provided) @@ -45,7 +47,7 @@ const EMAIL_UID_KEY: &str = "email"; /// Handles the given pepper request, returning the pepper base, pepper and account address pub async fn handle_pepper_request( - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, jwk_cache: JWKCache, cached_resources: CachedResources, jwt: String, @@ -94,16 +96,22 @@ pub async fn handle_pepper_request( ) .await?; - // Create the pepper base using the vuf private key and the pepper input - let pepper_base = create_pepper_base(vuf_private_key, &pepper_input)?; + // Derive the pepper base, pepper bytes and account address. + // Note: we do this using spawn_blocking() to avoid blocking the async runtime. + let (pepper_base, derived_pepper_bytes, address) = tokio::task::spawn_blocking(move || { + // Start the derivation timer + let derivation_start_time = Instant::now(); - // Derive the pepper using the verified derivation path and the pepper base - let verified_derivation_path = get_verified_derivation_path(derivation_path)?; - let derived_pepper = derive_pepper(&verified_derivation_path, &pepper_base)?; - let derived_pepper_bytes = derived_pepper.to_bytes().to_vec(); + // Derive the pepper and account address + let derivation_result = + derive_pepper_and_account_address(vuf_private_key, derivation_path, &pepper_input); - // Create the account address - let address = create_account_address(&pepper_input, &derived_pepper)?; + // Update the derivation metrics + metrics::update_pepper_derivation_metrics(derivation_result.is_ok(), derivation_start_time); + + derivation_result + }) + .await??; // Return the pepper base, derived pepper and address Ok((pepper_base, derived_pepper_bytes, address)) @@ -134,7 +142,7 @@ fn create_account_address( /// Creates the pepper base using the VUF private key and the pepper input fn create_pepper_base( - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, pepper_input: &PepperInput, ) -> Result, PepperServiceError> { // Serialize the pepper input using BCS @@ -147,7 +155,7 @@ fn create_pepper_base( // Generate the pepper base and proof using the VUF let (pepper_base, vuf_proof) = - vuf::bls12381_g1_bls::Bls12381G1Bls::eval(vuf_private_key, &input_bytes).map_err( + vuf::bls12381_g1_bls::Bls12381G1Bls::eval(vuf_private_key.deref(), &input_bytes).map_err( |error| { PepperServiceError::InternalError(format!( "Failed to evaluate bls12381_g1_bls VUF: {}", @@ -234,6 +242,26 @@ fn derive_pepper( Ok(derived_pepper) } +/// Derives the pepper base, pepper bytes and account address +fn derive_pepper_and_account_address( + vuf_private_key: Arc, + derivation_path: Option, + pepper_input: &PepperInput, +) -> Result<(Vec, Vec, AccountAddress), PepperServiceError> { + // Create the pepper base using the vuf private key and the pepper input + let pepper_base = create_pepper_base(vuf_private_key, pepper_input)?; + + // Derive the pepper using the verified derivation path and the pepper base + let verified_derivation_path = get_verified_derivation_path(derivation_path)?; + let derived_pepper = derive_pepper(&verified_derivation_path, &pepper_base)?; + let derived_pepper_bytes = derived_pepper.to_bytes().to_vec(); + + // Create the account address + let address = create_account_address(pepper_input, &derived_pepper)?; + + Ok((pepper_base, derived_pepper_bytes, address)) +} + /// Returns the on-chain keyless configuration from the cached resources fn get_on_chain_keyless_configuration( cached_resources: CachedResources, @@ -524,7 +552,7 @@ mod tests { // Verify the pepper base, derived pepper and account address verify_base_pepper_and_address_generation( - &vuf_private_key, + vuf_private_key, &pepper_input, TEST_SUB_PEPPER_BASE_HEX, TEST_SUB_DERIVED_PEPPER_HEX, @@ -560,7 +588,7 @@ mod tests { // Verify the pepper base, derived pepper and account address verify_base_pepper_and_address_generation( - &vuf_private_key, + vuf_private_key, &pepper_input, TEST_EMAIL_PEPPER_BASE_HEX, TEST_EMAIL_DERIVED_PEPPER_HEX, @@ -595,7 +623,7 @@ mod tests { // Verify the pepper base, derived pepper and account address verify_base_pepper_and_address_generation( - &vuf_private_key, + vuf_private_key, &pepper_input, TEST_SUB_PEPPER_BASE_HEX, TEST_SUB_DERIVED_PEPPER_HEX, @@ -635,7 +663,7 @@ mod tests { // Verify the pepper base, derived pepper and account address verify_base_pepper_and_address_generation( - &vuf_private_key, + vuf_private_key, &pepper_input, TEST_SUB_OVERRIDE_PEPPER_BASE_HEX, TEST_SUB_OVERRIDE_DERIVED_PEPPER_HEX, @@ -661,31 +689,26 @@ mod tests { } /// Returns the test VUF private key from the given seed - fn get_test_vuf_private_key(seed: [u8; 32]) -> ark_bls12_381::Fr { + fn get_test_vuf_private_key(seed: [u8; 32]) -> Arc { // Derive the VUF private key from the seed let mut sha3_hasher = sha3::Sha3_512::new(); sha3_hasher.update(seed); - ark_bls12_381::Fr::from_be_bytes_mod_order(sha3_hasher.finalize().as_slice()) + let vuf_private_key = + ark_bls12_381::Fr::from_be_bytes_mod_order(sha3_hasher.finalize().as_slice()); + Arc::new(vuf_private_key) } /// Verifies the generated pepper base, derived pepper and account address against the expected values fn verify_base_pepper_and_address_generation( - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, pepper_input: &PepperInput, expected_pepper_base_hex: &str, expected_derived_pepper_hex: &str, expected_account_address: &str, ) { - // Create the pepper base - let pepper_base = create_pepper_base(vuf_private_key, pepper_input).unwrap(); - - // Derive the pepper using the verified derivation path and the pepper base - let verified_derivation_path = get_verified_derivation_path(None).unwrap(); - let derived_pepper = derive_pepper(&verified_derivation_path, &pepper_base).unwrap(); - let derived_pepper_bytes = derived_pepper.to_bytes().to_vec(); - - // Create the account address - let address = create_account_address(pepper_input, &derived_pepper).unwrap(); + // Derive the pepper base, pepper and account address + let (pepper_base, derived_pepper_bytes, address) = + derive_pepper_and_account_address(vuf_private_key, None, pepper_input).unwrap(); // Verify the pepper base, derived pepper and account address assert_eq!(hex::encode(pepper_base), expected_pepper_base_hex); diff --git a/keyless/pepper/service/src/error.rs b/keyless/pepper/service/src/error.rs index ff73f14f0f339..755701b8dfa9f 100644 --- a/keyless/pepper/service/src/error.rs +++ b/keyless/pepper/service/src/error.rs @@ -14,3 +14,9 @@ pub enum PepperServiceError { #[error("Unexpected error: {0}")] UnexpectedError(String), } + +impl From for PepperServiceError { + fn from(error: tokio::task::JoinError) -> PepperServiceError { + PepperServiceError::UnexpectedError(format!("JoinError: {:?}", error)) + } +} diff --git a/keyless/pepper/service/src/main.rs b/keyless/pepper/service/src/main.rs index fa1b646b4ad46..fef78f6c6df2a 100644 --- a/keyless/pepper/service/src/main.rs +++ b/keyless/pepper/service/src/main.rs @@ -15,9 +15,9 @@ use aptos_keyless_pepper_service::{ metrics::DEFAULT_METRICS_SERVER_PORT, request_handler, request_handler::DEFAULT_PEPPER_SERVICE_PORT, - vuf_pub_key, + utils, vuf_pub_key, }; -use aptos_logger::{info, warn}; +use aptos_logger::{error, info, warn}; use clap::Parser; use hyper::{ service::{make_service_fn, service_fn}, @@ -37,6 +37,11 @@ struct Args { #[arg(long)] account_recovery_managers: Vec, + /// Disable asynchronous updates to the account recovery database. + /// By default, async updates are enabled to avoid blocking request handlers. + #[arg(long)] + disable_async_db_updates: bool, // Defaults to false if not provided + /// The Firestore database ID (required to connect to Firestore). /// Only required if not running in local development mode. #[arg( @@ -126,26 +131,32 @@ async fn main() { ); // Create the account recovery database - let account_recovery_db: Arc = if args - .local_development_mode - { - warn!("Running in local development mode! Using a test account recovery database!"); - Arc::new(TestAccountRecoveryDB::new()) - } else { - let google_project_id = args.google_project_id.expect( - "Google Project ID must be provided when not running in local development mode!", - ); - let firestore_database_id = args.firestore_database_id.expect( + let account_recovery_db: Arc = + if args.local_development_mode { + warn!("Running in local development mode! Using a test account recovery database!"); + Arc::new(TestAccountRecoveryDB::new()) + } else { + let google_project_id = args.google_project_id.expect( + "Google Project ID must be provided when not running in local development mode!", + ); + let firestore_database_id = args.firestore_database_id.expect( "Firestore Database ID must be provided when not running in local development mode!", ); - Arc::new(FirestoreAccountRecoveryDB::new(google_project_id, firestore_database_id).await) - }; + Arc::new( + FirestoreAccountRecoveryDB::new( + google_project_id, + firestore_database_id, + args.disable_async_db_updates, + ) + .await, + ) + }; // Start the JWK fetchers let jwk_cache = jwk_fetcher::start_jwk_fetchers(); // Start the pepper service - let vuf_keypair = Arc::new((vuf_public_key, vuf_private_key)); + let vuf_keypair = Arc::new((vuf_public_key, Arc::new(vuf_private_key))); start_pepper_service( args.pepper_service_port, vuf_keypair, @@ -179,7 +190,7 @@ fn start_metrics_server() { // Starts the pepper service async fn start_pepper_service( pepper_service_port: u16, - vuf_keypair: Arc<(String, ark_bls12_381::Fr)>, + vuf_keypair: Arc<(String, Arc)>, jwk_cache: JWKCache, cached_resources: CachedResources, account_recovery_managers: Arc, @@ -201,8 +212,11 @@ async fn start_pepper_service( async move { Ok::<_, Infallible>(service_fn(move |request| { - // Get the request start time, method and request path + // Start the request timer let request_start_time = Instant::now(); + + // Get the request origin, method and request path + let request_origin = utils::get_request_origin(&request); let request_method = request.method().clone(); let request_path = request.uri().path().to_owned(); @@ -226,14 +240,36 @@ async fn start_pepper_service( ) .await; - // Update the request handling metrics - if let Ok(response) = &result { - metrics::update_request_handling_metrics( - &request_path, - request_method, - response.status(), - request_start_time, - ); + // Update the request handling metrics and logs + match &result { + Ok(response) => { + // Update the request handling metrics + metrics::update_request_handling_metrics( + &request_path, + request_method.clone(), + response.status(), + request_start_time, + ); + + // If the response was not successful, log the request details + if !response.status().is_success() { + warn!( + "Handled request with non-successful response! Request origin: {:?}, \ + request path: {:?}, request method: {:?}, response status: {:?}", + request_origin, + request_path, + request_method, + response.status() + ); + } + }, + Err(error) => { + error!( + "Error occurred when handling request! Request origin: {:?}, \ + request path: {:?}, request method: {:?}, Error: {:?}", + request_origin, request_path, request_method, error + ); + }, } result diff --git a/keyless/pepper/service/src/metrics.rs b/keyless/pepper/service/src/metrics.rs index c04b099a71b59..872b7a6139d95 100644 --- a/keyless/pepper/service/src/metrics.rs +++ b/keyless/pepper/service/src/metrics.rs @@ -25,6 +25,17 @@ static LATENCY_BUCKETS: Lazy> = Lazy::new(|| { .unwrap() }); +// Histogram for tracking time taken to write to the account recovery database +static ACCOUNT_RECOVERY_DB_WRITE_SECONDS: Lazy = Lazy::new(|| { + register_histogram_vec!( + "keyless_pepper_service_account_recovery_db_write_seconds", + "Time taken to write to the account recovery database", + &["succeeded"], + LATENCY_BUCKETS.clone() + ) + .unwrap() +}); + // Histogram for tracking time taken to fetch external resources static EXTERNAL_RESOURCE_FETCH_SECONDS: Lazy = Lazy::new(|| { register_histogram_vec!( @@ -47,6 +58,17 @@ static JWK_FETCH_SECONDS: Lazy = Lazy::new(|| { .unwrap() }); +// Histogram for tracking time taken to derive peppers +static PEPPER_DERIVATION_SECONDS: Lazy = Lazy::new(|| { + register_histogram_vec!( + "keyless_pepper_service_pepper_derivation_seconds", + "Time taken to derive peppers", + &["succeeded"], + LATENCY_BUCKETS.clone() + ) + .unwrap() +}); + // Histogram for tracking time taken to handle pepper service requests static REQUEST_HANDLING_SECONDS: Lazy = Lazy::new(|| { register_histogram_vec!( @@ -80,6 +102,17 @@ pub async fn handle_metrics_request( Ok(response) } +/// Updates the account recovery DB write metrics with the given data +pub fn update_account_recovery_db_metrics(succeeded: bool, write_start_time: Instant) { + // Calculate the elapsed time + let elapsed = write_start_time.elapsed(); + + // Update the metrics + ACCOUNT_RECOVERY_DB_WRITE_SECONDS + .with_label_values(&[&succeeded.to_string()]) + .observe(elapsed.as_secs_f64()); +} + /// Updates the external resource fetch metrics with the given data pub fn update_external_resource_fetch_metrics( resource_name: &str, @@ -98,6 +131,17 @@ pub fn update_jwk_fetch_metrics(issuer: &str, succeeded: bool, elapsed: Duration .observe(elapsed.as_secs_f64()); } +/// Updates the pepper derivation metrics with the given data +pub fn update_pepper_derivation_metrics(succeeded: bool, derivation_start_time: Instant) { + // Calculate the elapsed time + let elapsed = derivation_start_time.elapsed(); + + // Update the metrics + PEPPER_DERIVATION_SECONDS + .with_label_values(&[&succeeded.to_string()]) + .observe(elapsed.as_secs_f64()); +} + /// Updates the request handling metrics with the given data pub fn update_request_handling_metrics( request_endpoint: &str, diff --git a/keyless/pepper/service/src/request_handler.rs b/keyless/pepper/service/src/request_handler.rs index 30125db1eadc7..c539d48de5dac 100644 --- a/keyless/pepper/service/src/request_handler.rs +++ b/keyless/pepper/service/src/request_handler.rs @@ -10,6 +10,7 @@ use crate::{ }, error::PepperServiceError, external_resources::{jwk_fetcher::JWKCache, resource_fetcher::CachedResources}, + utils, }; use aptos_build_info::build_information; use aptos_keyless_pepper_common::BadPepperRequestError; @@ -57,10 +58,6 @@ pub const ALL_PATHS: [&str; 9] = [ const CONTENT_TYPE_JSON: &str = "application/json"; const CONTENT_TYPE_TEXT: &str = "text/plain"; -// Origin header constants -const MISSING_ORIGIN_STRING: &str = ""; // Default to empty string if origin header is missing -const ORIGIN_HEADER: &str = "origin"; - // Useful message constants const METHOD_NOT_ALLOWED_MESSAGE: &str = "The request method is not allowed for the requested path!"; @@ -70,7 +67,7 @@ const UNEXPECTED_ERROR_MESSAGE: &str = "An unexpected error was encountered!"; async fn call_dedicated_request_handler( origin: String, request: Request, - vuf_private_key: &ark_bls12_381::Fr, + vuf_private_key: Arc, jwk_cache: JWKCache, cached_resources: CachedResources, request_handler: &TRequestHandler, @@ -291,16 +288,6 @@ fn generate_options_response(origin: String) -> Result, Infallibl Ok(response) } -/// Extracts the origin header from the request -fn get_request_origin(request: &Request) -> String { - request - .headers() - .get(ORIGIN_HEADER) - .and_then(|header_value| header_value.to_str().ok()) - .unwrap_or(MISSING_ORIGIN_STRING) - .to_owned() -} - /// Generates a text response with the given status code and body string fn generate_text_response( origin: String, @@ -317,14 +304,14 @@ fn generate_text_response( /// Handles the given request and returns a response pub async fn handle_request( request: Request, - vuf_keypair: Arc<(String, ark_bls12_381::Fr)>, + vuf_keypair: Arc<(String, Arc)>, jwk_cache: JWKCache, cached_resources: CachedResources, account_recovery_managers: Arc, account_recovery_db: Arc, ) -> Result, Infallible> { // Get the request origin - let origin = get_request_origin(&request); + let origin = utils::get_request_origin(&request); // Handle any OPTIONS requests let request_method = request.method(); @@ -374,7 +361,7 @@ pub async fn handle_request( return call_dedicated_request_handler( origin, request, - vuf_priv_key, + vuf_priv_key.clone(), jwk_cache, cached_resources, &V0DelegatedFetchHandler, @@ -387,7 +374,7 @@ pub async fn handle_request( return call_dedicated_request_handler( origin, request, - vuf_priv_key, + vuf_priv_key.clone(), jwk_cache, cached_resources, &V0FetchHandler, @@ -400,7 +387,7 @@ pub async fn handle_request( return call_dedicated_request_handler( origin, request, - vuf_priv_key, + vuf_priv_key.clone(), jwk_cache, cached_resources, &V0SignatureHandler, @@ -413,7 +400,7 @@ pub async fn handle_request( return call_dedicated_request_handler( origin, request, - vuf_priv_key, + vuf_priv_key.clone(), jwk_cache, cached_resources, &V0VerifyHandler, diff --git a/keyless/pepper/service/src/tests/pepper_request.rs b/keyless/pepper/service/src/tests/pepper_request.rs index f841ff6cec78b..580ed5ced7e55 100644 --- a/keyless/pepper/service/src/tests/pepper_request.rs +++ b/keyless/pepper/service/src/tests/pepper_request.rs @@ -25,7 +25,6 @@ use aptos_types::{ }; use std::{ collections::HashMap, - ops::Deref, sync::Arc, time::{SystemTime, UNIX_EPOCH}, }; @@ -33,8 +32,7 @@ use std::{ #[tokio::test] async fn request_ephemeral_public_key_expired() { // Generate a VUF private key - let vuf_keypair = utils::create_vuf_public_private_keypair(); - let (_, vuf_private_key) = vuf_keypair.deref(); + let (_, vuf_private_key) = utils::create_vuf_public_private_keypair(); // Create a JWK cache and resource cache let jwk_cache = Arc::new(Mutex::new(HashMap::new())); @@ -80,8 +78,7 @@ async fn request_ephemeral_public_key_expired() { #[tokio::test] async fn request_ephemeral_public_key_expiry_too_large() { // Generate a VUF private key - let vuf_keypair = utils::create_vuf_public_private_keypair(); - let (_, vuf_private_key) = vuf_keypair.deref(); + let (_, vuf_private_key) = utils::create_vuf_public_private_keypair(); // Create a JWK cache and resource cache let jwk_cache = Arc::new(Mutex::new(HashMap::new())); @@ -127,8 +124,7 @@ async fn request_ephemeral_public_key_expiry_too_large() { #[tokio::test] async fn request_invalid_oath_nonce() { // Generate a VUF private key - let vuf_keypair = utils::create_vuf_public_private_keypair(); - let (_, vuf_private_key) = vuf_keypair.deref(); + let (_, vuf_private_key) = utils::create_vuf_public_private_keypair(); // Create a JWK cache and resource cache let jwk_cache = Arc::new(Mutex::new(HashMap::new())); @@ -174,8 +170,7 @@ async fn request_invalid_oath_nonce() { #[tokio::test] async fn request_invalid_jwt() { // Generate a VUF private key - let vuf_keypair = utils::create_vuf_public_private_keypair(); - let (_, vuf_private_key) = vuf_keypair.deref(); + let (_, vuf_private_key) = utils::create_vuf_public_private_keypair(); // Create a JWK cache and resource cache let jwk_cache = Arc::new(Mutex::new(HashMap::new())); @@ -208,8 +203,7 @@ async fn request_invalid_jwt() { #[tokio::test] async fn request_invalid_jwt_signature() { // Generate a VUF private key - let vuf_keypair = utils::create_vuf_public_private_keypair(); - let (_, vuf_private_key) = vuf_keypair.deref(); + let (_, vuf_private_key) = utils::create_vuf_public_private_keypair(); // Create a JWK cache and resource cache let jwk_cache = Arc::new(Mutex::new(HashMap::new())); @@ -270,8 +264,7 @@ async fn request_invalid_jwt_signature() { #[tokio::test] async fn request_max_exp_data_secs_overflowed() { // Generate a VUF private key - let vuf_keypair = utils::create_vuf_public_private_keypair(); - let (_, vuf_private_key) = vuf_keypair.deref(); + let (_, vuf_private_key) = utils::create_vuf_public_private_keypair(); // Create a JWK cache and resource cache let jwk_cache = Arc::new(Mutex::new(HashMap::new())); diff --git a/keyless/pepper/service/src/tests/request_handler.rs b/keyless/pepper/service/src/tests/request_handler.rs index 8f773faeb47b2..faed0ed38a655 100644 --- a/keyless/pepper/service/src/tests/request_handler.rs +++ b/keyless/pepper/service/src/tests/request_handler.rs @@ -305,8 +305,8 @@ async fn test_get_vuf_pub_key_request() { let response_vuf_public_key = get_public_key_from_json(&body_string); // Get the expected public key from the keypair - let (vuf_public_key_json, _) = vuf_keypair.deref(); - let vuf_public_key = get_public_key_from_json(vuf_public_key_json); + let (vuf_public_key_json, _) = vuf_keypair; + let vuf_public_key = get_public_key_from_json(&vuf_public_key_json); // Verify the public key is correct assert_eq!(response_vuf_public_key, vuf_public_key); @@ -508,7 +508,7 @@ async fn send_request_to_path( method: Method, endpoint: &str, body: Body, - vuf_keypair: Option>, + vuf_keypair: Option<(String, Arc)>, jwk_cache: Option, cached_resources: Option, ) -> Response { @@ -526,7 +526,8 @@ async fn send_request_to_path( .unwrap(); // Get or create a VUF public private keypair - let vuf_keypair = vuf_keypair.unwrap_or_else(utils::create_vuf_public_private_keypair); + let vuf_keypair = + Arc::new(vuf_keypair.unwrap_or_else(utils::create_vuf_public_private_keypair)); // Get or create a JWK cache let jwk_cache = jwk_cache.unwrap_or_else(|| Arc::new(Mutex::new(HashMap::new()))); diff --git a/keyless/pepper/service/src/tests/utils.rs b/keyless/pepper/service/src/tests/utils.rs index ca86a127a327b..378631d080019 100644 --- a/keyless/pepper/service/src/tests/utils.rs +++ b/keyless/pepper/service/src/tests/utils.rs @@ -39,7 +39,7 @@ pub async fn advance_time_secs(time_service: TimeService, seconds: u64) { } /// Generates a random VUF public and private keypair for testing purposes -pub fn create_vuf_public_private_keypair() -> Arc<(String, ark_bls12_381::Fr)> { +pub fn create_vuf_public_private_keypair() -> (String, Arc) { // Generate a random VUF seed let private_key_seed = rand::random::<[u8; 32]>(); @@ -63,7 +63,7 @@ pub fn create_vuf_public_private_keypair() -> Arc<(String, ark_bls12_381::Fr)> { // Transform the public key object to a pretty JSON string let vuf_public_key_string = serde_json::to_string_pretty(&pepper_vuf_public_key).unwrap(); - Arc::new((vuf_public_key_string, vuf_private_key)) + (vuf_public_key_string, Arc::new(vuf_private_key)) } /// Returns an empty account managers and overrides instance diff --git a/keyless/pepper/service/src/utils.rs b/keyless/pepper/service/src/utils.rs index e4d68291b6ab2..e44ead54b7ebb 100644 --- a/keyless/pepper/service/src/utils.rs +++ b/keyless/pepper/service/src/utils.rs @@ -2,11 +2,16 @@ // SPDX-License-Identifier: Apache-2.0 use anyhow::{anyhow, ensure}; +use hyper::{Body, Request}; use reqwest::Client; use std::time::Duration; // Timeout for client requests -const CLIENT_REQUEST_TIMEOUT_SECS: u64 = 10; +const CLIENT_REQUEST_TIMEOUT_SECS: u64 = 15; + +// Origin header constants +const MISSING_ORIGIN_STRING: &str = ""; // Default to empty string if origin header is missing +const ORIGIN_HEADER: &str = "origin"; /// Creates and returns a reqwest HTTP client with a timeout pub fn create_request_client() -> Client { @@ -16,6 +21,16 @@ pub fn create_request_client() -> Client { .expect("Failed to build the request client!") } +/// Extracts the origin header from the request +pub fn get_request_origin(request: &Request) -> String { + request + .headers() + .get(ORIGIN_HEADER) + .and_then(|header_value| header_value.to_str().ok()) + .unwrap_or(MISSING_ORIGIN_STRING) + .to_owned() +} + /// Converts a hex-encoded string (with "0x" prefix) to a byte vector pub fn unhexlify_api_bytes(api_output: &str) -> anyhow::Result> { // Verify the input format