use core::mem::size_of;
use core::ops::{Deref, DerefMut};
use arrayref::{array_refs, mut_array_refs};
use solana_account_view::{AccountView, Ref};
use solana_program_error::ProgramError;
use solana_pubkey::Pubkey;
use crate::{Role, SettlementAccount, SettlementError};
pub const DISCRIMINATOR: u8 = SettlementAccount::SettlementState.discriminator();
const WIDTH_DISCRIMINATOR: usize = size_of::<u8>();
pub const WIDTH_PUBKEY: usize = size_of::<Pubkey>();
pub const WIDTH_HEADER: usize = WIDTH_DISCRIMINATOR + 2 * WIDTH_PUBKEY;
struct HeaderSlots<'a> {
discriminator: &'a [u8; WIDTH_DISCRIMINATOR],
manager: &'a [u8; WIDTH_PUBKEY],
reclaim_authority: &'a [u8; WIDTH_PUBKEY],
}
struct HeaderSlotsMut<'a> {
discriminator: &'a mut [u8; WIDTH_DISCRIMINATOR],
manager: &'a mut [u8; WIDTH_PUBKEY],
reclaim_authority: &'a mut [u8; WIDTH_PUBKEY],
}
fn header_slots(header: &[u8; WIDTH_HEADER]) -> HeaderSlots<'_> {
let (discriminator, manager, reclaim_authority) =
array_refs![header, WIDTH_DISCRIMINATOR, WIDTH_PUBKEY, WIDTH_PUBKEY];
HeaderSlots {
discriminator,
manager,
reclaim_authority,
}
}
fn header_slots_mut(header: &mut [u8; WIDTH_HEADER]) -> HeaderSlotsMut<'_> {
let (discriminator, manager, reclaim_authority) =
mut_array_refs![header, WIDTH_DISCRIMINATOR, WIDTH_PUBKEY, WIDTH_PUBKEY];
HeaderSlotsMut {
discriminator,
manager,
reclaim_authority,
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StateInitArgs {
pub manager: Pubkey,
pub reclaim_authority: Pubkey,
}
pub struct StateAccount<T>(T);
impl<T: Deref<Target = [u8]>> StateAccount<T> {
pub fn attach(bytes: T) -> Result<Self, ProgramError> {
let header = bytes
.first_chunk::<WIDTH_HEADER>()
.ok_or(ProgramError::InvalidAccountData)?;
if header_slots(header).discriminator != &[DISCRIMINATOR] {
return Err(ProgramError::InvalidAccountData);
}
Ok(Self(bytes))
}
fn header(&self) -> &[u8; WIDTH_HEADER] {
self.0
.first_chunk::<WIDTH_HEADER>()
.expect("header length is guaranteed by any constructor of `StateAccount`")
}
pub fn authority(&self, role: Role) -> Pubkey {
let slots = header_slots(self.header());
let holder = match role {
Role::Manager => slots.manager,
Role::ReclaimAuthority => slots.reclaim_authority,
};
Pubkey::new_from_array(*holder)
}
fn solver_region(&self) -> &[[u8; WIDTH_PUBKEY]] {
let region = self.0.get(WIDTH_HEADER..).unwrap_or_default();
let (solvers, _partial) = region.as_chunks::<WIDTH_PUBKEY>();
solvers
}
pub fn solver_search(&self, solver: &Pubkey) -> Result<usize, usize> {
let seek = solver.to_bytes();
self.solver_region().binary_search(&seek)
}
pub fn is_solver(&self, solver: &Pubkey) -> bool {
self.solver_search(solver).is_ok()
}
pub fn solvers(&self) -> impl Iterator<Item = Pubkey> + '_ {
self.solver_region()
.iter()
.map(|raw| Pubkey::new_from_array(*raw))
}
pub fn grown_len(&self) -> Result<usize, ProgramError> {
self.0
.len()
.checked_add(WIDTH_PUBKEY)
.ok_or(ProgramError::ArithmeticOverflow)
}
}
impl<'a> StateAccount<Ref<'a, [u8]>> {
pub fn from_account(account: &'a AccountView) -> Result<Self, ProgramError> {
Self::attach(account.try_borrow()?)
}
}
impl<T: DerefMut<Target = [u8]>> StateAccount<T> {
pub fn initialize(mut bytes: T, args: &StateInitArgs) -> Result<Self, ProgramError> {
{
let slots = header_slots_mut(
bytes
.first_chunk_mut::<WIDTH_HEADER>()
.ok_or(ProgramError::AccountDataTooSmall)?,
);
*slots.discriminator = [DISCRIMINATOR];
*slots.manager = args.manager.to_bytes();
*slots.reclaim_authority = args.reclaim_authority.to_bytes();
}
Ok(Self(bytes))
}
fn header_mut(&mut self) -> &mut [u8; WIDTH_HEADER] {
let bytes: &mut [u8] = &mut self.0;
bytes
.first_chunk_mut::<WIDTH_HEADER>()
.expect("header length is guaranteed by `new`")
}
pub fn set_authority(&mut self, role: Role, new: &Pubkey) {
let slots = header_slots_mut(self.header_mut());
let holder = match role {
Role::Manager => slots.manager,
Role::ReclaimAuthority => slots.reclaim_authority,
};
*holder = new.to_bytes();
}
pub fn insert_solver(&mut self, solver: &Pubkey) -> Result<(), ProgramError> {
let (_spare, occupied) = self
.solver_region()
.split_last()
.expect("account grown by one solver slot before insertion");
let index = match occupied.binary_search(&solver.to_bytes()) {
Ok(_) => return Err(SettlementError::SolverAlreadyExists.into()),
Err(index) => index,
};
let data: &mut [u8] = &mut self.0;
let occupied_end = data
.len()
.checked_sub(WIDTH_PUBKEY)
.expect("account grown by one solver slot before insertion");
let gap = WIDTH_HEADER
.checked_add(
index
.checked_mul(WIDTH_PUBKEY)
.expect("insertion index bound by account length"),
)
.expect("insertion offset bound by account length");
let gap_end = gap
.checked_add(WIDTH_PUBKEY)
.expect("insertion slot bound by account length");
data.copy_within(gap..occupied_end, gap_end);
data[gap..gap_end].copy_from_slice(&solver.to_bytes());
Ok(())
}
}
#[cfg(any(test, feature = "test-fixtures"))]
pub mod fixtures {
use proptest::prelude::*;
use solana_pubkey::Pubkey;
use super::{StateAccount, StateInitArgs, WIDTH_HEADER};
pub fn state_account_bytes(header: &StateInitArgs, solvers: &[Pubkey]) -> Vec<u8> {
let mut sorted = solvers.to_vec();
sorted.sort();
let mut bytes = vec![0u8; WIDTH_HEADER];
StateAccount::initialize(&mut bytes[..], header).expect("header fits");
for solver in &sorted {
bytes.extend_from_slice(&solver.to_bytes());
}
bytes
}
pub fn arb_init_params() -> impl Strategy<Value = StateInitArgs> {
(any::<[u8; 32]>(), any::<[u8; 32]>()).prop_map(|(manager, reclaim_authority)| {
StateInitArgs {
manager: Pubkey::new_from_array(manager),
reclaim_authority: Pubkey::new_from_array(reclaim_authority),
}
})
}
}
#[cfg(test)]
mod tests {
use std::sync::LazyLock;
use super::*;
use crate::fixtures::pubkey_from_seed;
const DISCRIMINATOR_OFFSET: usize = 0;
static SAMPLE_INIT_ARGS: LazyLock<StateInitArgs> = LazyLock::new(|| StateInitArgs {
manager: pubkey_from_seed("SAMPLE_INIT_ARGS's sample manager"),
reclaim_authority: pubkey_from_seed("SAMPLE_INIT_ARGS's sample reclaim authority"),
});
fn header_bytes() -> [u8; WIDTH_HEADER] {
let mut bytes = [0u8; WIDTH_HEADER];
StateAccount::initialize(&mut bytes[..], &SAMPLE_INIT_ARGS).expect("header fits");
bytes
}
fn state_bytes(solvers: &[Pubkey]) -> Vec<u8> {
super::fixtures::state_account_bytes(&SAMPLE_INIT_ARGS, solvers)
}
#[test]
fn header_has_the_canonical_wire_layout() {
assert_eq!(WIDTH_HEADER, 65);
let bytes = header_bytes();
assert_eq!(bytes[0], SettlementAccount::SettlementState.discriminator());
assert_eq!(&bytes[1..33], &SAMPLE_INIT_ARGS.manager.to_bytes()[..]);
assert_eq!(
&bytes[33..65],
&SAMPLE_INIT_ARGS.reclaim_authority.to_bytes()[..]
);
}
#[test]
fn reads_role_holders_from_the_header() {
let bytes = header_bytes();
let state = StateAccount::attach(&bytes[..]).expect("valid header");
assert_eq!(state.authority(Role::Manager), SAMPLE_INIT_ARGS.manager);
assert_eq!(
state.authority(Role::ReclaimAuthority),
SAMPLE_INIT_ARGS.reclaim_authority
);
}
#[test]
fn new_rejects_wrong_discriminator() {
let mut bytes = header_bytes();
bytes[DISCRIMINATOR_OFFSET] = 0xff;
assert_eq!(
StateAccount::attach(&bytes[..]).err(),
Some(ProgramError::InvalidAccountData),
);
}
#[test]
fn new_rejects_too_short_buffer() {
let bytes = header_bytes();
assert_eq!(
StateAccount::attach(&bytes[..WIDTH_HEADER - 1]).err(),
Some(ProgramError::InvalidAccountData),
);
}
fn assert_set_authority_updates_only(target: Role) {
let new_holder = pubkey_from_seed("set_authority's new holder");
let mut bytes = header_bytes();
let before = Role::ALL.map(|role| {
StateAccount::attach(&bytes[..])
.expect("valid header")
.authority(role)
});
StateAccount::attach(&mut bytes[..])
.expect("valid header")
.set_authority(target, &new_holder);
let state = StateAccount::attach(&bytes[..]).expect("valid header");
for (role, prior) in Role::ALL.into_iter().zip(before) {
let expected = if role == target { new_holder } else { prior };
assert_eq!(state.authority(role), expected);
}
}
macro_rules! set_authority_test {
($name:ident: $role:expr) => {
#[test]
fn $name() {
assert_set_authority_updates_only($role);
}
};
}
set_authority_test!(set_authority_updates_only_manager: Role::Manager);
set_authority_test!(set_authority_updates_only_reclaim_authority: Role::ReclaimAuthority);
#[test]
fn new_accepts_a_longer_account_and_reads_the_header() {
let mut bytes = header_bytes().to_vec();
bytes.push(0x42);
let state = StateAccount::attach(&bytes[..]).expect("header with trailing bytes is valid");
assert_eq!(state.authority(Role::Manager), SAMPLE_INIT_ARGS.manager);
assert_eq!(
state.authority(Role::ReclaimAuthority),
SAMPLE_INIT_ARGS.reclaim_authority
);
}
fn sample_solvers() -> (Vec<u8>, [Pubkey; 3]) {
let solvers = [
pubkey_from_seed("sample_solver's solver a"),
pubkey_from_seed("sample_solver's solver b"),
pubkey_from_seed("sample_solver's solver c"),
];
let bytes = state_bytes(&solvers);
let mut sorted = solvers;
sorted.sort();
(bytes, sorted)
}
#[test]
fn solvers_is_empty_without_any_stored() {
let bytes = state_bytes(&[]);
let state = StateAccount::attach(&bytes[..]).expect("valid header");
assert_eq!(state.solvers().count(), 0);
}
#[test]
fn solvers_lists_stored_solvers_sorted() {
let (bytes, sorted) = sample_solvers();
let state = StateAccount::attach(&bytes[..]).expect("valid header");
assert_eq!(state.solvers().collect::<Vec<_>>(), sorted);
}
#[test]
fn solver_search_on_empty_list() {
let bytes = state_bytes(&[]);
let state = StateAccount::attach(&bytes[..]).expect("valid header");
let absent = pubkey_from_seed("absent solver");
assert_eq!(state.solver_search(&absent), Err(0));
}
#[test]
fn solver_search_finds_present() {
let (bytes, sorted) = sample_solvers();
let state = StateAccount::attach(&bytes[..]).expect("valid header");
for (index, solver) in sorted.iter().enumerate() {
assert_eq!(state.solver_search(solver), Ok(index));
}
}
#[test]
fn solver_search_locates_absent() {
let (bytes, sorted) = sample_solvers();
let state = StateAccount::attach(&bytes[..]).expect("valid header");
let absent = pubkey_from_seed("absent solver");
let slot = state.solver_search(&absent).expect_err("absent solver");
assert_eq!(slot, sorted.partition_point(|s| s < &absent));
}
#[test]
fn initialize_rejects_too_small_buffer() {
let mut bytes = [0u8; WIDTH_HEADER - 1];
assert_eq!(
StateAccount::initialize(&mut bytes[..], &SAMPLE_INIT_ARGS).err(),
Some(ProgramError::AccountDataTooSmall),
);
}
mod proptest {
use ::proptest::prelude::*;
use super::*;
proptest! {
#[test]
fn is_solver_reflects_membership(
header in fixtures::arb_init_params(),
raw_solvers in prop::collection::btree_set(any::<[u8; 32]>(), 0..50),
raw_absent in any::<[u8; 32]>(),
) {
prop_assume!(!raw_solvers.contains(&raw_absent));
let stored: Vec<Pubkey> =
raw_solvers.into_iter().map(Pubkey::new_from_array).collect();
let absent = Pubkey::new_from_array(raw_absent);
let bytes = fixtures::state_account_bytes(&header, &stored);
let state = StateAccount::attach(&bytes[..]).expect("valid header");
for solver in &stored {
prop_assert!(state.is_solver(solver));
}
prop_assert!(!state.is_solver(&absent));
}
#[test]
fn account_encode_roundtrip(header in fixtures::arb_init_params()) {
let mut bytes = [0u8; WIDTH_HEADER];
StateAccount::initialize(&mut bytes[..], &header).expect("header fits");
let state = StateAccount::attach(&bytes[..]).expect("valid header");
let StateInitArgs { manager, reclaim_authority } = header;
prop_assert_eq!(state.authority(Role::Manager), manager);
prop_assert_eq!(state.authority(Role::ReclaimAuthority), reclaim_authority);
}
#[test]
fn insert_solver_inserts_an_absent_solver(
header in fixtures::arb_init_params(),
raw_solvers in prop::collection::btree_set(any::<[u8; 32]>(), 0..50),
raw_new in any::<[u8; 32]>(),
) {
prop_assume!(!raw_solvers.contains(&raw_new));
let stored: Vec<Pubkey> =
raw_solvers.into_iter().map(Pubkey::new_from_array).collect();
let new = Pubkey::new_from_array(raw_new);
let mut bytes = fixtures::state_account_bytes(&header, &stored);
let grown_len = StateAccount::attach(&bytes[..])
.expect("valid header")
.grown_len()
.expect("grown length fits");
prop_assert_eq!(grown_len, bytes.len().strict_add(WIDTH_PUBKEY));
bytes.resize(grown_len, 0);
StateAccount::attach(&mut bytes[..])
.expect("valid header")
.insert_solver(&new)
.expect("absent solver inserts");
let mut expected = stored;
expected.push(new);
expected.sort();
let state = StateAccount::attach(&bytes[..]).expect("valid header");
prop_assert_eq!(state.solvers().collect::<Vec<_>>(), expected);
}
#[test]
fn insert_solver_rejects_an_existing_solver(
header in fixtures::arb_init_params(),
raw_solvers in prop::collection::btree_set(any::<[u8; 32]>(), 1..50),
pick in any::<prop::sample::Index>(),
) {
let stored: Vec<Pubkey> =
raw_solvers.into_iter().map(Pubkey::new_from_array).collect();
let existing = stored[pick.index(stored.len())];
let mut bytes = fixtures::state_account_bytes(&header, &stored);
bytes.resize(bytes.len().strict_add(WIDTH_PUBKEY), 0);
prop_assert_eq!(
StateAccount::attach(&mut bytes[..])
.expect("valid header")
.insert_solver(&existing),
Err(SettlementError::SolverAlreadyExists.into()),
);
let state = StateAccount::attach(&bytes[..]).expect("valid header");
prop_assert_eq!(state.solvers().take(stored.len()).collect::<Vec<_>>(), stored);
}
}
}
}