diff --git a/addin-vesting/program/src/error.rs b/addin-vesting/program/src/error.rs index d631f5f..e88d400 100644 --- a/addin-vesting/program/src/error.rs +++ b/addin-vesting/program/src/error.rs @@ -59,6 +59,9 @@ pub enum VestingError { #[error("InvalidPercentage")] InvalidPercentage, + + #[error("Invalid TokenOwnerRecord")] + InvalidTokenOwnerRecord, } impl From for ProgramError { diff --git a/addin-vesting/program/src/lib.rs b/addin-vesting/program/src/lib.rs index e7459cd..4b8c5cf 100644 --- a/addin-vesting/program/src/lib.rs +++ b/addin-vesting/program/src/lib.rs @@ -6,5 +6,6 @@ pub mod instruction; pub mod state; pub mod voter_weight; pub mod max_voter_weight; +pub mod token_owner_record; pub mod processor; diff --git a/addin-vesting/program/src/processor.rs b/addin-vesting/program/src/processor.rs index 0ffa0ac..a3409a7 100644 --- a/addin-vesting/program/src/processor.rs +++ b/addin-vesting/program/src/processor.rs @@ -37,6 +37,9 @@ use crate::{ create_max_voter_weight_record, get_max_voter_weight_record_data_checked, }, + token_owner_record::{ + get_token_owner_record_data_if_exists, + }, }; pub struct Processor {} @@ -275,7 +278,7 @@ impl Processor { let realm_data = get_realm_data(governance_account.key, realm_account)?; realm_data.assert_is_valid_governing_token_mint(&vesting_record.mint)?; - let owner_record_data = get_token_owner_record_data_for_seeds( + let owner_record_optional_data = get_token_owner_record_data_if_exists( governance_account.key, owner_record_account, &get_token_owner_record_address_seeds( @@ -284,7 +287,9 @@ impl Processor { vesting_owner_account.key, ), )?; - owner_record_data.assert_can_withdraw_governing_tokens()?; + if let Some(owner_record_data) = owner_record_optional_data { + owner_record_data.assert_can_withdraw_governing_tokens()?; + } let mut voter_weight_record = get_voter_weight_record_data_checked( program_id, @@ -363,7 +368,7 @@ impl Processor { let realm_data = get_realm_data(governance_account.key, realm_account)?; realm_data.assert_is_valid_governing_token_mint(&vesting_record.mint)?; - let owner_record_data = get_token_owner_record_data_for_seeds( + let owner_record_optional_data = get_token_owner_record_data_if_exists( governance_account.key, owner_record_account, &get_token_owner_record_address_seeds( @@ -372,7 +377,9 @@ impl Processor { vesting_owner_account.key, ), )?; - owner_record_data.assert_can_withdraw_governing_tokens()?; + if let Some(owner_record_data) = owner_record_optional_data { + owner_record_data.assert_can_withdraw_governing_tokens()?; + } let mut voter_weight_record = get_voter_weight_record_data_checked( program_id, diff --git a/addin-vesting/program/src/token_owner_record.rs b/addin-vesting/program/src/token_owner_record.rs new file mode 100644 index 0000000..f385944 --- /dev/null +++ b/addin-vesting/program/src/token_owner_record.rs @@ -0,0 +1,41 @@ +use crate::error::VestingError; +use solana_program::{ + account_info::AccountInfo, + program_error::ProgramError, + pubkey::Pubkey, + system_program, +}; + +use spl_governance::state::{ + token_owner_record::{ + TokenOwnerRecordV2, + get_token_owner_record_data_for_seeds, + }, +}; + +pub fn get_token_owner_record_data_if_exists( + program_id: &Pubkey, + token_owner_record_info: &AccountInfo, + token_owner_record_seeds: &[&[u8]], +) -> Result, ProgramError> { + let (token_owner_record_address, _) = + Pubkey::find_program_address(token_owner_record_seeds, program_id); + + if token_owner_record_address != *token_owner_record_info.key { + return Err(VestingError::InvalidTokenOwnerRecord.into()); + } + + if token_owner_record_info.data_is_empty() { + if *token_owner_record_info.owner != system_program::id() { + return Err(VestingError::InvalidTokenOwnerRecord.into()); + } + + Ok(None) + } else { + Ok(Some(get_token_owner_record_data_for_seeds( + program_id, + token_owner_record_info, + token_owner_record_seeds)? + )) + } +} diff --git a/addin-vesting/program/tests/functional.rs b/addin-vesting/program/tests/functional.rs index 4011d1b..a0fe6c4 100644 --- a/addin-vesting/program/tests/functional.rs +++ b/addin-vesting/program/tests/functional.rs @@ -208,10 +208,10 @@ async fn test_token_vesting_with_realm() { let destination_account = Keypair::new(); let destination_token_account = Keypair::new(); + let destination_delegate = Keypair::new(); let new_destination_account = Keypair::new(); let new_destination_token_account = Keypair::new(); - let new_destination_delegate = Keypair::new(); let vesting_token_account = Keypair::new(); let (vesting_account_key,_) = Pubkey::find_program_address(&[&vesting_token_account.pubkey().as_ref()], &program_id); @@ -303,13 +303,6 @@ async fn test_token_vesting_with_realm() { &mint.pubkey(), &payer.pubkey(), ), - create_token_owner_record( - &governance_id, - &realm_address, - &new_destination_account.pubkey(), - &mint.pubkey(), - &payer.pubkey(), - ), ]; let mut create_realm_transaction = Transaction::new_with_payer( &create_realm_instructions, @@ -348,79 +341,79 @@ async fn test_token_vesting_with_realm() { banks_client.process_transaction(deposit_transaction).await.unwrap(); - let init_new_destination_instructions = [ - create_voter_weight_record( - &program_id, - &new_destination_account.pubkey(), - &payer.pubkey(), + let set_governance_delegate_instructions = [ + set_governance_delegate( &governance_id, + &destination_account.pubkey(), &realm_address, &mint.pubkey(), - ).unwrap(), + &destination_account.pubkey(), + &Some(destination_delegate.pubkey()), + ), ]; - let mut init_new_destination_transaction = Transaction::new_with_payer( - &init_new_destination_instructions, + let mut set_governance_delegate_transaction = Transaction::new_with_payer( + &set_governance_delegate_instructions, Some(&payer.pubkey()), ); - init_new_destination_transaction.partial_sign(&[&payer], recent_blockhash); - banks_client.process_transaction(init_new_destination_transaction).await.unwrap(); + set_governance_delegate_transaction.partial_sign(&[&payer, &destination_account], recent_blockhash); + banks_client.process_transaction(set_governance_delegate_transaction).await.unwrap(); - let change_owner_instructions = [ - change_owner_with_realm( + let set_vote_percentage_instructions = [ + set_vote_percentage_with_realm( &program_id, &vesting_token_account.pubkey(), &destination_account.pubkey(), - &new_destination_account.pubkey(), + &destination_delegate.pubkey(), &governance_id, &realm_address, &mint.pubkey(), + 30*100, ).unwrap(), ]; - let mut change_owner_transaction = Transaction::new_with_payer( - &change_owner_instructions, + let mut set_vote_percentage_transaction = Transaction::new_with_payer( + &set_vote_percentage_instructions, Some(&payer.pubkey()), ); - change_owner_transaction.partial_sign(&[&payer, &destination_account], recent_blockhash); - banks_client.process_transaction(change_owner_transaction).await.unwrap(); + set_vote_percentage_transaction.partial_sign(&[&payer, &destination_delegate], recent_blockhash); + banks_client.process_transaction(set_vote_percentage_transaction).await.unwrap(); - let set_governance_delegate_instructions = [ - set_governance_delegate( - &governance_id, + let init_new_destination_instructions = [ + create_voter_weight_record( + &program_id, &new_destination_account.pubkey(), + &payer.pubkey(), + &governance_id, &realm_address, &mint.pubkey(), - &new_destination_account.pubkey(), - &Some(new_destination_delegate.pubkey()), - ), + ).unwrap(), ]; - let mut set_governance_delegate_transaction = Transaction::new_with_payer( - &set_governance_delegate_instructions, + let mut init_new_destination_transaction = Transaction::new_with_payer( + &init_new_destination_instructions, Some(&payer.pubkey()), ); - set_governance_delegate_transaction.partial_sign(&[&payer, &new_destination_account], recent_blockhash); - banks_client.process_transaction(set_governance_delegate_transaction).await.unwrap(); + init_new_destination_transaction.partial_sign(&[&payer], recent_blockhash); + banks_client.process_transaction(init_new_destination_transaction).await.unwrap(); - let set_vote_percentage_instructions = [ - set_vote_percentage_with_realm( + let change_owner_instructions = [ + change_owner_with_realm( &program_id, &vesting_token_account.pubkey(), + &destination_account.pubkey(), &new_destination_account.pubkey(), - &new_destination_delegate.pubkey(), &governance_id, &realm_address, &mint.pubkey(), - 30*100, ).unwrap(), ]; - let mut set_vote_percentage_transaction = Transaction::new_with_payer( - &set_vote_percentage_instructions, + let mut change_owner_transaction = Transaction::new_with_payer( + &change_owner_instructions, Some(&payer.pubkey()), ); - set_vote_percentage_transaction.partial_sign(&[&payer, &new_destination_delegate], recent_blockhash); - banks_client.process_transaction(set_vote_percentage_transaction).await.unwrap(); + change_owner_transaction.partial_sign(&[&payer, &destination_account], recent_blockhash); + banks_client.process_transaction(change_owner_transaction).await.unwrap(); let voter_weight_record_address2 = get_voter_weight_record_address(