Skip to content

Commit

Permalink
pmonitor: do trial decryption in genesis block
Browse files Browse the repository at this point in the history
  • Loading branch information
@redshiftzero
redshiftzero committed Sep 15, 2024
1 parent 87e0e8a commit 47a2c8a
Show file tree
Hide file tree
Showing 2 changed files with 163 additions and 115 deletions.
95 changes: 47 additions & 48 deletions crates/bin/pmonitor/src/genesis.rs
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,10 @@
use std::{
collections::{BTreeMap, HashMap},
str::FromStr,
};
use std::{collections::BTreeMap, str::FromStr};

use penumbra_app::genesis::AppState;
use penumbra_asset::STAKING_TOKEN_ASSET_ID;
use penumbra_compact_block::{CompactBlock, StatePayload};
use penumbra_keys::FullViewingKey;
use penumbra_num::Amount;
use penumbra_shielded_pool::Note;
use penumbra_shielded_pool::{Note, NotePayload};
use penumbra_stake::{
rate::{BaseRateData, RateData},
DelegationToken,
Expand All @@ -26,78 +23,80 @@ pub struct FilteredGenesisBlock {
/// initial balances of the relevant addresses.
///
/// Assumption: There are no swaps or nullifiers in the genesis block.
pub fn scan_genesis_block(
genesis_app_state: AppState,
pub async fn scan_genesis_block(
CompactBlock {
height,
state_payloads,
..
}: CompactBlock,
fvks: Vec<FullViewingKey>,
) -> anyhow::Result<FilteredGenesisBlock> {
assert_eq!(height, 0);

let mut notes = BTreeMap::new();
let mut balances = BTreeMap::new();

let genesis_data = genesis_app_state
.content()
.expect("genesis app state should have content");
// We'll use the allocations from the genesis state.
let shielded_pool_content = &genesis_data.shielded_pool_content;

// Calculate the rate data for each validator in the initial validator set.
let base_rate = BaseRateData {
epoch_index: 0,
base_reward_rate: 0u128.into(),
base_exchange_rate: 1_0000_0000u128.into(),
};
let rate_data_map: HashMap<DelegationToken, RateData> = genesis_data
.stake_content
.validators
.iter()
.map(|validator| {
let identity_key = validator
.identity_key
.clone()
.expect("identity key should be present")
.try_into()
.expect("should be a valid identity key");
let rate_data = RateData {
identity_key,
validator_reward_rate: 0u128.into(),
validator_exchange_rate: base_rate.base_exchange_rate,
};
(DelegationToken::from(identity_key), rate_data)
})
.collect();

// We proceed one FVK at a time.
for fvk in fvks {
// Trial-decrypt a note with our a specific viewing key
let trial_decrypt_note =
|note_payload: NotePayload| -> tokio::task::JoinHandle<Option<Note>> {
let fvk2 = fvk.clone();
tokio::spawn(async move { note_payload.trial_decrypt(&fvk2) })
};

// Trial-decrypt the notes in this block, keeping track of the ones that were meant for the FVK
// we're monitoring.
let mut note_decryptions = Vec::new();

// We only care about notes, so we're ignoring swaps and rolled-up commitments.
for payload in state_payloads.iter() {
if let StatePayload::Note { note, .. } = payload {
note_decryptions.push(trial_decrypt_note((**note).clone()));
}
}

let mut notes_for_this_fvk = BTreeMap::new();
for allocation in &shielded_pool_content.allocations {
if fvk.incoming().views_address(&allocation.address) {
let note =
Note::from_allocation(allocation.clone()).expect("should be a valid note");
for decryption in note_decryptions {
if let Some(note) = decryption
.await
.expect("able to join tokio note decryption handle")
{
notes_for_this_fvk.insert(note.commit(), note.clone());

// Balance is expected to be in the staking or delegation token
let allocation_value = allocation.value();
if allocation_value.asset_id == *STAKING_TOKEN_ASSET_ID {
let note_value = note.value();
if note_value.asset_id == *STAKING_TOKEN_ASSET_ID {
balances
.entry(fvk.to_string())
.and_modify(|existing_amount| *existing_amount += allocation.amount())
.or_insert(allocation.amount());
.and_modify(|existing_amount| *existing_amount += note.amount())
.or_insert(note.amount());
} else if let Ok(delegation_token) =
DelegationToken::from_str(&allocation_value.asset_id.to_string())
DelegationToken::from_str(&note_value.asset_id.to_string())
{
// We need to convert the amount to the UM-equivalent amount
let rate_data = rate_data_map
.get(&delegation_token)
.expect("should be rate data for each validator");
let um_equivalent_balance = rate_data.unbonded_amount(allocation.amount());
let rate_data = RateData {
identity_key: delegation_token.validator(),
validator_reward_rate: 0u128.into(),
validator_exchange_rate: base_rate.base_exchange_rate,
};
let um_equivalent_balance = rate_data.unbonded_amount(note.amount());

balances
.entry(fvk.to_string())
.and_modify(|existing_amount| *existing_amount += um_equivalent_balance)
.or_insert(um_equivalent_balance);
} else {
tracing::warn!(
"ignoring note with unrecognized asset id: {}",
allocation_value.asset_id
"ignoring note with unknown asset id: {}",
note_value.asset_id
);
}
}
Expand Down
183 changes: 116 additions & 67 deletions crates/bin/pmonitor/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,18 +10,27 @@ use std::str::FromStr;
use url::Url;

use pcli::config::PcliConfig;
use penumbra_compact_block::CompactBlock;
use penumbra_keys::FullViewingKey;
use penumbra_num::Amount;
use penumbra_proto::box_grpc_svc;
use penumbra_proto::view::v1::{
view_service_client::ViewServiceClient, view_service_server::ViewServiceServer,
};
use penumbra_proto::{
core::component::compact_block::v1::CompactBlockRequest,
penumbra::core::component::compact_block::v1::query_service_client::QueryServiceClient as CompactBlockQueryServiceClient,
};
use penumbra_stake::rate::{BaseRateData, RateData};
use penumbra_stake::DelegationToken;
use penumbra_view::{ViewClient, ViewServer};

mod genesis;

// The maximum size of a compact block, in bytes (12MB).
const MAX_CB_SIZE_BYTES: usize = 12 * 1024 * 1024;

// The name of the view database file
const VIEW_FILE_NAME: &str = "pcli-view.sqlite";

#[tokio::main]
Expand Down Expand Up @@ -75,6 +84,7 @@ pub enum Command {
}

impl Opt {
/// Set up the view service for a given wallet.
pub async fn view(
&self,
path: Utf8PathBuf,
Expand All @@ -98,6 +108,7 @@ impl Opt {
Ok(view_service)
}

/// Sync that wallet to the latest block height.
pub async fn sync(
&self,
view_service: &mut ViewServiceClient<box_grpc_svc::BoxGrpcService>,
Expand Down Expand Up @@ -134,6 +145,23 @@ impl Opt {
Ok(())
}

/// Fetch the genesis compact block
pub async fn fetch_genesis_compact_block(&self, grpc_url: Url) -> Result<CompactBlock> {
let height = 0;
let mut client = CompactBlockQueryServiceClient::connect(grpc_url.to_string())
.await
.unwrap()
.max_decoding_message_size(MAX_CB_SIZE_BYTES);
let compact_block = client
.compact_block(CompactBlockRequest { height })
.await?
.into_inner()
.compact_block
.expect("response has compact block");
compact_block.try_into()
}

/// Execute the specified command.
pub async fn exec(&self) -> Result<()> {
let opt = self;
match &opt.cmd {
Expand Down Expand Up @@ -203,82 +231,103 @@ impl Opt {
base_exchange_rate: 1_0000_0000u128.into(),
};

// First, we need to sync each wallet to the latest block height.
// Parse all the wallets to get the FVKs
let mut fvks = Vec::new();
let mut configs = Vec::new();
let mut paths = Vec::new();
for entry in fs::read_dir(&opt.home)? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
println!("Syncing wallet: {}", path.to_str().unwrap());

let utf8_path =
Utf8PathBuf::from_path_buf(path).expect("should be valid utf8");
let config = PcliConfig::load(utf8_path.join("config.toml"))?;
let mut view_client = self
.view(utf8_path, config.full_viewing_key.clone(), config.grpc_url)
.await?;
// todo: do this in parallel
self.sync(&mut view_client).await?;
println!("Wallet synced successfully");

let notes = view_client.unspent_notes_by_asset_and_address().await?;
let mut total_um_equivalent_amount = Amount::from(0u64);
for (asset_id, map) in notes.iter() {
if *asset_id == *STAKING_TOKEN_ASSET_ID {
let total_amount = map
.iter()
.map(|(_, spendable_notes)| {
spendable_notes
.iter()
.map(|spendable_note| spendable_note.note.amount())
.sum::<Amount>()
})
.sum::<Amount>();
total_um_equivalent_amount += total_amount;
} else if let Ok(delegation_token) =
DelegationToken::from_str(&asset_id.to_string())
{
let total_amount = map
.iter()
.map(|(_, spendable_notes)| {
spendable_notes
.iter()
.map(|spendable_note| spendable_note.note.amount())
.sum::<Amount>()
})
.sum::<Amount>();

// We need to convert the amount to the UM-equivalent amount using the appropriate rate data
let dummy_rate_data = RateData {
identity_key: delegation_token.validator(),
validator_reward_rate: 0u128.into(),
validator_exchange_rate: dummy_base_rate.base_exchange_rate,
};
let um_equivalent_balance =
dummy_rate_data.unbonded_amount(total_amount);
total_um_equivalent_amount += um_equivalent_balance;
configs.push(config.clone());
fvks.push(config.full_viewing_key);
paths.push(utf8_path);
}
}

let genesis_compact_block = self
.fetch_genesis_compact_block(configs[0].grpc_url.clone())
.await?;
let genesis_filtered_block =
genesis::scan_genesis_block(genesis_compact_block, fvks).await?;

// Sync each wallet to the latest block height and check the balances.
for (config, path) in configs.iter().zip(paths.iter()) {
println!("Syncing wallet: {}", path.to_string());

let mut view_client = self
.view(
path.clone(),
config.full_viewing_key.clone(),
config.grpc_url.clone(),
)
.await?;

// todo: do this in parallel
self.sync(&mut view_client).await?;
println!("Wallet synced successfully");

let notes = view_client.unspent_notes_by_asset_and_address().await?;
let mut total_um_equivalent_amount = Amount::from(0u64);
for (asset_id, map) in notes.iter() {
if *asset_id == *STAKING_TOKEN_ASSET_ID {
let total_amount = map
.iter()
.map(|(_, spendable_notes)| {
spendable_notes
.iter()
.map(|spendable_note| spendable_note.note.amount())
.sum::<Amount>()
})
.sum::<Amount>();
total_um_equivalent_amount += total_amount;
} else if let Ok(delegation_token) =
DelegationToken::from_str(&asset_id.to_string())
{
let total_amount = map
.iter()
.map(|(_, spendable_notes)| {
spendable_notes
.iter()
.map(|spendable_note| spendable_note.note.amount())
.sum::<Amount>()
})
.sum::<Amount>();

// We need to convert the amount to the UM-equivalent amount using the appropriate rate data
let dummy_rate_data = RateData {
identity_key: delegation_token.validator(),
validator_reward_rate: 0u128.into(),
validator_exchange_rate: dummy_base_rate.base_exchange_rate,
};
}

println!("FVK: {:?}", config.full_viewing_key);
// todo: calculate the expected um equivalent balance from calling the genesis scanning method
let genesis_um_equivalent_amount = Amount::from(0u64);
println!(
"Genesis UM-equivalent balance: {:?}",
genesis_um_equivalent_amount
);
println!(
"Current UM-equivalent balance: {:?}",
total_um_equivalent_amount
);
let um_equivalent_balance =
dummy_rate_data.unbonded_amount(total_amount);
total_um_equivalent_amount += um_equivalent_balance;
};
}

println!("FVK: {:?}", config.full_viewing_key);
let genesis_um_equivalent_amount = genesis_filtered_block
.balances
.get(&config.full_viewing_key.to_string())
.expect("wallet must have genesis allocation");
println!(
"Genesis UM-equivalent balance: {:?}",
genesis_um_equivalent_amount
);
println!(
"Current UM-equivalent balance: {:?}",
total_um_equivalent_amount
);

// Let the user know if the balance is unexpected or not
if total_um_equivalent_amount < genesis_um_equivalent_amount {
println!(
"✘ Unexpected balance! Balance is less than the genesis balance"
);
} else {
println!("✅ Expected balance! Balance is greater than or equal to the genesis balance");
}
// Let the user know if the balance is unexpected or not
if total_um_equivalent_amount < *genesis_um_equivalent_amount {
println!("✘ Unexpected balance! Balance is less than the genesis balance");
} else {
println!("✅ Expected balance! Balance is greater than or equal to the genesis balance");
}
}
Ok(())
Expand Down

0 comments on commit 47a2c8a

Please sign in to comment.