Skip to content

Commit

Permalink
fix: validate system accounts (#20)
Browse files Browse the repository at this point in the history 
Fixes HAL-003

> ## Description
> In order to participate in a token locking process, a user needs to
first invoke the process_init entry point.
> 
> This entry point will create an account with the necessary size to
store the information provided in the next step of the process,
process_create .
> 
> The current version of this program is not checking that the provided
System Program and Rent Sysvar accounts corresponds to the correct
accounts, provided by the Solana Rust SDK.
> 
> Although there was not risk identified due to this lack of
restrictions, it is a good practice to implement this validations to
prevent any unexpected attack vector that might appear later.
  • Loading branch information
@wei3erHase
wei3erHase authored Oct 4, 2024
1 parent d4f9644 commit 7e7cde4
Show file tree
Hide file tree
Showing 2 changed files with 102 additions and 23 deletions.
123 changes: 101 additions & 22 deletions program/src/processor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,8 @@ use solana_program::{
pubkey::Pubkey,
rent::Rent,
system_instruction::create_account,
sysvar::{clock::Clock, Sysvar},
system_program,
sysvar::{clock, rent, Sysvar},
};

use num_traits::FromPrimitive;
Expand All @@ -37,9 +38,21 @@ impl Processor {
let payer = next_account_info(accounts_iter)?;
let vesting_account = next_account_info(accounts_iter)?;

// Validate that the system program account is correct
if *system_program_account.key != system_program::ID {
msg!("Invalid system program account");
return Err(ProgramError::InvalidArgument);
}

// Validate that the rent sysvar account is correct
if *rent_sysvar_account.key != rent::ID {
msg!("Invalid rent sysvar account");
return Err(ProgramError::InvalidArgument);
}

let rent = Rent::from_account_info(rent_sysvar_account)?;

// Find the non reversible public key for the vesting contract via the seed
// Create and validate the vesting account key with the provided seed
let vesting_account_key = Pubkey::create_program_address(&[&seeds], &program_id).unwrap();
if vesting_account_key != *vesting_account.key {
msg!("Provided vesting account is invalid");
Expand All @@ -48,6 +61,7 @@ impl Processor {

let state_size = VestingSchedule::LEN + VestingScheduleHeader::LEN;

// Create the vesting account creation instruction
let init_vesting_account = create_account(
&payer.key,
&vesting_account_key,
Expand All @@ -56,6 +70,7 @@ impl Processor {
&program_id,
);

// Invoke the vesting account creation instruction
invoke_signed(
&init_vesting_account,
&[
Expand Down Expand Up @@ -84,64 +99,84 @@ impl Processor {
let source_token_account_owner = next_account_info(accounts_iter)?;
let source_token_account = next_account_info(accounts_iter)?;

// Validate the SPL Token Program account
if spl_token_account.key != &spl_token::id() {
msg!("The provided spl token program account is invalid");
return Err(ProgramError::InvalidArgument);
}

// Validate the Clock Sysvar account
if *clock_sysvar_account.key != clock::ID {
msg!("Invalid clock sysvar account");
return Err(ProgramError::InvalidArgument);
}

// Get and validate vesting account key from the seeds
let vesting_account_key = Pubkey::create_program_address(&[&seeds], program_id)?;
if vesting_account_key != *vesting_account.key {
msg!("Provided vesting account is invalid");
msg!("Invalid vesting account key");
return Err(ProgramError::InvalidArgument);
}

// Validate that the source token account owner is a signer
if !source_token_account_owner.is_signer {
msg!("Source token account owner should be a signer.");
msg!("Source token account owner should be a signer");
return Err(ProgramError::InvalidArgument);
}

// Validate that the vesting account is owned by the program
if *vesting_account.owner != *program_id {
msg!("Program should own vesting account");
msg!("Vesting account is not owned by this program");
return Err(ProgramError::InvalidArgument);
}

// Verifying that no SVC was already created with this seed
// Validate that the vesting account is not already initialized
let is_initialized =
vesting_account.try_borrow_data()?[VestingScheduleHeader::LEN - 1] == 1;

if is_initialized {
msg!("Cannot overwrite an existing vesting contract.");
msg!("Cannot overwrite an existing vesting contract");
return Err(ProgramError::InvalidArgument);
}

// Unpack the vesting token account and validate it
let vesting_token_account_data = Account::unpack(&vesting_token_account.data.borrow())?;

if vesting_token_account_data.owner != vesting_account_key {
msg!("The vesting token account should be owned by the vesting account.");
msg!("The vesting token account should be owned by the vesting account");
return Err(ProgramError::InvalidArgument);
}

// Validate that the vesting token account has no delegate
if vesting_token_account_data.delegate.is_some() {
msg!("The vesting token account should not have a delegate authority");
return Err(ProgramError::InvalidAccountData);
}

// Validate that the vesting token account has no close authority
if vesting_token_account_data.close_authority.is_some() {
msg!("The vesting token account should not have a close authority");
return Err(ProgramError::InvalidAccountData);
}

// Pack the vesting schedule header into the vesting account data
let state_header = VestingScheduleHeader {
destination_address: *source_token_account.key,
mint_address: *mint_address,
is_initialized: true,
};

// Validate that the schedule data is not corrupted
let mut data = vesting_account.data.borrow_mut();
if data.len() != VestingScheduleHeader::LEN + VestingSchedule::LEN {
return Err(ProgramError::InvalidAccountData);
}
state_header.pack_into_slice(&mut data);

let clock = Clock::from_account_info(&clock_sysvar_account)?;
let mut total_amount: u64 = 0;
// Retrieve the clock sysvar and validate schedule time delta
let clock = clock::Clock::from_account_info(&clock_sysvar_account)?;

// NOTE: validate time delta to be 0 (unlocked), or a set of predefined values (3 month, 6 months, ...)
let mut total_amount: u64 = 0;
let release_time;
match schedule.time_delta {
/* Valid time_delta values:
Expand All @@ -163,6 +198,7 @@ impl Processor {
}
}

// Pack the schedule data
let state_schedule = VestingSchedule {
release_time: release_time,
amount: schedule.amount,
Expand All @@ -174,11 +210,13 @@ impl Processor {
None => return Err(ProgramError::InvalidInstructionData), // Total amount overflows u64
}

// Validate that the source token account has sufficient funds
if Account::unpack(&source_token_account.data.borrow())?.amount < total_amount {
msg!("The source token account has insufficient funds.");
return Err(ProgramError::InsufficientFunds);
};
}

// Create the transfer instruction
let transfer_tokens_to_vesting_account = transfer(
spl_token_account.key,
source_token_account.key,
Expand All @@ -188,6 +226,7 @@ impl Processor {
total_amount,
)?;

// Invoke the transfer instruction
invoke(
&transfer_tokens_to_vesting_account,
&[
Expand All @@ -213,17 +252,26 @@ impl Processor {
let vesting_token_account = next_account_info(accounts_iter)?;
let destination_token_account = next_account_info(accounts_iter)?;

let vesting_account_key = Pubkey::create_program_address(&[&seeds], program_id)?;
if vesting_account_key != *vesting_account.key {
msg!("Invalid vesting account key");
// Validate the SPL Token Program account
if spl_token_account.key != &spl_token::id() {
msg!("The provided spl token program account is invalid");
return Err(ProgramError::InvalidArgument);
}

if spl_token_account.key != &spl_token::id() {
msg!("The provided spl token program account is invalid");
// Validate that the clock sysvar account is correct
if *clock_sysvar_account.key != clock::ID {
msg!("Invalid clock sysvar account");
return Err(ProgramError::InvalidArgument);
}

// Validate that the vesting account public key is derived from the seeds
let vesting_account_key = Pubkey::create_program_address(&[&seeds], program_id)?;
if vesting_account_key != *vesting_account.key {
msg!("Invalid vesting account key");
return Err(ProgramError::InvalidArgument);
}

// Validate that the destination token account is the correct one from the schedule header
let packed_state = &vesting_account.data;
let header_state =
VestingScheduleHeader::unpack(&packed_state.borrow()[..VestingScheduleHeader::LEN])?;
Expand All @@ -233,6 +281,7 @@ impl Processor {
return Err(ProgramError::InvalidArgument);
}

// Unpack the vesting token account and validate it is owned by the vesting account
let vesting_token_account_data = Account::unpack(&vesting_token_account.data.borrow())?;

if vesting_token_account_data.owner != vesting_account_key {
Expand All @@ -241,26 +290,30 @@ impl Processor {
}

// Unlock the schedules that have reached maturity
let clock = Clock::from_account_info(&clock_sysvar_account)?;
let clock = clock::Clock::from_account_info(&clock_sysvar_account)?;
let mut schedule = unpack_schedule(&packed_state.borrow()[VestingScheduleHeader::LEN..])?;

let mut total_amount_to_transfer = 0;

// Ensure the schedule has been initialized (release time should not be 0)
if schedule.release_time == 0 {
msg!("Should initialize withdrawal first");
return Err(ProgramError::InvalidArgument);
}

// Check if the release time has been reached and release the amount
if clock.unix_timestamp as u64 >= schedule.release_time {
total_amount_to_transfer += schedule.amount;
schedule.amount = 0;
}

// Validate that there is an amount to transfer
if total_amount_to_transfer == 0 {
msg!("Vesting contract has not yet reached release time");
return Err(ProgramError::InvalidArgument);
}

// Create a token transfer from instruction
let transfer_tokens_from_vesting_account = transfer(
&spl_token_account.key,
&vesting_token_account.key,
Expand All @@ -270,6 +323,7 @@ impl Processor {
total_amount_to_transfer,
)?;

// Invoke the transfer from instruction
invoke_signed(
&transfer_tokens_from_vesting_account,
&[
Expand All @@ -281,7 +335,7 @@ impl Processor {
&[&[&seeds]],
)?;

// Reset released amounts to 0. This makes the simple unlock safe with complex scheduling contracts
// Reset the unlocked amounts in the schedule to 0 to avoid re-using
pack_schedule_into_slice(
schedule,
&mut packed_state.borrow_mut()[VestingScheduleHeader::LEN..],
Expand All @@ -303,42 +357,67 @@ impl Processor {
let vesting_token_account = next_account_info(accounts_iter)?;
let destination_token_account = next_account_info(accounts_iter)?;

// Validate the SPL Token Program account
if spl_token_account.key != &spl_token::id() {
msg!("The provided SPL token program account is invalid");
return Err(ProgramError::InvalidArgument);
}

// Validate the Clock Sysvar account
if *clock_sysvar_account.key != clock::ID {
msg!("Invalid clock sysvar account");
return Err(ProgramError::InvalidArgument);
}

// Validate the vesting account key derived from seeds
let vesting_account_key = Pubkey::create_program_address(&[&seeds], program_id)?;
if vesting_account_key != *vesting_account.key {
msg!("Invalid vesting account key");
return Err(ProgramError::InvalidArgument);
}

// Validate the SPL Token Program account
if spl_token_account.key != &spl_token::id() {
msg!("The provided spl token program account is invalid");
return Err(ProgramError::InvalidArgument);
}

// Validate that the vesting account is owned by the program
if *vesting_account.owner != *program_id {
msg!("Vesting account is not owned by this program");
return Err(ProgramError::InvalidArgument);
}

// Unpack the vesting account's state
let packed_state = &vesting_account.data;
let header_state =
VestingScheduleHeader::unpack(&packed_state.borrow()[..VestingScheduleHeader::LEN])?;

// Validate that the destination token account matches the contract's stored destination address
if header_state.destination_address != *destination_token_account.key {
msg!("Contract destination account does not matched provided account");
return Err(ProgramError::InvalidArgument);
}

// Unpack the vesting token account and validate ownership by the vesting account
let vesting_token_account_data = Account::unpack(&vesting_token_account.data.borrow())?;

if vesting_token_account_data.owner != vesting_account_key {
msg!("The vesting token account should be owned by the vesting account.");
return Err(ProgramError::InvalidArgument);
}

// Unlock the schedules that have reached maturity
let clock = Clock::from_account_info(&clock_sysvar_account)?;
// Unpack the schedule data
let clock = clock::Clock::from_account_info(&clock_sysvar_account)?;
let mut schedule = unpack_schedule(&packed_state.borrow()[VestingScheduleHeader::LEN..])?;

// Check if the vesting contract has already been fully claimed
if schedule.amount == 0 {
msg!("Vesting contract already claimed");
return Err(ProgramError::InvalidArgument);
}

// Ensure the withdrawal is not already initialized
if schedule.release_time != 0 {
msg!("Shouldn't initialize withdrawal for already initialized schedule");
return Err(ProgramError::InvalidArgument);
Expand All @@ -347,7 +426,7 @@ impl Processor {
// Withdrawal period is 7 days = 7 * 86400 = 604_800
schedule.release_time = clock.unix_timestamp as u64 + 604_800;

// Reset released amounts to 0. This makes the simple unlock safe with complex scheduling contracts
// Pack the updated schedule back into the account data
pack_schedule_into_slice(
schedule,
&mut packed_state.borrow_mut()[VestingScheduleHeader::LEN..],
Expand Down
2 changes: 1 addition & 1 deletion program/tests/functional.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ async fn test_token_vesting() {

let schedule = Schedule {
amount: 100,
time_delta: 60,
time_delta: 7_776_000,
};

let test_instructions = [
Expand Down

0 comments on commit 7e7cde4

Please sign in to comment.