Skip to content

[Keyless] Performance, metric and logging improvements for the pepper service. #17856

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 14, 2025
Merged

[Keyless] Performance, metric and logging improvements for the pepper service. #17856

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ca73c68
[Keyless] Small improvements to logging and timeout.
JoshLind Oct 11, 2025
a1bfd71
[Keyless] Use spawn_blocking for crypto operations.
JoshLind Oct 12, 2025
7d237c2
[Keyless] Make Firestore DB update async.
JoshLind Oct 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion keyless/pepper/common/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ impl PepperV0VufPubKey {
}
}

#[derive(Debug, Deserialize, Serialize)]
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct PepperInput {
pub iss: String,
pub aud: String,
Expand Down
217 changes: 131 additions & 86 deletions keyless/pepper/service/src/accounts/account_recovery_db.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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,
Expand All @@ -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: {}",
Expand All @@ -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(())
Comment on lines +95 to +99
Copy link
Contributor

@zjma zjma Oct 14, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i recall the original design is that we should wait for the recovery database txn to be committed first, otherwise if a db txn failed and it is the only login attempt in the history, it would be impossible to recover that account. Therefore async db update shouldn't be an option. (@alinush can confirm?)

That said, I'm not sure if we are still keep the original design.

Copy link
Contributor Author

@JoshLind JoshLind Oct 14, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Aah, good point! 😄 My thinking was the opposite, i.e.,:

If our DB is unhealthy, we shouldn't block users from accessing their pepper (and thus their on-chain accounts). Instead, we emit an error (so we can raise an alarm) and save the pepper input to our logs, so that we recover it manually (if we ever need to).

That way, for a pepper to be totally unrecoverable, we'd need: (1) the DB to be unhealthy (e.g., input never saved); (2) our logs to be missing (e.g., logs to be wiped); and (3) the dApp info to be lost/disabled.

Note: for (3), if any other users have logged into the same dApp, we should be able to take the "brute force approach" and find the input for the user more manually, right?

All in all, my preference is to still allow users to access their accounts (while we figure out what's going wrong with the DB). What do you think? 😄

}
}
}
Expand Down Expand Up @@ -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))
},
}
}
10 changes: 5 additions & 5 deletions keyless/pepper/service/src/dedicated_handlers/handlers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ pub trait HandlerTrait<TRequest, TResponse>: 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<ark_bls12_381::Fr>,
jwk_cache: JWKCache,
cached_resources: CachedResources,
request: TRequest,
Expand All @@ -56,7 +56,7 @@ pub struct V0DelegatedFetchHandler;
impl HandlerTrait<PepperRequestV2, PepperResponse> for V0DelegatedFetchHandler {
async fn handle_request(
&self,
vuf_private_key: &ark_bls12_381::Fr,
vuf_private_key: Arc<ark_bls12_381::Fr>,
jwk_cache: JWKCache,
cached_resources: CachedResources,
request: PepperRequestV2,
Expand Down Expand Up @@ -109,7 +109,7 @@ pub struct V0FetchHandler;
impl HandlerTrait<PepperRequest, PepperResponse> for V0FetchHandler {
async fn handle_request(
&self,
vuf_private_key: &ark_bls12_381::Fr,
vuf_private_key: Arc<ark_bls12_381::Fr>,
jwk_cache: JWKCache,
cached_resources: CachedResources,
request: PepperRequest,
Expand Down Expand Up @@ -158,7 +158,7 @@ pub struct V0SignatureHandler;
impl HandlerTrait<PepperRequest, SignatureResponse> for V0SignatureHandler {
async fn handle_request(
&self,
vuf_private_key: &ark_bls12_381::Fr,
vuf_private_key: Arc<ark_bls12_381::Fr>,
jwk_cache: JWKCache,
cached_resources: CachedResources,
request: PepperRequest,
Expand Down Expand Up @@ -207,7 +207,7 @@ pub struct V0VerifyHandler;
impl HandlerTrait<VerifyRequest, VerifyResponse> for V0VerifyHandler {
async fn handle_request(
&self,
_: &ark_bls12_381::Fr,
_: Arc<ark_bls12_381::Fr>,
jwk_cache: JWKCache,
cached_resources: CachedResources,
request: VerifyRequest,
Expand Down
Loading
Loading