Skip to content

Commit

Permalink
Implement token account balance assertion
Browse files Browse the repository at this point in the history
  • Loading branch information
@Jac0xb
Jac0xb committed Dec 18, 2023
1 parent fc623ae commit fa77f78
Show file tree
Hide file tree
Showing 16 changed files with 384 additions and 197 deletions.
4 changes: 2 additions & 2 deletions macros/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ use proc_macro::TokenStream;
// use proc_macro::TokenStream;
use quote::quote;
use syn::{
parse_macro_input, punctuated::Punctuated, token::Comma, Attribute, Data, DataStruct,
DeriveInput, Field, Fields,
parse_macro_input, punctuated::Punctuated, token::Comma, Data, DataStruct, DeriveInput, Field,
Fields,
};
// use syn::{parse_macro_input, DeriveInput};

Expand Down
2 changes: 1 addition & 1 deletion programs/lighthouse/program/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ pub mod lighthouse {
ctx: Context<'_, '_, '_, 'info, AssertV1<'info>>,
assertions: Vec<Assertion>,
logical_expression: Option<Vec<Expression>>,
options: Option<Config>,
options: Option<AssertionConfig>,
) -> Result<()> {
processor::assert(ctx, assertions, logical_expression, options)
}
Expand Down
236 changes: 103 additions & 133 deletions programs/lighthouse/program/src/processor/v1/assert.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
use std::collections::BTreeSet;

use anchor_lang::prelude::*;
use borsh::{BorshDeserialize, BorshSerialize};
use solana_program::program_pack::Pack;

use crate::error::ProgramError;
use crate::structs::{AccountInfoDataField, Assertion, AssertionState, Expression};
use crate::utils::print_result;
use crate::utils::print_assertion_result;

#[derive(Accounts)]
pub struct AssertV1<'info> {
Expand All @@ -14,200 +13,171 @@ pub struct AssertV1<'info> {
}

#[derive(BorshDeserialize, BorshSerialize, Debug)]
pub struct Config {
pub struct AssertionConfig {
pub verbose: bool,
}

pub fn assert<'info>(
ctx: Context<'_, '_, '_, 'info, AssertV1<'info>>,
assertions: Vec<Assertion>,
logical_expression: Option<Vec<Expression>>,
options: Option<Config>,
config: Option<AssertionConfig>,
) -> Result<()> {
let remaining_accounts = &mut ctx.remaining_accounts.iter();

let verbose = options.map(|options| options.verbose).unwrap_or(false);
let mut assertion_state = AssertionState::new(assertions.clone(), logical_expression)?;

for (i, assertion) in assertions.into_iter().enumerate() {
let mut assertion_result = false;

match assertion {
let assertion_result: Result<bool> = match &assertion {
Assertion::AccountOwnedBy(pubkey, operator) => {
let account = remaining_accounts.next().unwrap();
assertion_result = account.owner.key().eq(&pubkey);
let result = account.owner.key().eq(pubkey);

let value_str = account.owner.key().to_string();
let expected_value_str = pubkey.to_string();

if verbose {
print_result(assertion_result, i, operator, value_str, expected_value_str);
}
print_assertion_result(
&config,
assertion.format(),
result,
i,
operator,
value_str,
expected_value_str,
);

Ok(result)
}
Assertion::Memory(cache_offset, operator, memory_value) => {
let cache = ctx.accounts.cache.as_ref().unwrap(); // TODO: Graceful error handling
let cache_data = cache.try_borrow_data()?; // TODO: Graceful error handling

let (value_str, expected_value_str, result) = memory_value
.deserialize_and_compare(cache_data, (cache_offset + 8) as usize, &operator)?;

assertion_result = result;

if verbose {
print_result(assertion_result, i, operator, value_str, expected_value_str);
}
.deserialize_and_compare(cache_data, (cache_offset + 8) as usize, operator)?;

print_assertion_result(
&config,
assertion.format(),
result,
i,
operator,
value_str,
expected_value_str,
);

Ok(result)
}
Assertion::AccountData(account_offset, operator, memory_value) => {
let account = remaining_accounts.next().unwrap();
let account_data = account.try_borrow_data()?;

let (value_str, expected_value_str, result) = memory_value
.deserialize_and_compare(account_data, account_offset as usize, &operator)?;

assertion_result = result;

if verbose {
print_result(assertion_result, i, operator, value_str, expected_value_str);
}
.deserialize_and_compare(account_data, (*account_offset) as usize, operator)?;

print_assertion_result(
&config,
assertion.format(),
result,
i,
operator,
value_str,
expected_value_str,
);

Ok(result)
}
Assertion::AccountBalance(balance_value, operator) => {
let account = remaining_accounts.next().unwrap();
let result = operator.evaluate(&**account.try_borrow_lamports()?, balance_value);

let value_str = account.get_lamports().to_string();
let expected_value_str = balance_value.to_string();

print_assertion_result(
&config,
assertion.format(),
result,
i,
operator,
value_str,
expected_value_str,
);

Ok(result)
}
Assertion::AccountBalance(expected_balance, operator) => {
Assertion::TokenAccountBalance(balance_value, operator) => {
let account = remaining_accounts.next().unwrap();

assertion_result =
operator.evaluate(&**account.try_borrow_lamports()?, &expected_balance);

if verbose {
print_result(
assertion_result,
i,
operator,
account.get_lamports().to_string(),
expected_balance.to_string(),
);
if account.owner.eq(&spl_associated_token_account::id()) {
return Err(ProgramError::InvalidAccount.into());
}
}
Assertion::TokenAccountBalance(_, _) => {
return Err(ProgramError::Unimplemented.into());

let token_account =
spl_token::state::Account::unpack_from_slice(&account.try_borrow_data()?)?;

let result = operator.evaluate(&token_account.amount, balance_value);

let value_str = token_account.amount.to_string();
let expected_value_str = balance_value.to_string();

print_assertion_result(
&config,
assertion.format(),
result,
i,
operator,
value_str,
expected_value_str,
);

Ok(result)
}
Assertion::AccountInfo(account_info_fields, operator) => {
let account = remaining_accounts.next().unwrap();
let operator_result = true;

for account_info_field in account_info_fields {
match account_info_field {
let operator_result = match account_info_field {
AccountInfoDataField::Key(pubkey) => {
assertion_result = operator.evaluate(&account.key(), &pubkey);
operator.evaluate(&account.key(), pubkey)
}
AccountInfoDataField::Owner(pubkey) => {
assertion_result = operator.evaluate(account.owner, &pubkey);
operator.evaluate(account.owner, pubkey)
}
AccountInfoDataField::Lamports(lamports) => {
assertion_result =
operator.evaluate(&account.get_lamports(), &lamports);
operator.evaluate(&account.get_lamports(), lamports)
}
AccountInfoDataField::DataLength(data_length) => {
assertion_result =
operator.evaluate(&(account.data_len() as u64), &data_length);
operator.evaluate(&(account.data_len() as u64), data_length)
}
AccountInfoDataField::Executable(executable) => {
assertion_result = operator.evaluate(&account.executable, &executable);
operator.evaluate(&account.executable, executable)
}
AccountInfoDataField::IsSigner(is_signer) => {
assertion_result = operator.evaluate(&account.is_signer, &is_signer);
operator.evaluate(&account.is_signer, is_signer)
}
AccountInfoDataField::IsWritable(is_writable) => {
assertion_result =
operator.evaluate(&account.is_writable, &is_writable);
operator.evaluate(&account.is_writable, is_writable)
}
AccountInfoDataField::RentEpoch(rent_epoch) => {
assertion_result =
operator.evaluate(&account.rent_epoch as &u64, &rent_epoch);
operator.evaluate(&account.rent_epoch as &u64, rent_epoch)
}
};

if !operator_result {
break;
}
}
}
}

assertion_state.record_result(i, assertion_result)?;

// assertion_results.push(assertion_result);
Ok(operator_result)
}
};

// if (logical_expression.is_none()
// || !logically_dependent_assertions
// .as_ref()
// .unwrap()
// .contains(&(i as u8)))
// && !assertion_result
// {
// return Err(ProgramError::AssertionFailed.into());
// }
assertion_state.record_result(i, assertion_result?)?;
}

msg!("assertion_state: {:?}", assertion_state);
assertion_state.evaluate()?;

// if let Some(logical_expressions) = &logical_expression {
// for logical_expression in logical_expressions {
// match logical_expression {
// Expression::And(assertion_indexes) => {
// let mut result = true;

// for assertion_index in assertion_indexes {
// result = result && assertion_results[*assertion_index as usize];
// }

// if verbose {
// msg!(
// "{} Expression::And -> {:?} {}",
// if result {
// "[✅] SUCCESS"
// } else {
// "[❌] FAIL "
// },
// result,
// assertion_indexes
// .iter()
// .map(|i| format!("[{}]", i))
// .collect::<Vec<String>>()
// .join(" AND ")
// );
// }

// if !result {
// return Err(ProgramError::AssertionFailed.into());
// }
// }
// Expression::Or(assertion_indexes) => {
// let mut result = false;

// for assertion_index in assertion_indexes {
// result = result || assertion_results[*assertion_index as usize];
// }

// if verbose {
// msg!(
// "{} Expression::Or -> {:?} {}",
// if result {
// "[✅] SUCCESS"
// } else {
// "[❌] FAIL "
// },
// result,
// assertion_indexes
// .iter()
// .map(|i| format!("[{}]", i))
// .collect::<Vec<String>>()
// .join(" OR ")
// );
// }

// if !result {
// return Err(ProgramError::AssertionFailed.into());
// }
// }
// }
// }
// }

Ok(())
}

Expand Down
63 changes: 63 additions & 0 deletions programs/lighthouse/program/src/structs/assert/assertion.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
use anchor_lang::prelude::{
borsh,
borsh::{BorshDeserialize, BorshSerialize},
};
use solana_program::pubkey::Pubkey;

use crate::structs::{operator::Operator, AccountInfoDataField, DataValue};

#[derive(BorshDeserialize, BorshSerialize, Debug, Clone)]
pub enum Assertion {
// memory offset, assertion
Memory(u16, Operator, DataValue),

AccountInfo(Vec<AccountInfoDataField>, Operator),

// account data offset, borsh type, operator
AccountData(u16, Operator, DataValue),

// balance, operator
AccountBalance(u64, Operator),

AccountOwnedBy(Pubkey, Operator),

// token balance, operator
TokenAccountBalance(u64, Operator),
// TODO
// IsSigner,
}

impl Assertion {
pub fn format(&self) -> String {
match self {
Assertion::Memory(offset, operator, value) => {
format!(
"Memory[{}] {} {}",
offset,
operator.format(),
value.format()
)
}
Assertion::AccountData(offset, operator, value) => {
format!(
"AccountData[{}] {} {}",
offset,
operator.format(),
value.format()
)
}
Assertion::AccountBalance(balance, operator) => {
format!("AccountBalance {} {}", balance, operator.format())
}
Assertion::AccountOwnedBy(pubkey, operator) => {
format!("AccountOwnedBy {} {}", pubkey, operator.format())
}
Assertion::TokenAccountBalance(balance, operator) => {
format!("TokenAccountBalance {} {}", balance, operator.format())
}
Assertion::AccountInfo(fields, operator) => {
format!("AccountInfo {:?} {}", fields, operator.format())
}
}
}
}
Loading

0 comments on commit fa77f78

Please sign in to comment.