More object-oriented version of Vest (#6310)
This commit is contained in:
		@@ -14,7 +14,7 @@ use solana_sdk::{
 | 
				
			|||||||
    pubkey::Pubkey,
 | 
					    pubkey::Pubkey,
 | 
				
			||||||
};
 | 
					};
 | 
				
			||||||
 | 
					
 | 
				
			||||||
fn parse_date_account(
 | 
					fn verify_date_account(
 | 
				
			||||||
    keyed_account: &mut KeyedAccount,
 | 
					    keyed_account: &mut KeyedAccount,
 | 
				
			||||||
    expected_pubkey: &Pubkey,
 | 
					    expected_pubkey: &Pubkey,
 | 
				
			||||||
) -> Result<Date<Utc>, InstructionError> {
 | 
					) -> Result<Date<Utc>, InstructionError> {
 | 
				
			||||||
@@ -22,7 +22,7 @@ fn parse_date_account(
 | 
				
			|||||||
        return Err(InstructionError::IncorrectProgramId);
 | 
					        return Err(InstructionError::IncorrectProgramId);
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    let account = parse_account(keyed_account, expected_pubkey)?;
 | 
					    let account = verify_account(keyed_account, expected_pubkey)?;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    let config_data =
 | 
					    let config_data =
 | 
				
			||||||
        get_config_data(&account.data).map_err(|_| InstructionError::InvalidAccountData)?;
 | 
					        get_config_data(&account.data).map_err(|_| InstructionError::InvalidAccountData)?;
 | 
				
			||||||
@@ -32,7 +32,7 @@ fn parse_date_account(
 | 
				
			|||||||
    Ok(date_config.date_time.date())
 | 
					    Ok(date_config.date_time.date())
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
fn parse_account<'a>(
 | 
					fn verify_account<'a>(
 | 
				
			||||||
    keyed_account: &'a mut KeyedAccount,
 | 
					    keyed_account: &'a mut KeyedAccount,
 | 
				
			||||||
    expected_pubkey: &Pubkey,
 | 
					    expected_pubkey: &Pubkey,
 | 
				
			||||||
) -> Result<&'a mut Account, InstructionError> {
 | 
					) -> Result<&'a mut Account, InstructionError> {
 | 
				
			||||||
@@ -43,7 +43,7 @@ fn parse_account<'a>(
 | 
				
			|||||||
    Ok(keyed_account.account)
 | 
					    Ok(keyed_account.account)
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
fn parse_signed_account<'a>(
 | 
					fn verify_signed_account<'a>(
 | 
				
			||||||
    keyed_account: &'a mut KeyedAccount,
 | 
					    keyed_account: &'a mut KeyedAccount,
 | 
				
			||||||
    expected_pubkey: &Pubkey,
 | 
					    expected_pubkey: &Pubkey,
 | 
				
			||||||
) -> Result<&'a mut Account, InstructionError> {
 | 
					) -> Result<&'a mut Account, InstructionError> {
 | 
				
			||||||
@@ -51,7 +51,7 @@ fn parse_signed_account<'a>(
 | 
				
			|||||||
        return Err(InstructionError::MissingRequiredSignature);
 | 
					        return Err(InstructionError::MissingRequiredSignature);
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    parse_account(keyed_account, expected_pubkey)
 | 
					    verify_account(keyed_account, expected_pubkey)
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
pub fn process_instruction(
 | 
					pub fn process_instruction(
 | 
				
			||||||
@@ -60,66 +60,66 @@ pub fn process_instruction(
 | 
				
			|||||||
    data: &[u8],
 | 
					    data: &[u8],
 | 
				
			||||||
) -> Result<(), InstructionError> {
 | 
					) -> Result<(), InstructionError> {
 | 
				
			||||||
    let keyed_accounts_iter = &mut keyed_accounts.iter_mut();
 | 
					    let keyed_accounts_iter = &mut keyed_accounts.iter_mut();
 | 
				
			||||||
 | 
					    let contract_account = &mut next_keyed_account(keyed_accounts_iter)?.account;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    let instruction = deserialize(data).map_err(|_| InstructionError::InvalidInstructionData)?;
 | 
					    let instruction = deserialize(data).map_err(|_| InstructionError::InvalidInstructionData)?;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    match instruction {
 | 
					    let mut vest_state = if let VestInstruction::InitializeAccount {
 | 
				
			||||||
        VestInstruction::InitializeAccount {
 | 
					        terminator_pubkey,
 | 
				
			||||||
 | 
					        payee_pubkey,
 | 
				
			||||||
 | 
					        start_date_time,
 | 
				
			||||||
 | 
					        date_pubkey,
 | 
				
			||||||
 | 
					        total_lamports,
 | 
				
			||||||
 | 
					    } = instruction
 | 
				
			||||||
 | 
					    {
 | 
				
			||||||
 | 
					        VestState {
 | 
				
			||||||
            terminator_pubkey,
 | 
					            terminator_pubkey,
 | 
				
			||||||
            payee_pubkey,
 | 
					            payee_pubkey,
 | 
				
			||||||
            start_date_time,
 | 
					            start_date_time,
 | 
				
			||||||
            date_pubkey,
 | 
					            date_pubkey,
 | 
				
			||||||
            total_lamports,
 | 
					            total_lamports,
 | 
				
			||||||
        } => {
 | 
					            redeemed_lamports: 0,
 | 
				
			||||||
            let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					 | 
				
			||||||
            let contract_account = &mut contract_keyed_account.account;
 | 
					 | 
				
			||||||
            let vest_state = VestState {
 | 
					 | 
				
			||||||
                terminator_pubkey,
 | 
					 | 
				
			||||||
                payee_pubkey,
 | 
					 | 
				
			||||||
                start_date_time,
 | 
					 | 
				
			||||||
                date_pubkey,
 | 
					 | 
				
			||||||
                total_lamports,
 | 
					 | 
				
			||||||
                redeemed_lamports: 0,
 | 
					 | 
				
			||||||
            };
 | 
					 | 
				
			||||||
            vest_state.serialize(&mut contract_account.data)
 | 
					 | 
				
			||||||
        }
 | 
					        }
 | 
				
			||||||
 | 
					    } else {
 | 
				
			||||||
 | 
					        VestState::deserialize(&contract_account.data)?
 | 
				
			||||||
 | 
					    };
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    match instruction {
 | 
				
			||||||
 | 
					        VestInstruction::InitializeAccount { .. } => {}
 | 
				
			||||||
        VestInstruction::SetPayee(payee_pubkey) => {
 | 
					        VestInstruction::SetPayee(payee_pubkey) => {
 | 
				
			||||||
            let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					            verify_signed_account(
 | 
				
			||||||
            let old_payee_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					                next_keyed_account(keyed_accounts_iter)?,
 | 
				
			||||||
            let contract_account = &mut contract_keyed_account.account;
 | 
					                &vest_state.payee_pubkey,
 | 
				
			||||||
            let mut vest_state = VestState::deserialize(&contract_account.data)?;
 | 
					            )?;
 | 
				
			||||||
            parse_signed_account(old_payee_keyed_account, &vest_state.payee_pubkey)?;
 | 
					 | 
				
			||||||
            vest_state.payee_pubkey = payee_pubkey;
 | 
					            vest_state.payee_pubkey = payee_pubkey;
 | 
				
			||||||
            vest_state.serialize(&mut contract_account.data)
 | 
					 | 
				
			||||||
        }
 | 
					        }
 | 
				
			||||||
        VestInstruction::RedeemTokens => {
 | 
					        VestInstruction::RedeemTokens => {
 | 
				
			||||||
            let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					            let current_date = verify_date_account(
 | 
				
			||||||
            let date_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					                next_keyed_account(keyed_accounts_iter)?,
 | 
				
			||||||
            let payee_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					                &vest_state.date_pubkey,
 | 
				
			||||||
            let contract_account = &mut contract_keyed_account.account;
 | 
					            )?;
 | 
				
			||||||
            let mut vest_state = VestState::deserialize(&contract_account.data)?;
 | 
					            let payee_account = verify_account(
 | 
				
			||||||
            let current_date = parse_date_account(date_keyed_account, &vest_state.date_pubkey)?;
 | 
					                next_keyed_account(keyed_accounts_iter)?,
 | 
				
			||||||
            let payee_account = parse_account(payee_keyed_account, &vest_state.payee_pubkey)?;
 | 
					                &vest_state.payee_pubkey,
 | 
				
			||||||
 | 
					            )?;
 | 
				
			||||||
            vest_state.redeem_tokens(contract_account, current_date, payee_account);
 | 
					            vest_state.redeem_tokens(contract_account, current_date, payee_account);
 | 
				
			||||||
            vest_state.serialize(&mut contract_account.data)
 | 
					 | 
				
			||||||
        }
 | 
					        }
 | 
				
			||||||
        VestInstruction::Terminate => {
 | 
					        VestInstruction::Terminate => {
 | 
				
			||||||
            let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					            let terminator_account = verify_signed_account(
 | 
				
			||||||
            let terminator_keyed_account = next_keyed_account(keyed_accounts_iter)?;
 | 
					                next_keyed_account(keyed_accounts_iter)?,
 | 
				
			||||||
 | 
					                &vest_state.terminator_pubkey,
 | 
				
			||||||
 | 
					            )?;
 | 
				
			||||||
            let payee_keyed_account = keyed_accounts_iter.next();
 | 
					            let payee_keyed_account = keyed_accounts_iter.next();
 | 
				
			||||||
            let contract_account = &mut contract_keyed_account.account;
 | 
					 | 
				
			||||||
            let mut vest_state = VestState::deserialize(&contract_account.data)?;
 | 
					 | 
				
			||||||
            let terminator_account =
 | 
					 | 
				
			||||||
                parse_signed_account(terminator_keyed_account, &vest_state.terminator_pubkey)?;
 | 
					 | 
				
			||||||
            let payee_account = if let Some(payee_keyed_account) = payee_keyed_account {
 | 
					            let payee_account = if let Some(payee_keyed_account) = payee_keyed_account {
 | 
				
			||||||
                &mut payee_keyed_account.account
 | 
					                &mut payee_keyed_account.account
 | 
				
			||||||
            } else {
 | 
					            } else {
 | 
				
			||||||
                terminator_account
 | 
					                terminator_account
 | 
				
			||||||
            };
 | 
					            };
 | 
				
			||||||
            vest_state.terminate(contract_account, payee_account);
 | 
					            vest_state.terminate(contract_account, payee_account);
 | 
				
			||||||
            vest_state.serialize(&mut contract_account.data)
 | 
					 | 
				
			||||||
        }
 | 
					        }
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    vest_state.serialize(&mut contract_account.data)
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
#[cfg(test)]
 | 
					#[cfg(test)]
 | 
				
			||||||
@@ -233,7 +233,7 @@ mod tests {
 | 
				
			|||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    #[test]
 | 
					    #[test]
 | 
				
			||||||
    fn test_parse_account_unauthorized() {
 | 
					    fn test_verify_account_unauthorized() {
 | 
				
			||||||
        // Ensure client can't sneak in with an untrusted date account.
 | 
					        // Ensure client can't sneak in with an untrusted date account.
 | 
				
			||||||
        let date_pubkey = Pubkey::new_rand();
 | 
					        let date_pubkey = Pubkey::new_rand();
 | 
				
			||||||
        let mut account = Account::new(1, 0, &solana_config_api::id());
 | 
					        let mut account = Account::new(1, 0, &solana_config_api::id());
 | 
				
			||||||
@@ -241,68 +241,68 @@ mod tests {
 | 
				
			|||||||
 | 
					
 | 
				
			||||||
        let mallory_pubkey = Pubkey::new_rand(); // <-- Attack! Not the expected account.
 | 
					        let mallory_pubkey = Pubkey::new_rand(); // <-- Attack! Not the expected account.
 | 
				
			||||||
        assert_eq!(
 | 
					        assert_eq!(
 | 
				
			||||||
            parse_account(&mut keyed_account, &mallory_pubkey).unwrap_err(),
 | 
					            verify_account(&mut keyed_account, &mallory_pubkey).unwrap_err(),
 | 
				
			||||||
            VestError::Unauthorized.into()
 | 
					            VestError::Unauthorized.into()
 | 
				
			||||||
        );
 | 
					        );
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    #[test]
 | 
					    #[test]
 | 
				
			||||||
    fn test_parse_signed_account_missing_signature() {
 | 
					    fn test_verify_signed_account_missing_signature() {
 | 
				
			||||||
        // Ensure client can't sneak in with an unsigned account.
 | 
					        // Ensure client can't sneak in with an unsigned account.
 | 
				
			||||||
        let date_pubkey = Pubkey::new_rand();
 | 
					        let date_pubkey = Pubkey::new_rand();
 | 
				
			||||||
        let mut account = Account::new(1, 0, &solana_config_api::id());
 | 
					        let mut account = Account::new(1, 0, &solana_config_api::id());
 | 
				
			||||||
        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account); // <-- Attack! Unsigned transaction.
 | 
					        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account); // <-- Attack! Unsigned transaction.
 | 
				
			||||||
 | 
					
 | 
				
			||||||
        assert_eq!(
 | 
					        assert_eq!(
 | 
				
			||||||
            parse_signed_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
					            verify_signed_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
				
			||||||
            InstructionError::MissingRequiredSignature.into()
 | 
					            InstructionError::MissingRequiredSignature.into()
 | 
				
			||||||
        );
 | 
					        );
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    #[test]
 | 
					    #[test]
 | 
				
			||||||
    fn test_parse_date_account_incorrect_program_id() {
 | 
					    fn test_verify_date_account_incorrect_program_id() {
 | 
				
			||||||
        // Ensure client can't sneak in with a non-Config account.
 | 
					        // Ensure client can't sneak in with a non-Config account.
 | 
				
			||||||
        let date_pubkey = Pubkey::new_rand();
 | 
					        let date_pubkey = Pubkey::new_rand();
 | 
				
			||||||
        let mut account = Account::new(1, 0, &id()); // <-- Attack! Pass Vest account where Config account is expected.
 | 
					        let mut account = Account::new(1, 0, &id()); // <-- Attack! Pass Vest account where Config account is expected.
 | 
				
			||||||
        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
					        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
				
			||||||
        assert_eq!(
 | 
					        assert_eq!(
 | 
				
			||||||
            parse_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
					            verify_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
				
			||||||
            InstructionError::IncorrectProgramId
 | 
					            InstructionError::IncorrectProgramId
 | 
				
			||||||
        );
 | 
					        );
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    #[test]
 | 
					    #[test]
 | 
				
			||||||
    fn test_parse_date_account_uninitialized_config() {
 | 
					    fn test_verify_date_account_uninitialized_config() {
 | 
				
			||||||
        // Ensure no panic when `get_config_data()` returns an error.
 | 
					        // Ensure no panic when `get_config_data()` returns an error.
 | 
				
			||||||
        let date_pubkey = Pubkey::new_rand();
 | 
					        let date_pubkey = Pubkey::new_rand();
 | 
				
			||||||
        let mut account = Account::new(1, 0, &solana_config_api::id()); // <-- Attack! Zero space.
 | 
					        let mut account = Account::new(1, 0, &solana_config_api::id()); // <-- Attack! Zero space.
 | 
				
			||||||
        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
					        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
				
			||||||
        assert_eq!(
 | 
					        assert_eq!(
 | 
				
			||||||
            parse_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
					            verify_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
				
			||||||
            InstructionError::InvalidAccountData
 | 
					            InstructionError::InvalidAccountData
 | 
				
			||||||
        );
 | 
					        );
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    #[test]
 | 
					    #[test]
 | 
				
			||||||
    fn test_parse_date_account_invalid_date_config() {
 | 
					    fn test_verify_date_account_invalid_date_config() {
 | 
				
			||||||
        // Ensure no panic when `deserialize::<DateConfig>()` returns an error.
 | 
					        // Ensure no panic when `deserialize::<DateConfig>()` returns an error.
 | 
				
			||||||
        let date_pubkey = Pubkey::new_rand();
 | 
					        let date_pubkey = Pubkey::new_rand();
 | 
				
			||||||
        let mut account = Account::new(1, 1, &solana_config_api::id()); // Attack! 1 byte, enough to sneak by `get_config_data()`, but not DateConfig deserialize.
 | 
					        let mut account = Account::new(1, 1, &solana_config_api::id()); // Attack! 1 byte, enough to sneak by `get_config_data()`, but not DateConfig deserialize.
 | 
				
			||||||
        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
					        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
				
			||||||
        assert_eq!(
 | 
					        assert_eq!(
 | 
				
			||||||
            parse_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
					            verify_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
				
			||||||
            InstructionError::InvalidAccountData
 | 
					            InstructionError::InvalidAccountData
 | 
				
			||||||
        );
 | 
					        );
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    #[test]
 | 
					    #[test]
 | 
				
			||||||
    fn test_parse_date_account_deserialize() {
 | 
					    fn test_verify_date_account_deserialize() {
 | 
				
			||||||
        // Ensure no panic when `deserialize::<DateConfig>()` returns an error.
 | 
					        // Ensure no panic when `deserialize::<DateConfig>()` returns an error.
 | 
				
			||||||
        let date_pubkey = Pubkey::new_rand();
 | 
					        let date_pubkey = Pubkey::new_rand();
 | 
				
			||||||
        let mut account = Account::new(1, 1, &solana_config_api::id()); // Attack! 1 byte, enough to sneak by `get_config_data()`, but not DateConfig deserialize.
 | 
					        let mut account = Account::new(1, 1, &solana_config_api::id()); // Attack! 1 byte, enough to sneak by `get_config_data()`, but not DateConfig deserialize.
 | 
				
			||||||
        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
					        let mut keyed_account = KeyedAccount::new(&date_pubkey, false, &mut account);
 | 
				
			||||||
        assert_eq!(
 | 
					        assert_eq!(
 | 
				
			||||||
            parse_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
					            verify_date_account(&mut keyed_account, &date_pubkey).unwrap_err(),
 | 
				
			||||||
            InstructionError::InvalidAccountData
 | 
					            InstructionError::InvalidAccountData
 | 
				
			||||||
        );
 | 
					        );
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 
 | 
				
			|||||||
		Reference in New Issue
	
	Block a user