Add processing vesting into realm #5

addin-vesting/program/Cargo.toml

@@ -29,6 +29,7 @@ spl-associated-token-account = { version = "1.0.2", features = ["no-entrypoint"]
 #spl-governance-tools = { version = "0.1.2" }
 spl-governance = { path="/mnt/working/solana/solana-program-library.git/governance/program", features = ["no-entrypoint"] }
 spl-governance-tools = { path="/mnt/working/solana/solana-program-library.git/governance/tools", version = "0.1.2" }
+spl-governance-addin-api = { path="/mnt/working/solana/solana-program-library.git/governance/addin-api", version = "0.1.1" }
 arbitrary = { version = "0.4", features = ["derive"], optional = true }
 honggfuzz = { version = "0.5", optional = true }
 

addin-vesting/program/src/instruction.rs

@@ -1,4 +1,8 @@
-use crate::state::VestingSchedule;
+use crate::state::{
+    VestingSchedule,
+    get_voter_weight_record_address,
+    get_max_voter_weight_record_address,
+};
 
 use solana_program::{
     instruction::{AccountMeta, Instruction},
@@ -9,6 +13,7 @@ use solana_program::{
 };
 
 use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
+use spl_governance::state::token_owner_record::get_token_owner_record_address;
 
 #[cfg(feature = "fuzz")]
 use arbitrary::Arbitrary;
@@ -169,6 +174,51 @@ pub fn deposit(
     })
 }
 
+/// Creates a `Deposit` instruction to create and initialize the vesting token account
+/// inside the Realm
+pub fn deposit_with_realm(
+    program_id: &Pubkey,
+    token_program_id: &Pubkey,
+    seeds: [u8; 32],
+    vesting_token_account: &Pubkey,
+    source_token_owner: &Pubkey,
+    source_token_account: &Pubkey,
+    vesting_owner: &Pubkey,
+    payer: &Pubkey,
+    schedules: Vec<VestingSchedule>,
+    governance_id: &Pubkey,
+    realm: &Pubkey,
+    mint: &Pubkey,
+) -> Result<Instruction, ProgramError> {
+    let vesting_account = Pubkey::create_program_address(&[&seeds], program_id)?;
+    let voting_weight_record_account = get_voter_weight_record_address(program_id, realm, mint, vesting_owner);
+    let max_voting_weight_record_account = get_max_voter_weight_record_address(program_id, realm, mint);
+    let accounts = vec![
+        AccountMeta::new_readonly(system_program::id(), false),
+        AccountMeta::new_readonly(*token_program_id, false),
+        AccountMeta::new(vesting_account, false),
+        AccountMeta::new(*vesting_token_account, false),
+        AccountMeta::new_readonly(*source_token_owner, true),
+        AccountMeta::new(*source_token_account, false),
+        AccountMeta::new_readonly(*vesting_owner, false),
+        AccountMeta::new_readonly(*payer, true),
+        AccountMeta::new_readonly(sysvar::rent::id(), false),
+
+        AccountMeta::new_readonly(*governance_id, false),
+        AccountMeta::new_readonly(*realm, false),
+        AccountMeta::new(voting_weight_record_account, false),
+        AccountMeta::new(max_voting_weight_record_account, false),
+    ];
+
+    let instruction = VestingInstruction::Deposit { seeds, schedules };
+
+    Ok(Instruction {
+        program_id: *program_id,
+        accounts,
+        data: instruction.try_to_vec().unwrap(),
+    })
+}
+
 /// Creates a `Withdraw` instruction
 pub fn withdraw(
     program_id: &Pubkey,
@@ -197,6 +247,46 @@ pub fn withdraw(
     })
 }
 
+/// Creates a `Withdraw` instruction with realm
+pub fn withdraw_with_realm(
+    program_id: &Pubkey,
+    token_program_id: &Pubkey,
+    seeds: [u8; 32],
+    vesting_token_account: &Pubkey,
+    destination_token_account: &Pubkey,
+    vesting_owner: &Pubkey,
+    governance_id: &Pubkey,
+    realm: &Pubkey,
+    mint: &Pubkey,
+) -> Result<Instruction, ProgramError> {
+    let vesting_account = Pubkey::create_program_address(&[&seeds], program_id)?;
+    let owner_record_account = get_token_owner_record_address(governance_id, realm, mint, vesting_owner);
+    let voting_weight_record_account = get_voter_weight_record_address(program_id, realm, mint, vesting_owner);
+    let max_voting_weight_record_account = get_max_voter_weight_record_address(program_id, realm, mint);
+    let accounts = vec![
+        AccountMeta::new_readonly(*token_program_id, false),
+        AccountMeta::new_readonly(sysvar::clock::id(), false),
+        AccountMeta::new(vesting_account, false),
+        AccountMeta::new(*vesting_token_account, false),
+        AccountMeta::new(*destination_token_account, false),
+        AccountMeta::new_readonly(*vesting_owner, true),
+
+        AccountMeta::new_readonly(*governance_id, false),
+        AccountMeta::new_readonly(*realm, false),
+        AccountMeta::new_readonly(owner_record_account, false),
+        AccountMeta::new(voting_weight_record_account, false),
+        AccountMeta::new(max_voting_weight_record_account, false),
+    ];
+
+    let instruction = VestingInstruction::Withdraw { seeds };
+
+    Ok(Instruction {
+        program_id: *program_id,
+        accounts,
+        data: instruction.try_to_vec().unwrap(),
+    })
+}
+
 /// Creates a `Withdraw` instruction
 pub fn change_owner(
     program_id: &Pubkey,
@@ -220,6 +310,41 @@ pub fn change_owner(
     })
 }
 
+/// Creates a `Withdraw` instruction
+pub fn change_owner_with_realm(
+    program_id: &Pubkey,
+    seeds: [u8; 32],
+    vesting_owner: &Pubkey,
+    new_vesting_owner: &Pubkey,
+    governance_id: &Pubkey,
+    realm: &Pubkey,
+    mint: &Pubkey,
+) -> Result<Instruction, ProgramError> {
+    let vesting_account = Pubkey::create_program_address(&[&seeds], program_id)?;
+    let current_owner_record_account = get_token_owner_record_address(governance_id, realm, mint, vesting_owner);
+    let current_voting_weight_record_account = get_voter_weight_record_address(program_id, realm, mint, vesting_owner);
+    let new_voting_weight_record_account = get_voter_weight_record_address(program_id, realm, mint, new_vesting_owner);
+    let accounts = vec![
+        AccountMeta::new(vesting_account, false),
+        AccountMeta::new(*vesting_owner, true),
+        AccountMeta::new(*new_vesting_owner, false),
+
+        AccountMeta::new_readonly(*governance_id, false),
+        AccountMeta::new_readonly(*realm, false),
+        AccountMeta::new_readonly(current_owner_record_account, false),
+        AccountMeta::new(current_voting_weight_record_account, false),
+        AccountMeta::new(new_voting_weight_record_account, false),
+    ];
+
+    let instruction = VestingInstruction::ChangeOwner { seeds };
+
+    Ok(Instruction {
+        program_id: *program_id,
+        accounts,
+        data: instruction.try_to_vec().unwrap(),
+    })
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

addin-vesting/program/src/processor.rs

@@ -20,11 +20,18 @@ use spl_governance_tools::account::{
     get_account_data,
     create_and_serialize_account_signed,
 };
+use spl_governance_addin_api::{
+    voter_weight::VoterWeightRecord,
+    max_voter_weight::MaxVoterWeightRecord,
+};
 
 use crate::{
     error::VestingError,
     instruction::VestingInstruction,
-    state::{VestingAccountType, VestingRecord, VestingSchedule},
+    state::{
+        VestingAccountType, VestingRecord, VestingSchedule,
+        get_voter_weight_record_seeds, get_max_voter_weight_record_seeds,
+    },
 };
 
 pub struct Processor {}
@@ -49,6 +56,15 @@ impl Processor {
         let payer_account = next_account_info(accounts_iter)?;
         let rent_sysvar_info = next_account_info(accounts_iter)?;
 
+        let realm_info = if let Some(governance) = accounts_iter.next() {
+            let realm = next_account_info(accounts_iter)?;
+            let voter_weight = next_account_info(accounts_iter)?;
+            let max_voter_weight = next_account_info(accounts_iter)?;
+            Some((governance, realm, voter_weight, max_voter_weight,))
+        } else {
+            None
+        };
+
         let vesting_account_key = Pubkey::create_program_address(&[&seeds], program_id)?;
         if vesting_account_key != *vesting_account.key {
             msg!("Provided vesting account is invalid");
@@ -92,7 +108,7 @@ impl Processor {
             account_type: VestingAccountType::VestingRecord,
             owner: *vesting_owner_account.key,
             mint: vesting_token_account_data.mint,
-            realm: None,
+            realm: realm_info.map(|v| *v.1.key),
             schedule: schedules
         };
         create_and_serialize_account_signed::<VestingRecord>(
@@ -128,6 +144,64 @@ impl Processor {
                 source_token_account_owner.clone(),
             ],
         )?;
+
+        if let Some((_governance, realm_account, voter_weight_record_account, max_voter_weight_record_account)) = realm_info {
+            if voter_weight_record_account.data_is_empty() {
+                let voter_weight_record = VoterWeightRecord {
+                    account_discriminator: VoterWeightRecord::ACCOUNT_DISCRIMINATOR,
+                    realm: *realm_account.key,
+                    governing_token_mint: vesting_token_account_data.mint,
+                    governing_token_owner: *vesting_owner_account.key,
+                    voter_weight: total_amount,
+                    voter_weight_expiry: None,
+                    weight_action: None,
+                    weight_action_target: None,
+                    reserved: [0u8; 8],
+                };
+                create_and_serialize_account_signed::<VoterWeightRecord>(
+                    payer_account,
+                    voter_weight_record_account,
+                    &voter_weight_record,
+                    &get_voter_weight_record_seeds(realm_account.key, &vesting_token_account_data.mint, vesting_owner_account.key),
+                    program_id,
+                    system_program_account,
+                    rent
+                )?;
+            } else {
+                let mut voter_weight_record = get_account_data::<VoterWeightRecord>(&program_id, voter_weight_record_account)?;
+                // TODO check realm & voter_weight_record
+                let voter_weight = &mut voter_weight_record.voter_weight;
+                *voter_weight = voter_weight.checked_add(total_amount).ok_or_else(|| ProgramError::InvalidInstructionData)?;
+                voter_weight_record.serialize(&mut *voter_weight_record_account.data.borrow_mut())?;
+            }
+
+            if max_voter_weight_record_account.data_is_empty() {
+                let max_voter_weight_record = MaxVoterWeightRecord {
+                    account_discriminator: MaxVoterWeightRecord::ACCOUNT_DISCRIMINATOR,
+                    realm: *realm_account.key,
+                    governing_token_mint: vesting_token_account_data.mint,
+                    max_voter_weight: total_amount,
+                    max_voter_weight_expiry: None,
+                    reserved: [0u8; 8],
+                };
+                create_and_serialize_account_signed::<MaxVoterWeightRecord>(
+                    payer_account,
+                    max_voter_weight_record_account,
+                    &max_voter_weight_record,
+                    &get_max_voter_weight_record_seeds(realm_account.key, &vesting_token_account_data.mint),
+                    program_id,
+                    system_program_account,
+                    rent
+                )?;
+            } else {
+                let mut max_voter_weight_record = get_account_data::<MaxVoterWeightRecord>(&program_id, max_voter_weight_record_account)?;
+                // TODO check realm & max_voter_weight_record
+                let max_voter_weight = &mut max_voter_weight_record.max_voter_weight;
+                *max_voter_weight = max_voter_weight.checked_add(total_amount).ok_or_else(|| ProgramError::InvalidInstructionData)?;
+                max_voter_weight_record.serialize(&mut *max_voter_weight_record_account.data.borrow_mut())?;
+            }
+        }
+
         Ok(())
     }
 
@@ -145,6 +219,16 @@ impl Processor {
         let destination_token_account = next_account_info(accounts_iter)?;
         let vesting_owner_account = next_account_info(accounts_iter)?;
 
+        let realm_info = if let Some(governance) = accounts_iter.next() {
+            let realm = next_account_info(accounts_iter)?;
+            let owner_record = next_account_info(accounts_iter)?;
+            let voter_weight = next_account_info(accounts_iter)?;
+            let max_voter_weight = next_account_info(accounts_iter)?;
+            Some((governance, realm, owner_record, voter_weight, max_voter_weight,))
+        } else {
+            None
+        };
+
         let vesting_account_key = Pubkey::create_program_address(&[&seeds], program_id)?;
         if vesting_account_key != *vesting_account.key {
             msg!("Invalid vesting account key");
@@ -211,6 +295,25 @@ impl Processor {
         // Reset released amounts to 0. This makes the simple unlock safe with complex scheduling contracts
         vesting_record.serialize(&mut *vesting_account.data.borrow_mut())?;
 
+        if let Some(expected_realm_account) = vesting_record.realm {
+            let (governance, realm, owner_record, voter_weight_record_account, max_voter_weight_record_account) = realm_info.ok_or_else(|| {
+                msg!("Vesting contract expects realm accounts");
+                ProgramError::InvalidArgument
+            })?;
+            // TODO check realm, owner_record, voter_weight and max_voter_weight correct
+            // TODO check that owner can withdraw tokens
+
+            let mut voter_weight_record = get_account_data::<VoterWeightRecord>(&program_id, voter_weight_record_account)?;
+            let voter_weight = &mut voter_weight_record.voter_weight;
+            *voter_weight = voter_weight.checked_sub(total_amount_to_transfer).ok_or_else(|| ProgramError::InvalidInstructionData)?;
+            voter_weight_record.serialize(&mut *voter_weight_record_account.data.borrow_mut())?;
+
+            let mut max_voter_weight_record = get_account_data::<MaxVoterWeightRecord>(&program_id, max_voter_weight_record_account)?;
+            let max_voter_weight = &mut max_voter_weight_record.max_voter_weight;
+            *max_voter_weight = max_voter_weight.checked_sub(total_amount_to_transfer).ok_or_else(|| ProgramError::InvalidInstructionData)?;
+            max_voter_weight_record.serialize(&mut *max_voter_weight_record_account.data.borrow_mut())?;
+        }
+
         Ok(())
     }
 
@@ -224,6 +327,15 @@ impl Processor {
         let vesting_account = next_account_info(accounts_iter)?;
         let vesting_owner_account = next_account_info(accounts_iter)?;
         let new_vesting_owner_account = next_account_info(accounts_iter)?;
+        let realm_info = if let Some(governance) = accounts_iter.next() {
+            let realm = next_account_info(accounts_iter)?;
+            let current_owner_record = next_account_info(accounts_iter)?;
+            let current_voter_weight = next_account_info(accounts_iter)?;
+            let new_voter_weight = next_account_info(accounts_iter)?;
+            Some((governance, realm, current_owner_record, current_voter_weight, new_voter_weight,))
+        } else {
+            None
+        };
 
         msg!("Change owner {} -> {}", vesting_owner_account.key, new_vesting_owner_account.key);
 
@@ -245,9 +357,35 @@ impl Processor {
             return Err(ProgramError::InvalidArgument);
         }
 
+        let mut total_amount = 0;
+        for s in vesting_record.schedule.iter_mut() {
+            total_amount += s.amount;
+        }
+
         vesting_record.owner = *new_vesting_owner_account.key;
         vesting_record.serialize(&mut *vesting_account.data.borrow_mut())?;
 
+        if let Some(expected_realm_account) = vesting_record.realm {
+            let (governance, realm, current_owner_record, current_voter_weight_record_account, new_voter_weight_record_account) = realm_info.ok_or_else(|| {
+                msg!("Vesting contract expects realm accounts");
+                ProgramError::InvalidArgument
+            })?;
+            // TODO check realm, owner_record, current_voter_weight and new_voter_weight correct
+            // TODO check that owner can withdraw tokens
+
+            let mut current_voter_weight_record = get_account_data::<VoterWeightRecord>(&program_id, current_voter_weight_record_account)?;
+            let current_voter_weight = &mut current_voter_weight_record.voter_weight;
+            *current_voter_weight = current_voter_weight.checked_sub(total_amount).ok_or_else(|| ProgramError::InvalidInstructionData)?;
+            current_voter_weight_record.serialize(&mut *current_voter_weight_record_account.data.borrow_mut())?;
+
+            // TODO create new_voter_weight_record if missing
+            let mut new_voter_weight_record = get_account_data::<VoterWeightRecord>(&program_id, new_voter_weight_record_account)?;
+            let new_voter_weight = &mut new_voter_weight_record.voter_weight;
+            *new_voter_weight = new_voter_weight.checked_add(total_amount).ok_or_else(|| ProgramError::InvalidInstructionData)?;
+            new_voter_weight_record.serialize(&mut *new_voter_weight_record_account.data.borrow_mut())?;
+
+        }
+
         Ok(())
     }