Allow missing TokenOwnerRecord for Withdraw and ChangeOwner instructi…

…ons neonevm#23 (neonevm#24)

Co-authored-by: Semen Medvedev <sm@neonlabs.org>

addin-vesting/program/src/error.rs

@@ -59,6 +59,9 @@ pub enum VestingError {
 
     #[error("InvalidPercentage")]
     InvalidPercentage,
+
+    #[error("Invalid TokenOwnerRecord")]
+    InvalidTokenOwnerRecord,
 }
 
 impl From<VestingError> for ProgramError {

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;

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,

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<Option<TokenOwnerRecordV2>, 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)?
+        ))
+    }
+}

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(