blockworks-foundation
/
pyth-crosschain
mirror of https://github.com/blockworks-foundation/pyth-crosschain.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[Solana] Load balance treasury (#1267) 

* Checkpoint

* Checkpoint

* Cleanup

* Checkpoint, debug

* Go

* Checkpoint

* Fix

* Add new error and test

* Cleanup

* Add another test

* Keep adding errors

* Another test

* Add comment

* More

* Do it

* Again

* Nice

* Checkpoint

* Checkpoint

* Nice
This commit is contained in:
guibescos 2024-02-01 16:27:35 +00:00 committed by GitHub
parent f5b78e5a8c
commit b2c90a310d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 142 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
target_chains/solana/cli/src/main.rs
View File

@ -16,9 +16,13 @@ use {
Cli,
},
pyth_solana_receiver::{
sdk::deserialize_accumulator_update_data,
sdk::{
deserialize_accumulator_update_data,
DEFAULT_TREASURY_ID,
},
state::config::DataSource,
PostUpdatesAtomicParams,
PostUpdatesParams,
},
pythnet_sdk::wire::v1::MerklePriceUpdate,
serde_wormhole::RawMessage,
@ -256,6 +260,7 @@ pub fn process_post_price_update_atomic(
price_update_keypair.pubkey(),
*wormhole,
header.guardian_set_index,
DEFAULT_TREASURY_ID,
)
.to_account_metas(None);
@ -266,6 +271,7 @@ pub fn process_post_price_update_atomic(
params: PostUpdatesAtomicParams {
merkle_price_update: merkle_price_update.clone(),
vaa: serde_wormhole::to_vec(&(header, body)).unwrap(),
treasury_id: DEFAULT_TREASURY_ID,
},
}
.data(),
@ -467,7 +473,10 @@ pub fn process_write_encoded_vaa_and_post_price_update(
program_id: pyth_solana_receiver::id(),
accounts: post_update_accounts,
data: pyth_solana_receiver::instruction::PostUpdates {
price_update: merkle_price_update.clone(),
params: PostUpdatesParams {
merkle_price_update: merkle_price_update.clone(),
treasury_id: DEFAULT_TREASURY_ID,
},
}
.data(),
};

17
target_chains/solana/programs/pyth-solana-receiver/src/lib.rs
View File

@ -191,7 +191,7 @@ pub mod pyth_solana_receiver {
/// Post a price update using an encoded_vaa account and a MerklePriceUpdate calldata.
/// This should be called after the client has already verified the Vaa via the Wormhole contract.
/// Check out target_chains/solana/cli/src/main.rs for an example of how to do this.
pub fn post_updates(ctx: Context<PostUpdates>, price_update: MerklePriceUpdate) -> Result<()> {
pub fn post_updates(ctx: Context<PostUpdates>, params: PostUpdatesParams) -> Result<()> {
let config = &ctx.accounts.config;
let payer: &Signer<'_> = &ctx.accounts.payer;
let encoded_vaa = VaaAccount::load(&ctx.accounts.encoded_vaa)?;
@ -212,7 +212,7 @@ pub mod pyth_solana_receiver {
price_update_account,
&vaa_components,
encoded_vaa.try_payload()?.as_ref(),
&price_update,
&params.merkle_price_update,
)?;
Ok(())
@ -259,6 +259,7 @@ pub struct AcceptGovernanceAuthorityTransfer<'info> {
}
#[derive(Accounts)]
#[instruction(params: PostUpdatesParams)]
pub struct PostUpdates<'info> {
#[account(mut)]
pub payer: Signer<'info>,
@ -267,7 +268,7 @@ pub struct PostUpdates<'info> {
pub encoded_vaa: AccountInfo<'info>,
#[account(seeds = [CONFIG_SEED.as_ref()], bump)]
pub config: Account<'info, Config>,
#[account(seeds = [TREASURY_SEED.as_ref()], bump)]
#[account(seeds = [TREASURY_SEED.as_ref(), &[params.treasury_id]], bump)]
/// CHECK: This is just a PDA controlled by the program. There is currently no way to withdraw funds from it.
#[account(mut)]
pub treasury: AccountInfo<'info>,
@ -279,6 +280,7 @@ pub struct PostUpdates<'info> {
}
#[derive(Accounts)]
#[instruction(params: PostUpdatesAtomicParams)]
pub struct PostUpdatesAtomic<'info> {
#[account(mut)]
pub payer: Signer<'info>,
@ -289,7 +291,7 @@ pub struct PostUpdatesAtomic<'info> {
pub guardian_set: AccountInfo<'info>,
#[account(seeds = [CONFIG_SEED.as_ref()], bump)]
pub config: Account<'info, Config>,
#[account(mut, seeds = [TREASURY_SEED.as_ref()], bump)]
#[account(mut, seeds = [TREASURY_SEED.as_ref(), &[params.treasury_id]], bump)]
/// CHECK: This is just a PDA controlled by the program. There is currently no way to withdraw funds from it.
pub treasury: AccountInfo<'info>,
/// The contraint is such that either the price_update_account is uninitialized or the payer is the write_authority.
@ -311,6 +313,13 @@ pub struct ReclaimRent<'info> {
pub struct PostUpdatesAtomicParams {
pub vaa: Vec<u8>,
pub merkle_price_update: MerklePriceUpdate,
pub treasury_id: u8,
}
#[derive(Debug, AnchorSerialize, AnchorDeserialize, Clone)]
pub struct PostUpdatesParams {
pub merkle_price_update: MerklePriceUpdate,
pub treasury_id: u8,
}

23
target_chains/solana/programs/pyth-solana-receiver/src/sdk.rs
View File

@ -7,6 +7,7 @@ use {
DataSource,
},
PostUpdatesAtomicParams,
PostUpdatesParams,
CONFIG_SEED,
ID,
TREASURY_SEED,
@ -25,6 +26,9 @@ use {
wormhole_core_bridge_solana::state::GuardianSet,
};
pub const DEFAULT_TREASURY_ID: u8 = 0;
pub const SECONDARY_TREASURY_ID: u8 = 1;
impl accounts::Initialize {
pub fn populate(payer: &Pubkey) -> Self {
let config = get_config_address();
@ -42,9 +46,10 @@ impl accounts::PostUpdatesAtomic {
price_update_account: Pubkey,
wormhole_address: Pubkey,
guardian_set_index: u32,
treasury_id: u8,
) -> Self {
let config = get_config_address();
let treasury = get_treasury_address();
let treasury = get_treasury_address(treasury_id);
let guardian_set = get_guardian_set_address(wormhole_address, guardian_set_index);
@ -62,7 +67,7 @@ impl accounts::PostUpdatesAtomic {
impl accounts::PostUpdates {
pub fn populate(payer: Pubkey, encoded_vaa: Pubkey, price_update_account: Pubkey) -> Self {
let config = get_config_address();
let treasury = get_treasury_address();
let treasury = get_treasury_address(DEFAULT_TREASURY_ID);
accounts::PostUpdates {
payer,
encoded_vaa,
@ -122,7 +127,10 @@ impl instruction::PostUpdates {
program_id: ID,
accounts: post_update_accounts,
data: instruction::PostUpdates {
price_update: merkle_price_update,
params: PostUpdatesParams {
merkle_price_update,
treasury_id: DEFAULT_TREASURY_ID,
},
}
.data(),
}
@ -138,12 +146,14 @@ impl instruction::PostUpdatesAtomic {
guardian_set_index: u32,
vaa: Vec<u8>,
merkle_price_update: MerklePriceUpdate,
treasury_id: u8,
) -> Instruction {
let post_update_accounts = accounts::PostUpdatesAtomic::populate(
payer,
price_update_account,
wormhole_address,
guardian_set_index,
treasury_id,
)
.to_account_metas(None);
Instruction {
@ -153,6 +163,7 @@ impl instruction::PostUpdatesAtomic {
params: PostUpdatesAtomicParams {
vaa,
merkle_price_update,
treasury_id,
},
}
.data(),
@ -252,8 +263,10 @@ impl instruction::ReclaimRent {
}
pub fn get_treasury_address() -> Pubkey {
Pubkey::find_program_address(&[TREASURY_SEED.as_ref()], &ID).0
// There is one treasury for each u8 value
// This is to load balance the write load
pub fn get_treasury_address(treasury_id: u8) -> Pubkey {
Pubkey::find_program_address(&[TREASURY_SEED.as_ref(), &[treasury_id]], &ID).0
}
pub fn get_config_address() -> Pubkey {

17
target_chains/solana/programs/pyth-solana-receiver/tests/common/mod.rs
View File

@ -8,6 +8,8 @@ use {
get_config_address,
get_guardian_set_address,
get_treasury_address,
DEFAULT_TREASURY_ID,
SECONDARY_TREASURY_ID,
},
state::config::{
Config,
@ -200,10 +202,20 @@ pub async fn setup_pyth_receiver(
assert_eq!(config_account, initial_config);
program_simulator
.airdrop(&get_treasury_address(), Rent::default().minimum_balance(0))
.airdrop(
&get_treasury_address(DEFAULT_TREASURY_ID),
Rent::default().minimum_balance(0),
)
.await
.unwrap();
program_simulator
.airdrop(
&get_treasury_address(SECONDARY_TREASURY_ID),
Rent::default().minimum_balance(0),
)
.await
.unwrap();
ProgramTestFixtures {
program_simulator,
@ -215,9 +227,10 @@ pub async fn setup_pyth_receiver(
pub async fn assert_treasury_balance(
program_simulator: &mut ProgramSimulator,
expected_balance: u64,
treasury_id: u8,
) {
let treasury_balance = program_simulator
.get_balance(get_treasury_address())
.get_balance(get_treasury_address(treasury_id))
.await
.unwrap();

30
target_chains/solana/programs/pyth-solana-receiver/tests/test_post_price_update_from_vaa.rs
View File

@ -16,7 +16,10 @@ use {
SetDataSources,
SetFee,
},
sdk::deserialize_accumulator_update_data,
sdk::{
deserialize_accumulator_update_data,
DEFAULT_TREASURY_ID,
},
state::{
config::DataSource,
price_update::{
@ -82,9 +85,10 @@ async fn test_invalid_wormhole_message() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa,
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
None
)
.await
.unwrap_err()
@ -126,6 +130,7 @@ async fn test_invalid_update_message() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa,
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -167,7 +172,7 @@ async fn test_post_price_update_from_vaa() {
)
.await;
assert_treasury_balance(&mut program_simulator, 0).await;
assert_treasury_balance(&mut program_simulator, 0, DEFAULT_TREASURY_ID).await;
let poster = program_simulator.get_funded_keypair().await.unwrap();
let price_update_keypair = Keypair::new();
@ -183,6 +188,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates2[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -204,6 +210,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[2].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -242,6 +249,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -279,6 +287,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -315,6 +324,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID,
),
&vec![&poster, &price_update_keypair],
None,
@ -322,7 +332,7 @@ async fn test_post_price_update_from_vaa() {
.await
.unwrap();
assert_treasury_balance(&mut program_simulator, 1).await;
assert_treasury_balance(&mut program_simulator, 1, DEFAULT_TREASURY_ID).await;
let mut price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())
@ -361,6 +371,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[1].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -371,7 +382,7 @@ async fn test_post_price_update_from_vaa() {
into_transaction_error(ReceiverError::InsufficientFunds)
);
assert_treasury_balance(&mut program_simulator, 1).await;
assert_treasury_balance(&mut program_simulator, 1, DEFAULT_TREASURY_ID).await;
price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())
@ -411,6 +422,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[1].clone(),
DEFAULT_TREASURY_ID,
),
&vec![&poster, &price_update_keypair],
None,
@ -418,7 +430,12 @@ async fn test_post_price_update_from_vaa() {
.await
.unwrap();
assert_treasury_balance(&mut program_simulator, LAMPORTS_PER_SOL + 1).await;
assert_treasury_balance(
&mut program_simulator,
LAMPORTS_PER_SOL + 1,
DEFAULT_TREASURY_ID,
)
.await;
price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())
@ -448,6 +465,7 @@ async fn test_post_price_update_from_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster_2, &price_update_keypair],
None,

11
target_chains/solana/programs/pyth-solana-receiver/tests/test_post_updates.rs
View File

@ -14,7 +14,10 @@ use {
PostUpdates,
ReclaimRent,
},
sdk::deserialize_accumulator_update_data,
sdk::{
deserialize_accumulator_update_data,
DEFAULT_TREASURY_ID,
},
state::price_update::{
PriceUpdateV1,
VerificationLevel,
@ -55,7 +58,7 @@ async fn test_post_updates() {
)
.await;
assert_treasury_balance(&mut program_simulator, 0).await;
assert_treasury_balance(&mut program_simulator, 0, DEFAULT_TREASURY_ID).await;
let poster = program_simulator.get_funded_keypair().await.unwrap();
let price_update_keypair = Keypair::new();
@ -75,7 +78,7 @@ async fn test_post_updates() {
.await
.unwrap();
assert_treasury_balance(&mut program_simulator, 1).await;
assert_treasury_balance(&mut program_simulator, 1, DEFAULT_TREASURY_ID).await;
let mut price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())
@ -107,7 +110,7 @@ async fn test_post_updates() {
.await
.unwrap();
assert_treasury_balance(&mut program_simulator, 2).await;
assert_treasury_balance(&mut program_simulator, 2, DEFAULT_TREASURY_ID).await;
price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())

57
target_chains/solana/programs/pyth-solana-receiver/tests/test_post_updates_atomic.rs
View File

@ -15,6 +15,8 @@ use {
sdk::{
deserialize_accumulator_update_data,
get_guardian_set_address,
DEFAULT_TREASURY_ID,
SECONDARY_TREASURY_ID,
},
state::price_update::{
PriceUpdateV1,
@ -62,7 +64,7 @@ async fn test_post_updates_atomic() {
let poster = program_simulator.get_funded_keypair().await.unwrap();
let price_update_keypair = Keypair::new();
assert_treasury_balance(&mut program_simulator, 0).await;
assert_treasury_balance(&mut program_simulator, 0, DEFAULT_TREASURY_ID).await;
// post one update atomically
program_simulator
@ -74,6 +76,7 @@ async fn test_post_updates_atomic() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID,
),
&vec![&poster, &price_update_keypair],
None,
@ -81,7 +84,7 @@ async fn test_post_updates_atomic() {
.await
.unwrap();
assert_treasury_balance(&mut program_simulator, 1).await;
assert_treasury_balance(&mut program_simulator, 1, DEFAULT_TREASURY_ID).await;
let mut price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())
@ -108,6 +111,7 @@ async fn test_post_updates_atomic() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[1].clone(),
DEFAULT_TREASURY_ID,
),
&vec![&poster, &price_update_keypair],
None,
@ -115,7 +119,8 @@ async fn test_post_updates_atomic() {
.await
.unwrap();
assert_treasury_balance(&mut program_simulator, 2).await;
assert_treasury_balance(&mut program_simulator, 2, DEFAULT_TREASURY_ID).await;
assert_treasury_balance(&mut program_simulator, 0, SECONDARY_TREASURY_ID).await;
price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())
@ -131,6 +136,41 @@ async fn test_post_updates_atomic() {
Message::PriceFeedMessage(price_update_account.price_message),
feed_2
);
// use another treasury account
program_simulator
.process_ix_with_default_compute_limit(
PostUpdatesAtomic::populate(
poster.pubkey(),
price_update_keypair.pubkey(),
BRIDGE_ID,
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
SECONDARY_TREASURY_ID,
),
&vec![&poster, &price_update_keypair],
None,
)
.await
.unwrap();
assert_treasury_balance(&mut program_simulator, 2, DEFAULT_TREASURY_ID).await;
assert_treasury_balance(&mut program_simulator, 1, SECONDARY_TREASURY_ID).await;
price_update_account = program_simulator
.get_anchor_account_data::<PriceUpdateV1>(price_update_keypair.pubkey())
.await
.unwrap();
assert_eq!(price_update_account.write_authority, poster.pubkey());
assert_eq!(
price_update_account.verification_level,
VerificationLevel::Partial(5)
);
assert_eq!(
Message::PriceFeedMessage(price_update_account.price_message),
feed_1
);
}
#[tokio::test]
@ -162,6 +202,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa_buffer_copy,
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -187,6 +228,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa_wrong_num_signatures.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -211,6 +253,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
serde_wormhole::to_vec(&vaa_copy).unwrap(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -234,6 +277,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
serde_wormhole::to_vec(&vaa_copy).unwrap(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -257,6 +301,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
serde_wormhole::to_vec(&vaa_copy).unwrap(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -281,6 +326,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
serde_wormhole::to_vec(&vaa_copy).unwrap(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -304,6 +350,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
serde_wormhole::to_vec(&vaa_copy).unwrap(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -328,6 +375,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
serde_wormhole::to_vec(&vaa_copy).unwrap(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -345,6 +393,7 @@ async fn test_post_updates_atomic_wrong_vaa() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID,
);
let wrong_guardian_set = get_guardian_set_address(BRIDGE_ID, 1);
@ -388,6 +437,7 @@ async fn test_post_updates_atomic_wrong_setup() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,
@ -415,6 +465,7 @@ async fn test_post_updates_atomic_wrong_setup() {
DEFAULT_GUARDIAN_SET_INDEX,
vaa.clone(),
merkle_price_updates[0].clone(),
DEFAULT_TREASURY_ID
),
&vec![&poster, &price_update_keypair],
None,