diff --git a/interface/src/lib.rs b/interface/src/lib.rs index 231ba644..7d678566 100644 --- a/interface/src/lib.rs +++ b/interface/src/lib.rs @@ -7,18 +7,20 @@ use solana_sdk::{ }; pub use swig; use swig::actions::{ - add_authority_v1::AddAuthorityV1Args, create_session_v1::CreateSessionV1Args, - create_sub_account_v1::CreateSubAccountV1Args, create_v1::CreateV1Args, - remove_authority_v1::RemoveAuthorityV1Args, sub_account_sign_v1::SubAccountSignV1Args, - toggle_sub_account_v1::ToggleSubAccountV1Args, + add_authority_v1::AddAuthorityV1Args, add_authorization_lock_v1::AddAuthorizationLockV1Args, + create_session_v1::CreateSessionV1Args, create_sub_account_v1::CreateSubAccountV1Args, + create_v1::CreateV1Args, remove_authority_v1::RemoveAuthorityV1Args, + remove_authorization_lock_v1::RemoveAuthorizationLockV1Args, + sub_account_sign_v1::SubAccountSignV1Args, toggle_sub_account_v1::ToggleSubAccountV1Args, withdraw_from_sub_account_v1::WithdrawFromSubAccountV1Args, }; pub use swig_compact_instructions::*; use swig_state_x::{ action::{ - all::All, manage_authority::ManageAuthority, program::Program, program_scope::ProgramScope, - sol_limit::SolLimit, sol_recurring_limit::SolRecurringLimit, stake_all::StakeAll, - stake_limit::StakeLimit, stake_recurring_limit::StakeRecurringLimit, + all::All, manage_authority::ManageAuthority, + manage_authorization_locks::ManageAuthorizationLocks, program::Program, + program_scope::ProgramScope, sol_limit::SolLimit, sol_recurring_limit::SolRecurringLimit, + stake_all::StakeAll, stake_limit::StakeLimit, stake_recurring_limit::StakeRecurringLimit, sub_account::SubAccount, token_limit::TokenLimit, token_recurring_limit::TokenRecurringLimit, Action, Permission, }, @@ -39,6 +41,7 @@ pub enum ClientAction { ProgramScope(ProgramScope), All(All), ManageAuthority(ManageAuthority), + ManageAuthorizationLocks(ManageAuthorizationLocks), SubAccount(SubAccount), StakeLimit(StakeLimit), StakeRecurringLimit(StakeRecurringLimit), @@ -60,6 +63,10 @@ impl ClientAction { ClientAction::ProgramScope(_) => (Permission::ProgramScope, ProgramScope::LEN), ClientAction::All(_) => (Permission::All, All::LEN), ClientAction::ManageAuthority(_) => (Permission::ManageAuthority, ManageAuthority::LEN), + ClientAction::ManageAuthorizationLocks(_) => ( + Permission::ManageAuthorizationLocks, + ManageAuthorizationLocks::LEN, + ), ClientAction::SubAccount(_) => (Permission::SubAccount, SubAccount::LEN), ClientAction::StakeLimit(_) => (Permission::StakeLimit, StakeLimit::LEN), ClientAction::StakeRecurringLimit(_) => { @@ -86,6 +93,7 @@ impl ClientAction { ClientAction::ProgramScope(action) => action.into_bytes(), ClientAction::All(action) => action.into_bytes(), ClientAction::ManageAuthority(action) => action.into_bytes(), + ClientAction::ManageAuthorizationLocks(action) => action.into_bytes(), ClientAction::SubAccount(action) => action.into_bytes(), ClientAction::StakeLimit(action) => action.into_bytes(), ClientAction::StakeRecurringLimit(action) => action.into_bytes(), @@ -969,3 +977,65 @@ impl ToggleSubAccountInstruction { }) } } + +pub struct AddAuthorizationLockInstruction; + +impl AddAuthorizationLockInstruction { + pub fn new( + swig_account: Pubkey, + authority: Pubkey, + payer: Pubkey, + acting_role_id: u32, + token_mint: [u8; 32], + amount: u64, + expiry_slot: u64, + ) -> anyhow::Result { + let accounts = vec![ + AccountMeta::new(swig_account, false), + AccountMeta::new(payer, true), + AccountMeta::new_readonly(system_program::ID, false), + AccountMeta::new_readonly(authority, true), + ]; + + let args = AddAuthorizationLockV1Args::new(acting_role_id, token_mint, amount, expiry_slot); + let args_bytes = args + .into_bytes() + .map_err(|e| anyhow::anyhow!("Failed to serialize args {:?}", e))?; + + Ok(Instruction { + program_id: program_id(), + accounts, + data: [args_bytes, &[3]].concat(), + }) + } +} + +pub struct RemoveAuthorizationLockInstruction; + +impl RemoveAuthorizationLockInstruction { + pub fn new( + swig_account: Pubkey, + authority: Pubkey, + payer: Pubkey, + acting_role_id: u32, + lock_index: u32, + ) -> anyhow::Result { + let accounts = vec![ + AccountMeta::new(swig_account, false), + AccountMeta::new(payer, true), + AccountMeta::new_readonly(system_program::ID, false), + AccountMeta::new_readonly(authority, true), + ]; + + let args = RemoveAuthorizationLockV1Args::new(acting_role_id, lock_index); + let args_bytes = args + .into_bytes() + .map_err(|e| anyhow::anyhow!("Failed to serialize args {:?}", e))?; + + Ok(Instruction { + program_id: program_id(), + accounts, + data: [args_bytes, &[3]].concat(), + }) + } +} diff --git a/program/src/actions/add_authorization_lock_v1.rs b/program/src/actions/add_authorization_lock_v1.rs new file mode 100644 index 00000000..52ff2099 --- /dev/null +++ b/program/src/actions/add_authorization_lock_v1.rs @@ -0,0 +1,347 @@ +/// Module for adding authorization locks to Swig accounts. +/// Authorization locks pre-authorize token spending up to a specific amount +/// and expiry slot, providing a mechanism for payment preauthorizations. +use no_padding::NoPadding; +use pinocchio::{ + account_info::AccountInfo, + program_error::ProgramError, + sysvars::{clock::Clock, rent::Rent, Sysvar}, + ProgramResult, +}; +use pinocchio_pubkey::pubkey; +use pinocchio_system::instructions::Transfer; +use swig_assertions::*; +use swig_state_x::{ + action::{ + all::All, manage_authorization_locks::ManageAuthorizationLocks, sol_limit::SolLimit, + sol_recurring_limit::SolRecurringLimit, token_limit::TokenLimit, + token_recurring_limit::TokenRecurringLimit, + }, + role::Position, + swig::{AuthorizationLock, Swig, SwigBuilder, SwigWithRoles}, + Discriminator, IntoBytes, SwigAuthenticateError, SwigStateError, Transmutable, TransmutableMut, +}; + +use crate::{ + error::SwigError, + instruction::{ + accounts::{AddAuthorizationLockV1Accounts, Context}, + SwigInstruction, + }, +}; + +/// Arguments for adding an authorization lock to a Swig wallet. +/// +/// # Fields +/// * `instruction` - The instruction type identifier +/// * `token_mint` - The mint of the token to lock +/// * `amount` - The maximum amount that can be spent +/// * `expiry_slot` - The slot when this lock expires +/// * `acting_role_id` - ID of the role performing the operation +#[derive(Debug, NoPadding)] +#[repr(C, align(8))] +pub struct AddAuthorizationLockV1Args { + instruction: SwigInstruction, + _padding: [u8; 2], // Reduced padding for new field + pub acting_role_id: u32, + pub token_mint: [u8; 32], + pub amount: u64, + pub expiry_slot: u64, +} + +impl AddAuthorizationLockV1Args { + /// Creates a new instance of AddAuthorizationLockV1Args. + /// + /// # Arguments + /// * `acting_role_id` - ID of the role performing the operation + /// * `token_mint` - The mint of the token to lock + /// * `amount` - The maximum amount that can be spent + /// * `expiry_slot` - The slot when this lock expires + pub fn new(acting_role_id: u32, token_mint: [u8; 32], amount: u64, expiry_slot: u64) -> Self { + Self { + instruction: SwigInstruction::AddAuthorizationLockV1, + _padding: [0; 2], + acting_role_id, + token_mint, + amount, + expiry_slot, + } + } +} + +impl Transmutable for AddAuthorizationLockV1Args { + const LEN: usize = core::mem::size_of::(); +} + +impl IntoBytes for AddAuthorizationLockV1Args { + fn into_bytes(&self) -> Result<&[u8], ProgramError> { + Ok(unsafe { core::slice::from_raw_parts(self as *const Self as *const u8, Self::LEN) }) + } +} + +/// Structured data for the add authorization lock instruction. +pub struct AddAuthorizationLockV1<'a> { + pub args: &'a AddAuthorizationLockV1Args, + data_payload: &'a [u8], + authority_payload: &'a [u8], +} + +impl<'a> AddAuthorizationLockV1<'a> { + /// Parses the instruction data bytes into an AddAuthorizationLockV1 + /// instance. + pub fn from_instruction_bytes(data: &'a [u8]) -> Result { + if data.len() < AddAuthorizationLockV1Args::LEN { + return Err(SwigError::InvalidSwigSignInstructionDataTooShort.into()); + } + + let (inst, authority_payload) = data.split_at(AddAuthorizationLockV1Args::LEN); + let args = unsafe { AddAuthorizationLockV1Args::load_unchecked(inst)? }; + + Ok(Self { + args, + data_payload: inst, + authority_payload, + }) + } +} + +/// Adds an authorization lock to a Swig wallet. +/// +/// This function: +/// 1. Validates the acting role's permissions (All or ManageAuthorizationLocks) +/// 2. Authenticates the request +/// 3. Validates the Swig account and lock parameters +/// 4. Reallocates the account to accommodate the new lock +/// 5. Adds the authorization lock to the end of the account data +/// +/// # Arguments +/// * `ctx` - The account context for the operation +/// * `data` - Raw instruction data bytes +/// * `all_accounts` - All accounts involved in the operation +/// +/// # Returns +/// * `ProgramResult` - Success or error status +#[inline(always)] +pub fn add_authorization_lock_v1( + ctx: Context, + data: &[u8], + all_accounts: &[AccountInfo], +) -> ProgramResult { + check_stack_height(1, SwigError::Cpi)?; + + let add_lock = AddAuthorizationLockV1::from_instruction_bytes(data)?; + + // Get current slot to validate expiry + let clock = Clock::get()?; + if add_lock.args.expiry_slot <= clock.slot { + return Err(SwigError::InvalidAuthorizationLockExpiry.into()); + } + + let swig_account_data = unsafe { ctx.accounts.swig.borrow_mut_data_unchecked() }; + if unsafe { *swig_account_data.get_unchecked(0) } != Discriminator::SwigAccount as u8 { + return Err(SwigError::InvalidSwigAccountDiscriminator.into()); + } + + // Authentication and permission checking - consolidate loads + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account_data).unwrap(); + let role = swig_with_roles.get_role(add_lock.args.acting_role_id)?; + + // Get existing authorization locks for this role using a smaller array to avoid + // stack overflow + const MAX_LOCKS: usize = 10; // Smaller bound to prevent stack overflow + let (existing_locks, _count) = swig_with_roles + .get_authorization_locks_by_role::(add_lock.args.acting_role_id)?; + + // Convert Option array to Vec of actual locks + let existing_locks_vec: Vec = existing_locks + .iter() + .filter_map(|opt_lock| *opt_lock) + .collect(); + + // TODO need to merge in fix for getting all roles because of action boundary + // cursor positions + let (swig_header, swig_roles) = unsafe { swig_account_data.split_at_mut_unchecked(Swig::LEN) }; + let swig = unsafe { Swig::load_mut_unchecked(swig_header)? }; + let acting_role = Swig::get_mut_role(add_lock.args.acting_role_id, swig_roles)?; + if acting_role.is_none() { + return Err(SwigError::InvalidAuthorityNotFoundByRoleId.into()); + } + let acting_role = acting_role.unwrap(); + + // Authenticate the caller + let slot = clock.slot; + if acting_role.authority.session_based() { + acting_role.authority.authenticate_session( + all_accounts, + add_lock.authority_payload, + add_lock.data_payload, + slot, + )?; + } else { + acting_role.authority.authenticate( + all_accounts, + add_lock.authority_payload, + add_lock.data_payload, + slot, + )?; + } + + // Check permissions: must have All or ManageAuthorizationLocks + let all = acting_role.get_action::(&[])?; + let manage_auth_locks = acting_role.get_action::(&[])?; + + if all.is_none() && manage_auth_locks.is_none() { + return Err(SwigAuthenticateError::PermissionDeniedMissingPermission.into()); + } + + // Validate the new lock against existing token limits + validate_authorization_lock_against_limits( + &acting_role, + add_lock.args.token_mint, + add_lock.args.amount, + &existing_locks_vec, + )?; + + // Re-borrow data after authentication + let swig_account_data = unsafe { ctx.accounts.swig.borrow_mut_data_unchecked() }; + let (swig_header, remaining_data) = + unsafe { swig_account_data.split_at_mut_unchecked(Swig::LEN) }; + + // Find the end of roles data to determine where authorization locks start + let mut roles_end = 0; + let mut cursor = 0; + for _i in 0..swig.roles { + if cursor + Position::LEN > remaining_data.len() { + return Err(SwigStateError::InvalidRoleData.into()); + } + let position = + unsafe { Position::load_unchecked(&remaining_data[cursor..cursor + Position::LEN])? }; + cursor = position.boundary() as usize; + roles_end = cursor; + } + + // Calculate required space for new authorization lock + let new_lock_size = AuthorizationLock::LEN; + let current_auth_locks_size = swig.authorization_locks as usize * AuthorizationLock::LEN; + let required_total_size = Swig::LEN + roles_end + current_auth_locks_size + new_lock_size; + + // Check if we need to reallocate + let current_size = ctx.accounts.swig.data_len(); + if required_total_size > current_size { + // Reallocate account + ctx.accounts.swig.realloc(required_total_size, false)?; + let rent = Rent::get()?; + let rent_required = rent.minimum_balance(required_total_size); + let current_lamports = ctx.accounts.swig.lamports(); + if rent_required > current_lamports { + let additional_rent = rent_required - current_lamports; + Transfer { + from: ctx.accounts.payer, + to: ctx.accounts.swig, + lamports: additional_rent, + } + .invoke()?; + } + } + + // Re-borrow data after potential reallocation + let swig_account_data = unsafe { ctx.accounts.swig.borrow_mut_data_unchecked() }; + + // Write the new lock at the end of the authorization locks section using + // zero-copy + let auth_locks_start = Swig::LEN + roles_end; + let new_lock_offset = auth_locks_start + current_auth_locks_size; + + // Zero-copy: write directly to the account data buffer + let lock_slice = &mut swig_account_data[new_lock_offset..new_lock_offset + new_lock_size]; + let new_lock = unsafe { &mut *(lock_slice.as_mut_ptr() as *mut AuthorizationLock) }; + + // Initialize the lock fields directly in memory + new_lock.token_mint = add_lock.args.token_mint; + new_lock.amount = add_lock.args.amount; + new_lock.expiry_slot = add_lock.args.expiry_slot; + new_lock.role_id = add_lock.args.acting_role_id; + new_lock._padding = [0; 4]; + + // Update the authorization locks count in the header + let (swig_header, _) = unsafe { swig_account_data.split_at_mut_unchecked(Swig::LEN) }; + let swig = unsafe { Swig::load_mut_unchecked(swig_header)? }; + swig.authorization_locks += 1; + + Ok(()) +} + +/// Validates that the new authorization lock doesn't exceed existing token +/// limits for the role. +/// +/// This function checks if adding the new authorization lock would cause the +/// total authorization locks for the token to exceed any existing token limits +/// (simple or recurring) that the acting role has for that specific token. +/// For SOL (wrapped SOL mint), it also checks against SOL limits. +/// +/// # Arguments +/// * `acting_role` - The role that is creating the authorization lock +/// * `token_mint` - The mint of the token +/// * `new_lock_amount` - The amount of the new authorization lock +/// * `existing_locks` - All existing authorization locks for this role and +/// token +/// +/// # Returns +/// * `Ok(())` - If the new lock is within limits +/// * `Err(ProgramError)` - If the new lock would exceed limits +fn validate_authorization_lock_against_limits<'a>( + acting_role: &'a swig_state_x::role::RoleMut<'a>, + token_mint: [u8; 32], + new_lock_amount: u64, + existing_locks: &[AuthorizationLock], +) -> ProgramResult { + // Wrapped SOL mint address + const WRAPPED_SOL_MINT: [u8; 32] = pubkey!("So11111111111111111111111111111111111111112"); + + // Calculate total existing authorization lock amount for this token + let existing_total = existing_locks + .iter() + .filter(|lock| lock.token_mint == token_mint) + .map(|lock| lock.amount) + .sum::(); + + let total_with_new_lock = existing_total.saturating_add(new_lock_amount); + + // Check if this is the wrapped SOL mint + if token_mint == WRAPPED_SOL_MINT { + // Check against SOL limits first + if let Ok(Some(sol_limit)) = acting_role.get_action::(&[]) { + if total_with_new_lock > sol_limit.amount { + return Err(SwigAuthenticateError::PermissionDeniedInsufficientBalance.into()); + } + } + + // Check against recurring SOL limit + if let Ok(Some(sol_recurring_limit)) = acting_role.get_action::(&[]) { + if total_with_new_lock > sol_recurring_limit.current_amount { + return Err(SwigAuthenticateError::PermissionDeniedInsufficientBalance.into()); + } + } + } else { + // Check token limits for non-SOL tokens + let mint_data = &token_mint[..]; + + // Check against simple token limit + if let Ok(Some(token_limit)) = acting_role.get_action::(mint_data) { + if total_with_new_lock > token_limit.current_amount { + return Err(SwigAuthenticateError::PermissionDeniedInsufficientBalance.into()); + } + } + + // Check against recurring token limit + if let Ok(Some(token_recurring_limit)) = + acting_role.get_action::(mint_data) + { + if total_with_new_lock > token_recurring_limit.limit { + return Err(SwigAuthenticateError::PermissionDeniedInsufficientBalance.into()); + } + } + } + + Ok(()) +} diff --git a/program/src/actions/mod.rs b/program/src/actions/mod.rs index 6d7ef273..0cc94201 100644 --- a/program/src/actions/mod.rs +++ b/program/src/actions/mod.rs @@ -6,10 +6,12 @@ //! instruction's business logic. pub mod add_authority_v1; +pub mod add_authorization_lock_v1; pub mod create_session_v1; pub mod create_sub_account_v1; pub mod create_v1; pub mod remove_authority_v1; +pub mod remove_authorization_lock_v1; pub mod sign_v1; pub mod sub_account_sign_v1; pub mod toggle_sub_account_v1; @@ -19,15 +21,17 @@ use num_enum::FromPrimitive; use pinocchio::{account_info::AccountInfo, msg, program_error::ProgramError, ProgramResult}; use self::{ - add_authority_v1::*, create_session_v1::*, create_sub_account_v1::*, create_v1::*, - remove_authority_v1::*, sign_v1::*, sub_account_sign_v1::*, toggle_sub_account_v1::*, + add_authority_v1::*, add_authorization_lock_v1::*, create_session_v1::*, + create_sub_account_v1::*, create_v1::*, remove_authority_v1::*, + remove_authorization_lock_v1::*, sign_v1::*, sub_account_sign_v1::*, toggle_sub_account_v1::*, withdraw_from_sub_account_v1::*, }; use crate::{ instruction::{ accounts::{ - AddAuthorityV1Accounts, CreateSessionV1Accounts, CreateSubAccountV1Accounts, - CreateV1Accounts, RemoveAuthorityV1Accounts, SignV1Accounts, SubAccountSignV1Accounts, + AddAuthorityV1Accounts, AddAuthorizationLockV1Accounts, CreateSessionV1Accounts, + CreateSubAccountV1Accounts, CreateV1Accounts, RemoveAuthorityV1Accounts, + RemoveAuthorizationLockV1Accounts, SignV1Accounts, SubAccountSignV1Accounts, ToggleSubAccountV1Accounts, WithdrawFromSubAccountV1Accounts, }, SwigInstruction, @@ -45,6 +49,8 @@ use crate::{ /// * `accounts` - List of accounts involved in the instruction /// * `account_classification` - Classification of each account's type and role /// * `data` - Raw instruction data +/// * `authorization_lock_cache` - Optional cache of authorization locks for +/// performance /// /// # Returns /// * `ProgramResult` - Success or error status @@ -53,6 +59,7 @@ pub fn process_action( accounts: &[AccountInfo], account_classification: &[AccountClassification], data: &[u8], + authorization_lock_cache: Option<&crate::util::AuthorizationLockCache>, ) -> ProgramResult { if data.len() < 2 { return Err(ProgramError::InvalidInstructionData); @@ -61,7 +68,12 @@ pub fn process_action( let ix = SwigInstruction::from_primitive(discriminator); match ix { SwigInstruction::CreateV1 => process_create_v1(accounts, data), - SwigInstruction::SignV1 => process_sign_v1(accounts, account_classification, data), + SwigInstruction::SignV1 => process_sign_v1( + accounts, + account_classification, + data, + authorization_lock_cache, + ), SwigInstruction::AddAuthorityV1 => process_add_authority_v1(accounts, data), SwigInstruction::RemoveAuthorityV1 => process_remove_authority_v1(accounts, data), SwigInstruction::CreateSessionV1 => process_create_session_v1(accounts, data), @@ -73,6 +85,12 @@ pub fn process_action( process_sub_account_sign_v1(accounts, account_classification, data) }, SwigInstruction::ToggleSubAccountV1 => process_toggle_sub_account_v1(accounts, data), + SwigInstruction::AddAuthorizationLockV1 => { + process_add_authorization_lock_v1(accounts, data) + }, + SwigInstruction::RemoveAuthorizationLockV1 => { + process_remove_authorization_lock_v1(accounts, data) + }, } } @@ -91,9 +109,16 @@ fn process_sign_v1( accounts: &[AccountInfo], account_classification: &[AccountClassification], data: &[u8], + authorization_lock_cache: Option<&crate::util::AuthorizationLockCache>, ) -> ProgramResult { let account_ctx = SignV1Accounts::context(accounts)?; - sign_v1(account_ctx, accounts, data, account_classification) + sign_v1( + account_ctx, + accounts, + data, + account_classification, + authorization_lock_cache, + ) } /// Processes an AddAuthorityV1 instruction. @@ -159,3 +184,19 @@ fn process_toggle_sub_account_v1(accounts: &[AccountInfo], data: &[u8]) -> Progr let account_ctx = ToggleSubAccountV1Accounts::context(accounts)?; toggle_sub_account_v1(account_ctx, data, accounts) } + +/// Processes an AddAuthorizationLockV1 instruction. +/// +/// Adds an authorization lock to the wallet. +fn process_add_authorization_lock_v1(accounts: &[AccountInfo], data: &[u8]) -> ProgramResult { + let account_ctx = AddAuthorizationLockV1Accounts::context(accounts)?; + add_authorization_lock_v1(account_ctx, data, accounts) +} + +/// Processes a RemoveAuthorizationLockV1 instruction. +/// +/// Removes an authorization lock from the wallet. +fn process_remove_authorization_lock_v1(accounts: &[AccountInfo], data: &[u8]) -> ProgramResult { + let account_ctx = RemoveAuthorizationLockV1Accounts::context(accounts)?; + remove_authorization_lock_v1(account_ctx, data, accounts) +} diff --git a/program/src/actions/remove_authorization_lock_v1.rs b/program/src/actions/remove_authorization_lock_v1.rs new file mode 100644 index 00000000..df874250 --- /dev/null +++ b/program/src/actions/remove_authorization_lock_v1.rs @@ -0,0 +1,261 @@ +/// Module for removing authorization locks from Swig accounts. +/// Authorization locks can be removed by authorities with proper permissions, +/// which helps manage payment preauthorizations by revoking them when needed. +use no_padding::NoPadding; +use pinocchio::{ + account_info::AccountInfo, + msg, + program_error::ProgramError, + sysvars::{clock::Clock, Sysvar}, + ProgramResult, +}; +use swig_assertions::*; +use swig_state_x::{ + action::{all::All, manage_authorization_locks::ManageAuthorizationLocks}, + role::Position, + swig::{AuthorizationLock, Swig, SwigBuilder}, + Discriminator, IntoBytes, SwigAuthenticateError, SwigStateError, Transmutable, TransmutableMut, +}; + +use crate::{ + error::SwigError, + instruction::{ + accounts::{Context, RemoveAuthorizationLockV1Accounts}, + SwigInstruction, + }, +}; + +/// Arguments for removing an authorization lock from a Swig wallet. +/// +/// # Fields +/// * `instruction` - The instruction type identifier +/// * `acting_role_id` - ID of the role performing the operation +/// * `lock_index` - Index of the authorization lock to remove +#[derive(Debug, NoPadding)] +#[repr(C, align(8))] +pub struct RemoveAuthorizationLockV1Args { + instruction: SwigInstruction, + _padding: [u8; 6], // Adjusted padding for proper alignment + pub acting_role_id: u32, + pub lock_index: u32, +} + +impl RemoveAuthorizationLockV1Args { + /// Creates a new instance of RemoveAuthorizationLockV1Args. + /// + /// # Arguments + /// * `acting_role_id` - ID of the role performing the operation + /// * `lock_index` - Index of the authorization lock to remove + pub fn new(acting_role_id: u32, lock_index: u32) -> Self { + Self { + instruction: SwigInstruction::RemoveAuthorizationLockV1, + _padding: [0; 6], + acting_role_id, + lock_index, + } + } +} + +impl Transmutable for RemoveAuthorizationLockV1Args { + const LEN: usize = core::mem::size_of::(); +} + +impl IntoBytes for RemoveAuthorizationLockV1Args { + fn into_bytes(&self) -> Result<&[u8], ProgramError> { + Ok(unsafe { core::slice::from_raw_parts(self as *const Self as *const u8, Self::LEN) }) + } +} + +/// Structured data for the remove authorization lock instruction. +pub struct RemoveAuthorizationLockV1<'a> { + pub args: &'a RemoveAuthorizationLockV1Args, + data_payload: &'a [u8], + authority_payload: &'a [u8], +} + +impl<'a> RemoveAuthorizationLockV1<'a> { + /// Parses the instruction data bytes into a RemoveAuthorizationLockV1 + /// instance. + pub fn from_instruction_bytes(data: &'a [u8]) -> Result { + if data.len() < RemoveAuthorizationLockV1Args::LEN { + return Err(SwigError::InvalidSwigSignInstructionDataTooShort.into()); + } + + let (inst, authority_payload) = data.split_at(RemoveAuthorizationLockV1Args::LEN); + let args = unsafe { RemoveAuthorizationLockV1Args::load_unchecked(inst)? }; + + Ok(Self { + args, + data_payload: inst, + authority_payload, + }) + } +} + +/// Removes an authorization lock from a Swig wallet. +/// +/// This function: +/// 1. Validates the acting role's permissions (All or ManageAuthorizationLocks) +/// 2. Authenticates the request +/// 3. Validates role ownership (both All and ManageAuthorizationLocks can only +/// remove own locks) +/// 4. Validates the Swig account and lock index +/// 5. Removes the authorization lock by shifting remaining locks down +/// 6. Updates the authorization lock count +/// +/// # Arguments +/// * `ctx` - The account context for the operation +/// * `data` - Raw instruction data bytes +/// * `all_accounts` - All accounts involved in the operation +/// +/// # Returns +/// * `ProgramResult` - Success or error status +#[inline(always)] +pub fn remove_authorization_lock_v1( + ctx: Context, + data: &[u8], + all_accounts: &[AccountInfo], +) -> ProgramResult { + check_stack_height(1, SwigError::Cpi)?; + + let remove_lock = RemoveAuthorizationLockV1::from_instruction_bytes(data)?; + + // Get current slot for authentication + let clock = Clock::get()?; + + let swig_account_data = unsafe { ctx.accounts.swig.borrow_mut_data_unchecked() }; + if unsafe { *swig_account_data.get_unchecked(0) } != Discriminator::SwigAccount as u8 { + return Err(SwigError::InvalidSwigAccountDiscriminator.into()); + } + + // Authentication and permission checking + let (swig_header, swig_roles) = unsafe { swig_account_data.split_at_mut_unchecked(Swig::LEN) }; + let swig = unsafe { Swig::load_mut_unchecked(swig_header)? }; + let acting_role = Swig::get_mut_role(remove_lock.args.acting_role_id, swig_roles)?; + if acting_role.is_none() { + return Err(SwigError::InvalidAuthorityNotFoundByRoleId.into()); + } + let acting_role = acting_role.unwrap(); + + // Authenticate the caller + let slot = clock.slot; + if acting_role.authority.session_based() { + acting_role.authority.authenticate_session( + all_accounts, + remove_lock.authority_payload, + remove_lock.data_payload, + slot, + )?; + } else { + acting_role.authority.authenticate( + all_accounts, + remove_lock.authority_payload, + remove_lock.data_payload, + slot, + )?; + } + + // Check permissions: must have All or ManageAuthorizationLocks + let has_all_permission = acting_role.get_action::(&[])?.is_some(); + let has_manage_auth_locks_permission = acting_role + .get_action::(&[])? + .is_some(); + + if !has_all_permission && !has_manage_auth_locks_permission { + return Err(SwigAuthenticateError::PermissionDeniedMissingPermission.into()); + } + + // Re-borrow data after authentication + let swig_account_data = unsafe { ctx.accounts.swig.borrow_mut_data_unchecked() }; + let (swig_header, remaining_data) = + unsafe { swig_account_data.split_at_mut_unchecked(Swig::LEN) }; + let swig = unsafe { Swig::load_mut_unchecked(swig_header)? }; + + // Validate that we have authorization locks to remove + if swig.authorization_locks == 0 { + return Err(SwigError::InvalidAuthorizationLockIndex.into()); + } + + // Validate lock index + if remove_lock.args.lock_index >= swig.authorization_locks as u32 { + return Err(SwigError::InvalidAuthorizationLockIndex.into()); + } + + // Find the end of roles data to determine where authorization locks start + let mut roles_end = 0; + let mut cursor = 0; + for _i in 0..swig.roles { + if cursor + Position::LEN > remaining_data.len() { + return Err(SwigStateError::InvalidRoleData.into()); + } + let position = + unsafe { Position::load_unchecked(&remaining_data[cursor..cursor + Position::LEN])? }; + cursor = position.boundary() as usize; + roles_end = cursor; + } + + let auth_locks_start = roles_end; + let lock_size = AuthorizationLock::LEN; + let total_locks = swig.authorization_locks as usize; + let lock_index = remove_lock.args.lock_index as usize; + + // Calculate positions + let lock_to_remove_start = auth_locks_start + (lock_index * lock_size); + let lock_to_remove_end = lock_to_remove_start + lock_size; + let locks_after_start = lock_to_remove_end; + let locks_after_end = auth_locks_start + (total_locks * lock_size); + + // Validate role ownership - both All and ManageAuthorizationLocks permissions + // can only remove locks created by their own role + if lock_to_remove_start + lock_size <= remaining_data.len() { + // Zero-copy: cast the raw bytes directly to a reference + let lock = unsafe { + &*(remaining_data[lock_to_remove_start..lock_to_remove_end].as_ptr() + as *const AuthorizationLock) + }; + if lock.role_id != remove_lock.args.acting_role_id { + msg!( + "Permission denied: Role {} cannot remove authorization lock created by role {}", + remove_lock.args.acting_role_id, + lock.role_id + ); + return Err(SwigAuthenticateError::PermissionDeniedMissingPermission.into()); + } + } + + // Log the lock being removed for debugging + if lock_to_remove_start + lock_size <= remaining_data.len() { + // Zero-copy: cast the raw bytes directly to a reference + let lock = unsafe { + &*(remaining_data[lock_to_remove_start..lock_to_remove_end].as_ptr() + as *const AuthorizationLock) + }; + } + + // Shift all locks after the removed lock down by one position + if lock_index < total_locks - 1 { + let locks_after_count = total_locks - lock_index - 1; + let move_size = locks_after_count * lock_size; + + // Safety check: ensure we don't go out of bounds + if locks_after_end <= remaining_data.len() + && lock_to_remove_start + move_size <= remaining_data.len() + { + // Use copy_within to safely move the data + let source_start = locks_after_start; + let source_end = locks_after_end; + let dest_start = lock_to_remove_start; + + remaining_data.copy_within(source_start..source_end, dest_start); + } else { + return Err(SwigError::InvalidAuthorizationLockIndex.into()); + } + } + + // Update the authorization locks count in the header + let (swig_header, _) = unsafe { swig_account_data.split_at_mut_unchecked(Swig::LEN) }; + let swig = unsafe { Swig::load_mut_unchecked(swig_header)? }; + swig.authorization_locks -= 1; + + Ok(()) +} diff --git a/program/src/actions/sign_v1.rs b/program/src/actions/sign_v1.rs index bea7ac5d..76ac7242 100644 --- a/program/src/actions/sign_v1.rs +++ b/program/src/actions/sign_v1.rs @@ -13,7 +13,7 @@ use pinocchio::{ sysvars::{clock::Clock, Sysvar}, ProgramResult, }; -use pinocchio_pubkey::from_str; +use pinocchio_pubkey::{from_str, pubkey}; use swig_assertions::*; use swig_compact_instructions::InstructionIterator; use swig_state_x::{ @@ -30,7 +30,7 @@ use swig_state_x::{ }, authority::AuthorityType, role::RoleMut, - swig::{swig_account_signer, Swig}, + swig::{swig_account_signer, AuthorizationLock, Swig, SwigWithRoles}, Discriminator, IntoBytes, SwigAuthenticateError, Transmutable, TransmutableMut, }; @@ -137,6 +137,8 @@ impl<'a> SignV1<'a> { /// * `all_accounts` - All accounts involved in the transaction /// * `data` - Raw signing instruction data /// * `account_classifiers` - Classifications for involved accounts +/// * `authorization_lock_cache` - Optional cache of authorization locks for +/// performance /// /// # Returns /// * `ProgramResult` - Success or error status @@ -146,6 +148,7 @@ pub fn sign_v1( all_accounts: &[AccountInfo], data: &[u8], account_classifiers: &[AccountClassification], + authorization_lock_cache: Option<&crate::util::AuthorizationLockCache>, ) -> ProgramResult { check_stack_height(1, SwigError::Cpi)?; // KEEP remove since we enfoce swig is owned in lib.rs @@ -227,25 +230,89 @@ pub fn sign_v1( if lamports > ¤t_lamports { let amount_diff = lamports - current_lamports; - { - if let Some(action) = RoleMut::get_action_mut::(actions, &[])? - { - action.run(amount_diff)?; - continue; - }; + // Wrapped SOL mint address + const WRAPPED_SOL_MINT: [u8; 32] = + pubkey!("So11111111111111111111111111111111111111112"); + + // Check authorization locks first for SOL spending using the cache + let mut total_authorized_amount = 0u64; + let mut has_active_locks = false; + + if let Some(cache) = authorization_lock_cache { + // Use the cached authorization locks for better performance + let locks = cache.get_locks_for_role_and_mint( + sign_v1.args.role_id, + &WRAPPED_SOL_MINT, + ); + for auth_lock in locks { + // Only check non-expired locks for this role + if auth_lock.expiry_slot > slot { + has_active_locks = true; + total_authorized_amount = + total_authorized_amount.saturating_add(auth_lock.amount); + } + } + } else { + // Fallback to the original method if cache is not available + let swig_account_data = + unsafe { ctx.accounts.swig.borrow_data_unchecked() }; + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account_data)?; + + let _: Result<(), ProgramError> = swig_with_roles + .for_each_authorization_lock_by_mint( + &WRAPPED_SOL_MINT, + |auth_lock| { + // Only check non-expired locks for this role + if auth_lock.expiry_slot > slot + && auth_lock.role_id == sign_v1.args.role_id + { + has_active_locks = true; + total_authorized_amount = total_authorized_amount + .saturating_add(auth_lock.amount); + } + Ok(()) + }, + ); } - { - if let Some(action) = - RoleMut::get_action_mut::(actions, &[])? + + // If there are active authorization locks, check against the total + if has_active_locks { + if amount_diff > total_authorized_amount { + return Err( + SwigAuthenticateError::PermissionDeniedMissingPermission.into(), + ); + } else { + // This spending is within the combined authorization lock limits + matched = true; + } + } + + // If not covered by authorization locks, check regular SOL permissions + if !matched { { - action.run(amount_diff, slot)?; - continue; - }; + if let Some(action) = + RoleMut::get_action_mut::(actions, &[])? + { + action.run(amount_diff)?; + continue; + }; + } + { + if let Some(action) = + RoleMut::get_action_mut::(actions, &[])? + { + action.run(amount_diff, slot)?; + continue; + }; + } + return Err( + SwigAuthenticateError::PermissionDeniedMissingPermission.into() + ); } - return Err(SwigAuthenticateError::PermissionDeniedMissingPermission.into()); } }, AccountClassification::SwigTokenAccount { balance } => { + let mut matched = false; let data = unsafe { &all_accounts.get_unchecked(index).borrow_data_unchecked() }; let mint = unsafe { data.get_unchecked(0..32) }; @@ -259,6 +326,8 @@ pub fn sign_v1( .map_err(|_| ProgramError::InvalidAccountData)? }); + let diff = balance - current_token_balance; + if delegate != [0u8; 4] || close_authority != [0u8; 4] { return Err( SwigAuthenticateError::PermissionDeniedTokenAccountDelegatePresent @@ -278,8 +347,60 @@ pub fn sign_v1( ); } - if balance > ¤t_token_balance { - let diff = balance - current_token_balance; + // Convert mint slice to array for comparison + let mint_array: [u8; 32] = mint + .try_into() + .map_err(|_| ProgramError::InvalidAccountData)?; + + // Check authorization locks first using the cache + let mut total_authorized_amount = 0u64; + let mut has_active_locks = false; + + if let Some(cache) = authorization_lock_cache { + // Use the cached authorization locks for better performance + // For tokens, we check ALL authorization locks across all roles (different + // from SOL) + let locks = cache.get_all_locks_for_mint(&mint_array); + for auth_lock in locks { + // Only check non-expired locks + if auth_lock.expiry_slot > slot { + has_active_locks = true; + total_authorized_amount = + total_authorized_amount.saturating_add(auth_lock.amount); + } + } + } else { + // Fallback to the original method if cache is not available + let swig_account_data = + unsafe { ctx.accounts.swig.borrow_data_unchecked() }; + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account_data)?; + + let _: Result<(), ProgramError> = swig_with_roles + .for_each_authorization_lock_by_mint(&mint_array, |auth_lock| { + // Only check non-expired locks + if auth_lock.expiry_slot > slot { + has_active_locks = true; + total_authorized_amount = + total_authorized_amount.saturating_add(auth_lock.amount); + } + Ok(()) + }); + } + + // If there are active authorization locks, check against the total + if has_active_locks { + if diff > total_authorized_amount { + return Err( + SwigAuthenticateError::PermissionDeniedMissingPermission.into() + ); + } else { + // This spending is within the combined authorization lock limits + matched = true; + } + } + + // If not covered by authorization locks, check regular token permissions + if !matched || balance > ¤t_token_balance { { if let Some(action) = RoleMut::get_action_mut::(actions, mint)? @@ -293,11 +414,12 @@ pub fn sign_v1( RoleMut::get_action_mut::(actions, mint)? { action.run(diff)?; + // matched = true; continue; }; } - return Err(SwigAuthenticateError::PermissionDeniedMissingPermission.into()); } + return Err(SwigAuthenticateError::PermissionDeniedMissingPermission.into()); }, AccountClassification::SwigStakeAccount { state, balance } => { // Get current stake balance from account data @@ -416,5 +538,21 @@ pub fn sign_v1( } } + // Clean up expired authorization locks at the end of the transaction + // Note: We don't reallocate the account to keep it simple and avoid potential + // issues. The unused space at the end of the account is acceptable. + { + let mut swig_account_data = unsafe { ctx.accounts.swig.borrow_mut_data_unchecked() }; + if unsafe { *swig_account_data.get_unchecked(0) } == Discriminator::SwigAccount as u8 { + let (swig_header, rest) = + unsafe { swig_account_data.split_at_mut_unchecked(Swig::LEN) }; + let swig = unsafe { Swig::load_mut_unchecked(swig_header)? }; + let _removed_count = + SwigWithRoles::remove_expired_authorization_locks_mut(swig, rest, slot)?; + // Account size remains the same - unused space at the end is + // acceptable + } + } + Ok(()) } diff --git a/program/src/actions/withdraw_from_sub_account_v1.rs b/program/src/actions/withdraw_from_sub_account_v1.rs index 5d2e59a3..66c9e184 100644 --- a/program/src/actions/withdraw_from_sub_account_v1.rs +++ b/program/src/actions/withdraw_from_sub_account_v1.rs @@ -191,7 +191,6 @@ pub fn withdraw_from_sub_account_v1( amount, token_program: token_account_program_owner, }; - msg!("amount: {}", amount); let role_id_bytes = sub_account.role_id.to_le_bytes(); let bump_byte = [sub_account.bump]; let seeds = sub_account_signer(&swig.id, &role_id_bytes, &bump_byte); diff --git a/program/src/error.rs b/program/src/error.rs index afe4aff8..4dbf43f1 100644 --- a/program/src/error.rs +++ b/program/src/error.rs @@ -105,6 +105,12 @@ pub enum SwigError { InvalidSwigTokenAccountOwner, /// Invalid program scope balance field configuration InvalidProgramScopeBalanceFields, + /// Authorization lock expiry slot is invalid + InvalidAuthorizationLockExpiry, + /// Authorization lock insufficient balance + AuthorizationLockInsufficientBalance, + /// Invalid authorization lock index + InvalidAuthorizationLockIndex, } /// Implements conversion from SwigError to ProgramError. diff --git a/program/src/instruction.rs b/program/src/instruction.rs index c39bcf43..9218edbb 100644 --- a/program/src/instruction.rs +++ b/program/src/instruction.rs @@ -128,4 +128,26 @@ pub enum SwigInstruction { #[account(1, signer, name="payer", desc="the payer")] #[account(2, writable, name="sub_account", desc="the sub account to toggle enabled state")] ToggleSubAccountV1 = 10, + + /// Adds an authorization lock to the wallet. + /// + /// Required accounts: + /// 1. `[writable, signer]` Swig wallet account + /// 2. `[writable, signer]` Payer account + /// 3. System program account + #[account(0, writable, signer, name="swig", desc="the swig smart wallet")] + #[account(1, writable, signer, name="payer", desc="the payer")] + #[account(2, name="system_program", desc="the system program")] + AddAuthorizationLockV1 = 11, + + /// Removes an authorization lock from the wallet. + /// + /// Required accounts: + /// 1. `[writable, signer]` Swig wallet account + /// 2. `[writable, signer]` Payer account + /// 3. System program account + #[account(0, writable, signer, name="swig", desc="the swig smart wallet")] + #[account(1, writable, signer, name="payer", desc="the payer")] + #[account(2, name="system_program", desc="the system program")] + RemoveAuthorizationLockV1 = 12, } diff --git a/program/src/lib.rs b/program/src/lib.rs index 8f1f60d4..dafba303 100644 --- a/program/src/lib.rs +++ b/program/src/lib.rs @@ -34,7 +34,7 @@ use swig_state_x::{ swig::{Swig, SwigWithRoles}, AccountClassification, Discriminator, StakeAccountState, Transmutable, }; -use util::{read_program_scope_account_balance, ProgramScopeCache}; +use util::{read_program_scope_account_balance, AuthorizationLockCache, ProgramScopeCache}; /// Program ID for the Swig wallet program declare_id!("swigDk8JezhiAVde8k6NMwxpZfgGm2NNuMe1KYCmUjP"); @@ -118,21 +118,23 @@ unsafe fn execute( } // Create program scope cache if first account is a valid Swig account - let program_scope_cache = if index > 0 { + let (program_scope_cache, authorization_lock_cache) = if index > 0 { let first_account = accounts[0].assume_init_ref(); if first_account.owner() == &crate::ID { let data = first_account.borrow_data_unchecked(); if data.len() >= Swig::LEN && *data.get_unchecked(0) == Discriminator::SwigAccount as u8 { - ProgramScopeCache::load_from_swig(data) + let program_scope_cache = ProgramScopeCache::load_from_swig(data); + let authorization_lock_cache = AuthorizationLockCache::load_from_swig(data); + (program_scope_cache, authorization_lock_cache) } else { - None + (None, None) } } else { - None + (None, None) } } else { - None + (None, None) }; // Process remaining accounts using the cache @@ -161,6 +163,7 @@ unsafe fn execute( core::slice::from_raw_parts(accounts.as_ptr() as _, index), core::slice::from_raw_parts(account_classification.as_ptr() as _, index), instruction, + authorization_lock_cache.as_ref(), )?; Ok(()) } diff --git a/program/src/util/mod.rs b/program/src/util/mod.rs index fda032d5..4035e77d 100644 --- a/program/src/util/mod.rs +++ b/program/src/util/mod.rs @@ -21,9 +21,9 @@ use swig_state_x::{ program_scope::{NumericType, ProgramScope}, Action, Permission, }, - constants::PROGRAM_SCOPE_BYTE_SIZE, + constants::{AUTHORIZATION_LOCK_BYTE_SIZE, PROGRAM_SCOPE_BYTE_SIZE}, read_numeric_field, - swig::{Swig, SwigWithRoles}, + swig::{AuthorizationLock, Swig, SwigWithRoles}, Transmutable, }; @@ -158,6 +158,132 @@ impl ProgramScopeCache { } } +/// Cache for authorization lock information to optimize lookups. +/// +/// This struct maintains a mapping of role IDs to their authorization locks. +/// It helps avoid repeated parsing of authorization lock data from the Swig +/// account during transaction processing. +pub(crate) struct AuthorizationLockCache { + /// Maps role_id to Vec<(mint, AuthorizationLock)> for efficient lookups + locks_by_role: Vec<(u32, Vec<([u8; 32], AuthorizationLock)>)>, +} + +impl AuthorizationLockCache { + /// Creates a new empty authorization lock cache. + /// + /// Initializes with a reasonable capacity to avoid frequent reallocations. + pub(crate) fn new() -> Self { + Self { + locks_by_role: Vec::with_capacity(8), // Reasonable initial capacity for roles + } + } + + /// Loads authorization lock information from a Swig account's data. + /// + /// This function parses the Swig account data to extract all authorization + /// locks and builds a cache for efficient lookup by role. + /// + /// # Arguments + /// * `data` - Raw Swig account data + /// + /// # Returns + /// * `Option` - The populated cache if successful, None if data is + /// invalid + pub(crate) fn load_from_swig(data: &[u8]) -> Option { + if data.len() < Swig::LEN { + return None; + } + + let swig_with_roles = SwigWithRoles::from_bytes(data).ok()?; + let mut cache = Self::new(); + + // Iterate through all authorization locks using the zero-copy iterator + let _: Result<(), ProgramError> = + swig_with_roles.for_each_authorization_lock(|auth_lock| { + let role_id = auth_lock.role_id; + let mint = auth_lock.token_mint; + + // Find existing role entry or create new one + if let Some((_, locks)) = cache + .locks_by_role + .iter_mut() + .find(|(id, _)| *id == role_id) + { + locks.push((mint, *auth_lock)); + } else { + cache + .locks_by_role + .push((role_id, vec![(mint, *auth_lock)])); + } + + Ok(()) + }); + + Some(cache) + } + + /// Gets all authorization locks for a specific role and mint. + /// + /// # Arguments + /// * `role_id` - The role ID to look up locks for + /// * `mint` - The token mint to filter by + /// + /// # Returns + /// * `Vec<&AuthorizationLock>` - All matching authorization locks + pub(crate) fn get_locks_for_role_and_mint( + &self, + role_id: u32, + mint: &[u8; 32], + ) -> Vec<&AuthorizationLock> { + self.locks_by_role + .iter() + .find(|(id, _)| *id == role_id) + .map(|(_, locks)| { + locks + .iter() + .filter(|(lock_mint, _)| lock_mint == mint) + .map(|(_, lock)| lock) + .collect() + }) + .unwrap_or_default() + } + + /// Gets all authorization locks for a specific role. + /// + /// # Arguments + /// * `role_id` - The role ID to look up locks for + /// + /// # Returns + /// * `Vec<&AuthorizationLock>` - All authorization locks for the role + pub(crate) fn get_locks_for_role(&self, role_id: u32) -> Vec<&AuthorizationLock> { + self.locks_by_role + .iter() + .find(|(id, _)| *id == role_id) + .map(|(_, locks)| locks.iter().map(|(_, lock)| lock).collect()) + .unwrap_or_default() + } + + /// Gets all authorization locks for a specific mint across all roles. + /// + /// # Arguments + /// * `mint` - The token mint to filter by + /// + /// # Returns + /// * `Vec<&AuthorizationLock>` - All matching authorization locks from all + /// roles + pub(crate) fn get_all_locks_for_mint(&self, mint: &[u8; 32]) -> Vec<&AuthorizationLock> { + let mut result = Vec::new(); + for (_, role_locks) in &self.locks_by_role { + for (lock_mint, auth_lock) in role_locks { + if lock_mint == mint { + result.push(auth_lock); + } + } + } + result + } +} + /// Reads a numeric balance from an account's data based on a `ProgramScope` /// configuration. /// diff --git a/program/tests/authorization_lock_permissions_test.rs b/program/tests/authorization_lock_permissions_test.rs new file mode 100644 index 00000000..a52144cb --- /dev/null +++ b/program/tests/authorization_lock_permissions_test.rs @@ -0,0 +1,591 @@ +#![cfg(not(feature = "program_scope_test"))] + +mod common; +use common::*; +use litesvm_token::spl_token; +use solana_sdk::{ + clock::Clock, + message::{v0, VersionedMessage}, + program_pack::Pack, + pubkey::Pubkey, + signature::Keypair, + signer::Signer, + transaction::VersionedTransaction, +}; +use swig_interface::{AuthorityConfig, ClientAction}; +use swig_state_x::{ + action::{ + all::All, manage_authorization_locks::ManageAuthorizationLocks, token_limit::TokenLimit, + }, + swig::{swig_account_seeds, AuthorizationLock, Swig, SwigWithRoles}, + IntoBytes, Transmutable, +}; + +/// Test that validates only authorities with "All" or +/// "ManageAuthorizationLocks" permissions can add authorization locks, while +/// others are denied. +#[test_log::test] +fn test_authorization_lock_permission_enforcement() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); // Will have All permissions + let manage_auth_locks_authority = Keypair::new(); // Will have ManageAuthorizationLocks permission + let token_authority = Keypair::new(); // Will have only token limit permissions (should be denied) + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&manage_auth_locks_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&token_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + println!("=== AUTHORIZATION LOCK PERMISSION ENFORCEMENT TEST ==="); + println!( + "Testing that only authorities with proper permissions can manage authorization locks" + ); + println!(); + + // Add authority with ManageAuthorizationLocks permission + let manage_auth_locks_action = + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}); + let add_manage_auth_locks_authority_ix = + swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, // Acting role ID (swig_authority has All permissions) + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &manage_auth_locks_authority.pubkey().to_bytes(), + }, + vec![manage_auth_locks_action], + ) + .unwrap(); + + let add_manage_auth_locks_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_manage_auth_locks_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_manage_auth_locks_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_manage_auth_locks_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_manage_auth_locks_authority_result = context + .svm + .send_transaction(add_manage_auth_locks_authority_tx); + assert!( + add_manage_auth_locks_authority_result.is_ok(), + "Adding ManageAuthorizationLocks authority should succeed" + ); + println!("✅ Added authority with ManageAuthorizationLocks permission (role ID: 1)"); + + // Add authority with limited token permissions (should NOT be able to manage + // auth locks) + let token_action = ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 1000, + }); + + let add_token_authority_ix = + swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, // Acting role ID (swig_authority has All permissions) + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &token_authority.pubkey().to_bytes(), + }, + vec![token_action], + ) + .unwrap(); + + let add_token_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_token_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_token_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_token_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_token_authority_result = context.svm.send_transaction(add_token_authority_tx); + assert!( + add_token_authority_result.is_ok(), + "Adding token authority should succeed" + ); + println!("✅ Added authority with TokenLimit permission only (role ID: 2)"); + println!(); + + // Test parameters for authorization lock + let current_slot = context.svm.get_sysvar::().slot; + let lock_amount = 500u64; + let expiry_slot = current_slot + 1000; + + println!("TEST SCENARIOS:"); + println!(" 1. Authority with All permission (role 0): Should PASS"); + println!(" 2. Authority with ManageAuthorizationLocks permission (role 1): Should PASS"); + println!(" 3. Authority with TokenLimit permission only (role 2): Should FAIL"); + println!(); + + // Test 1: Authority with All permission should succeed + println!("EXECUTING TEST SCENARIO 1:"); + println!("Adding authorization lock using authority with All permission (role 0)"); + + let add_lock_all_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority (All permission) + mint_pubkey.to_bytes(), + lock_amount, + expiry_slot, + ) + .unwrap(); + + let add_lock_all_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_all_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_all_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_all_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock_all_result = context.svm.send_transaction(add_lock_all_tx); + assert!( + add_lock_all_result.is_ok(), + "Authority with All permission should be able to add authorization lock" + ); + println!( + "✅ Scenario 1 PASSED: Authority with All permission successfully added authorization lock" + ); + + // Verify the lock was added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 1); + println!( + " → Verified: Authorization lock count = {}", + swig_with_roles.state.authorization_locks + ); + println!(); + + // Test 2: Authority with ManageAuthorizationLocks permission should succeed + println!("EXECUTING TEST SCENARIO 2:"); + println!( + "Adding authorization lock using authority with ManageAuthorizationLocks permission (role \ + 1)" + ); + + let add_lock_manage_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + manage_auth_locks_authority.pubkey(), + context.default_payer.pubkey(), + 1, // acting_role_id: manage_auth_locks_authority (ManageAuthorizationLocks permission) + mint_pubkey.to_bytes(), + lock_amount + 100, // Different amount + expiry_slot, + ) + .unwrap(); + + let add_lock_manage_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_manage_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_manage_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_manage_message), + &[&context.default_payer, &manage_auth_locks_authority], + ) + .unwrap(); + + let add_lock_manage_result = context.svm.send_transaction(add_lock_manage_tx); + assert!( + add_lock_manage_result.is_ok(), + "Authority with ManageAuthorizationLocks permission should be able to add authorization \ + lock" + ); + println!( + "✅ Scenario 2 PASSED: Authority with ManageAuthorizationLocks permission successfully \ + added authorization lock" + ); + + // Verify the second lock was added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 2); + println!( + " → Verified: Authorization lock count = {}", + swig_with_roles.state.authorization_locks + ); + println!(); + + // Test 3: Authority with only TokenLimit permission should fail + println!("EXECUTING TEST SCENARIO 3:"); + println!( + "Attempting to add authorization lock using authority with TokenLimit permission only \ + (role 2)" + ); + println!("Expected: FAIL (insufficient permissions)"); + + let add_lock_token_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + token_authority.pubkey(), + context.default_payer.pubkey(), + 2, // acting_role_id: token_authority (TokenLimit permission only) + mint_pubkey.to_bytes(), + lock_amount + 200, // Different amount + expiry_slot, + ) + .unwrap(); + + let add_lock_token_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_token_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_token_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_token_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let add_lock_token_result = context.svm.send_transaction(add_lock_token_tx); + assert!( + add_lock_token_result.is_err(), + "Authority with only TokenLimit permission should NOT be able to add authorization lock" + ); + println!("✅ Scenario 3 PASSED: Authority with TokenLimit permission was correctly denied"); + + // Verify the lock count didn't change + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!( + swig_with_roles.state.authorization_locks, 2, + "Lock count should remain unchanged after failed attempt" + ); + println!( + " → Verified: Authorization lock count remains = {}", + swig_with_roles.state.authorization_locks + ); + println!(); + + println!("FINAL VERIFICATION:"); + let (auth_locks, count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + println!("Total authorization locks in account: {}", count); + for i in 0..count { + if let Some(lock) = auth_locks[i] { + println!( + "Lock {}: mint={:?}, amount={}, expiry_slot={}", + i, lock.token_mint, lock.amount, lock.expiry_slot + ); + } + } + + println!(); + println!("✅ PERMISSION ENFORCEMENT TEST COMPLETED SUCCESSFULLY!"); + println!( + "✅ Only authorities with All or ManageAuthorizationLocks permissions can manage \ + authorization locks" + ); + println!("✅ Authorities with insufficient permissions are properly denied"); + println!("======================================================"); +} + +/// Test that validates role ID validation and non-existent role handling +#[test_log::test] +fn test_authorization_lock_invalid_role_handling() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + println!("=== AUTHORIZATION LOCK INVALID ROLE HANDLING TEST ==="); + println!("Testing behavior with invalid role IDs"); + println!(); + + let current_slot = context.svm.get_sysvar::().slot; + let lock_amount = 500u64; + let expiry_slot = current_slot + 1000; + + // Test 1: Using non-existent role ID should fail + println!("TEST SCENARIO 1:"); + println!("Attempting to add authorization lock using non-existent role ID (999)"); + println!("Expected: FAIL (role not found)"); + + let add_lock_invalid_role_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 999, // Non-existent role ID + mint_pubkey.to_bytes(), + lock_amount, + expiry_slot, + ) + .unwrap(); + + let add_lock_invalid_role_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_invalid_role_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_invalid_role_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_invalid_role_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock_invalid_role_result = context.svm.send_transaction(add_lock_invalid_role_tx); + assert!( + add_lock_invalid_role_result.is_err(), + "Using non-existent role ID should fail" + ); + println!("✅ Scenario 1 PASSED: Non-existent role ID was correctly rejected"); + println!(); + + // Test 2: Valid role ID (0) should succeed + println!("TEST SCENARIO 2:"); + println!("Adding authorization lock using valid role ID (0)"); + println!("Expected: PASS (role 0 has All permissions)"); + + let add_lock_valid_role_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // Valid role ID + mint_pubkey.to_bytes(), + lock_amount, + expiry_slot, + ) + .unwrap(); + + let add_lock_valid_role_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_valid_role_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_valid_role_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_valid_role_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock_valid_role_result = context.svm.send_transaction(add_lock_valid_role_tx); + assert!( + add_lock_valid_role_result.is_ok(), + "Using valid role ID should succeed" + ); + println!("✅ Scenario 2 PASSED: Valid role ID successfully added authorization lock"); + + // Verify the lock was added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 1); + println!( + " → Verified: Authorization lock count = {}", + swig_with_roles.state.authorization_locks + ); + + println!(); + println!("✅ INVALID ROLE HANDLING TEST COMPLETED SUCCESSFULLY!"); + println!("======================================================"); +} + +/// Test that validates expired authorization lock rejection +#[test_log::test] +fn test_authorization_lock_expiry_validation() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + println!("=== AUTHORIZATION LOCK EXPIRY VALIDATION TEST ==="); + println!("Testing that already-expired authorization locks are rejected"); + println!(); + + let current_slot = context.svm.get_sysvar::().slot; + let lock_amount = 500u64; + + println!("Current slot: {}", current_slot); + println!(); + + // Test 1: Try to add an authorization lock that has already expired + let expired_slot = if current_slot > 0 { + current_slot - 1 + } else { + 0 + }; + + println!("TEST SCENARIO 1:"); + println!( + "Attempting to add authorization lock with expired slot: {}", + expired_slot + ); + println!("Expected: FAIL (already expired)"); + + let add_expired_lock_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id + mint_pubkey.to_bytes(), + lock_amount, + expired_slot, + ) + .unwrap(); + + let add_expired_lock_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_expired_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_expired_lock_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_expired_lock_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_expired_lock_result = context.svm.send_transaction(add_expired_lock_tx); + assert!( + add_expired_lock_result.is_err(), + "Adding expired authorization lock should fail" + ); + println!("✅ Scenario 1 PASSED: Expired authorization lock was correctly rejected"); + println!(); + + // Test 2: Add a valid authorization lock with future expiry + let future_slot = current_slot + 1000; + + println!("TEST SCENARIO 2:"); + println!( + "Adding authorization lock with future expiry slot: {}", + future_slot + ); + println!("Expected: PASS (valid expiry)"); + + let add_valid_lock_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id + mint_pubkey.to_bytes(), + lock_amount, + future_slot, + ) + .unwrap(); + + let add_valid_lock_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_valid_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_valid_lock_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_valid_lock_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_valid_lock_result = context.svm.send_transaction(add_valid_lock_tx); + assert!( + add_valid_lock_result.is_ok(), + "Adding valid authorization lock should succeed" + ); + println!("✅ Scenario 2 PASSED: Valid authorization lock successfully added"); + + // Verify the lock was added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 1); + println!( + " → Verified: Authorization lock count = {}", + swig_with_roles.state.authorization_locks + ); + + println!(); + println!("✅ EXPIRY VALIDATION TEST COMPLETED SUCCESSFULLY!"); + println!("======================================================"); +} diff --git a/program/tests/authorization_lock_role_tracking_test.rs b/program/tests/authorization_lock_role_tracking_test.rs new file mode 100644 index 00000000..ab4878f0 --- /dev/null +++ b/program/tests/authorization_lock_role_tracking_test.rs @@ -0,0 +1,342 @@ +#![cfg(not(feature = "program_scope_test"))] + +mod common; +use common::*; +use litesvm_token::spl_token; +use solana_sdk::{ + clock::Clock, + message::{v0, VersionedMessage}, + program_pack::Pack, + pubkey::Pubkey, + signature::Keypair, + signer::Signer, + transaction::VersionedTransaction, +}; +use swig_interface::{AuthorityConfig, ClientAction}; +use swig_state_x::{ + action::{ + all::All, manage_authorization_locks::ManageAuthorizationLocks, token_limit::TokenLimit, + }, + swig::{swig_account_seeds, AuthorizationLock, Swig, SwigWithRoles}, + IntoBytes, Transmutable, +}; + +/// Test that validates authorization locks track the role ID that created them +/// and that we can retrieve locks by role ID. +#[test_log::test] +fn test_authorization_lock_role_tracking() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts - simplified to test role tracking with just 2 roles + let swig_authority = Keypair::new(); // Will have All permissions (role 0) + let manage_auth_locks_authority = Keypair::new(); // Will have ManageAuthorizationLocks permission (role 1) + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&manage_auth_locks_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + println!("=== AUTHORIZATION LOCK ROLE TRACKING TEST ==="); + println!( + "Testing that authorization locks track creator role IDs and can be retrieved by role" + ); + println!(); + + // Add authority with ManageAuthorizationLocks permission + let manage_auth_locks_action = + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}); + let add_manage_auth_locks_authority_ix = + swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, // Acting role ID (swig_authority has All permissions) + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &manage_auth_locks_authority.pubkey().to_bytes(), + }, + vec![manage_auth_locks_action], + ) + .unwrap(); + + let add_manage_auth_locks_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_manage_auth_locks_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_manage_auth_locks_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_manage_auth_locks_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_manage_auth_locks_authority_result = context + .svm + .send_transaction(add_manage_auth_locks_authority_tx); + assert!( + add_manage_auth_locks_authority_result.is_ok(), + "Adding ManageAuthorizationLocks authority should succeed" + ); + println!("✅ Added authority with ManageAuthorizationLocks permission (role ID: 1)"); + println!(); + + // Test parameters for authorization locks + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + + // Expire blockhash + context.svm.expire_blockhash(); + + // Add authorization locks from different roles + println!("ADDING AUTHORIZATION LOCKS FROM DIFFERENT ROLES:"); + + // Lock 1: Added by role 0 (swig_authority with All permission) + println!("Adding lock 1 from role 0 (All permission): 500 tokens"); + let add_lock1_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority (All permission) + mint_pubkey.to_bytes(), + 500, + expiry_slot, + ) + .unwrap(); + + let add_lock1_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock1_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock1_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock1_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock1_result = context.svm.send_transaction(add_lock1_tx); + assert!( + add_lock1_result.is_ok(), + "Adding lock from role 0 should succeed" + ); + + // Expire blockhash + context.svm.expire_blockhash(); + + // Lock 2: Added by role 1 (ManageAuthorizationLocks permission) + println!("Adding lock 2 from role 1 (ManageAuthorizationLocks permission): 300 tokens"); + let add_lock2_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + manage_auth_locks_authority.pubkey(), + context.default_payer.pubkey(), + 1, // acting_role_id: manage_auth_locks_authority + mint_pubkey.to_bytes(), + 300, + expiry_slot, + ) + .unwrap(); + + let add_lock2_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock2_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock2_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock2_message), + &[&context.default_payer, &manage_auth_locks_authority], + ) + .unwrap(); + + let add_lock2_result = context.svm.send_transaction(add_lock2_tx); + assert!( + add_lock2_result.is_ok(), + "Adding lock from role 1 should succeed" + ); + + // Expire blockhash + context.svm.expire_blockhash(); + + // Lock 3: Another lock from role 0 + println!("Adding lock 3 from role 0 (All permission): 100 tokens"); + let add_lock3_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority + mint_pubkey.to_bytes(), + 100, + expiry_slot, + ) + .unwrap(); + + let add_lock3_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock3_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock3_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock3_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock3_result = context.svm.send_transaction(add_lock3_tx); + assert!( + add_lock3_result.is_ok(), + "Adding third lock from role 0 should succeed" + ); + + // Lock 4: Another lock from role 1 + println!("Adding lock 4 from role 1 (ManageAuthorizationLocks permission): 200 tokens"); + + // Expire blockhash + context.svm.expire_blockhash(); + + let add_lock4_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + manage_auth_locks_authority.pubkey(), + context.default_payer.pubkey(), + 1, // acting_role_id: manage_auth_locks_authority + mint_pubkey.to_bytes(), + 200, + expiry_slot, + ) + .unwrap(); + + let add_lock4_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock4_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock4_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock4_message), + &[&context.default_payer, &manage_auth_locks_authority], + ) + .unwrap(); + + let add_lock4_result = context.svm.send_transaction(add_lock4_tx); + assert!( + add_lock4_result.is_ok(), + "Adding second lock from role 1 should succeed" + ); + + println!("✅ Added 4 authorization locks from different roles"); + println!(); + + // Verify all locks and their role IDs + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 4); + + println!("VERIFICATION - All authorization locks in account:"); + let (all_locks, total_count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + println!("Total authorization locks: {}", total_count); + for i in 0..total_count { + if let Some(lock) = all_locks[i] { + println!( + "Lock {}: role_id={}, amount={}, mint={:?}", + i, lock.role_id, lock.amount, lock.token_mint + ); + } + } + println!(); + + // Test getting locks by role ID + println!("TESTING GET LOCKS BY ROLE ID:"); + + // Get locks for role 0 (should have 2 locks: 500 and 100 tokens) + println!("Getting locks for role 0:"); + let (role0_locks, role0_count) = swig_with_roles + .get_authorization_locks_by_role::<10>(0) + .unwrap(); + println!("Found {} locks for role 0", role0_count); + assert_eq!(role0_count, 2, "Role 0 should have 2 locks"); + + let mut role0_amounts = Vec::new(); + for i in 0..role0_count { + if let Some(lock) = role0_locks[i] { + println!(" Role 0 Lock {}: amount={}", i, lock.amount); + assert_eq!(lock.role_id, 0, "Lock should belong to role 0"); + role0_amounts.push(lock.amount); + } + } + // Should have 500 and 100 token locks + role0_amounts.sort(); + assert_eq!( + role0_amounts, + vec![100, 500], + "Role 0 should have 100 and 500 token locks" + ); + + // Get locks for role 1 (should have 2 locks: 300 and 200 tokens) + println!("Getting locks for role 1:"); + let (role1_locks, role1_count) = swig_with_roles + .get_authorization_locks_by_role::<10>(1) + .unwrap(); + println!("Found {} locks for role 1", role1_count); + assert_eq!(role1_count, 2, "Role 1 should have 2 locks"); + + let mut role1_amounts = Vec::new(); + for i in 0..role1_count { + if let Some(lock) = role1_locks[i] { + println!(" Role 1 Lock {}: amount={}", i, lock.amount); + assert_eq!(lock.role_id, 1, "Lock should belong to role 1"); + role1_amounts.push(lock.amount); + } + } + // Should have 300 and 200 token locks + role1_amounts.sort(); + assert_eq!( + role1_amounts, + vec![200, 300], + "Role 1 should have 200 and 300 token locks" + ); + + // Get locks for role 999 (should have 0 locks) + println!("Getting locks for role 999 (non-existent):"); + let (role999_locks, role999_count) = swig_with_roles + .get_authorization_locks_by_role::<10>(999) + .unwrap(); + println!("Found {} locks for role 999", role999_count); + assert_eq!(role999_count, 0, "Role 999 should have no locks"); + + println!(); + println!("✅ ROLE TRACKING VERIFICATION COMPLETE!"); + println!("✅ Authorization locks correctly track creator role IDs"); + println!("✅ get_authorization_locks_by_role function works correctly"); + println!("✅ Role 0: 2 locks (500, 100 tokens)"); + println!("✅ Role 1: 2 locks (300, 200 tokens)"); + println!("✅ Non-existent roles return 0 locks"); + println!("================================================================"); +} diff --git a/program/tests/authorization_lock_sol_limits_test.rs b/program/tests/authorization_lock_sol_limits_test.rs new file mode 100644 index 00000000..15a766e8 --- /dev/null +++ b/program/tests/authorization_lock_sol_limits_test.rs @@ -0,0 +1,177 @@ +#![cfg(not(feature = "program_scope_test"))] + +mod common; +use common::*; +use pinocchio_pubkey::pubkey; +use solana_sdk::{ + clock::Clock, + message::{v0, VersionedMessage}, + signature::Keypair, + signer::Signer, + transaction::VersionedTransaction, +}; +use swig_interface::{AddAuthorizationLockInstruction, AuthorityConfig, ClientAction}; +use swig_state_x::{ + action::{manage_authorization_locks::ManageAuthorizationLocks, sol_limit::SolLimit}, + swig::{AuthorizationLock, SwigWithRoles}, + IntoBytes, Transmutable, +}; + +/// Test that validates authorization locks respect simple SOL limits. +/// +/// This test creates a role with a SOL limit of 1000 lamports, then: +/// 1. Successfully adds an authorization lock for 800 lamports (within limit) +/// 2. Fails to add another authorization lock for 300 lamports (would exceed +/// limit: 800 + 300 = 1100 > 1000) +#[test_log::test] +fn test_authorization_lock_respects_simple_sol_limit() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Create swig account with root authority (All permissions) + let swig_id = [1u8; 32]; + let (swig_pubkey, _) = create_swig_ed25519(&mut context, &swig_authority, swig_id).unwrap(); + + // Add a SOL authority with SOL limit permission (1000 lamports) + let sol_limit_amount = 1000u64; + let sol_authority = Keypair::new(); + context + .svm + .airdrop(&sol_authority.pubkey(), 10_000_000_000) + .unwrap(); + + let sol_authority_actions = vec![ + ClientAction::SolLimit(SolLimit { + amount: sol_limit_amount, + }), + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}), + ]; + + add_authority_with_ed25519_root( + &mut context, + &swig_pubkey, + &swig_authority, + AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: sol_authority.pubkey().as_ref(), + }, + sol_authority_actions, + ) + .unwrap(); + + // Get role ID for the SOL authority + let swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + let sol_role_id = swig_with_roles + .lookup_role_id(sol_authority.pubkey().as_ref()) + .unwrap() + .unwrap(); + + // Wrapped SOL mint address (So11111111111111111111111111111111111111112) + let wrapped_sol_mint = pubkey!("So11111111111111111111111111111111111111112"); + + // Test 1: Successfully add authorization lock for 800 lamports (within limit) + let auth_lock_amount_1 = 800u64; + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + + let add_auth_lock_ix_1 = AddAuthorizationLockInstruction::new( + swig_pubkey, + sol_authority.pubkey(), + context.default_payer.pubkey(), + sol_role_id, + wrapped_sol_mint, + auth_lock_amount_1, + expiry_slot, + ) + .unwrap(); + + let msg_1 = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_auth_lock_ix_1], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx_1 = VersionedTransaction::try_new( + VersionedMessage::V0(msg_1), + &[ + context.default_payer.insecure_clone(), + sol_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result_1 = context.svm.send_transaction(tx_1); + assert!( + result_1.is_ok(), + "First authorization lock (800 lamports) should succeed: {:?}", + result_1 + ); + + // Test 2: Fail to add authorization lock for 300 lamports (would exceed limit: + // 800 + 300 = 1100 > 1000) + let auth_lock_amount_2 = 300u64; + + let add_auth_lock_ix_2 = AddAuthorizationLockInstruction::new( + swig_pubkey, + sol_authority.pubkey(), + context.default_payer.pubkey(), + sol_role_id, + wrapped_sol_mint, + auth_lock_amount_2, + expiry_slot, + ) + .unwrap(); + + let msg_2 = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_auth_lock_ix_2], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx_2 = VersionedTransaction::try_new( + VersionedMessage::V0(msg_2), + &[ + context.default_payer.insecure_clone(), + sol_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result_2 = context.svm.send_transaction(tx_2); + assert!( + result_2.is_err(), + "Second authorization lock (300 lamports) should fail due to exceeding limit" + ); + + // Verify final state: should have exactly 1 authorization lock for 800 lamports + let final_swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let final_swig_with_roles = SwigWithRoles::from_bytes(&final_swig_account.data).unwrap(); + + let (auth_locks, count) = final_swig_with_roles + .get_authorization_locks_by_role::<10>(sol_role_id) + .unwrap(); + + assert_eq!(count, 1, "Should have exactly 1 authorization lock"); + + let total_locked: u64 = auth_locks + .iter() + .filter_map(|opt_lock| *opt_lock) + .filter(|lock| lock.token_mint == wrapped_sol_mint) + .map(|lock| lock.amount) + .sum(); + + assert_eq!(total_locked, 800, "Total locked amount should be 800"); +} diff --git a/program/tests/authorization_lock_test.rs b/program/tests/authorization_lock_test.rs new file mode 100644 index 00000000..ff934589 --- /dev/null +++ b/program/tests/authorization_lock_test.rs @@ -0,0 +1,1443 @@ +#![cfg(not(feature = "program_scope_test"))] + +mod common; +use common::*; +use litesvm_token::spl_token; +use solana_sdk::{ + clock::Clock, + message::{v0, VersionedMessage}, + program_pack::Pack, + pubkey::Pubkey, + signature::Keypair, + signer::Signer, + transaction::VersionedTransaction, +}; +use swig_interface::{AuthorityConfig, ClientAction}; +use swig_state_x::{ + swig::{swig_account_seeds, AuthorizationLock, Swig, SwigWithRoles}, + IntoBytes, Transmutable, +}; + +/// Test that validates creating a swig, adding an authorization lock, and then +/// trying to spend over the authorization lock limit should fail, but spending +/// within the limit should succeed. +#[test_log::test] +fn test_authorization_lock_enforcement() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); // This will have All permissions + let token_authority = Keypair::new(); // This will have only token limit permissions + let recipient = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&token_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&recipient.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + // Setup token accounts + let swig_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &swig, + &context.default_payer, + ) + .unwrap(); + + let recipient_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &recipient.pubkey(), + &context.default_payer, + ) + .unwrap(); + + // Mint tokens to swig account + let initial_token_amount = 1000; + mint_to( + &mut context.svm, + &mint_pubkey, + &context.default_payer, + &swig_ata, + initial_token_amount, + ) + .unwrap(); + + // Add token authority with limited token permissions FIRST + use swig_state_x::action::token_limit::TokenLimit; + let token_action = ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 1000, // Allow up to 1000 tokens + }); + + let add_authority_ix = swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, // Acting role ID (swig_authority has All permissions) + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &token_authority.pubkey().to_bytes(), + }, + vec![token_action], + ) + .unwrap(); + + let add_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_authority_result = context.svm.send_transaction(add_authority_tx); + assert!( + add_authority_result.is_ok(), + "Adding token authority should succeed" + ); + + // Add authorization lock for 500 tokens with a future expiry AFTER adding the + // authority + let current_slot = context.svm.get_sysvar::().slot; + let lock_amount = 500u64; + let expiry_slot = current_slot + 1000; // Far in the future + + let add_lock_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority has All permissions (role 0) + mint_pubkey.to_bytes(), + lock_amount, + expiry_slot, + ) + .unwrap(); + + let add_lock_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock_result = context.svm.send_transaction(add_lock_tx); + assert!( + add_lock_result.is_ok(), + "Adding authorization lock should succeed" + ); + + // Verify the authorization lock was added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 1); + + let (auth_locks, count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + assert_eq!(count, 1); + + println!("=== AUTHORIZATION LOCK ENFORCEMENT TEST ==="); + println!("Authorization locks count: {}", count); + for i in 0..count { + if let Some(lock) = auth_locks[i] { + println!( + "Lock {}: mint={:?}, amount={}, expiry_slot={}", + i, lock.token_mint, lock.amount, lock.expiry_slot + ); + } + } + println!("Token authority limit: 1000 tokens"); + println!("Test scenarios:"); + println!(" - Over limit (600 tokens): Should FAIL (exceeds 500 auth lock)"); + println!(" - Within limit (400 tokens): Should PASS (within 500 auth lock)"); + println!(" - Exact limit (500 tokens): Should PASS (equals 500 auth lock)"); + println!("==============================================="); + + let first_lock = auth_locks[0].unwrap(); + assert_eq!(first_lock.token_mint, mint_pubkey.to_bytes()); + assert_eq!(first_lock.amount, lock_amount); + assert_eq!(first_lock.expiry_slot, expiry_slot); + + // Test 1: Try to transfer more than the authorization lock limit (600 tokens) + // This should fail because it exceeds the authorization lock + let over_limit_amount = 600; + + let over_limit_transfer_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata, + &recipient_ata, + &swig, + &[], + over_limit_amount, + ) + .unwrap(); + + let over_limit_sign_ix = swig_interface::SignInstruction::new_ed25519( + swig, + context.default_payer.pubkey(), + token_authority.pubkey(), + over_limit_transfer_ix, + 1, // authority role id (token_authority is role 1) + ) + .unwrap(); + + let over_limit_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[over_limit_sign_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let over_limit_tx = VersionedTransaction::try_new( + VersionedMessage::V0(over_limit_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let over_limit_result = context.svm.send_transaction(over_limit_tx); + assert!( + over_limit_result.is_err(), + "Transfer over authorization lock limit should fail" + ); + + // Test 2: Transfer within the authorization lock limit (400 tokens) + // This should succeed + let within_limit_amount = 400; + + let within_limit_transfer_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata, + &recipient_ata, + &swig, + &[], + within_limit_amount, + ) + .unwrap(); + + let within_limit_sign_ix = swig_interface::SignInstruction::new_ed25519( + swig, + context.default_payer.pubkey(), + token_authority.pubkey(), + within_limit_transfer_ix, + 1, // authority role id (token_authority is role 1) + ) + .unwrap(); + + let within_limit_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[within_limit_sign_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let within_limit_tx = VersionedTransaction::try_new( + VersionedMessage::V0(within_limit_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let within_limit_result = context.svm.send_transaction(within_limit_tx); + println!("{}", within_limit_result.unwrap().pretty_logs()); + // assert!(within_limit_result.is_ok(), "Transfer within authorization lock + // limit should succeed"); + + // Verify the token transfer actually happened + let recipient_token_account = context.svm.get_account(&recipient_ata).unwrap(); + let recipient_balance = spl_token::state::Account::unpack(&recipient_token_account.data) + .unwrap() + .amount; + assert_eq!(recipient_balance, within_limit_amount); + + // Verify swig balance decreased + let swig_token_account = context.svm.get_account(&swig_ata).unwrap(); + let swig_balance = spl_token::state::Account::unpack(&swig_token_account.data) + .unwrap() + .amount; + assert_eq!(swig_balance, initial_token_amount - within_limit_amount); + + // Test 3: Try to transfer exactly the authorization lock limit (500 tokens + // remaining) This should succeed + let exact_limit_amount = 500; + + let exact_limit_transfer_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata, + &recipient_ata, + &swig, + &[], + exact_limit_amount, + ) + .unwrap(); + + let exact_limit_sign_ix = swig_interface::SignInstruction::new_ed25519( + swig, + context.default_payer.pubkey(), + token_authority.pubkey(), + exact_limit_transfer_ix, + 1, // authority role id (token_authority is role 1) + ) + .unwrap(); + + let exact_limit_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[exact_limit_sign_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let exact_limit_tx = VersionedTransaction::try_new( + VersionedMessage::V0(exact_limit_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let exact_limit_result = context.svm.send_transaction(exact_limit_tx); + assert!( + exact_limit_result.is_ok(), + "Transfer of exact authorization lock limit should succeed" + ); + + // Verify final balances + let final_recipient_account = context.svm.get_account(&recipient_ata).unwrap(); + let final_recipient_balance = spl_token::state::Account::unpack(&final_recipient_account.data) + .unwrap() + .amount; + assert_eq!( + final_recipient_balance, + within_limit_amount + exact_limit_amount + ); + + let final_swig_account = context.svm.get_account(&swig_ata).unwrap(); + let final_swig_balance = spl_token::state::Account::unpack(&final_swig_account.data) + .unwrap() + .amount; + assert_eq!( + final_swig_balance, + initial_token_amount - within_limit_amount - exact_limit_amount + ); +} + +/// Test authorization lock expiry behavior +#[test_log::test] +fn test_authorization_lock_expiry() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + let recipient = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&recipient.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + // Setup token accounts + let swig_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &swig, + &context.default_payer, + ) + .unwrap(); + + let recipient_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &recipient.pubkey(), + &context.default_payer, + ) + .unwrap(); + + // Mint tokens to swig account + let initial_token_amount = 1000; + mint_to( + &mut context.svm, + &mint_pubkey, + &context.default_payer, + &swig_ata, + initial_token_amount, + ) + .unwrap(); + + // Test 1: Try to add an authorization lock that has already expired + let current_slot = context.svm.get_sysvar::().slot; + let expired_slot = if current_slot > 0 { + current_slot - 1 + } else { + 0 + }; // Already expired + + println!("=== AUTHORIZATION LOCK EXPIRY TEST ==="); + println!("Current slot: {}", current_slot); + println!("Trying to add lock with expired slot: {}", expired_slot); + println!("Expected result: FAIL (expired authorization lock should be rejected)"); + println!("=========================================="); + + let add_expired_lock_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority has All permissions + mint_pubkey.to_bytes(), + 500, + expired_slot, + ) + .unwrap(); + + let add_expired_lock_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_expired_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_expired_lock_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_expired_lock_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_expired_lock_result = context.svm.send_transaction(add_expired_lock_tx); + assert!( + add_expired_lock_result.is_err(), + "Adding expired authorization lock should fail" + ); +} + +/// Test that expired authorization locks are automatically removed during sign +/// operations +#[test_log::test] +fn test_expired_authorization_lock_cleanup() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + let token_authority = Keypair::new(); + let recipient = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&token_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&recipient.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + // Setup token accounts + let swig_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &swig, + &context.default_payer, + ) + .unwrap(); + + let recipient_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &recipient.pubkey(), + &context.default_payer, + ) + .unwrap(); + + // Mint tokens to swig account + let initial_token_amount = 1000; + mint_to( + &mut context.svm, + &mint_pubkey, + &context.default_payer, + &swig_ata, + initial_token_amount, + ) + .unwrap(); + + // Add token authority + use swig_state_x::action::token_limit::TokenLimit; + let token_action = ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 1000, + }); + + let add_authority_ix = swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &token_authority.pubkey().to_bytes(), + }, + vec![token_action], + ) + .unwrap(); + + let add_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_authority_result = context.svm.send_transaction(add_authority_tx); + assert!( + add_authority_result.is_ok(), + "Adding token authority should succeed" + ); + + // Add authorization lock that will expire soon + let current_slot = context.svm.get_sysvar::().slot; + let short_expiry_slot = current_slot + 5; // Will expire soon + + let add_lock_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority has All permissions + mint_pubkey.to_bytes(), + 500, + short_expiry_slot, + ) + .unwrap(); + + let add_lock_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock_result = context.svm.send_transaction(add_lock_tx); + assert!( + add_lock_result.is_ok(), + "Adding authorization lock should succeed" + ); + + // Verify the authorization lock was added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 1); + + println!("=== EXPIRED AUTHORIZATION LOCK CLEANUP TEST ==="); + println!("Current slot: {}", current_slot); + println!("Lock expiry slot: {}", short_expiry_slot); + println!( + "Authorization locks before expiry: {}", + swig_with_roles.state.authorization_locks + ); + + // Display the lock details + let (auth_locks, count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + for i in 0..count { + if let Some(lock) = auth_locks[i] { + println!( + "Lock {}: mint={:?}, amount={}, expiry_slot={}", + i, lock.token_mint, lock.amount, lock.expiry_slot + ); + println!( + " → Lock will expire at slot {}, current slot is {}", + lock.expiry_slot, current_slot + ); + } + } + + // Advance time past the expiry slot + context.svm.warp_to_slot(short_expiry_slot + 10); + let new_current_slot = context.svm.get_sysvar::().slot; + println!(); + println!("TIME WARP:"); + println!("New current slot after warp: {}", new_current_slot); + println!("Lock expiry slot: {}", short_expiry_slot); + println!( + "Lock is now {} slots expired", + new_current_slot - short_expiry_slot + ); + println!("Expected: Lock should be removed during next sign operation"); + println!(); + + // Perform a token transfer that will trigger the cleanup + let transfer_amount = 100; + println!("PERFORMING SIGN OPERATION:"); + println!("Transfer amount: {} tokens", transfer_amount); + println!("This will trigger expired lock cleanup..."); + + let transfer_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata, + &recipient_ata, + &swig, + &[], + transfer_amount, + ) + .unwrap(); + + let sign_ix = swig_interface::SignInstruction::new_ed25519( + swig, + context.default_payer.pubkey(), + token_authority.pubkey(), + transfer_ix, + 1, // authority role id + ) + .unwrap(); + + let sign_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[sign_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let sign_tx = VersionedTransaction::try_new( + VersionedMessage::V0(sign_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let sign_result = context.svm.send_transaction(sign_tx); + assert!(sign_result.is_ok(), "Sign operation should succeed"); + println!("✅ Sign operation completed successfully"); + println!(); + + // Verify that the expired authorization lock was removed + let swig_account_after = context.svm.get_account(&swig).unwrap(); + let swig_with_roles_after = SwigWithRoles::from_bytes(&swig_account_after.data).unwrap(); + + println!("CLEANUP RESULTS:"); + println!( + "Authorization locks after cleanup: {}", + swig_with_roles_after.state.authorization_locks + ); + println!( + "Authorization locks before cleanup: {}", + swig_with_roles.state.authorization_locks + ); + println!( + "Locks removed: {}", + swig_with_roles.state.authorization_locks - swig_with_roles_after.state.authorization_locks + ); + + // Display remaining locks (should be none) + let (remaining_locks, remaining_count) = swig_with_roles_after + .get_authorization_locks_for_test::<10>() + .unwrap(); + println!("Remaining locks: {}", remaining_count); + for i in 0..remaining_count { + if let Some(lock) = remaining_locks[i] { + println!( + " Remaining Lock {}: mint={:?}, amount={}, expiry_slot={}", + i, lock.token_mint, lock.amount, lock.expiry_slot + ); + } + } + + println!("Expected: 0 (expired lock should be removed)"); + println!("============================================"); + + assert_eq!( + swig_with_roles_after.state.authorization_locks, 0, + "Expired authorization lock should have been removed" + ); + + // Verify the token transfer still succeeded + let recipient_token_account = context.svm.get_account(&recipient_ata).unwrap(); + let recipient_balance = spl_token::state::Account::unpack(&recipient_token_account.data) + .unwrap() + .amount; + assert_eq!(recipient_balance, transfer_amount); + println!( + "✅ Token transfer verification: {} tokens successfully transferred", + transfer_amount + ); +} + +/// Test that multiple authorization locks work correctly +#[test_log::test] +fn test_multiple_authorization_locks() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); // This will have All permissions + let token_authority = Keypair::new(); // This will have only token limit permissions + let recipient = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&token_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&recipient.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup two different token mints + let mint1_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + let mint2_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + // Setup token accounts for both mints + let swig_ata1 = setup_ata( + &mut context.svm, + &mint1_pubkey, + &swig, + &context.default_payer, + ) + .unwrap(); + + let swig_ata2 = setup_ata( + &mut context.svm, + &mint2_pubkey, + &swig, + &context.default_payer, + ) + .unwrap(); + + let recipient_ata1 = setup_ata( + &mut context.svm, + &mint1_pubkey, + &recipient.pubkey(), + &context.default_payer, + ) + .unwrap(); + + let recipient_ata2 = setup_ata( + &mut context.svm, + &mint2_pubkey, + &recipient.pubkey(), + &context.default_payer, + ) + .unwrap(); + + // Mint tokens to both swig accounts + let initial_token_amount = 1000; + mint_to( + &mut context.svm, + &mint1_pubkey, + &context.default_payer, + &swig_ata1, + initial_token_amount, + ) + .unwrap(); + + mint_to( + &mut context.svm, + &mint2_pubkey, + &context.default_payer, + &swig_ata2, + initial_token_amount, + ) + .unwrap(); + + // Add authorization locks for both tokens + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + + println!("=== MULTIPLE AUTHORIZATION LOCKS TEST ==="); + println!("Setting up authorization locks:"); + println!("Current slot: {}", current_slot); + println!("Expiry slot: {} (+1000 slots)", expiry_slot); + println!("Mint1: {:?}", mint1_pubkey.to_bytes()); + println!("Mint2: {:?}", mint2_pubkey.to_bytes()); + println!(); + + // Lock 1: 300 tokens for mint1 + println!("Adding Lock 1 for mint1: 300 tokens"); + let add_lock1_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority has All permissions + mint1_pubkey.to_bytes(), + 300, + expiry_slot, + ) + .unwrap(); + + // Lock 2: 400 tokens for mint2 + println!("Adding Lock 2 for mint2: 400 tokens"); + let add_lock2_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority has All permissions + mint2_pubkey.to_bytes(), + 400, + expiry_slot, + ) + .unwrap(); + + // Add both locks + for lock_ix in [add_lock1_ix, add_lock2_ix] { + let message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx = VersionedTransaction::try_new( + VersionedMessage::V0(message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let result = context.svm.send_transaction(tx); + assert!(result.is_ok(), "Adding authorization lock should succeed"); + } + println!("✅ Both authorization locks added successfully"); + println!(); + + // Verify both locks were added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 2); + + let (all_auth_locks, count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + println!("VERIFICATION - Authorization locks in account:"); + println!("Total authorization locks count: {}", count); + for i in 0..count { + if let Some(lock) = all_auth_locks[i] { + println!( + "Lock {}: mint={:?}, amount={}, expiry_slot={}", + i, lock.token_mint, lock.amount, lock.expiry_slot + ); + // Check which mint this corresponds to + if lock.token_mint == mint1_pubkey.to_bytes() { + println!(" → This is the MINT1 lock (300 tokens)"); + } else if lock.token_mint == mint2_pubkey.to_bytes() { + println!(" → This is the MINT2 lock (400 tokens)"); + } + } + } + println!(); + println!("AUTHORITY SETUP:"); + println!("Token authority limit: 250 tokens for mint1 ONLY"); + println!("Swig authority: unlimited permissions (All)"); + println!(); + println!("TEST SCENARIOS:"); + println!( + " 1. Transfer 200 tokens of mint1 using swig_authority: Should PASS (within 300 auth \ + lock)" + ); + println!( + " 2. Transfer 350 tokens of mint2 using swig_authority: Should PASS (within 400 auth \ + lock)" + ); + println!( + " 3. Transfer 400 tokens of mint1 using token_authority: Should FAIL (exceeds 300 auth \ + lock)" + ); + println!("============================================="); + + // Add token authority with limited token permissions for mint1 + use swig_state_x::action::token_limit::TokenLimit; + let token_action = ClientAction::TokenLimit(TokenLimit { + token_mint: mint1_pubkey.to_bytes(), + current_amount: 250, // Allow up to 250 tokens for mint1 (less than the 300 auth lock) + }); + + let add_authority_ix = swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, // Acting role ID (swig_authority has All permissions) + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &token_authority.pubkey().to_bytes(), + }, + vec![token_action], + ) + .unwrap(); + + let add_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_authority_result = context.svm.send_transaction(add_authority_tx); + assert!( + add_authority_result.is_ok(), + "Adding token authority should succeed" + ); + + // Test transfers within each lock's limits + println!(); + println!("EXECUTING TEST SCENARIO 1:"); + println!("Transfer 200 tokens of mint1 using swig_authority (All permissions)"); + println!("Expected: PASS (200 < 300 auth lock limit)"); + + let transfer1_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata1, + &recipient_ata1, + &swig, + &[], + 200, + ) + .unwrap(); + + let sign1_ix = swig_interface::SignInstruction::new_ed25519( + swig, + swig_authority.pubkey(), + swig_authority.pubkey(), + transfer1_ix, + 0, + ) + .unwrap(); + + let message1 = v0::Message::try_compile( + &swig_authority.pubkey(), + &[sign1_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx1 = + VersionedTransaction::try_new(VersionedMessage::V0(message1), &[&swig_authority]).unwrap(); + + let result1 = context.svm.send_transaction(tx1); + assert!( + result1.is_ok(), + "Transfer of mint1 within lock limit should succeed" + ); + println!("✅ Scenario 1 PASSED: 200 tokens of mint1 transferred successfully"); + + // Transfer 350 tokens of mint2 (within 400 limit) + println!(); + println!("EXECUTING TEST SCENARIO 2:"); + println!("Transfer 350 tokens of mint2 using swig_authority (All permissions)"); + println!("Expected: PASS (350 < 400 auth lock limit)"); + + let transfer2_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata2, + &recipient_ata2, + &swig, + &[], + 350, + ) + .unwrap(); + + let sign2_ix = swig_interface::SignInstruction::new_ed25519( + swig, + swig_authority.pubkey(), + swig_authority.pubkey(), + transfer2_ix, + 0, + ) + .unwrap(); + + let message2 = v0::Message::try_compile( + &swig_authority.pubkey(), + &[sign2_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx2 = + VersionedTransaction::try_new(VersionedMessage::V0(message2), &[&swig_authority]).unwrap(); + + let result2 = context.svm.send_transaction(tx2); + assert!( + result2.is_ok(), + "Transfer of mint2 within lock limit should succeed" + ); + println!("✅ Scenario 2 PASSED: 350 tokens of mint2 transferred successfully"); + + // Test transfer that exceeds mint1 lock (400 tokens, exceeds 300 limit) + println!(); + println!("EXECUTING TEST SCENARIO 3:"); + println!("Transfer 400 tokens of mint1 using token_authority (250 token limit)"); + println!("Expected: FAIL (400 > 300 auth lock limit, AND 400 > 250 token limit)"); + println!( + "Note: This tests that authorization locks are enforced even with limited authorities" + ); + + let over_limit_transfer_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata1, + &recipient_ata1, + &swig, + &[], + 400, + ) + .unwrap(); + + let over_limit_sign_ix = swig_interface::SignInstruction::new_ed25519( + swig, + context.default_payer.pubkey(), + token_authority.pubkey(), + over_limit_transfer_ix, + 1, // authority role id (token_authority is role 1) + ) + .unwrap(); + + let over_limit_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[over_limit_sign_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let over_limit_tx = VersionedTransaction::try_new( + VersionedMessage::V0(over_limit_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let over_limit_result = context.svm.send_transaction(over_limit_tx); + assert!( + over_limit_result.is_err(), + "Transfer exceeding mint1 lock limit should fail" + ); + println!("✅ Scenario 3 PASSED: Transfer correctly rejected (exceeds auth lock limit)"); + println!(); + println!("FINAL VERIFICATION:"); + + // Check final state + let final_swig_account = context.svm.get_account(&swig).unwrap(); + let final_swig_with_roles = SwigWithRoles::from_bytes(&final_swig_account.data).unwrap(); + let (final_locks, final_count) = final_swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + + println!("Authorization locks still present: {}", final_count); + println!("All locks should still be active (none expired)"); + for i in 0..final_count { + if let Some(lock) = final_locks[i] { + let current_test_slot = context.svm.get_sysvar::().slot; + let expires_in = if lock.expiry_slot > current_test_slot { + lock.expiry_slot - current_test_slot + } else { + 0 + }; + println!( + "Lock {}: amount={}, expires in {} slots", + i, lock.amount, expires_in + ); + } + } + println!("✅ Multiple authorization locks test completed successfully!"); + println!("============================================="); +} + +/// Test that multiple authorization locks for the same token mint are combined +#[test_log::test] +fn test_combined_authorization_locks_same_mint() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + let token_authority = Keypair::new(); + let recipient = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&token_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&recipient.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + // Setup token accounts + let swig_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &swig, + &context.default_payer, + ) + .unwrap(); + + let recipient_ata = setup_ata( + &mut context.svm, + &mint_pubkey, + &recipient.pubkey(), + &context.default_payer, + ) + .unwrap(); + + // Mint tokens to swig account + let initial_token_amount = 1000; + mint_to( + &mut context.svm, + &mint_pubkey, + &context.default_payer, + &swig_ata, + initial_token_amount, + ) + .unwrap(); + + // Add token authority with limited permissions + use swig_state_x::action::token_limit::TokenLimit; + let token_action = ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 220, // Allow up to 150 tokens (less than combined auth locks) + }); + + let add_authority_ix = swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &token_authority.pubkey().to_bytes(), + }, + vec![token_action], + ) + .unwrap(); + + let add_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_authority_result = context.svm.send_transaction(add_authority_tx); + assert!( + add_authority_result.is_ok(), + "Adding token authority should succeed" + ); + + // Add multiple authorization locks for the SAME mint + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + + println!("=== COMBINED AUTHORIZATION LOCKS TEST ==="); + println!("Setting up multiple authorization locks for the SAME mint:"); + println!("Current slot: {}", current_slot); + println!("Expiry slot: {} (+1000 slots)", expiry_slot); + println!("Mint: {:?}", mint_pubkey.to_bytes()); + println!(); + + // Lock 1: 100 tokens for the same mint + println!("Adding Lock 1 for mint: 100 tokens"); + let add_lock1_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority has All permissions + mint_pubkey.to_bytes(), + 100, + expiry_slot, + ) + .unwrap(); + + // Lock 2: 120 tokens for the same mint + println!("Adding Lock 2 for SAME mint: 120 tokens"); + let add_lock2_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority has All permissions + mint_pubkey.to_bytes(), + 120, + expiry_slot, + ) + .unwrap(); + + // Add both locks + for (i, lock_ix) in [add_lock1_ix, add_lock2_ix].iter().enumerate() { + let message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[lock_ix.clone()], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx = VersionedTransaction::try_new( + VersionedMessage::V0(message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let result = context.svm.send_transaction(tx); + assert!( + result.is_ok(), + "Adding authorization lock {} should succeed", + i + 1 + ); + println!("{}", result.unwrap().pretty_logs()); + } + println!("✅ Both authorization locks added successfully"); + println!(); + + // Verify both locks were added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 2); + + let (all_auth_locks, count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + println!("VERIFICATION - Authorization locks in account:"); + println!("Total authorization locks count: {}", count); + let mut total_amount = 0u64; + for i in 0..count { + if let Some(lock) = all_auth_locks[i] { + println!( + "Lock {}: mint={:?}, amount={}, expiry_slot={}", + i, lock.token_mint, lock.amount, lock.expiry_slot + ); + total_amount += lock.amount; + } + } + println!(); + println!("COMBINED AUTHORIZATION:"); + println!("Individual locks: 100 + 120 = {} tokens", total_amount); + println!("Token authority limit: 150 tokens"); + println!( + "Expected behavior: Combined auth locks should allow up to {} tokens", + total_amount + ); + println!(); + + // Test scenarios + println!("TEST SCENARIOS:"); + println!(" 1. Transfer 200 tokens: Should PASS (200 < 220 combined auth locks)"); + println!(" 2. Transfer 250 tokens: Should FAIL (250 > 220 combined auth locks)"); + println!("============================================="); + println!(); + + // Test 1: Transfer 200 tokens (within combined limit of 220) + println!("EXECUTING TEST SCENARIO 1:"); + println!("Transfer 200 tokens using token_authority (150 token limit)"); + println!("Expected: PASS (200 < 220 combined auth lock limit)"); + + let transfer_amount = 200; + let transfer_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata, + &recipient_ata, + &swig, + &[], + transfer_amount, + ) + .unwrap(); + + let sign_ix = swig_interface::SignInstruction::new_ed25519( + swig, + context.default_payer.pubkey(), + token_authority.pubkey(), + transfer_ix, + 1, // authority role id (token_authority is role 1) + ) + .unwrap(); + + let sign_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[sign_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let sign_tx = VersionedTransaction::try_new( + VersionedMessage::V0(sign_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let sign_result = context.svm.send_transaction(sign_tx); + println!("{}", sign_result.unwrap().pretty_logs()); + // assert!( + // sign_result.is_ok(), + // "Transfer within combined authorization lock limit should succeed" + // ); + println!( + "✅ Scenario 1 PASSED: 200 tokens transferred successfully (combined auth locks worked)" + ); + + // Verify the token transfer actually happened + let recipient_token_account = context.svm.get_account(&recipient_ata).unwrap(); + let recipient_balance = spl_token::state::Account::unpack(&recipient_token_account.data) + .unwrap() + .amount; + assert_eq!(recipient_balance, transfer_amount); + + // Test 2: Transfer 250 tokens (exceeds combined limit of 220) + println!(); + println!("EXECUTING TEST SCENARIO 2:"); + println!("Transfer 250 tokens using token_authority (150 token limit)"); + println!("Expected: FAIL (250 > 220 combined auth lock limit)"); + + let over_limit_amount = 250; + let over_limit_transfer_ix = spl_token::instruction::transfer( + &spl_token::ID, + &swig_ata, + &recipient_ata, + &swig, + &[], + over_limit_amount, + ) + .unwrap(); + + let over_limit_sign_ix = swig_interface::SignInstruction::new_ed25519( + swig, + context.default_payer.pubkey(), + token_authority.pubkey(), + over_limit_transfer_ix, + 1, // authority role id + ) + .unwrap(); + + let over_limit_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[over_limit_sign_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let over_limit_tx = VersionedTransaction::try_new( + VersionedMessage::V0(over_limit_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let over_limit_result = context.svm.send_transaction(over_limit_tx); + assert!( + over_limit_result.is_err(), + "Transfer exceeding combined authorization lock limit should fail" + ); + println!( + "✅ Scenario 2 PASSED: Transfer correctly rejected (exceeds combined auth lock limit)" + ); + + println!(); + println!("FINAL VERIFICATION:"); + println!("✅ Combined authorization locks working correctly!"); + println!("✅ Multiple locks for same mint are properly summed (100 + 120 = 220)"); + println!("✅ Transfers within combined limit (200) succeed"); + println!("✅ Transfers exceeding combined limit (250) are rejected"); + println!("============================================="); +} diff --git a/program/tests/authorization_lock_token_limits_test.rs b/program/tests/authorization_lock_token_limits_test.rs new file mode 100644 index 00000000..06cd7cce --- /dev/null +++ b/program/tests/authorization_lock_token_limits_test.rs @@ -0,0 +1,179 @@ +#![cfg(not(feature = "program_scope_test"))] + +mod common; +use common::*; +use litesvm_token::spl_token; +use solana_sdk::{ + clock::Clock, + message::{v0, VersionedMessage}, + program_pack::Pack, + pubkey::Pubkey, + signature::Keypair, + signer::Signer, + transaction::VersionedTransaction, +}; +use swig_interface::{AuthorityConfig, ClientAction}; +use swig_state_x::{ + action::{manage_authorization_locks::ManageAuthorizationLocks, token_limit::TokenLimit}, + swig::{swig_account_seeds, AuthorizationLock, Swig, SwigWithRoles}, + IntoBytes, Transmutable, +}; + +/// Test that validates authorization locks respect simple token limits. +/// +/// This test creates a role with a TokenLimit of 1000 tokens, then: +/// 1. Successfully adds an authorization lock for 800 tokens (within limit) +/// 2. Fails to add another authorization lock for 300 tokens (would exceed +/// limit: 800 + 300 = 1100 > 1000) +#[test_log::test] +fn test_authorization_lock_respects_simple_token_limit() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + let token_authority = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&token_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account with root authority (All permissions) + let swig_id = [1u8; 32]; + let (swig_pubkey, _) = create_swig_ed25519(&mut context, &swig_authority, swig_id).unwrap(); + + // Add a token authority with TokenLimit permission (1000 tokens) + let token_limit_amount = 1000u64; + let token_authority_actions = vec![ + ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: token_limit_amount, + }), + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}), + ]; + + add_authority_with_ed25519_root( + &mut context, + &swig_pubkey, + &swig_authority, + AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: token_authority.pubkey().as_ref(), + }, + token_authority_actions, + ) + .unwrap(); + + // Get role ID for the token authority + let swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + let token_role_id = swig_with_roles + .lookup_role_id(token_authority.pubkey().as_ref()) + .unwrap() + .unwrap(); + + // Test 1: Successfully add authorization lock for 800 tokens (within limit) + let auth_lock_amount_1 = 800u64; + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + + let add_auth_lock_ix_1 = swig_interface::AddAuthorizationLockInstruction::new( + swig_pubkey, + token_authority.pubkey(), + context.default_payer.pubkey(), + token_role_id, + mint_pubkey.to_bytes(), + auth_lock_amount_1, + expiry_slot, + ) + .unwrap(); + + let msg_1 = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_auth_lock_ix_1], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx_1 = VersionedTransaction::try_new( + VersionedMessage::V0(msg_1), + &[ + context.default_payer.insecure_clone(), + token_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result_1 = context.svm.send_transaction(tx_1); + assert!( + result_1.is_ok(), + "First authorization lock (800 tokens) should succeed: {:?}", + result_1 + ); + + // Test 2: Fail to add authorization lock for 300 tokens (would exceed limit: + // 800 + 300 = 1100 > 1000) + let auth_lock_amount_2 = 300u64; + + let add_auth_lock_ix_2 = swig_interface::AddAuthorizationLockInstruction::new( + swig_pubkey, + token_authority.pubkey(), + context.default_payer.pubkey(), + token_role_id, + mint_pubkey.to_bytes(), + auth_lock_amount_2, + expiry_slot, + ) + .unwrap(); + + let msg_2 = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_auth_lock_ix_2], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx_2 = VersionedTransaction::try_new( + VersionedMessage::V0(msg_2), + &[ + context.default_payer.insecure_clone(), + token_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result_2 = context.svm.send_transaction(tx_2); + assert!( + result_2.is_err(), + "Second authorization lock (300 tokens) should fail due to exceeding limit" + ); + + // Verify final state: should have exactly 1 authorization lock for 800 tokens + let final_swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let final_swig_with_roles = SwigWithRoles::from_bytes(&final_swig_account.data).unwrap(); + + let (auth_locks, count) = final_swig_with_roles + .get_authorization_locks_by_role::<10>(token_role_id) + .unwrap(); + + assert_eq!(count, 1, "Should have exactly 1 authorization lock"); + + let total_locked: u64 = auth_locks + .iter() + .filter_map(|opt_lock| *opt_lock) + .filter(|lock| lock.token_mint == mint_pubkey.to_bytes()) + .map(|lock| lock.amount) + .sum(); + + assert_eq!(total_locked, 800, "Total locked amount should be 800"); +} diff --git a/program/tests/remove_authorization_lock_role_ownership_test.rs b/program/tests/remove_authorization_lock_role_ownership_test.rs new file mode 100644 index 00000000..2ddf3d38 --- /dev/null +++ b/program/tests/remove_authorization_lock_role_ownership_test.rs @@ -0,0 +1,426 @@ +#![cfg(not(feature = "program_scope_test"))] + +mod common; +use common::*; +use solana_sdk::{ + clock::Clock, + message::{v0, VersionedMessage}, + signature::Keypair, + signer::Signer, + transaction::VersionedTransaction, +}; +use swig_interface::{ + AddAuthorizationLockInstruction, AuthorityConfig, ClientAction, + RemoveAuthorizationLockInstruction, +}; +use swig_state_x::{ + action::{manage_authorization_locks::ManageAuthorizationLocks, token_limit::TokenLimit}, + swig::SwigWithRoles, + IntoBytes, Transmutable, +}; + +/// Test that validates only the role that created an authorization lock can +/// remove it. +/// +/// This test: +/// 1. Creates two roles with ManageAuthorizationLocks permission +/// 2. Role A creates an authorization lock +/// 3. Role B tries to remove Role A's lock and fails +/// 4. Role A successfully removes its own lock +#[test_log::test] +fn test_remove_authorization_lock_role_ownership() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let root_authority = Keypair::new(); + let role_a_authority = Keypair::new(); + let role_b_authority = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&root_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&role_a_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&role_b_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account with root authority + let swig_id = [1u8; 32]; + let (swig_pubkey, _) = create_swig_ed25519(&mut context, &root_authority, swig_id).unwrap(); + + // Add Role A with token limit and manage authorization locks permission + let role_a_actions = vec![ + ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 1000u64, + }), + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}), + ]; + + add_authority_with_ed25519_root( + &mut context, + &swig_pubkey, + &root_authority, + AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: role_a_authority.pubkey().as_ref(), + }, + role_a_actions, + ) + .unwrap(); + + // Add Role B with token limit and manage authorization locks permission + let role_b_actions = vec![ + ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 1000u64, + }), + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}), + ]; + + add_authority_with_ed25519_root( + &mut context, + &swig_pubkey, + &root_authority, + AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: role_b_authority.pubkey().as_ref(), + }, + role_b_actions, + ) + .unwrap(); + + // Get role IDs + let swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + + let role_a_id = swig_with_roles + .lookup_role_id(role_a_authority.pubkey().as_ref()) + .unwrap() + .unwrap(); + + let role_b_id = swig_with_roles + .lookup_role_id(role_b_authority.pubkey().as_ref()) + .unwrap() + .unwrap(); + + // Role A creates an authorization lock + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + let auth_lock_amount = 500u64; + + let add_auth_lock_ix = AddAuthorizationLockInstruction::new( + swig_pubkey, + role_a_authority.pubkey(), + context.default_payer.pubkey(), + role_a_id, + mint_pubkey.to_bytes(), + auth_lock_amount, + expiry_slot, + ) + .unwrap(); + + let msg = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_auth_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx = VersionedTransaction::try_new( + VersionedMessage::V0(msg), + &[ + context.default_payer.insecure_clone(), + role_a_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result = context.svm.send_transaction(tx); + assert!( + result.is_ok(), + "Role A should successfully create authorization lock: {:?}", + result + ); + + // Verify the lock was created + let swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + let (auth_locks, count) = swig_with_roles + .get_authorization_locks_by_role::<10>(role_a_id) + .unwrap(); + assert_eq!( + count, 1, + "Should have exactly 1 authorization lock for Role A" + ); + + // Role B tries to remove Role A's authorization lock (should fail) + let remove_auth_lock_ix_b = RemoveAuthorizationLockInstruction::new( + swig_pubkey, + role_b_authority.pubkey(), + context.default_payer.pubkey(), + role_b_id, + 0, // First (and only) lock index + ) + .unwrap(); + + let msg_b = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_auth_lock_ix_b], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx_b = VersionedTransaction::try_new( + VersionedMessage::V0(msg_b), + &[ + context.default_payer.insecure_clone(), + role_b_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result_b = context.svm.send_transaction(tx_b); + assert!( + result_b.is_err(), + "Role B should NOT be able to remove Role A's authorization lock" + ); + + // Verify the lock is still there + let swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + let (auth_locks, count) = swig_with_roles + .get_authorization_locks_by_role::<10>(role_a_id) + .unwrap(); + assert_eq!( + count, 1, + "Authorization lock should still exist after failed removal" + ); + + // Role A removes its own authorization lock (should succeed) + let remove_auth_lock_ix_a = RemoveAuthorizationLockInstruction::new( + swig_pubkey, + role_a_authority.pubkey(), + context.default_payer.pubkey(), + role_a_id, + 0, // First (and only) lock index + ) + .unwrap(); + + let msg_a = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_auth_lock_ix_a], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx_a = VersionedTransaction::try_new( + VersionedMessage::V0(msg_a), + &[ + context.default_payer.insecure_clone(), + role_a_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result_a = context.svm.send_transaction(tx_a); + assert!( + result_a.is_ok(), + "Role A should successfully remove its own authorization lock: {:?}", + result_a + ); + + // Verify the lock was removed + let final_swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let final_swig_with_roles = SwigWithRoles::from_bytes(&final_swig_account.data).unwrap(); + let (final_locks, final_count) = final_swig_with_roles + .get_authorization_locks_by_role::<10>(role_a_id) + .unwrap(); + assert_eq!(final_count, 0, "Authorization lock should be removed"); +} + +/// Test that validates a role with All permissions can only remove +/// authorization locks they created. +/// +/// This test: +/// 1. Creates a role A with limited permissions that creates an authorization +/// lock +/// 2. Creates a role B with All permissions +/// 3. Role B fails to remove Role A's lock (can only remove own locks) +#[test_log::test] +fn test_remove_authorization_lock_all_permission_override() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let root_authority = Keypair::new(); + let role_a_authority = Keypair::new(); + let role_b_authority = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&root_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&role_a_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&role_b_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account with root authority + let swig_id = [1u8; 32]; + let (swig_pubkey, _) = create_swig_ed25519(&mut context, &root_authority, swig_id).unwrap(); + + // Add Role A with limited permissions + let role_a_actions = vec![ + ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 1000u64, + }), + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}), + ]; + + add_authority_with_ed25519_root( + &mut context, + &swig_pubkey, + &root_authority, + AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: role_a_authority.pubkey().as_ref(), + }, + role_a_actions, + ) + .unwrap(); + + // Add Role B with All permissions + let role_b_actions = vec![ClientAction::All(swig_state_x::action::all::All {})]; + + add_authority_with_ed25519_root( + &mut context, + &swig_pubkey, + &root_authority, + AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: role_b_authority.pubkey().as_ref(), + }, + role_b_actions, + ) + .unwrap(); + + // Get role IDs + let swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + + let role_a_id = swig_with_roles + .lookup_role_id(role_a_authority.pubkey().as_ref()) + .unwrap() + .unwrap(); + + let role_b_id = swig_with_roles + .lookup_role_id(role_b_authority.pubkey().as_ref()) + .unwrap() + .unwrap(); + + // Role A creates an authorization lock + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + let auth_lock_amount = 500u64; + + let add_auth_lock_ix = AddAuthorizationLockInstruction::new( + swig_pubkey, + role_a_authority.pubkey(), + context.default_payer.pubkey(), + role_a_id, + mint_pubkey.to_bytes(), + auth_lock_amount, + expiry_slot, + ) + .unwrap(); + + let msg = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_auth_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx = VersionedTransaction::try_new( + VersionedMessage::V0(msg), + &[ + context.default_payer.insecure_clone(), + role_a_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result = context.svm.send_transaction(tx); + assert!( + result.is_ok(), + "Role A should successfully create authorization lock: {:?}", + result + ); + + // Role B (with All permissions) tries to remove Role A's authorization lock + // (should fail) + let remove_auth_lock_ix_b = RemoveAuthorizationLockInstruction::new( + swig_pubkey, + role_b_authority.pubkey(), + context.default_payer.pubkey(), + role_b_id, + 0, // First (and only) lock index + ) + .unwrap(); + + let msg_b = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_auth_lock_ix_b], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx_b = VersionedTransaction::try_new( + VersionedMessage::V0(msg_b), + &[ + context.default_payer.insecure_clone(), + role_b_authority.insecure_clone(), + ], + ) + .unwrap(); + + let result_b = context.svm.send_transaction(tx_b); + assert!( + result_b.is_err(), + "Role B with All permissions should NOT be able to remove Role A's authorization lock" + ); + + // Verify the lock was NOT removed + let final_swig_account = context.svm.get_account(&swig_pubkey).unwrap(); + let final_swig_with_roles = SwigWithRoles::from_bytes(&final_swig_account.data).unwrap(); + let (final_locks, final_count) = final_swig_with_roles + .get_authorization_locks_by_role::<10>(role_a_id) + .unwrap(); + assert_eq!( + final_count, 1, + "Authorization lock should still exist after failed removal attempt" + ); +} diff --git a/program/tests/remove_authorization_lock_test.rs b/program/tests/remove_authorization_lock_test.rs new file mode 100644 index 00000000..83ed477a --- /dev/null +++ b/program/tests/remove_authorization_lock_test.rs @@ -0,0 +1,702 @@ +#![cfg(not(feature = "program_scope_test"))] + +mod common; +use common::*; +use litesvm_token::spl_token; +use solana_sdk::{ + clock::Clock, + message::{v0, VersionedMessage}, + program_pack::Pack, + pubkey::Pubkey, + signature::Keypair, + signer::Signer, + transaction::VersionedTransaction, +}; +use swig_interface::{AuthorityConfig, ClientAction}; +use swig_state_x::{ + action::{ + all::All, manage_authorization_locks::ManageAuthorizationLocks, token_limit::TokenLimit, + }, + swig::{swig_account_seeds, AuthorizationLock, Swig, SwigWithRoles}, + IntoBytes, Transmutable, +}; + +/// Test that validates only authorities with "All" or +/// "ManageAuthorizationLocks" permissions can remove authorization locks, while +/// others are denied. +#[test_log::test] +fn test_remove_authorization_lock_permission_enforcement() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); // Will have All permissions + let manage_auth_locks_authority = Keypair::new(); // Will have ManageAuthorizationLocks permission + let token_authority = Keypair::new(); // Will have only token limit permissions (should be denied) + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&manage_auth_locks_authority.pubkey(), 10_000_000_000) + .unwrap(); + context + .svm + .airdrop(&token_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + println!("=== REMOVE AUTHORIZATION LOCK PERMISSION ENFORCEMENT TEST ==="); + println!( + "Testing that only authorities with proper permissions can remove authorization locks" + ); + println!(); + + // Add authority with ManageAuthorizationLocks permission + let manage_auth_locks_action = + ClientAction::ManageAuthorizationLocks(ManageAuthorizationLocks {}); + let add_manage_auth_locks_authority_ix = + swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, // Acting role ID (swig_authority has All permissions) + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &manage_auth_locks_authority.pubkey().to_bytes(), + }, + vec![manage_auth_locks_action], + ) + .unwrap(); + + let add_manage_auth_locks_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_manage_auth_locks_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_manage_auth_locks_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_manage_auth_locks_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_manage_auth_locks_authority_result = context + .svm + .send_transaction(add_manage_auth_locks_authority_tx); + assert!( + add_manage_auth_locks_authority_result.is_ok(), + "Adding ManageAuthorizationLocks authority should succeed" + ); + println!("✅ Added authority with ManageAuthorizationLocks permission (role ID: 1)"); + + // Add authority with limited token permissions (should NOT be able to manage + // auth locks) + let token_action = ClientAction::TokenLimit(TokenLimit { + token_mint: mint_pubkey.to_bytes(), + current_amount: 1000, + }); + + let add_token_authority_ix = + swig_interface::AddAuthorityInstruction::new_with_ed25519_authority( + swig, + context.default_payer.pubkey(), + swig_authority.pubkey(), + 0, // Acting role ID (swig_authority has All permissions) + swig_interface::AuthorityConfig { + authority_type: swig_state_x::authority::AuthorityType::Ed25519, + authority: &token_authority.pubkey().to_bytes(), + }, + vec![token_action], + ) + .unwrap(); + + let add_token_authority_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_token_authority_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_token_authority_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_token_authority_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_token_authority_result = context.svm.send_transaction(add_token_authority_tx); + assert!( + add_token_authority_result.is_ok(), + "Adding token authority should succeed" + ); + println!("✅ Added authority with TokenLimit permission only (role ID: 2)"); + println!(); + + // Add multiple authorization locks to test with + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + + // Add first lock using manage_auth_locks_authority (role 1) + let add_lock1_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + manage_auth_locks_authority.pubkey(), + context.default_payer.pubkey(), + 1, // acting_role_id: manage_auth_locks_authority + mint_pubkey.to_bytes(), + 500, + expiry_slot, + ) + .unwrap(); + + // Add second lock using manage_auth_locks_authority (role 1) + let add_lock2_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + manage_auth_locks_authority.pubkey(), + context.default_payer.pubkey(), + 1, // acting_role_id: manage_auth_locks_authority + mint_pubkey.to_bytes(), + 300, + expiry_slot, + ) + .unwrap(); + + // Add third lock using swig_authority (role 0) + let add_lock3_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority + mint_pubkey.to_bytes(), + 200, + expiry_slot, + ) + .unwrap(); + + // Add all three locks + for (i, lock_ix) in [add_lock1_ix, add_lock2_ix, add_lock3_ix] + .iter() + .enumerate() + { + let authority = if i < 2 { + &manage_auth_locks_authority + } else { + &swig_authority + }; + + let message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[lock_ix.clone()], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let tx = VersionedTransaction::try_new( + VersionedMessage::V0(message), + &[&context.default_payer, authority], + ) + .unwrap(); + + let result = context.svm.send_transaction(tx); + assert!( + result.is_ok(), + "Adding authorization lock {} should succeed", + i + 1 + ); + } + println!("✅ Added 3 authorization locks for testing removal"); + + // Expire blockhash before verification and removal operations + context.svm.expire_blockhash(); + + // Verify all locks were added + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 3); + + let (auth_locks, count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + println!(); + println!("INITIAL STATE - Authorization locks in account:"); + println!("Total authorization locks count: {}", count); + for i in 0..count { + if let Some(lock) = auth_locks[i] { + println!( + "Lock {}: mint={:?}, amount={}, expiry_slot={}", + i, lock.token_mint, lock.amount, lock.expiry_slot + ); + } + } + println!(); + + println!("TEST SCENARIOS:"); + println!( + " 1. Remove lock using authority with All permission (role 0) - removing another role's \ + lock: Should FAIL" + ); + println!( + " 2. Remove lock using authority with ManageAuthorizationLocks permission (role 1) - \ + removing their own lock: Should PASS" + ); + println!( + " 3. Remove lock using authority with TokenLimit permission only (role 2): Should FAIL" + ); + println!(); + + // Test 1: Authority with All permission should FAIL when trying to remove + // another role's lock + println!("EXECUTING TEST SCENARIO 1:"); + println!( + "Attempting to remove authorization lock at index 1 using authority with All permission \ + (role 0) - removing another role's lock" + ); + println!("Expected: FAIL (can only remove own locks)"); + + let remove_lock_all_ix = swig_interface::RemoveAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id: swig_authority (All permission) + 1, // Remove the middle lock (index 1, created by role 1) + ) + .unwrap(); + + let remove_lock_all_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_lock_all_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let remove_lock_all_tx = VersionedTransaction::try_new( + VersionedMessage::V0(remove_lock_all_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let remove_lock_all_result = context.svm.send_transaction(remove_lock_all_tx); + assert!( + remove_lock_all_result.is_err(), + "Authority with All permission should NOT be able to remove another role's authorization \ + lock" + ); + println!( + "✅ Scenario 1 PASSED: Authority with All permission correctly denied from removing \ + another role's lock" + ); + + // Verify the lock was NOT removed and count remains unchanged + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 3); + println!( + " → Verified: Authorization lock count remains = {}", + swig_with_roles.state.authorization_locks + ); + println!(); + + // Expire blockhash before next removal + context.svm.expire_blockhash(); + + // Test 2: Authority with ManageAuthorizationLocks permission should succeed + println!("EXECUTING TEST SCENARIO 2:"); + println!( + "Removing authorization lock at index 0 using authority with ManageAuthorizationLocks \ + permission (role 1) - removing their own lock" + ); + + let remove_lock_manage_ix = swig_interface::RemoveAuthorizationLockInstruction::new( + swig, + manage_auth_locks_authority.pubkey(), + context.default_payer.pubkey(), + 1, // acting_role_id: manage_auth_locks_authority (ManageAuthorizationLocks permission) + 0, // Remove the first lock (index 0, created by role 1) + ) + .unwrap(); + + let remove_lock_manage_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_lock_manage_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let remove_lock_manage_tx = VersionedTransaction::try_new( + VersionedMessage::V0(remove_lock_manage_message), + &[&context.default_payer, &manage_auth_locks_authority], + ) + .unwrap(); + + let remove_lock_manage_result = context.svm.send_transaction(remove_lock_manage_tx); + assert!( + remove_lock_manage_result.is_ok(), + "Authority with ManageAuthorizationLocks permission should be able to remove their own \ + authorization lock" + ); + println!( + "✅ Scenario 2 PASSED: Authority with ManageAuthorizationLocks permission successfully \ + removed their own authorization lock" + ); + + // Verify the lock was removed + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 2); + println!( + " → Verified: Authorization lock count = {}", + swig_with_roles.state.authorization_locks + ); + + // Should have locks with amounts 300 and 200 remaining (the 500 amount lock was + // removed) + let (final_locks, final_count) = swig_with_roles + .get_authorization_locks_for_test::<10>() + .unwrap(); + assert_eq!(final_count, 2); + assert_eq!(final_locks[0].unwrap().amount, 300); + assert_eq!(final_locks[1].unwrap().amount, 200); + println!(" → Verified: 300 and 200 amount locks remain (500 amount lock was removed)"); + println!(); + + // Expire blockhash before final test + context.svm.expire_blockhash(); + + // Test 3: Authority with only TokenLimit permission should fail + println!("EXECUTING TEST SCENARIO 3:"); + println!( + "Attempting to remove authorization lock using authority with TokenLimit permission only \ + (role 2)" + ); + println!("Expected: FAIL (insufficient permissions)"); + + let remove_lock_token_ix = swig_interface::RemoveAuthorizationLockInstruction::new( + swig, + token_authority.pubkey(), + context.default_payer.pubkey(), + 2, // acting_role_id: token_authority (TokenLimit permission only) + 0, // Try to remove the remaining lock + ) + .unwrap(); + + let remove_lock_token_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_lock_token_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let remove_lock_token_tx = VersionedTransaction::try_new( + VersionedMessage::V0(remove_lock_token_message), + &[&context.default_payer, &token_authority], + ) + .unwrap(); + + let remove_lock_token_result = context.svm.send_transaction(remove_lock_token_tx); + assert!( + remove_lock_token_result.is_err(), + "Authority with only TokenLimit permission should NOT be able to remove authorization lock" + ); + println!("✅ Scenario 3 PASSED: Authority with TokenLimit permission was correctly denied"); + + // Verify the lock count didn't change + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!( + swig_with_roles.state.authorization_locks, 2, + "Lock count should remain unchanged after failed attempt" + ); + println!( + " → Verified: Authorization lock count remains = {}", + swig_with_roles.state.authorization_locks + ); + println!(); + + println!("✅ REMOVE AUTHORIZATION LOCK PERMISSION ENFORCEMENT TEST COMPLETED SUCCESSFULLY!"); + println!("✅ Roles with All permissions can only remove their own authorization locks"); + println!( + "✅ Roles with ManageAuthorizationLocks permissions can only remove their own \ + authorization locks" + ); + println!("✅ Authorities with insufficient permissions are properly denied"); + println!("✅ Lock removal properly shifts remaining locks and updates count"); + println!("================================================================"); +} + +/// Test that validates edge cases and error conditions for remove authorization +/// lock +#[test_log::test] +fn test_remove_authorization_lock_edge_cases() { + let mut context = setup_test_context().unwrap(); + + // Setup accounts + let swig_authority = Keypair::new(); + + // Airdrop to participants + context + .svm + .airdrop(&swig_authority.pubkey(), 10_000_000_000) + .unwrap(); + + // Setup token mint + let mint_pubkey = setup_mint(&mut context.svm, &context.default_payer).unwrap(); + + // Create swig account + let id = rand::random::<[u8; 32]>(); + let (swig, _) = Pubkey::find_program_address(&swig_account_seeds(&id), &program_id()); + let swig_create_result = create_swig_ed25519(&mut context, &swig_authority, id); + assert!(swig_create_result.is_ok()); + + println!("=== REMOVE AUTHORIZATION LOCK EDGE CASES TEST ==="); + println!("Testing error conditions and edge cases"); + println!(); + + // Test 1: Try to remove lock when no locks exist + println!("TEST SCENARIO 1:"); + println!("Attempting to remove authorization lock when no locks exist"); + println!("Expected: FAIL (no locks to remove)"); + + let remove_empty_ix = swig_interface::RemoveAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id + 0, // lock_index + ) + .unwrap(); + + let remove_empty_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_empty_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let remove_empty_tx = VersionedTransaction::try_new( + VersionedMessage::V0(remove_empty_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let remove_empty_result = context.svm.send_transaction(remove_empty_tx); + assert!( + remove_empty_result.is_err(), + "Removing lock when no locks exist should fail" + ); + println!("✅ Scenario 1 PASSED: Correctly rejected removal when no locks exist"); + println!(); + + // Add a single lock for testing + let current_slot = context.svm.get_sysvar::().slot; + let expiry_slot = current_slot + 1000; + + let add_lock_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id + mint_pubkey.to_bytes(), + 100, + expiry_slot, + ) + .unwrap(); + + let add_lock_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock_result = context.svm.send_transaction(add_lock_tx); + assert!(add_lock_result.is_ok(), "Adding single lock should succeed"); + println!("✅ Added single authorization lock for edge case testing"); + + // Expire blockhash to avoid transaction replay issues + context.svm.expire_blockhash(); + + // Test 2: Try to remove lock with invalid index (out of bounds) + println!(); + println!("TEST SCENARIO 2:"); + println!( + "Attempting to remove authorization lock with invalid index (1 when only index 0 exists)" + ); + println!("Expected: FAIL (index out of bounds)"); + + let remove_invalid_ix = swig_interface::RemoveAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id + 1, // Invalid index (only 0 exists) + ) + .unwrap(); + + let remove_invalid_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_invalid_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let remove_invalid_tx = VersionedTransaction::try_new( + VersionedMessage::V0(remove_invalid_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let remove_invalid_result = context.svm.send_transaction(remove_invalid_tx); + assert!( + remove_invalid_result.is_err(), + "Removing lock with invalid index should fail" + ); + println!("✅ Scenario 2 PASSED: Correctly rejected removal with invalid index"); + println!(); + + // Expire blockhash before next transaction + context.svm.expire_blockhash(); + + // Test 3: Successfully remove the valid lock + println!("TEST SCENARIO 3:"); + println!("Removing the valid authorization lock at index 0"); + println!("Expected: PASS (valid removal)"); + + let remove_valid_ix = swig_interface::RemoveAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id + 0, // Valid index + ) + .unwrap(); + + let remove_valid_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_valid_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let remove_valid_tx = VersionedTransaction::try_new( + VersionedMessage::V0(remove_valid_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let remove_valid_result = context.svm.send_transaction(remove_valid_tx); + assert!( + remove_valid_result.is_ok(), + "Removing valid lock should succeed" + ); + println!("✅ Scenario 3 PASSED: Successfully removed valid authorization lock"); + + // Verify no locks remain + let swig_account = context.svm.get_account(&swig).unwrap(); + let swig_with_roles = SwigWithRoles::from_bytes(&swig_account.data).unwrap(); + assert_eq!(swig_with_roles.state.authorization_locks, 0); + println!( + " → Verified: No authorization locks remain (count = {})", + swig_with_roles.state.authorization_locks + ); + println!(); + + // Test 4: Try to remove lock with non-existent role ID + println!("TEST SCENARIO 4:"); + println!("Attempting to remove authorization lock using non-existent role ID (999)"); + println!("Expected: FAIL (role not found)"); + + // First add a lock again for this test + let add_lock_again_ix = swig_interface::AddAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 0, // acting_role_id + mint_pubkey.to_bytes(), + 150, + expiry_slot, + ) + .unwrap(); + + let add_lock_again_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[add_lock_again_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let add_lock_again_tx = VersionedTransaction::try_new( + VersionedMessage::V0(add_lock_again_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let add_lock_again_result = context.svm.send_transaction(add_lock_again_tx); + assert!( + add_lock_again_result.is_ok(), + "Adding lock for role test should succeed" + ); + + // Expire blockhash before next transaction + context.svm.expire_blockhash(); + + let remove_bad_role_ix = swig_interface::RemoveAuthorizationLockInstruction::new( + swig, + swig_authority.pubkey(), + context.default_payer.pubkey(), + 999, // Non-existent role ID + 0, // Valid lock index + ) + .unwrap(); + + let remove_bad_role_message = v0::Message::try_compile( + &context.default_payer.pubkey(), + &[remove_bad_role_ix], + &[], + context.svm.latest_blockhash(), + ) + .unwrap(); + + let remove_bad_role_tx = VersionedTransaction::try_new( + VersionedMessage::V0(remove_bad_role_message), + &[&context.default_payer, &swig_authority], + ) + .unwrap(); + + let remove_bad_role_result = context.svm.send_transaction(remove_bad_role_tx); + assert!( + remove_bad_role_result.is_err(), + "Using non-existent role ID should fail" + ); + println!("✅ Scenario 4 PASSED: Non-existent role ID was correctly rejected"); + + println!(); + println!("✅ REMOVE AUTHORIZATION LOCK EDGE CASES TEST COMPLETED SUCCESSFULLY!"); + println!("================================================================"); +} diff --git a/state-x/src/action/manage_authorization_locks.rs b/state-x/src/action/manage_authorization_locks.rs new file mode 100644 index 00000000..9863208f --- /dev/null +++ b/state-x/src/action/manage_authorization_locks.rs @@ -0,0 +1,50 @@ +//! Manage authorization locks action type. +//! +//! This module defines the ManageAuthorizationLocks action type which grants +//! permission to add and remove authorization locks within the Swig wallet +//! system. + +use no_padding::NoPadding; +use pinocchio::program_error::ProgramError; + +use super::{Actionable, Permission}; +use crate::{IntoBytes, Transmutable, TransmutableMut}; + +/// Represents permission to manage authorization locks. +/// +/// This action grants the authority to add and remove authorization locks +/// for any token mint. It's a powerful permission that should be granted +/// carefully as it allows control over payment preauthorization limits. +#[repr(C, align(8))] +#[derive(Debug, NoPadding)] +pub struct ManageAuthorizationLocks {} + +impl Transmutable for ManageAuthorizationLocks { + /// Size of the ManageAuthorizationLocks struct in bytes (empty struct) + const LEN: usize = 1; // Minimum size for empty struct +} + +impl TransmutableMut for ManageAuthorizationLocks {} + +impl IntoBytes for ManageAuthorizationLocks { + fn into_bytes(&self) -> Result<&[u8], ProgramError> { + let bytes = + unsafe { core::slice::from_raw_parts(self as *const Self as *const u8, Self::LEN) }; + Ok(bytes) + } +} + +impl<'a> Actionable<'a> for ManageAuthorizationLocks { + /// This action represents the ManageAuthorizationLocks permission type + const TYPE: Permission = Permission::ManageAuthorizationLocks; + /// Only one ManageAuthorizationLocks permission per role is needed + const REPEATABLE: bool = false; + + /// No specific data matching required for this permission. + /// + /// # Arguments + /// * `_data` - Unused data parameter + fn match_data(&self, _data: &[u8]) -> bool { + true // This permission applies globally, no specific data matching + } +} diff --git a/state-x/src/action/mod.rs b/state-x/src/action/mod.rs index 8e945a1f..6d3b5b89 100644 --- a/state-x/src/action/mod.rs +++ b/state-x/src/action/mod.rs @@ -7,6 +7,7 @@ pub mod all; pub mod manage_authority; +pub mod manage_authorization_locks; pub mod program; pub mod program_scope; pub mod sol_limit; @@ -19,6 +20,7 @@ pub mod token_limit; pub mod token_recurring_limit; use all::All; use manage_authority::ManageAuthority; +use manage_authorization_locks::ManageAuthorizationLocks; use no_padding::NoPadding; use pinocchio::program_error::ProgramError; use program::Program; @@ -130,6 +132,8 @@ pub enum Permission { StakeRecurringLimit = 11, /// Permission to perform all stake operations StakeAll = 12, + /// Permission to manage authorization locks + ManageAuthorizationLocks = 13, } impl TryFrom for Permission { @@ -139,7 +143,7 @@ impl TryFrom for Permission { fn try_from(value: u16) -> Result { match value { // SAFETY: `value` is guaranteed to be in the range of the enum variants. - 0..=12 => Ok(unsafe { core::mem::transmute::(value) }), + 0..=13 => Ok(unsafe { core::mem::transmute::(value) }), _ => Err(SwigStateError::PermissionLoadError.into()), } } @@ -196,6 +200,7 @@ impl ActionLoader { Permission::StakeLimit => StakeLimit::valid_layout(data), Permission::StakeRecurringLimit => StakeRecurringLimit::valid_layout(data), Permission::StakeAll => StakeAll::valid_layout(data), + Permission::ManageAuthorizationLocks => ManageAuthorizationLocks::valid_layout(data), _ => Ok(false), } } diff --git a/state-x/src/action/token_limit.rs b/state-x/src/action/token_limit.rs index 45859a34..b94e9fe2 100644 --- a/state-x/src/action/token_limit.rs +++ b/state-x/src/action/token_limit.rs @@ -5,7 +5,7 @@ //! to a particular token mint. use no_padding::NoPadding; -use pinocchio::program_error::ProgramError; +use pinocchio::{msg, program_error::ProgramError}; use super::{Actionable, Permission}; use crate::{IntoBytes, SwigAuthenticateError, Transmutable, TransmutableMut}; diff --git a/state-x/src/constants.rs b/state-x/src/constants.rs index dc2b63b6..c9041da1 100644 --- a/state-x/src/constants.rs +++ b/state-x/src/constants.rs @@ -6,3 +6,8 @@ /// This is used for memory allocation and validation when handling program /// scope actions. pub const PROGRAM_SCOPE_BYTE_SIZE: usize = 144; + +/// Size in bytes of an authorization lock data structure. +/// This is used for memory allocation and validation when handling +/// authorization lock actions. +pub const AUTHORIZATION_LOCK_BYTE_SIZE: usize = 56; diff --git a/state-x/src/swig.rs b/state-x/src/swig.rs index 311e67e8..86545a2f 100644 --- a/state-x/src/swig.rs +++ b/state-x/src/swig.rs @@ -103,6 +103,34 @@ impl IntoBytes for SwigSubAccount { } } +/// Represents an authorization lock for pre-authorizing token spending limits. +#[repr(C, align(8))] +#[derive(Debug, PartialEq, Copy, Clone, NoPadding)] +pub struct AuthorizationLock { + /// Token mint public key that this lock applies to + pub token_mint: [u8; 32], + /// Maximum amount that can be spent + pub amount: u64, + /// Slot number when this lock expires + pub expiry_slot: u64, + /// Role ID that created this authorization lock + pub role_id: u32, + /// Padding to ensure struct has no padding + pub _padding: [u8; 4], +} + +impl Transmutable for AuthorizationLock { + const LEN: usize = core::mem::size_of::(); +} + +impl TransmutableMut for AuthorizationLock {} + +impl IntoBytes for AuthorizationLock { + fn into_bytes(&self) -> Result<&[u8], ProgramError> { + Ok(unsafe { core::slice::from_raw_parts(self as *const Self as *const u8, Self::LEN) }) + } +} + /// Builder for constructing and modifying Swig accounts. pub struct SwigBuilder<'a> { /// Buffer for role data @@ -342,6 +370,10 @@ pub struct Swig { pub role_counter: u32, /// Amount of lamports reserved for rent pub reserved_lamports: u64, + /// Number of authorization locks in this account + pub authorization_locks: u16, + /// Reserved bytes for future use + _reserved: [u8; 6], } impl Swig { @@ -354,6 +386,8 @@ impl Swig { roles: 0, role_counter: 0, reserved_lamports, + authorization_locks: 0, + _reserved: [0; 6], } } @@ -428,12 +462,12 @@ impl IntoBytes for Swig { } } -/// Wrapper structure for a Swig account with its roles. +/// Wrapper structure for a Swig account with its roles and authorization locks. pub struct SwigWithRoles<'a> { /// Reference to the Swig account state pub state: &'a Swig, - /// Raw bytes containing role data - roles: &'a [u8], + /// Raw bytes containing all data after the Swig header (roles + auth locks) + data: &'a [u8], } impl<'a> SwigWithRoles<'a> { @@ -444,19 +478,44 @@ impl<'a> SwigWithRoles<'a> { } let state = unsafe { Swig::load_unchecked(&bytes[..Swig::LEN])? }; - let roles = &bytes[Swig::LEN..]; + let data = &bytes[Swig::LEN..]; - Ok(SwigWithRoles { state, roles }) + Ok(SwigWithRoles { state, data }) + } + + /// Gets the roles data slice from the combined data. + fn roles_data(&self) -> &[u8] { + let mut cursor = 0; + for _i in 0..self.state.roles { + if cursor + Position::LEN > self.data.len() { + return &[]; + } + let position = + unsafe { Position::load_unchecked(&self.data[cursor..cursor + Position::LEN]) }; + if let Ok(pos) = position { + cursor = pos.boundary() as usize; + } else { + return &[]; + } + } + &self.data[..cursor] + } + + /// Gets the authorization locks data slice from the combined data. + fn authorization_locks_data(&self) -> &[u8] { + let roles_end = self.roles_data().len(); + &self.data[roles_end..] } /// Looks up a role ID by authority data. pub fn lookup_role_id(&'a self, authority_data: &'a [u8]) -> Result, ProgramError> { + let roles = self.roles_data(); let mut cursor = 0; for _i in 0..self.state.roles { let offset = cursor + Position::LEN; let position = - unsafe { Position::load_unchecked(self.roles.get_unchecked(cursor..offset))? }; + unsafe { Position::load_unchecked(roles.get_unchecked(cursor..offset))? }; let auth_type = position.authority_type()?; let auth_len = position.authority_length() as usize; @@ -464,22 +523,22 @@ impl<'a> SwigWithRoles<'a> { let authority: &dyn AuthorityInfo = match auth_type { AuthorityType::Ed25519 => unsafe { ED25519Authority::load_unchecked( - self.roles.get_unchecked(offset..offset + auth_len), + roles.get_unchecked(offset..offset + auth_len), )? }, AuthorityType::Ed25519Session => unsafe { Ed25519SessionAuthority::load_unchecked( - self.roles.get_unchecked(offset..offset + auth_len), + roles.get_unchecked(offset..offset + auth_len), )? }, AuthorityType::Secp256k1 => unsafe { Secp256k1Authority::load_unchecked( - self.roles.get_unchecked(offset..offset + auth_len), + roles.get_unchecked(offset..offset + auth_len), )? }, AuthorityType::Secp256k1Session => unsafe { Secp256k1SessionAuthority::load_unchecked( - self.roles.get_unchecked(offset..offset + auth_len), + roles.get_unchecked(offset..offset + auth_len), )? }, @@ -502,31 +561,32 @@ impl<'a> SwigWithRoles<'a> { /// Gets a reference to a role by ID. pub fn get_role(&'a self, id: u32) -> Result>, ProgramError> { + let roles = self.roles_data(); let mut cursor = 0; for _i in 0..self.state.roles { let offset = cursor + Position::LEN; let position = - unsafe { Position::load_unchecked(self.roles.get_unchecked(cursor..offset))? }; + unsafe { Position::load_unchecked(roles.get_unchecked(cursor..offset))? }; if position.id() == id { let authority: &dyn AuthorityInfo = match position.authority_type()? { AuthorityType::Ed25519 => unsafe { - ED25519Authority::load_unchecked(self.roles.get_unchecked( + ED25519Authority::load_unchecked(roles.get_unchecked( offset..offset + position.authority_length() as usize, ))? }, AuthorityType::Ed25519Session => unsafe { - Ed25519SessionAuthority::load_unchecked(self.roles.get_unchecked( + Ed25519SessionAuthority::load_unchecked(roles.get_unchecked( offset..offset + position.authority_length() as usize, ))? }, AuthorityType::Secp256k1 => unsafe { - Secp256k1Authority::load_unchecked(self.roles.get_unchecked( + Secp256k1Authority::load_unchecked(roles.get_unchecked( offset..offset + position.authority_length() as usize, ))? }, AuthorityType::Secp256k1Session => unsafe { - Secp256k1SessionAuthority::load_unchecked(self.roles.get_unchecked( + Secp256k1SessionAuthority::load_unchecked(roles.get_unchecked( offset..offset + position.authority_length() as usize, ))? }, @@ -537,7 +597,7 @@ impl<'a> SwigWithRoles<'a> { position, authority, actions: unsafe { - self.roles.get_unchecked( + roles.get_unchecked( offset + position.authority_length() as usize ..position.boundary() as usize, ) @@ -603,8 +663,254 @@ impl<'a> SwigWithRoles<'a> { } None } + + /// Iterates over all authorization locks from the account, calling the + /// provided function for each lock. This is zero-copy - passes direct + /// references to the raw lock data. + pub fn for_each_authorization_lock(&self, mut f: F) -> Result<(), E> + where + F: FnMut(&AuthorizationLock) -> Result<(), E>, + E: From, + { + let auth_locks_data = self.authorization_locks_data(); + + let expected_size = self.state.authorization_locks as usize * AuthorizationLock::LEN; + if auth_locks_data.len() < expected_size { + return Err(ProgramError::InvalidAccountData.into()); + } + + let mut cursor = 0; + for _i in 0..self.state.authorization_locks { + if cursor + AuthorizationLock::LEN > auth_locks_data.len() { + break; + } + // Zero-copy: cast the raw bytes directly to a reference + let lock = unsafe { + &*(auth_locks_data[cursor..cursor + AuthorizationLock::LEN].as_ptr() + as *const AuthorizationLock) + }; + f(lock)?; + cursor += AuthorizationLock::LEN; + } + + Ok(()) + } + + /// Iterates over authorization locks for a specific token mint, calling the + /// provided function for each matching lock. + pub fn for_each_authorization_lock_by_mint( + &self, + mint: &[u8; 32], + mut f: F, + ) -> Result<(), E> + where + F: FnMut(&AuthorizationLock) -> Result<(), E>, + E: From, + { + self.for_each_authorization_lock(|lock| { + if &lock.token_mint == mint { + f(lock) + } else { + Ok(()) + } + }) + } + + /// Helper method for tests to get authorization locks in a fixed-size + /// array. Only collects up to MAX_LOCKS authorization locks for testing + /// purposes. + pub fn get_authorization_locks_for_test( + &self, + ) -> Result<([Option; MAX_LOCKS], usize), ProgramError> { + let mut locks = [None; MAX_LOCKS]; + let mut count = 0; + + self.for_each_authorization_lock::<_, ProgramError>(|lock| { + if count < MAX_LOCKS { + locks[count] = Some(*lock); + count += 1; + } + Ok(()) + })?; + + Ok((locks, count)) + } + + /// Gets authorization locks created by a specific role ID. + /// Returns a tuple of (locks array, count) where count is the number of + /// locks found. + pub fn get_authorization_locks_by_role( + &self, + role_id: u32, + ) -> Result<([Option; MAX_LOCKS], usize), ProgramError> { + let mut locks = [None; MAX_LOCKS]; + let mut count = 0; + + self.for_each_authorization_lock::<_, ProgramError>(|lock| { + if lock.role_id == role_id && count < MAX_LOCKS { + locks[count] = Some(*lock); + count += 1; + } + Ok(()) + })?; + + Ok((locks, count)) + } + + /// Iterates over authorization locks for a specific role ID and applies a + /// function to each. This is useful for operations that need to process + /// locks without collecting them into an array. + pub fn for_each_authorization_lock_by_role(&self, role_id: u32, mut f: F) -> Result<(), E> + where + F: FnMut(&AuthorizationLock) -> Result<(), E>, + E: From, + { + self.for_each_authorization_lock::<_, E>(|lock| { + if lock.role_id == role_id { + f(lock) + } else { + Ok(()) + } + }) + } + + /// Gets a zero-copy reference to an authorization lock by index. + /// Returns None if the index is out of bounds. + pub fn get_authorization_lock_by_index(&self, index: usize) -> Option<&AuthorizationLock> { + if index >= self.state.authorization_locks as usize { + return None; + } + + let auth_locks_data = self.authorization_locks_data(); + let lock_offset = index * AuthorizationLock::LEN; + + if lock_offset + AuthorizationLock::LEN > auth_locks_data.len() { + return None; + } + + // Zero-copy: cast the raw bytes directly to a reference + unsafe { + Some( + &*(auth_locks_data[lock_offset..lock_offset + AuthorizationLock::LEN].as_ptr() + as *const AuthorizationLock), + ) + } + } + + /// Zero-copy iterator over authorization locks. + /// Returns an iterator that yields direct references to authorization locks + /// in memory. + pub fn authorization_locks_iter(&self) -> AuthorizationLockIterator { + AuthorizationLockIterator { + data: self.authorization_locks_data(), + count: self.state.authorization_locks as usize, + current: 0, + } + } + + /// Removes expired authorization locks from the account. + /// Takes mutable references to state and data to allow modification. + /// Returns the number of locks removed. + pub fn remove_expired_authorization_locks_mut( + state: &mut Swig, + data: &mut [u8], + current_slot: u64, + ) -> Result { + let auth_locks_count = state.authorization_locks; + + // Calculate where authorization locks start (after roles data) + let mut roles_cursor = 0; + for _i in 0..state.roles { + if roles_cursor + Position::LEN > data.len() { + break; + } + let position = unsafe { + Position::load_unchecked(&data[roles_cursor..roles_cursor + Position::LEN])? + }; + roles_cursor = position.boundary() as usize; + } + + let auth_locks_data = &mut data[roles_cursor..]; + let mut removed_count = 0u16; + let mut write_cursor = 0; + let mut read_cursor = 0; + + // Iterate through all authorization locks + for _i in 0..auth_locks_count { + if read_cursor + AuthorizationLock::LEN > auth_locks_data.len() { + break; + } + + // Zero-copy: cast the raw bytes directly to a reference + let lock = unsafe { + &*(auth_locks_data[read_cursor..read_cursor + AuthorizationLock::LEN].as_ptr() + as *const AuthorizationLock) + }; + + // If lock is not expired, copy it to the write position + if lock.expiry_slot > current_slot { + if write_cursor != read_cursor { + auth_locks_data.copy_within( + read_cursor..read_cursor + AuthorizationLock::LEN, + write_cursor, + ); + } + write_cursor += AuthorizationLock::LEN; + } else { + // Lock is expired, don't copy it (effectively removing it) + removed_count += 1; + } + + read_cursor += AuthorizationLock::LEN; + } + + // Update the authorization locks count + state.authorization_locks -= removed_count; + + Ok(removed_count) + } +} + +/// Zero-copy iterator over authorization locks. +/// Yields direct references to authorization locks stored in the account data. +pub struct AuthorizationLockIterator<'a> { + data: &'a [u8], + count: usize, + current: usize, +} + +impl<'a> Iterator for AuthorizationLockIterator<'a> { + type Item = &'a AuthorizationLock; + + fn next(&mut self) -> Option { + if self.current >= self.count { + return None; + } + + let lock_offset = self.current * AuthorizationLock::LEN; + + if lock_offset + AuthorizationLock::LEN > self.data.len() { + return None; + } + + // Zero-copy: cast the raw bytes directly to a reference + let lock = unsafe { + &*(self.data[lock_offset..lock_offset + AuthorizationLock::LEN].as_ptr() + as *const AuthorizationLock) + }; + + self.current += 1; + Some(lock) + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self.count - self.current; + (remaining, Some(remaining)) + } } +impl<'a> ExactSizeIterator for AuthorizationLockIterator<'a> {} + #[cfg(test)] mod tests { use super::*;