From 18954fa7655cceedb27592491488fe41ee63aabf Mon Sep 17 00:00:00 2001 From: Serguei <96189670+sekoje@users.noreply.github.com> Date: Thu, 16 Jun 2022 17:52:51 +0000 Subject: [PATCH] Fix Custodian ATA Verification (#44) Co-authored-by: skojenov Co-authored-by: Karl Kempe --- anchor-contributor/package.json | 54 +++--- .../anchor-contributor/src/context.rs | 40 +++-- .../programs/anchor-contributor/src/lib.rs | 159 ++++++++---------- .../anchor-contributor/src/state/sale.rs | 37 +++- 4 files changed, 163 insertions(+), 127 deletions(-) diff --git a/anchor-contributor/package.json b/anchor-contributor/package.json index 0590d2b0..6a52dc86 100644 --- a/anchor-contributor/package.json +++ b/anchor-contributor/package.json @@ -1,29 +1,29 @@ { - "scripts": { - "unit-test": "bash -ac '. test.env && cargo test'", - "integration-test": "bash -ac '. test.env && anchor test'", - "deploy-devnet": "bash migrations/deploy-devnet.sh", - "lint:fix": "prettier */*.js \"*/**/*{.js,.ts}\" -w", - "lint": "prettier */*.js \"*/**/*{.js,.ts}\" --check" - }, - "dependencies": { - "@certusone/wormhole-sdk": "^0.3.4", - "@project-serum/anchor": "^0.24.2", - "@solana/spl-token": "^0.2.0", - "byteify": "^2.0.10", - "elliptic": "^6.5.4", - "ethers": "^5.6.8", - "keccak256": "^1.0.6", - "web3-utils": "^1.7.3" - }, - "devDependencies": { - "@types/bn.js": "^5.1.0", - "@types/chai": "^4.3.0", - "@types/mocha": "^9.0.0", - "chai": "^4.3.4", - "mocha": "^9.0.3", - "prettier": "^2.6.2", - "ts-mocha": "^10.0.0", - "typescript": "^4.3.5" - } + "scripts": { + "unit-test": "bash -ac '. test.env && cargo test'", + "integration-test": "bash -ac '. test.env && anchor test'", + "deploy-devnet": "bash migrations/deploy-devnet.sh", + "lint:fix": "prettier */*.js \"*/**/*{.js,.ts}\" -w", + "lint": "prettier */*.js \"*/**/*{.js,.ts}\" --check" + }, + "dependencies": { + "@certusone/wormhole-sdk": "^0.3.4", + "@project-serum/anchor": "^0.24.2", + "@solana/spl-token": "^0.2.0", + "byteify": "^2.0.10", + "elliptic": "^6.5.4", + "ethers": "^5.6.8", + "keccak256": "^1.0.6", + "web3-utils": "^1.7.3" + }, + "devDependencies": { + "@types/bn.js": "^5.1.0", + "@types/chai": "^4.3.0", + "@types/mocha": "^9.0.0", + "chai": "^4.3.4", + "mocha": "^9.0.3", + "prettier": "^2.6.2", + "ts-mocha": "^10.0.0", + "typescript": "^4.3.5" + } } diff --git a/anchor-contributor/programs/anchor-contributor/src/context.rs b/anchor-contributor/programs/anchor-contributor/src/context.rs index 3d4a9bf1..8a0bc622 100644 --- a/anchor-contributor/programs/anchor-contributor/src/context.rs +++ b/anchor-contributor/programs/anchor-contributor/src/context.rs @@ -2,7 +2,10 @@ use anchor_lang::{ prelude::*, solana_program::sysvar::{clock, rent}, }; -use anchor_spl::token::{Mint, Token, TokenAccount}; +use anchor_spl::{ + associated_token::AssociatedToken, + token::{Mint, Token, TokenAccount}, +}; use crate::{ constants::*, @@ -77,17 +80,17 @@ pub struct InitSale<'info> { pub core_bridge_vaa: AccountInfo<'info>, pub sale_token_mint: Account<'info, Mint>, - #[ - account( - constraint = custodian_sale_token_acct.mint == sale_token_mint.key(), - constraint = custodian_sale_token_acct.owner == custodian.key(), - ) - ] + #[account( + associated_token::mint = sale_token_mint, + associated_token::authority = custodian, + )] + /// This must be an associated token account pub custodian_sale_token_acct: Account<'info, TokenAccount>, #[account(mut)] pub payer: Signer<'info>, pub system_program: Program<'info, System>, + pub associated_token_program: Program<'info, AssociatedToken>, } /// Context provides all accounts required for user to send contribution @@ -145,14 +148,17 @@ pub struct Contribute<'info> { #[account( mut, - constraint = custodian_token_acct.owner == custodian.key(), + associated_token::mint = buyer_token_acct.mint, + associated_token::authority = custodian, )] + /// This must be an associated token account pub custodian_token_acct: Account<'info, TokenAccount>, #[account(mut)] pub owner: Signer<'info>, pub system_program: Program<'info, System>, pub token_program: Program<'info, Token>, + pub associated_token_program: Program<'info, AssociatedToken>, } /// Context provides all accounts required to attest contributions. @@ -308,9 +314,10 @@ pub struct BridgeSealedContribution<'info> { /// CHECK: Check if owned by ATA Program #[account( mut, - constraint = custodian_token_acct.owner == custodian.key(), - constraint = custodian_token_acct.mint == accepted_mint.key() + associated_token::mint = accepted_mint, + associated_token::authority = custodian, )] + /// This must be an associated token account pub custodian_token_acct: Box>, #[account(mut)] @@ -321,6 +328,7 @@ pub struct BridgeSealedContribution<'info> { pub payer: Signer<'info>, pub token_program: Program<'info, Token>, pub system_program: Program<'info, System>, + pub associated_token_program: Program<'info, AssociatedToken>, #[account( constraint = token_bridge.key() == Custodian::token_bridge()? @@ -522,12 +530,14 @@ pub struct SealSale<'info> { pub core_bridge_vaa: AccountInfo<'info>, #[account( - constraint = custodian_sale_token_acct.mint == sale.sale_token_mint, - constraint = custodian_sale_token_acct.owner == custodian.key(), + associated_token::mint = sale.sale_token_mint, + associated_token::authority = custodian, )] + /// This must be an associated token account pub custodian_sale_token_acct: Account<'info, TokenAccount>, pub system_program: Program<'info, System>, + pub associated_token_program: Program<'info, AssociatedToken>, } /// Context provides all accounts required for user to claim his allocation @@ -575,9 +585,10 @@ pub struct ClaimAllocation<'info> { #[account( mut, - constraint = custodian_sale_token_acct.mint == sale.sale_token_mint, - constraint = custodian_sale_token_acct.owner == custodian.key(), + associated_token::mint = sale.sale_token_mint, + associated_token::authority = custodian, )] + /// This must be an associated token account pub custodian_sale_token_acct: Account<'info, TokenAccount>, #[account( @@ -591,6 +602,7 @@ pub struct ClaimAllocation<'info> { pub owner: Signer<'info>, pub system_program: Program<'info, System>, pub token_program: Program<'info, Token>, + pub associated_token_program: Program<'info, AssociatedToken>, } /// Context provides all accounts required for user to claim his refunds diff --git a/anchor-contributor/programs/anchor-contributor/src/lib.rs b/anchor-contributor/programs/anchor-contributor/src/lib.rs index 51a5e843..d98849b0 100644 --- a/anchor-contributor/programs/anchor-contributor/src/lib.rs +++ b/anchor-contributor/programs/anchor-contributor/src/lib.rs @@ -99,49 +99,67 @@ pub mod anchor_contributor { /// contribution amount and the contribution will be transferred from the buyer's /// associated token account to the custodian's associated token account. pub fn contribute(ctx: Context, amount: u64, kyc_signature: Vec) -> Result<()> { - // We need to use the buyer's associated token account to help us find the token index - // for this particular mint he wishes to contribute. - let sale = &ctx.accounts.sale; + // buyer's token account -> custodian's associated token account + // These references will be used throughout the instruction let buyer_token_acct = &ctx.accounts.buyer_token_acct; - let (idx, asset) = sale.get_total_info(&buyer_token_acct.mint)?; - let token_index = asset.token_index; + let custodian_token_acct = &ctx.accounts.custodian_token_acct; - // If the buyer account wasn't initialized before, we will do so here. This initializes - // the state for all of this buyer's contributions. - let buyer = &mut ctx.accounts.buyer; - if !buyer.initialized { - buyer.initialize(sale.totals.len()); - } - - // We verify the KYC signature by encoding specific details of this contribution the - // same way the KYC entity signed for the transaction. If we cannot recover the KYC's - // public key using ecdsa recovery, we cannot allow the contribution to continue. - // - // We also refer to the buyer (owner) of this instruction as the transfer_authority + // We refer to the buyer (owner) of this instruction as the transfer_authority // for the SPL transfer that will happen at the end of all the accounting processing. let transfer_authority = &ctx.accounts.owner; - sale.verify_kyc_authority( - token_index, - amount, - &transfer_authority.key(), - buyer.contributions[idx].amount, - &kyc_signature, - )?; + + // Find indices used for contribution accounting + let (idx, token_index) = { + let sale = &ctx.accounts.sale; + + // We need to use the buyer's associated token account to help us find the token index + // for this particular mint he wishes to contribute. + let (idx, asset) = sale.get_total_info(&buyer_token_acct.mint)?; + + // is this overkill to check given anchor constraints? + asset.verify_ata( + &custodian_token_acct.to_account_info(), + &ctx.accounts.custodian.key(), + )?; + + // If the buyer account wasn't initialized before, we will do so here. This initializes + // the state for all of this buyer's contributions. + let buyer = &mut ctx.accounts.buyer; + if !buyer.initialized { + buyer.initialize(sale.totals.len()); + } + + let token_index = asset.token_index; + + // We verify the KYC signature by encoding specific details of this contribution the + // same way the KYC entity signed for the transaction. If we cannot recover the KYC's + // public key using ecdsa recovery, we cannot allow the contribution to continue. + sale.verify_kyc_authority( + token_index, + amount, + &transfer_authority.key(), + buyer.contributions[idx].amount, + &kyc_signature, + )?; + + (idx, token_index) + }; // We need to grab the current block's timestamp and verify that the buyer is allowed // to contribute now. A user cannot contribute before the sale has started. If all the // sale checks pass, the Sale's total contributions uptick to reflect this buyer's // contribution. let clock = Clock::get()?; - let sale = &mut ctx.accounts.sale; - sale.update_total_contributions(clock.unix_timestamp, token_index, amount)?; + ctx.accounts + .sale + .update_total_contributions(clock.unix_timestamp, token_index, amount)?; // And we do the same with the Buyer account. - buyer.contribute(idx, amount)?; + ctx.accounts.buyer.contribute(idx, amount)?; // Finally transfer SPL tokens from the buyer's associated token account to the - // custodian's associated token account. - let custodian_token_acct = &ctx.accounts.custodian_token_acct; + // custodian's associated token account. Verify the custodian's associated + // token account token::transfer( CpiContext::new( ctx.accounts.token_program.to_account_info(), @@ -295,27 +313,16 @@ pub mod anchor_contributor { ); // We will mutate the buyer's accounting and state for each contributed mint. - for (total, custodian_token_acct) in izip!(totals, custodian_token_accts) { - // Verify the authority of the custodian's associated token account + for (asset, custodian_token_acct) in izip!(totals, custodian_token_accts) { + // re-derive custodian_token_acct address and check it. + // Verifies the authority and mint of the custodian's associated token account + let ata = asset.verify_ata(custodian_token_acct, &custodian.key())?; require!( - token::accessor::authority(&custodian_token_acct)? == custodian.key(), - ContributorError::InvalidAccount - ); - - // We need to verify that the mints are the same between the two - // associated token accounts. After which, we will use the sale account - // to find the correct index to reference in the buyer account. - require!( - token::accessor::mint(&custodian_token_acct)? == total.mint, - ContributorError::InvalidAccount - ); - // verify custodian ata has enough funds after conductor accounting - require!( - token::accessor::amount(&custodian_token_acct)? >= total.contributions, + ata.amount >= asset.contributions, ContributorError::InsufficientFunds ); - total.prepare_for_transfer(); + asset.prepare_for_transfer(); } // Finish instruction. @@ -338,8 +345,16 @@ pub mod anchor_contributor { require!(sale.is_sealed(), ContributorError::SaleNotSealed); let custodian_token_acct = &ctx.accounts.custodian_token_acct; + let accepted_mint = &ctx.accounts.accepted_mint; let (idx, asset) = sale.get_total_info(&accepted_mint.key())?; + + let custodian = &ctx.accounts.custodian; + + // is this overkill to check given anchor constraints? + asset.verify_ata(&custodian_token_acct.to_account_info(), &custodian.key())?; + + // check if asset is in the correct state after sealing the sale require!( asset.is_ready_for_transfer(), ContributorError::TransferNotAllowed @@ -353,7 +368,7 @@ pub mod anchor_contributor { token::Approve { to: custodian_token_acct.to_account_info(), delegate: authority_signer.to_account_info(), - authority: ctx.accounts.custodian.to_account_info(), + authority: custodian.to_account_info(), }, &[&[SEED_PREFIX_CUSTODIAN.as_bytes(), &[ctx.bumps["custodian"]]]], ), @@ -384,7 +399,7 @@ pub mod anchor_contributor { AccountMeta::new(ctx.accounts.payer.key(), true), AccountMeta::new_readonly(ctx.accounts.token_bridge_config.key(), false), AccountMeta::new(custodian_token_acct.key(), false), - AccountMeta::new_readonly(ctx.accounts.custodian.key(), true), + AccountMeta::new_readonly(custodian.key(), true), AccountMeta::new(accepted_mint.key(), false), AccountMeta::new_readonly(wrapped_meta.key(), false), AccountMeta::new_readonly(authority_signer.key(), false), @@ -522,27 +537,14 @@ pub mod anchor_contributor { // We will mutate the buyer's accounting and state for each contributed mint. let buyer = &mut ctx.accounts.buyer; - for (idx, (total, custodian_token_acct, buyer_token_acct)) in + for (idx, (asset, custodian_token_acct, buyer_token_acct)) in izip!(totals, custodian_token_accts, buyer_token_accts).enumerate() { // Verify the custodian's associated token account - require!( - token::accessor::authority(&custodian_token_acct)? == transfer_authority.key(), - ContributorError::InvalidAccount - ); - require!( - token::accessor::mint(&custodian_token_acct)? == total.mint, - ContributorError::InvalidAccount - ); - // And verify the buyer's associated token account - require!( - token::accessor::authority(&buyer_token_acct)? == owner.key(), - ContributorError::InvalidAccount - ); - require!( - token::accessor::mint(&buyer_token_acct)? == total.mint, - ContributorError::InvalidAccount - ); + asset.verify_ata(custodian_token_acct, &ctx.accounts.custodian.key())?; + + // And verify the buyer's token account + asset.verify_token_account(buyer_token_acct, &owner.key())?; // Now calculate the refund and transfer to the buyer's associated // token account if there is any amount to refund. @@ -695,31 +697,18 @@ pub mod anchor_contributor { // We will mutate the buyer's accounting and state for each contributed mint. let buyer = &mut ctx.accounts.buyer; - for (idx, (total, custodian_token_acct, buyer_token_acct)) in + for (idx, (asset, custodian_token_acct, buyer_token_acct)) in izip!(totals, custodian_token_accts, buyer_token_accts).enumerate() { // Verify the custodian's associated token account - require!( - token::accessor::authority(&custodian_token_acct)? == transfer_authority.key(), - ContributorError::InvalidAccount - ); - require!( - token::accessor::mint(&custodian_token_acct)? == total.mint, - ContributorError::InvalidAccount - ); - // And verify the buyer's associated token account - require!( - token::accessor::authority(&buyer_token_acct)? == owner.key(), - ContributorError::InvalidAccount - ); - require!( - token::accessor::mint(&buyer_token_acct)? == total.mint, - ContributorError::InvalidAccount - ); + asset.verify_ata(custodian_token_acct, &ctx.accounts.custodian.key())?; + + // And verify the buyer's token account + asset.verify_token_account(buyer_token_acct, &owner.key())?; // Now calculate the excess contribution and transfer to the // buyer's associated token account if there is any amount calculated. - let excess = buyer.claim_excess(idx, total)?; + let excess = buyer.claim_excess(idx, asset)?; if excess == 0 { continue; } diff --git a/anchor-contributor/programs/anchor-contributor/src/state/sale.rs b/anchor-contributor/programs/anchor-contributor/src/state/sale.rs index 7012c260..37d0104f 100644 --- a/anchor-contributor/programs/anchor-contributor/src/state/sale.rs +++ b/anchor-contributor/programs/anchor-contributor/src/state/sale.rs @@ -1,5 +1,8 @@ use anchor_lang::{prelude::*, solana_program::keccak}; -use anchor_spl::token::Mint; +use anchor_spl::{ + associated_token::get_associated_token_address, + token::{Mint, TokenAccount}, +}; use num::{bigint::BigUint, traits::ToPrimitive}; use num_derive::*; use std::{mem::size_of_val, u64}; @@ -101,6 +104,38 @@ impl AssetTotal { pub fn set_transferred(&mut self) { self.status = AssetStatus::TransferredToConductor; } + + pub fn verify_ata( + &self, + token_acct_info: &AccountInfo, + authority: &Pubkey, + ) -> Result { + require!( + get_associated_token_address(&authority, &self.mint) == token_acct_info.key(), + ContributorError::InvalidAccount + ); + + AssetTotal::deserialize_token_account(token_acct_info) + } + + pub fn verify_token_account( + &self, + token_acct_info: &AccountInfo, + owner: &Pubkey, + ) -> Result { + let token_acct = AssetTotal::deserialize_token_account(token_acct_info)?; + require!(token_acct.owner == *owner, ContributorError::InvalidAccount); + require!( + token_acct.mint == self.mint, + ContributorError::InvalidAccount + ); + Ok(token_acct) + } + + pub fn deserialize_token_account(token_acct_info: &AccountInfo) -> Result { + let mut bf: &[u8] = &token_acct_info.try_borrow_data()?; + TokenAccount::try_deserialize_unchecked(&mut bf) + } } impl Sale {