feat: implement governance plugin with improved security and maintain…

programs/registrar/src/deposit_governance_token.rs

@@ -0,0 +1,91 @@
+use crate::{
+    error::RegistrarError,
+    registrar_config::RegistrarConfig,
+    voter_weight_record::VoterWeightRecord,
+};
+use solana_program::{
+    account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, sysvar::clock::Clock,
+};
+use spl_token::instruction::transfer;
+
+#[derive(Accounts)]
+pub struct DepositGovernanceToken<'info> {
+    #[account(mut)]
+    pub voter_weight_record: Account<'info, VoterWeightRecord>,
+    #[account(mut)]
+    pub voter_token_account: AccountInfo<'info>,
+    #[account(mut)]
+    pub governance_token_mint: AccountInfo<'info>,
+    pub registrar_config: AccountInfo<'info>,
+    pub token_program: AccountInfo<'info>,
+    pub authority: AccountInfo<'info>,
+}
+
+pub fn deposit_governance_token(
+    ctx: Context<DepositGovernanceToken>,
+    amount: u64,
+) -> ProgramResult {
+    let registrar_config = RegistrarConfig::unpack_from_slice(&ctx.accounts.registrar_config.data.borrow())?;
+
+    // Check if the token is an accepted governance token
+    let token_mint = ctx.accounts.governance_token_mint.key();
+    if !registrar_config.accepted_tokens.contains(&token_mint) {
+        return Err(RegistrarError::InvalidArgument.into());
+    }
+
+    // Transfer tokens from the voter's account to the governance program's account
+    transfer_tokens(
+        &ctx.accounts.voter_token_account,
+        &ctx.accounts.governance_token_mint, // Replace with the governance program's token account
+        amount,
+        &ctx.accounts.token_program,
+        &ctx.accounts.authority,
+    )?;
+
+    // Update the VoterWeightRecord account
+    let voter_weight_record = &mut ctx.accounts.voter_weight_record;
+    let clock = Clock::get()?;
+
+    if voter_weight_record.last_deposit_or_withdrawal_slot == clock.slot {
+        return Err(RegistrarError::InvalidOperation.into());
+    }
+
+    let weight_increase = amount * registrar_config.weights[registrar_config.accepted_tokens.iter().position(|&token| token == token_mint).ok_or(RegistrarError::InvalidArgument)?];
+    voter_weight_record.weight = voter_weight_record.weight.checked_add(weight_increase).ok_or(RegistrarError::Overflow)?;
+    voter_weight_record.last_deposit_or_withdrawal_slot = clock.slot;
+    voter_weight_record.serialize(&mut ctx.accounts.voter_weight_record.data.borrow_mut()[..])?;
+
+    // Update the MaxVoterWeightRecord account
+    // ...
+
+    Ok(())
+}
+
+fn transfer_tokens(
+    source_account: &AccountInfo,
+    destination_account: &AccountInfo,
+    amount: u64,
+    token_program: &AccountInfo,
+    authority: &AccountInfo,
+) -> ProgramResult {
+    let transfer_instruction = transfer(
+        token_program.key,
+        source_account.key,
+        destination_account.key,
+        authority.key,
+        &[],
+        amount,
+    )?;
+
+    invoke(
+        &transfer_instruction,
+        &[
+            source_account.clone(),
+            destination_account.clone(),
+            authority.clone(),
+            token_program.clone(),
+        ],
+    )?;
+
+    Ok(())
+}

programs/registrar/src/error.rs

@@ -0,0 +1,15 @@
+#[derive(Debug, PartialEq, Eq)]
+pub enum RegistrarError {
+    InvalidArgument,
+    InvalidAccountData,
+    InvalidOperation,
+    Overflow,
+    InsufficientFunds,
+    // Add more error variants as needed
+}
+
+impl From<RegistrarError> for ProgramError {
+    fn from(error: RegistrarError) -> Self {
+        ProgramError::Custom(error as u32)
+    }
+}

programs/registrar/src/initialize_registrar.rs

@@ -0,0 +1,35 @@
+use crate::{error::RegistrarError, registrar_config::RegistrarConfig};
+use solana_program::{
+    account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, system_program,
+};
+
+#[derive(Accounts)]
+pub struct InitializeRegistrar<'info> {
+    #[account(init, payer = payer, space = RegistrarConfig::LEN)]
+    pub registrar_config: Account<'info, RegistrarConfig>,
+    #[account(mut)]
+    pub payer: AccountInfo<'info>,
+    pub system_program: Program<'info, System>,
+}
+
+pub fn initialize_registrar(
+    ctx: Context<InitializeRegistrar>,
+    accepted_tokens: Vec<Pubkey>,
+    weights: Vec<u64>,
+) -> ProgramResult {
+    if accepted_tokens.len() != weights.len() {
+        return Err(RegistrarError::InvalidArgument.into());
+    }
+
+    if accepted_tokens.len() > MAX_ACCEPTED_TOKENS {
+        return Err(RegistrarError::InvalidArgument.into());
+    }
+
+    let registrar_info = &mut ctx.accounts.registrar_config;
+    let config = RegistrarConfig {
+        accepted_tokens,
+        weights,
+    };
+    config.pack_into_slice(&mut registrar_info.data.borrow_mut());
+    Ok(())
+}

programs/registrar/src/lib.rs

@@ -0,0 +1,52 @@
+pub mod error;
+pub mod initialize_registrar;
+pub mod deposit_governance_token;
+pub mod withdraw_governance_token;
+pub mod registrar_config;
+pub mod voter_weight_record;
+pub mod max_voter_weight_record;
+
+use solana_program::{
+    account_info::AccountInfo, entrypoint, entrypoint::ProgramResult, pubkey::Pubkey,
+};
+
+entrypoint!(process_instruction);
+
+pub fn process_instruction(
+    program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    instruction_data: &[u8],
+) -> ProgramResult {
+    let instruction = RegistrarInstruction::try_from_slice(instruction_data)?;
+
+    match instruction {
+        RegistrarInstruction::InitializeRegistrar {
+            accepted_tokens,
+            weights,
+        } => {
+            let accounts = initialize_registrar::InitializeRegistrar::try_accounts(
+                program_id,
+                accounts,
+                &accepted_tokens,
+                &weights,
+            )?;
+            initialize_registrar::initialize_registrar(accounts, accepted_tokens, weights)
+        }
+        RegistrarInstruction::DepositGovernanceToken { amount } => {
+            let accounts = deposit_governance_token::DepositGovernanceToken::try_accounts(
+                program_id,
+                accounts,
+                &amount,
+            )?;
+            deposit_governance_token::deposit_governance_token(accounts, amount)
+        }
+        RegistrarInstruction::WithdrawGovernanceToken { amount } => {
+            let accounts = withdraw_governance_token::WithdrawGovernanceToken::try_accounts(
+                program_id,
+                accounts,
+                &amount,
+            )?;
+            withdraw_governance_token::withdraw_governance_token(accounts, amount)
+        }
+    }
+}

programs/registrar/src/max_voter_weight_record.rs

@@ -0,0 +1,25 @@
+use borsh::{BorshDeserialize, BorshSerialize};
+use solana_program::program_error::ProgramError;
+
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct MaxVoterWeightRecord {
+    pub max_weight: u64,
+}
+
+impl Sealed for MaxVoterWeightRecord {}
+
+impl Pack for MaxVoterWeightRecord {
+    const LEN: usize = 8;
+
+    fn unpack_from_slice(src: &[u8]) -> Result<Self, ProgramError> {
+        let max_weight = src.get(..8).ok_or(ProgramError::InvalidAccountData)?.get_u64();
+        Ok(MaxVoterWeightRecord { max_weight })
+    }
+
+    fn pack_into_slice(&self, dst: &mut [u8]) {
+        dst.get_mut(..8)
+            .ok_or(ProgramError::InvalidAccountData)
+            .unwrap()
+            .copy_from_slice(&self.max_weight.to_le_bytes());
+    }
+}

programs/registrar/src/registrar_config.rs

@@ -0,0 +1,34 @@
+use borsh::{BorshDeserialize, BorshSerialize};
+use solana_program::{program_error::ProgramError, pubkey::Pubkey};
+
+const MAX_ACCEPTED_TOKENS: usize = 10;
+
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct RegistrarConfig {
+    pub accepted_tokens: Vec<Pubkey>,
+    pub weights: Vec<u64>,
+}
+
+impl Sealed for RegistrarConfig {}
+
+impl Pack for RegistrarConfig {
+    const LEN: usize = 8 + (4 + 32) * MAX_ACCEPTED_TOKENS;
+
+    fn unpack_from_slice(src: &[u8]) -> Result<Self, ProgramError> {
+        let (accepted_tokens, weights) = array_refs![src, MAX_ACCEPTED_TOKENS; Pubkey, u64];
+        let accepted_tokens = accepted_tokens.to_vec();
+        let weights = weights.to_vec();
+        Ok(RegistrarConfig {
+            accepted_tokens,
+            weights,
+        })
+    }
+
+    fn pack_into_slice(&self, dst: &mut [u8]) {
+        let (accepted_tokens, weights) = array_mut_refs![dst, MAX_ACCEPTED_TOKENS; Pubkey, u64];
+        for (i, token) in self.accepted_tokens.iter().enumerate() {
+            accepted_tokens[i] = *token;
+            weights[i] = self.weights[i];
+        }
+    }
+}

programs/registrar/src/voter_weight_record.rs

@@ -0,0 +1,33 @@
+use borsh::{BorshDeserialize, BorshSerialize};
+use solana_program::{program_error::ProgramError, pubkey::Pubkey};
+
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct VoterWeightRecord {
+    pub voter: Pubkey,
+    pub weight: u64,
+    pub last_deposit_or_withdrawal_slot: u64,
+}
+
+impl Sealed for VoterWeightRecord {}
+
+impl Pack for VoterWeightRecord {
+    const LEN: usize = 8 + 32 + 8;
+
+    fn unpack_from_slice(src: &[u8]) -> Result<Self, ProgramError> {
+        let (voter, weight, last_deposit_or_withdrawal_slot) =
+            array_refs![src, 0, 32, 40; Pubkey, u64, u64];
+        Ok(VoterWeightRecord {
+            voter: *voter,
+            weight: *weight,
+            last_deposit_or_withdrawal_slot: *last_deposit_or_withdrawal_slot,
+        })
+    }
+
+    fn pack_into_slice(&self, dst: &mut [u8]) {
+        let (voter, weight, last_deposit_or_withdrawal_slot) =
+            array_mut_refs![dst, 0, 32, 40; Pubkey, u64, u64];
+        *voter = self.voter;
+        *weight = self.weight;
+        *last_deposit_or_withdrawal_slot = self.last_deposit_or_withdrawal_slot;
+    }
+}

programs/registrar/src/withdraw_governance_token.rs

@@ -0,0 +1,99 @@
+use crate::{
+    error::RegistrarError,
+    registrar_config::RegistrarConfig,
+    voter_weight_record::VoterWeightRecord,
+};
+use solana_program::{
+    account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, sysvar::clock::Clock,
+};
+use spl_token::instruction::transfer;
+
+#[derive(Accounts)]
+pub struct WithdrawGovernanceToken<'info> {
+    #[account(mut)]
+    pub voter_weight_record: Account<'info, VoterWeightRecord>,
+    #[account(mut)]
+    pub voter_token_account: AccountInfo<'info>,
+    pub governance_token_mint: AccountInfo<'info>,
+    pub registrar_config: AccountInfo<'info>,
+    pub token_program: AccountInfo<'info>,
+    pub spl_governance_program: AccountInfo<'info>,
+    pub authority: AccountInfo<'info>,
+}
+
+pub fn withdraw_governance_token(
+    ctx: Context<WithdrawGovernanceToken>,
+    amount: u64,
+) -> ProgramResult {
+    let voter_weight_record = &mut ctx.accounts.voter_weight_record;
+    let clock = Clock::get()?;
+
+    // Check if the withdrawal is allowed by the spl-governance program
+    // ...
+
+    // Check if the current slot is the same as the last deposit or withdrawal slot
+    if voter_weight_record.last_deposit_or_withdrawal_slot == clock.slot {
+        return Err(RegistrarError::InvalidOperation.into());
+    }
+
+    let registrar_config = RegistrarConfig::unpack_from_slice(&ctx.accounts.registrar_config.data.borrow())?;
+
+    // Check if the token is an accepted governance token
+    let token_mint = ctx.accounts.governance_token_mint.key();
+    if !registrar_config.accepted_tokens.contains(&token_mint) {
+        return Err(RegistrarError::InvalidArgument.into());
+    }
+
+    let weight_decrease = amount * registrar_config.weights[registrar_config.accepted_tokens.iter().position(|&token| token == token_mint).ok_or(RegistrarError::InvalidArgument)?];
+    if voter_weight_record.weight < weight_decrease {
+        return Err(RegistrarError::InsufficientFunds.into());
+    }
+
+    // Transfer tokens from the governance program's account to the voter's account
+    transfer_tokens(
+        &ctx.accounts.governance_token_mint, // Replace with the governance program's token account
+        &ctx.accounts.voter_token_account,
+        amount,
+        &ctx.accounts.token_program,
+        &ctx.accounts.authority,
+    )?;
+
+    // Update the VoterWeightRecord account
+    voter_weight_record.weight = voter_weight_record.weight.checked_sub(weight_decrease).ok_or(RegistrarError::Overflow)?;
+    voter_weight_record.last_deposit_or_withdrawal_slot = clock.slot;
+    voter_weight_record.serialize(&mut ctx.accounts.voter_weight_record.data.borrow_mut()[..])?;
+
+    // Update the MaxVoterWeightRecord account
+    // ...
+
+    Ok(())
+}
+
+fn transfer_tokens(
+    source_account: &AccountInfo,
+    destination_account: &AccountInfo,
+    amount: u64,
+    token_program: &AccountInfo,
+    authority: &AccountInfo,
+) -> ProgramResult {
+    let transfer_instruction = transfer(
+        token_program.key,
+        source_account.key,
+        destination_account.key,
+        authority.key,
+        &[],
+        amount,
+    )?;
+
+    invoke(
+        &transfer_instruction,
+        &[
+            source_account.clone(),
+            destination_account.clone(),
+            authority.clone(),
+            token_program.clone(),
+        ],
+    )?;
+
+    Ok(())
+}