From bff1868ccc7f4c66eedaec930e01cac9f3f30e13 Mon Sep 17 00:00:00 2001 From: Hardhat Chad Date: Mon, 26 Aug 2024 21:02:52 +0000 Subject: [PATCH] instruction macro and simplification --- Cargo.toml | 6 +- api/src/instruction.rs | 243 +++++++++++++++++--------------------- api/src/state/bus.rs | 8 +- api/src/state/config.rs | 4 +- api/src/state/mod.rs | 2 +- api/src/state/proof.rs | 4 +- api/src/state/treasury.rs | 4 +- program/src/claim.rs | 4 +- program/src/initialize.rs | 9 +- program/src/mine.rs | 4 +- program/src/open.rs | 4 +- program/src/stake.rs | 4 +- program/src/upgrade.rs | 4 +- utils/src/loaders.rs | 60 ++++++++++ utils/src/macros.rs | 19 ++- utils/src/traits.rs | 8 +- 16 files changed, 217 insertions(+), 170 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3ad3e3f..3005406 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,14 +19,12 @@ const-crypto = "0.1.0" drillx = { version = "2.0.0", features = ["solana"] } mpl-token-metadata = "4.1.2" num_enum = "0.7.2" -ore-api = { path = "api", version = "2.1.7" } -ore-utils = { path = "utils", features = ["spl"], version = "2.1.7" } +ore-api = { path = "api", version = "2.1.8" } +ore-utils = { path = "utils", features = ["spl"], version = "2.1.8" } solana-program = "^1.18" spl-token = { version = "^4", features = ["no-entrypoint"] } spl-associated-token-account = { version = "^2.3", features = [ "no-entrypoint" ] } static_assertions = "1.1.0" thiserror = "1.0.57" -# [patch.crates-io] -# drillx = { path = "../drillx/drillx" } diff --git a/api/src/instruction.rs b/api/src/instruction.rs index 0cdfa8d..0ee2457 100644 --- a/api/src/instruction.rs +++ b/api/src/instruction.rs @@ -1,18 +1,20 @@ use bytemuck::{Pod, Zeroable}; use drillx::Solution; use num_enum::TryFromPrimitive; -use ore_utils::instruction; +use ore_utils::*; use solana_program::{ instruction::{AccountMeta, Instruction}, pubkey::Pubkey, system_program, sysvar, }; -use crate::consts::*; +use crate::{ + consts::*, + state::{bus_pda, config_pda, proof_pda, treasury_pda}, +}; #[repr(u8)] #[derive(Clone, Copy, Debug, Eq, PartialEq, TryFromPrimitive)] -#[rustfmt::skip] pub enum OreInstruction { // User Claim = 0, @@ -28,15 +30,52 @@ pub enum OreInstruction { Initialize = 100, } -impl OreInstruction { - pub fn to_vec(&self) -> Vec { - vec![*self as u8] - } +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Claim { + pub amount: [u8; 8], } #[repr(C)] #[derive(Clone, Copy, Debug, Pod, Zeroable)] -pub struct InitializeArgs { +pub struct Close {} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Mine { + pub digest: [u8; 16], + pub nonce: [u8; 8], +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Open { + pub bump: u8, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Reset {} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Stake { + pub amount: [u8; 8], +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Update {} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Upgrade { + pub amount: [u8; 8], +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct Initialize { pub bus_0_bump: u8, pub bus_1_bump: u8, pub bus_2_bump: u8, @@ -51,43 +90,15 @@ pub struct InitializeArgs { pub treasury_bump: u8, } -#[repr(C)] -#[derive(Clone, Copy, Debug, Pod, Zeroable)] -pub struct OpenArgs { - pub bump: u8, -} - -#[repr(C)] -#[derive(Clone, Copy, Debug, Pod, Zeroable)] -pub struct MineArgs { - pub digest: [u8; 16], - pub nonce: [u8; 8], -} - -#[repr(C)] -#[derive(Clone, Copy, Debug, Pod, Zeroable)] -pub struct ClaimArgs { - pub amount: [u8; 8], -} - -#[repr(C)] -#[derive(Clone, Copy, Debug, Pod, Zeroable)] -pub struct StakeArgs { - pub amount: [u8; 8], -} - -#[repr(C)] -#[derive(Clone, Copy, Debug, Pod, Zeroable)] -pub struct UpgradeArgs { - pub amount: [u8; 8], -} - -instruction!(InitializeArgs); -instruction!(OpenArgs); -instruction!(MineArgs); -instruction!(ClaimArgs); -instruction!(StakeArgs); -instruction!(UpgradeArgs); +instruction!(OreInstruction, Claim); +instruction!(OreInstruction, Close); +instruction!(OreInstruction, Mine); +instruction!(OreInstruction, Open); +instruction!(OreInstruction, Reset); +instruction!(OreInstruction, Stake); +instruction!(OreInstruction, Update); +instruction!(OreInstruction, Upgrade); +instruction!(OreInstruction, Initialize); /// Builds an auth instruction. pub fn auth(proof: Pubkey) -> Instruction { @@ -100,7 +111,7 @@ pub fn auth(proof: Pubkey) -> Instruction { /// Builds a claim instruction. pub fn claim(signer: Pubkey, beneficiary: Pubkey, amount: u64) -> Instruction { - let proof = Pubkey::find_program_address(&[PROOF, signer.as_ref()], &crate::id()).0; + let proof = proof_pda(signer).0; let treasury_tokens = spl_associated_token_account::get_associated_token_address( &TREASURY_ADDRESS, &MINT_ADDRESS, @@ -115,40 +126,30 @@ pub fn claim(signer: Pubkey, beneficiary: Pubkey, amount: u64) -> Instruction { AccountMeta::new(treasury_tokens, false), AccountMeta::new_readonly(spl_token::id(), false), ], - data: [ - OreInstruction::Claim.to_vec(), - ClaimArgs { - amount: amount.to_le_bytes(), - } - .to_bytes() - .to_vec(), - ] - .concat(), + data: Claim { + amount: amount.to_le_bytes(), + } + .to_bytes(), } } /// Builds a close instruction. pub fn close(signer: Pubkey) -> Instruction { - let proof_pda = Pubkey::find_program_address(&[PROOF, signer.as_ref()], &crate::id()); + let proof = proof_pda(signer).0; Instruction { program_id: crate::id(), accounts: vec![ AccountMeta::new(signer, true), - AccountMeta::new(proof_pda.0, false), + AccountMeta::new(proof, false), AccountMeta::new_readonly(solana_program::system_program::id(), false), ], - data: OreInstruction::Close.to_vec(), + data: Close {}.to_bytes(), } } /// Builds a mine instruction. -pub fn mine( - signer: Pubkey, - proof_authority: Pubkey, - bus: Pubkey, - solution: Solution, -) -> Instruction { - let proof = Pubkey::find_program_address(&[PROOF, proof_authority.as_ref()], &crate::id()).0; +pub fn mine(signer: Pubkey, authority: Pubkey, bus: Pubkey, solution: Solution) -> Instruction { + let proof = proof_pda(authority).0; Instruction { program_id: crate::id(), accounts: vec![ @@ -159,22 +160,17 @@ pub fn mine( AccountMeta::new_readonly(sysvar::instructions::id(), false), AccountMeta::new_readonly(sysvar::slot_hashes::id(), false), ], - data: [ - OreInstruction::Mine.to_vec(), - MineArgs { - digest: solution.d, - nonce: solution.n, - } - .to_bytes() - .to_vec(), - ] - .concat(), + data: Mine { + digest: solution.d, + nonce: solution.n, + } + .to_bytes(), } } /// Builds an open instruction. pub fn open(signer: Pubkey, miner: Pubkey, payer: Pubkey) -> Instruction { - let proof_pda = Pubkey::find_program_address(&[PROOF, signer.as_ref()], &crate::id()); + let proof_pda = proof_pda(signer); Instruction { program_id: crate::id(), accounts: vec![ @@ -185,11 +181,7 @@ pub fn open(signer: Pubkey, miner: Pubkey, payer: Pubkey) -> Instruction { AccountMeta::new_readonly(solana_program::system_program::id(), false), AccountMeta::new_readonly(sysvar::slot_hashes::id(), false), ], - data: [ - OreInstruction::Open.to_vec(), - OpenArgs { bump: proof_pda.1 }.to_bytes().to_vec(), - ] - .concat(), + data: Open { bump: proof_pda.1 }.to_bytes(), } } @@ -217,13 +209,13 @@ pub fn reset(signer: Pubkey) -> Instruction { AccountMeta::new(treasury_tokens, false), AccountMeta::new_readonly(spl_token::id(), false), ], - data: OreInstruction::Reset.to_vec(), + data: Reset {}.to_bytes(), } } /// Build a stake instruction. pub fn stake(signer: Pubkey, sender: Pubkey, amount: u64) -> Instruction { - let proof = Pubkey::find_program_address(&[PROOF, signer.as_ref()], &crate::id()).0; + let proof = proof_pda(signer).0; let treasury_tokens = spl_associated_token_account::get_associated_token_address( &TREASURY_ADDRESS, &MINT_ADDRESS, @@ -237,21 +229,16 @@ pub fn stake(signer: Pubkey, sender: Pubkey, amount: u64) -> Instruction { AccountMeta::new(treasury_tokens, false), AccountMeta::new_readonly(spl_token::id(), false), ], - data: [ - OreInstruction::Stake.to_vec(), - StakeArgs { - amount: amount.to_le_bytes(), - } - .to_bytes() - .to_vec(), - ] - .concat(), + data: Stake { + amount: amount.to_le_bytes(), + } + .to_bytes(), } } // Build an update instruction. pub fn update(signer: Pubkey, miner: Pubkey) -> Instruction { - let proof = Pubkey::find_program_address(&[PROOF, signer.as_ref()], &crate::id()).0; + let proof = proof_pda(signer).0; Instruction { program_id: crate::id(), accounts: vec![ @@ -259,7 +246,7 @@ pub fn update(signer: Pubkey, miner: Pubkey) -> Instruction { AccountMeta::new_readonly(miner, false), AccountMeta::new(proof, false), ], - data: OreInstruction::Update.to_vec(), + data: Update {}.to_bytes(), } } @@ -276,33 +263,28 @@ pub fn upgrade(signer: Pubkey, beneficiary: Pubkey, sender: Pubkey, amount: u64) AccountMeta::new(TREASURY_ADDRESS, false), AccountMeta::new_readonly(spl_token::id(), false), ], - data: [ - OreInstruction::Upgrade.to_vec(), - UpgradeArgs { - amount: amount.to_le_bytes(), - } - .to_bytes() - .to_vec(), - ] - .concat(), + data: Upgrade { + amount: amount.to_le_bytes(), + } + .to_bytes(), } } /// Builds an initialize instruction. pub fn initialize(signer: Pubkey) -> Instruction { let bus_pdas = [ - Pubkey::find_program_address(&[BUS, &[0]], &crate::id()), - Pubkey::find_program_address(&[BUS, &[1]], &crate::id()), - Pubkey::find_program_address(&[BUS, &[2]], &crate::id()), - Pubkey::find_program_address(&[BUS, &[3]], &crate::id()), - Pubkey::find_program_address(&[BUS, &[4]], &crate::id()), - Pubkey::find_program_address(&[BUS, &[5]], &crate::id()), - Pubkey::find_program_address(&[BUS, &[6]], &crate::id()), - Pubkey::find_program_address(&[BUS, &[7]], &crate::id()), + bus_pda(0), + bus_pda(1), + bus_pda(2), + bus_pda(3), + bus_pda(4), + bus_pda(5), + bus_pda(6), + bus_pda(7), ]; - let config_pda = Pubkey::find_program_address(&[CONFIG], &crate::id()); + let config_pda = config_pda(); let mint_pda = Pubkey::find_program_address(&[MINT, MINT_NOISE.as_slice()], &crate::id()); - let treasury_pda = Pubkey::find_program_address(&[TREASURY], &crate::id()); + let treasury_pda = treasury_pda(); let treasury_tokens = spl_associated_token_account::get_associated_token_address(&treasury_pda.0, &mint_pda.0); let metadata_pda = Pubkey::find_program_address( @@ -336,25 +318,20 @@ pub fn initialize(signer: Pubkey) -> Instruction { AccountMeta::new_readonly(mpl_token_metadata::ID, false), AccountMeta::new_readonly(sysvar::rent::id(), false), ], - data: [ - OreInstruction::Initialize.to_vec(), - InitializeArgs { - bus_0_bump: bus_pdas[0].1, - bus_1_bump: bus_pdas[1].1, - bus_2_bump: bus_pdas[2].1, - bus_3_bump: bus_pdas[3].1, - bus_4_bump: bus_pdas[4].1, - bus_5_bump: bus_pdas[5].1, - bus_6_bump: bus_pdas[6].1, - bus_7_bump: bus_pdas[7].1, - config_bump: config_pda.1, - metadata_bump: metadata_pda.1, - mint_bump: mint_pda.1, - treasury_bump: treasury_pda.1, - } - .to_bytes() - .to_vec(), - ] - .concat(), + data: Initialize { + bus_0_bump: bus_pdas[0].1, + bus_1_bump: bus_pdas[1].1, + bus_2_bump: bus_pdas[2].1, + bus_3_bump: bus_pdas[3].1, + bus_4_bump: bus_pdas[4].1, + bus_5_bump: bus_pdas[5].1, + bus_6_bump: bus_pdas[6].1, + bus_7_bump: bus_pdas[7].1, + config_bump: config_pda.1, + metadata_bump: metadata_pda.1, + mint_bump: mint_pda.1, + treasury_bump: treasury_pda.1, + } + .to_bytes(), } } diff --git a/api/src/state/bus.rs b/api/src/state/bus.rs index 4ff5d8b..0d07296 100644 --- a/api/src/state/bus.rs +++ b/api/src/state/bus.rs @@ -4,7 +4,7 @@ use solana_program::pubkey::Pubkey; use crate::consts::BUS; -use super::AccountDiscriminator; +use super::OreAccount; /// Bus accounts are responsible for distributing mining rewards. There are 8 busses total /// to minimize write-lock contention and allow Solana to process mine instructions in parallel. @@ -26,8 +26,8 @@ pub struct Bus { } /// Fetch the PDA of a bus account. -pub fn bus_pda(id: u64) -> (Pubkey, u8) { - Pubkey::find_program_address(&[BUS, id.to_le_bytes().as_slice()], &crate::id()) +pub fn bus_pda(id: u8) -> (Pubkey, u8) { + Pubkey::find_program_address(&[BUS, &[id]], &crate::id()) } -account!(AccountDiscriminator, Bus); +account!(OreAccount, Bus); diff --git a/api/src/state/config.rs b/api/src/state/config.rs index 5933f36..35dded5 100644 --- a/api/src/state/config.rs +++ b/api/src/state/config.rs @@ -4,7 +4,7 @@ use solana_program::pubkey::Pubkey; use crate::consts::CONFIG; -use super::AccountDiscriminator; +use super::OreAccount; /// Config is a singleton account which manages program global variables. #[repr(C)] @@ -28,4 +28,4 @@ pub fn config_pda() -> (Pubkey, u8) { Pubkey::find_program_address(&[CONFIG], &crate::id()) } -account!(AccountDiscriminator, Config); +account!(OreAccount, Config); diff --git a/api/src/state/mod.rs b/api/src/state/mod.rs index e228cac..6803553 100644 --- a/api/src/state/mod.rs +++ b/api/src/state/mod.rs @@ -12,7 +12,7 @@ use num_enum::{IntoPrimitive, TryFromPrimitive}; #[repr(u8)] #[derive(Clone, Copy, Debug, Eq, PartialEq, IntoPrimitive, TryFromPrimitive)] -pub enum AccountDiscriminator { +pub enum OreAccount { Bus = 100, Config = 101, Proof = 102, diff --git a/api/src/state/proof.rs b/api/src/state/proof.rs index 03fb9b0..ba68660 100644 --- a/api/src/state/proof.rs +++ b/api/src/state/proof.rs @@ -4,7 +4,7 @@ use solana_program::pubkey::Pubkey; use crate::consts::PROOF; -use super::AccountDiscriminator; +use super::OreAccount; /// Proof accounts track a miner's current hash, claimable rewards, and lifetime stats. /// Every miner is allowed one proof account which is required by the program to mine or claim rewards. @@ -44,4 +44,4 @@ pub fn proof_pda(authority: Pubkey) -> (Pubkey, u8) { Pubkey::find_program_address(&[PROOF, authority.as_ref()], &crate::id()) } -account!(AccountDiscriminator, Proof); +account!(OreAccount, Proof); diff --git a/api/src/state/treasury.rs b/api/src/state/treasury.rs index e2dede5..20c34fc 100644 --- a/api/src/state/treasury.rs +++ b/api/src/state/treasury.rs @@ -4,7 +4,7 @@ use solana_program::pubkey::Pubkey; use crate::consts::TREASURY; -use super::AccountDiscriminator; +use super::OreAccount; /// Treasury is a singleton account which is the mint authority for the ORE token and the authority of /// the program's global token account. @@ -17,4 +17,4 @@ pub fn treasury_pda() -> (Pubkey, u8) { Pubkey::find_program_address(&[TREASURY], &crate::id()) } -account!(AccountDiscriminator, Treasury); +account!(OreAccount, Treasury); diff --git a/program/src/claim.rs b/program/src/claim.rs index fa16658..aedf620 100644 --- a/program/src/claim.rs +++ b/program/src/claim.rs @@ -1,4 +1,4 @@ -use ore_api::{consts::*, error::OreError, instruction::ClaimArgs, loaders::*, state::Proof}; +use ore_api::{consts::*, error::OreError, instruction::*, loaders::*, state::Proof}; use ore_utils::*; use solana_program::{ account_info::AccountInfo, entrypoint::ProgramResult, program_error::ProgramError, @@ -7,7 +7,7 @@ use solana_program::{ /// Claim distributes claimable ORE from the treasury to a miner. pub fn process_claim(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult { // Parse args. - let args = ClaimArgs::try_from_bytes(data)?; + let args = Claim::try_from_bytes(data)?; let amount = u64::from_le_bytes(args.amount); // Load accounts. diff --git a/program/src/initialize.rs b/program/src/initialize.rs index 46bd4a5..6759876 100644 --- a/program/src/initialize.rs +++ b/program/src/initialize.rs @@ -7,18 +7,15 @@ use ore_api::{ }; use ore_utils::*; use solana_program::{ - account_info::AccountInfo, - entrypoint::ProgramResult, - program_error::ProgramError, - program_pack::Pack, - system_program, {self, sysvar}, + self, account_info::AccountInfo, entrypoint::ProgramResult, program_error::ProgramError, + program_pack::Pack, system_program, sysvar, }; use spl_token::state::Mint; /// Initialize sets up the ORE program to begin mining. pub fn process_initialize(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult { // Parse args. - let args = InitializeArgs::try_from_bytes(data)?; + let args = Initialize::try_from_bytes(data)?; // Load accounts. let [signer, bus_0_info, bus_1_info, bus_2_info, bus_3_info, bus_4_info, bus_5_info, bus_6_info, bus_7_info, config_info, metadata_info, mint_info, treasury_info, treasury_tokens_info, system_program, token_program, associated_token_program, metadata_program, rent_sysvar] = diff --git a/program/src/mine.rs b/program/src/mine.rs index 3d1180a..9efd3dc 100644 --- a/program/src/mine.rs +++ b/program/src/mine.rs @@ -5,7 +5,7 @@ use ore_api::{ consts::*, error::OreError, event::MineEvent, - instruction::MineArgs, + instruction::Mine, loaders::*, state::{Bus, Config, Proof}, }; @@ -28,7 +28,7 @@ use solana_program::{ /// Mine validates hashes and increments a miner's collectable balance. pub fn process_mine(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult { // Parse args. - let args = MineArgs::try_from_bytes(data)?; + let args = Mine::try_from_bytes(data)?; // Load accounts. let [signer, bus_info, config_info, proof_info, instructions_sysvar, slot_hashes_sysvar] = diff --git a/program/src/open.rs b/program/src/open.rs index e42ecd2..4ded87c 100644 --- a/program/src/open.rs +++ b/program/src/open.rs @@ -1,6 +1,6 @@ use std::mem::size_of; -use ore_api::{consts::*, instruction::OpenArgs, state::Proof}; +use ore_api::{consts::*, instruction::Open, state::Proof}; use ore_utils::*; use solana_program::{ account_info::AccountInfo, @@ -16,7 +16,7 @@ use solana_program::{ /// Open creates a new proof account to track a miner's state. pub fn process_open(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult { // Parse args. - let args = OpenArgs::try_from_bytes(data)?; + let args = Open::try_from_bytes(data)?; // Load accounts. let [signer, miner_info, payer_info, proof_info, system_program, slot_hashes_info] = accounts diff --git a/program/src/stake.rs b/program/src/stake.rs index e82912a..e64f5b3 100644 --- a/program/src/stake.rs +++ b/program/src/stake.rs @@ -1,4 +1,4 @@ -use ore_api::{consts::*, instruction::StakeArgs, loaders::*, state::Proof}; +use ore_api::{consts::*, instruction::Stake, loaders::*, state::Proof}; use ore_utils::*; use solana_program::{ account_info::AccountInfo, clock::Clock, entrypoint::ProgramResult, @@ -8,7 +8,7 @@ use solana_program::{ /// Stake deposits ORE into a proof account to earn multiplier. pub fn process_stake(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult { // Parse args. - let args = StakeArgs::try_from_bytes(data)?; + let args = Stake::try_from_bytes(data)?; let amount = u64::from_le_bytes(args.amount); // Load accounts. diff --git a/program/src/upgrade.rs b/program/src/upgrade.rs index a2bf8f5..a167872 100644 --- a/program/src/upgrade.rs +++ b/program/src/upgrade.rs @@ -1,4 +1,4 @@ -use ore_api::{consts::*, error::OreError, instruction::StakeArgs}; +use ore_api::{consts::*, error::OreError, instruction::Stake}; use ore_utils::*; use solana_program::{ account_info::AccountInfo, entrypoint::ProgramResult, program_error::ProgramError, @@ -9,7 +9,7 @@ use spl_token::state::Mint; /// Upgrade allows a user to migrate a v1 token to a v2 token at a 1:1 exchange rate. pub fn process_upgrade(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult { // Parse args - let args = StakeArgs::try_from_bytes(data)?; + let args = Stake::try_from_bytes(data)?; let amount = u64::from_le_bytes(args.amount); // Load accounts diff --git a/utils/src/loaders.rs b/utils/src/loaders.rs index 365218f..72681fe 100644 --- a/utils/src/loaders.rs +++ b/utils/src/loaders.rs @@ -146,6 +146,30 @@ pub fn load_mint( Ok(()) } +/// Errors if: +/// - Owner is not SPL token program. +/// - Data is empty. +/// - Data cannot deserialize into a mint account. +/// - Expected to be writable, but is not. +#[cfg(feature = "spl")] +pub fn load_any_mint(info: &AccountInfo<'_>, is_writable: bool) -> Result<(), ProgramError> { + if info.owner.ne(&spl_token::id()) { + return Err(ProgramError::InvalidAccountOwner); + } + + if info.data_is_empty() { + return Err(ProgramError::UninitializedAccount); + } + + Mint::unpack(&info.data.borrow())?; + + if is_writable && !info.is_writable { + return Err(ProgramError::InvalidAccountData); + } + + Ok(()) +} + /// Errors if: /// - Owner is not SPL token program. /// - Data is empty. @@ -187,3 +211,39 @@ pub fn load_token_account( Ok(()) } + +/// Errors if: +/// - Owner is not SPL token program. +/// - Data is empty. +/// - Data cannot deserialize into a token account. +/// - Address does not match the expected associated token address. +/// - Expected to be writable, but is not. +#[cfg(feature = "spl")] +pub fn load_associated_token_account( + info: &AccountInfo<'_>, + owner: &Pubkey, + mint: &Pubkey, + is_writable: bool, +) -> Result<(), ProgramError> { + if info.owner.ne(&spl_token::id()) { + return Err(ProgramError::InvalidAccountOwner); + } + + if info.data_is_empty() { + return Err(ProgramError::UninitializedAccount); + } + + let account_data = info.data.borrow(); + let _ = spl_token::state::Account::unpack(&account_data)?; + + let address = spl_associated_token_account::get_associated_token_address(owner, mint); + if info.key.ne(&address) { + return Err(ProgramError::InvalidSeeds); + } + + if is_writable && !info.is_writable { + return Err(ProgramError::InvalidAccountData); + } + + Ok(()) +} diff --git a/utils/src/macros.rs b/utils/src/macros.rs index 0a2e064..1bda272 100644 --- a/utils/src/macros.rs +++ b/utils/src/macros.rs @@ -98,8 +98,23 @@ macro_rules! event { #[macro_export] macro_rules! instruction { - ($struct_name:ident) => { - $crate::impl_to_bytes!($struct_name); + ($discriminator_name:ident, $struct_name:ident) => { $crate::impl_instruction_from_bytes!($struct_name); + + impl $crate::Discriminator for $struct_name { + fn discriminator() -> u8 { + $discriminator_name::$struct_name as u8 + } + } + + impl $struct_name { + pub fn to_bytes(&self) -> Vec { + [ + [$discriminator_name::$struct_name as u8].to_vec(), + bytemuck::bytes_of(self).to_vec(), + ] + .concat() + } + } }; } diff --git a/utils/src/traits.rs b/utils/src/traits.rs index 32d1696..755a5fe 100644 --- a/utils/src/traits.rs +++ b/utils/src/traits.rs @@ -1,10 +1,10 @@ use solana_program::program_error::ProgramError; -pub trait Discriminator { - fn discriminator() -> u8; -} - pub trait AccountDeserialize { fn try_from_bytes(data: &[u8]) -> Result<&Self, ProgramError>; fn try_from_bytes_mut(data: &mut [u8]) -> Result<&mut Self, ProgramError>; } + +pub trait Discriminator { + fn discriminator() -> u8; +}