Skip to content

Commit

Permalink
feat: use checked arithmetic in vesting contract
Browse files Browse the repository at this point in the history
Description
============
- Introduced enhanced error handling.
- Introduced checked arithmetic for
  safer calculations.

Impact
======
This commit addresses issue use-ink#36
  • Loading branch information
0xf333 committed Aug 26, 2023
1 parent 616ac4c commit 76d18ad
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 61 deletions.
2 changes: 1 addition & 1 deletion vesting/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "vesting"
version = "0.1.0"
version = "0.1.5"
authors = ["Parity Technologies <admin@parity.io>"]
edition = "2021"

Expand Down
166 changes: 106 additions & 60 deletions vesting/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,24 @@ mod vesting_contract {
owner: AccountId,
}

/// Error for when the beneficiary is a zero address.
/// & Error for when the releasable balance is zero.
/// # Errors for the vesting contract:
/// - InvalidBeneficiary: when the beneficiary is a zero address.
/// - ZeroReleasableBalance: when the releasable balance is zero.
/// - Overflow: when an arithmetic operation resulted in an overflow.
/// - Underflow: when a substraction operation resulted in an underflow.
/// - DivisionByZero: when an attempt was made to divide by zero.
/// - TimestampError: in cases where an error related to timestamp
/// calculations occurs.
///
#[derive(Debug, PartialEq, Eq, scale::Encode, scale::Decode)]
#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
pub enum Error {
InvalidBeneficiary,
ZeroReleasableBalance,
Overflow,
Underflow,
DivisionByZero,
TimestampError,
}

/// To emit events when a release is made.
Expand All @@ -29,15 +40,16 @@ mod vesting_contract {
to: AccountId,
}

/// ## This is to set the following during contract deployment:
/// - beneficiary: the account that will receive the tokens
/// - duration_time: duration of the vesting period,
/// please note that this is in seconds
/// - start_time: the time (as Unix time) at which point
/// vesting starts
/// - owner: the account that can release the tokens
/// - releasable_balance: the initial amount of tokens vested
/// - released_balance: the initial amount of tokens released

/// # This is to set the following during contract deployment:
/// - beneficiary: the account that will receive the tokens
/// - duration_time: duration of the vesting period,
/// please note that this is in seconds
/// - start_time: the time (as Unix time) at which point
/// vesting starts
/// - owner: the account that can release the tokens
/// - releasable_balance: the initial amount of tokens vested
/// - released_balance: the initial amount of tokens released
///
/// # Note:
/// The beneficiary cannot be the zero address.
Expand All @@ -52,9 +64,11 @@ mod vesting_contract {
return Err(Error::InvalidBeneficiary);
}

// This is multiplied by 1000 to conform to the
// This is multiplied by 1000 to conform to the
// Timestamp fomat in ink.
let duration_time = duration_time_in_sec * 1000;
let duration_time = duration_time_in_sec
.checked_mul(1000)
.ok_or(Error::Overflow)?;

let start_time = Self::env().block_timestamp();
let owner = Self::env().caller();
Expand Down Expand Up @@ -105,18 +119,22 @@ mod vesting_contract {
/// This returns the time at which point
/// vesting ends.
#[ink(message)]
pub fn end_time(&self) -> Timestamp {
self.start_time() + self.duration_time()
pub fn end_time(&self) -> Result<Timestamp, Error> {
self.start_time()
.checked_add(self.duration_time())
.ok_or(Error::Overflow)
}

/// This returns the amount of time remaining
/// until the end of the vesting period.
#[ink(message)]
pub fn time_remaining(&self) -> Timestamp {
if self.time_now() < self.end_time() {
return self.end_time() - self.time_now();
pub fn time_remaining(&self) -> Result<Timestamp, Error> {
if self.time_now() < self.end_time()? {
self.end_time()?
.checked_sub(self.time_now())
.ok_or(Error::TimestampError)
} else {
return 0;
Ok(0)
}
}

Expand All @@ -130,27 +148,40 @@ mod vesting_contract {
/// This returns the amount of native token that
/// is currently available for release.
#[ink(message)]
pub fn releasable_balance(&self) -> Balance {
return self.vested_amount() as Balance - self.released_balance();
pub fn releasable_balance(&self) -> Result<Balance, Error> {
self.vested_amount()?
.checked_sub(self.released_balance())
.ok_or(Error::Underflow)
}

/// This calculates the amount that has already vested
/// but hasn't been released from the contract yet.
#[ink(message)]
pub fn vested_amount(&self) -> Balance {
return self.vesting_schedule(self.this_contract_balance(), self.time_now());
pub fn vested_amount(&self) -> Result<Balance, Error> {
let total_allocation = self.this_contract_balance()
.checked_add(self.released_balance)
.ok_or(Error::Overflow)?;

self.vesting_schedule(
total_allocation,
self.time_now()
)
}

/// This sends the releasable balance to the beneficiary.
/// This sends the releasable balance to the beneficiary
/// wallet address; no matter who triggers the release.
#[ink(message)]
pub fn release(&mut self) -> Result<(), Error> {
let releasable = self.releasable_balance();
let releasable = self.releasable_balance()?;

if releasable == 0 {
return Err(Error::ZeroReleasableBalance);
}

self.released_balance += releasable;
self.released_balance = self.released_balance
.checked_add(releasable)
.ok_or(Error::Overflow)?;

self.env()
.transfer(self.beneficiary, releasable)
.expect("Transfer failed during release");
Expand All @@ -170,18 +201,18 @@ mod vesting_contract {
/// released evenly over the vesting duration.
///
/// # Parameters:
/// - total_allocation: The total number of tokens
/// allocated for vesting.
/// - current_time: The current timestamp for which
/// we want to check the vested amount.
/// - total_allocation: The total number of tokens
/// allocated for vesting.
/// - current_time: The current timestamp for which
/// we want to check the vested amount.
///
/// # Returns:
/// - `0` if the current_time is before the vesting start time.
/// - total_allocation if the current_time is after the vesting
/// end time or at least equal to it.
/// - A prorated amount based on how much time has passed since
/// the start of the vesting period if the `current_time` is
/// during the vesting period.
/// - `0` if the current_time is before the vesting start time.
/// - total_allocation if the current_time is after the vesting
/// end time or at least equal to it.
/// - A prorated amount based on how much time has passed since
/// the start of the vesting period if the `current_time` is
/// during the vesting period.
///
/// # Example:
/// If the vesting duration is 200 seconds and 100 seconds have
Expand All @@ -192,17 +223,32 @@ mod vesting_contract {
&self,
total_allocation: Balance,
current_time: Timestamp,
) -> Balance {
) -> Result<Balance, Error> {
if current_time < self.start_time() {
return 0;
} else if current_time >= self.end_time() {
return total_allocation;
Ok(0)
} else if current_time >= self.end_time()? {
Ok(total_allocation)
} else {
return (total_allocation
* (current_time - self.start_time()) as Balance)
/ self.duration_time() as Balance;

let time_elapsed = current_time
.checked_sub(self.start_time())
.ok_or(Error::TimestampError)?;

// This represents the portion of the total allocation
// that has vested based on the time elapsed.
let vested_portion = total_allocation
.checked_mul(time_elapsed as u128)
.ok_or(Error::Overflow)?;

let total_vesting_duration = self.duration_time() as u128;

let vested_balance = vested_portion
.checked_div(total_vesting_duration)
.ok_or(Error::DivisionByZero)?;

Ok(vested_balance)
}
}
}
}

#[cfg(test)]
Expand All @@ -217,22 +263,22 @@ mod vesting_contract {
assert_eq!(contract.beneficiary(), AccountId::from([0x01; 32]));
assert_eq!(contract.duration_time(), 200 * 1000);
assert_eq!(contract.released_balance(), 0);
assert_eq!(contract.releasable_balance(), 0);
assert_eq!(contract.releasable_balance().unwrap(), 0);
}

/// There should be some time remaining before the vesting period ends.
#[ink::test]
fn time_remaining_works() {
let contract = VestingContract::new(AccountId::from([0x01; 32]), 200).unwrap();
assert!(contract.time_remaining() > 0);
assert!(contract.time_remaining().unwrap() > 0);
}

/// # Checking that tokens cannot be released before
/// the vesting period:
/// - Trying to release tokens before the vesting period
/// has ended, it will return an error.
/// - The released_balance should remain 0 since no tokens
/// were released.
/// - Trying to release tokens before the vesting period
/// has ended, it will return an error.
/// - The released_balance should remain 0 since no tokens
/// were released.
#[ink::test]
fn release_before_vesting_period_fails() {
let mut contract = VestingContract::new(AccountId::from([0x01; 32]), 200).unwrap();
Expand All @@ -242,11 +288,11 @@ mod vesting_contract {
}

/// # Checking if tokens can be released after the vesting period:
/// - Setting the duration_time to 0 to simulate the end of
/// the vesting period.
/// - And then simulate a deposit into the contract.
/// - After releasing, the released_balance should match the
/// amount we simulated as a deposit.
/// - Setting the duration_time to 0 to simulate the end of
/// the vesting period.
/// - And then simulate a deposit into the contract.
/// - After releasing, the released_balance should match the
/// amount we simulated as a deposit.
#[ink::test]
fn release_after_vesting_period_works() {
let mut contract = VestingContract::new(AccountId::from([0x01; 32]), 0).unwrap();
Expand All @@ -255,18 +301,18 @@ mod vesting_contract {
assert_eq!(contract.release(), Ok(()));
assert_eq!(contract.released_balance(), 1000000);
}

/// # Checking the vesting_schedule function for a specific behavior:
/// - Given a total allocation and a current time halfway through
/// the vesting period, the vested amount should be half of
/// the total allocation.
/// - Given a total allocation and a current time halfway through
/// the vesting period, the vested amount should be half of
/// the total allocation.
#[ink::test]
fn vesting_schedule_works() {
let contract = VestingContract::new(AccountId::from([0x01; 32]), 200).unwrap();

assert_eq!(
contract.vesting_schedule(1000, contract.start_time() + 100 * 1000),
500
Ok(500)
);
}
}
Expand Down

0 comments on commit 76d18ad

Please sign in to comment.