diff --git a/dlp-api/src/v2/args/mod.rs b/dlp-api/src/v2/args/mod.rs index 1c6dd317..d7ba57a9 100644 --- a/dlp-api/src/v2/args/mod.rs +++ b/dlp-api/src/v2/args/mod.rs @@ -6,9 +6,11 @@ mod register_operator; mod register_verifier; mod update_protocol_config; mod update_verifier_registry; +mod write_state_buffer; pub use init_protocol_config::*; pub use register_operator::*; pub use register_verifier::*; pub use update_protocol_config::*; pub use update_verifier_registry::*; +pub use write_state_buffer::*; diff --git a/dlp-api/src/v2/args/write_state_buffer.rs b/dlp-api/src/v2/args/write_state_buffer.rs new file mode 100644 index 00000000..ce6781bc --- /dev/null +++ b/dlp-api/src/v2/args/write_state_buffer.rs @@ -0,0 +1,15 @@ +use wheels::variable_offset_layout; + +#[derive(Clone, Debug, PartialEq, Eq)] +#[variable_offset_layout(buffer_offset = 1)] +pub struct WriteStateBufferArgs { + pub commit_id: u64, + + pub total_len: u32, + + /// Must equal bytes already written, unless retrying an exact old chunk. + pub offset: u32, + + #[flexible = 4] + pub chunk: Vec, +} diff --git a/dlp-api/src/v2/instruction.rs b/dlp-api/src/v2/instruction.rs index 03e19eb7..a18e8434 100644 --- a/dlp-api/src/v2/instruction.rs +++ b/dlp-api/src/v2/instruction.rs @@ -15,6 +15,12 @@ pub enum DlpV2Instruction { UpdateVerifierRegistry = 103, /// Updates global v2 config for future commitments. UpdateProtocolConfig = 104, + /// Writes full account-state bytes into a v2 state buffer. + /// + /// TODO (snawaz/optimization): we can split this into two instructions such that + /// InitStateBuffer takes more arguments and AppendStateBuffer takes as less as + /// possible. + WriteStateBuffer = 107, } impl DlpV2Instruction { diff --git a/dlp-api/src/v2/instruction_builder/mod.rs b/dlp-api/src/v2/instruction_builder/mod.rs index d93de000..128e03ff 100644 --- a/dlp-api/src/v2/instruction_builder/mod.rs +++ b/dlp-api/src/v2/instruction_builder/mod.rs @@ -3,9 +3,11 @@ mod register_operator; mod register_verifier; mod update_protocol_config; mod update_verifier_registry; +mod write_state_buffer; pub use init_protocol_config::*; pub use register_operator::*; pub use register_verifier::*; pub use update_protocol_config::*; pub use update_verifier_registry::*; +pub use write_state_buffer::*; diff --git a/dlp-api/src/v2/instruction_builder/write_state_buffer.rs b/dlp-api/src/v2/instruction_builder/write_state_buffer.rs new file mode 100644 index 00000000..cbe54287 --- /dev/null +++ b/dlp-api/src/v2/instruction_builder/write_state_buffer.rs @@ -0,0 +1,47 @@ +use solana_program::{ + instruction::{AccountMeta, Instruction}, + pubkey::Pubkey, +}; +use solana_sdk_ids::system_program; +use wheels::layout::Encodable; + +use crate::{ + compat::{Compatize, Modernize}, + v2::{ + pda::{protocol_config_pda, state_buffer_pda}, + DlpV2Instruction, WriteStateBufferArgs, + }, +}; + +/// Builds the instruction that writes full account-state bytes to a v2 buffer. +pub fn write_state_buffer( + payer: Pubkey, + authority: Pubkey, + account: Pubkey, + args: WriteStateBufferArgs, +) -> Instruction { + Instruction { + program_id: crate::id().modernize(), + accounts: vec![ + AccountMeta::new(payer, true), + AccountMeta::new_readonly(authority, true), + AccountMeta::new( + state_buffer_pda( + &account.compatize(), + args.commit_id, + &authority.compatize(), + ) + .modernize(), + false, + ), + AccountMeta::new_readonly(account, false), + AccountMeta::new_readonly(protocol_config_pda().modernize(), false), + AccountMeta::new_readonly(system_program::id(), false), + ], + data: [ + DlpV2Instruction::WriteStateBuffer.to_vec(), + args.encode().unwrap(), + ] + .concat(), + } +} diff --git a/dlp-api/src/v2/pda.rs b/dlp-api/src/v2/pda.rs index d446aab4..c783480b 100644 --- a/dlp-api/src/v2/pda.rs +++ b/dlp-api/src/v2/pda.rs @@ -4,6 +4,7 @@ pub const PROTOCOL_CONFIG_SEED: &[u8] = b"protocol-config"; pub const OPERATOR_BOND_SEED: &[u8] = b"operator-bond"; pub const VERIFIER_BOND_SEED: &[u8] = b"verifier-bond"; pub const VERIFIER_REGISTRY_SEED: &[u8] = b"verifier-registry"; +pub const STATE_BUFFER_SEED: &[u8] = b"state-buffer"; // TODO (snawaz): Precompute these addresses if PDA derivation becomes const-safe. @@ -30,3 +31,20 @@ pub fn verifier_bond_pda(verifier: &Pubkey) -> Pubkey { ) .0 } + +pub fn state_buffer_pda( + account: &Pubkey, + commit_id: u64, + authority: &Pubkey, +) -> Pubkey { + Pubkey::find_program_address( + &[ + STATE_BUFFER_SEED, + account.as_ref(), + &commit_id.to_le_bytes(), + authority.as_ref(), + ], + &crate::id(), + ) + .0 +} diff --git a/dlp-api/src/v2/state/mod.rs b/dlp-api/src/v2/state/mod.rs index ee5e1d3c..a92b8062 100644 --- a/dlp-api/src/v2/state/mod.rs +++ b/dlp-api/src/v2/state/mod.rs @@ -1,9 +1,11 @@ mod operator_bond; mod protocol_config; +mod state_buffer; mod verifier_bond; mod verifier_registry; pub use operator_bond::*; pub use protocol_config::*; +pub use state_buffer::*; pub use verifier_bond::*; pub use verifier_registry::*; diff --git a/dlp-api/src/v2/state/state_buffer.rs b/dlp-api/src/v2/state/state_buffer.rs new file mode 100644 index 00000000..594b25f9 --- /dev/null +++ b/dlp-api/src/v2/state/state_buffer.rs @@ -0,0 +1,81 @@ +use wheels::fixed_offset_layout; + +use crate::{compat::Pubkey, error::DlpError}; + +/// PDA: `["state-buffer", account, commit_id, authority]`. +/// Created by `WriteStateBuffer`. +/// Closed by `CloseTerminalAccounts` after finalize, cancel, or expiry. +#[derive(Clone, Debug, PartialEq, Eq)] +#[fixed_offset_layout(buffer_offset = 0)] +pub struct StateBuffer { + /// Account type marker. + pub discriminator: [u8; 8], + + /// Writer that owns this opened buffer. + /// + /// This is the operator identity for an operator commitment buffer, or the + /// challenger identity for a challenger dispute buffer. + pub authority: Pubkey, + + /// Delegated account whose payload is stored after this header. + pub account_pubkey: Pubkey, + + /// Flow-specific nonce that identifies this opened buffer. + pub commit_id: u64, + + /// Hash of the finalized payload. Zero until finalized. + pub data_hash: [u8; 32], + + /// Expected final byte length of the payload. + pub total_len: u32, + + /// Once true, buffer content cannot change except exact duplicate retries. + pub finalized: bool, + + /// Active payload bytes for `PostCommitment` or `RaiseChallenge`. + /// + /// This is not a serialized Solana account. Depending on the flow, it can + /// be the delegated account's complete `Account::data` bytes or an encoded + /// diff of those bytes. `payload.len()` is the written prefix, while + /// `payload.capacity()` is the allocated account-backed span and can be + /// shorter than `total_len` until later writes grow the buffer account. + #[extendable = 4] + pub payload: Vec, +} + +impl StateBuffer { + pub const DISCRIMINATOR: [u8; 8] = *b"v2sbuf00"; + + /// Maximum account data bytes a StateBuffer PDA may allocate. + pub const MAX_ACCOUNT_DATA_LEN: usize = 10 * 1024 * 1024; + + /// Maximum account data bytes a StateBuffer PDA may grow in one write. + pub const MAX_ACCOUNT_DATA_GROWTH_PER_WRITE: usize = 10_240; + + /// Offset of the extendable payload length header. + pub const PAYLOAD_LEN_HEADER_OFFSET: usize = Self::MIN_DATA_LEN; + + /// Byte length of the extendable payload length header. + pub const PAYLOAD_LEN_HEADER_LEN: usize = 4; + + /// Offset where payload bytes begin in account data. + pub const PAYLOAD_BYTES_OFFSET: usize = + Self::PAYLOAD_LEN_HEADER_OFFSET + Self::PAYLOAD_LEN_HEADER_LEN; + + /// Maximum payload bytes allocated when a StateBuffer PDA is created. + pub const MAX_INITIAL_PAYLOAD_LEN: usize = + Self::MAX_ACCOUNT_DATA_GROWTH_PER_WRITE - Self::PAYLOAD_BYTES_OFFSET; + + /// Maximum payload bytes accepted across all writes. + pub const MAX_TOTAL_PAYLOAD_LEN: u32 = + (Self::MAX_ACCOUNT_DATA_LEN - Self::PAYLOAD_BYTES_OFFSET) as u32; + + /// Returns the serialized data length needed for a payload capacity. + pub fn data_len_from_payload_capacity( + payload_capacity: usize, + ) -> Result { + Self::PAYLOAD_BYTES_OFFSET + .checked_add(payload_capacity) + .ok_or(DlpError::Overflow) + } +} diff --git a/src/processor/fast/utils/pda.rs b/src/processor/fast/utils/pda.rs index 08b4e1ec..7a330d3a 100644 --- a/src/processor/fast/utils/pda.rs +++ b/src/processor/fast/utils/pda.rs @@ -16,15 +16,33 @@ pub(crate) fn create_pda( pda_signers: &[Signer], payer: &AccountView, ) -> ProgramResult { - // Create the account manually or using the create instruction + create_pda_with_rent_exempt_lamports( + target_account, + owner, + space, + Rent::get()?.try_minimum_balance(space)?, + pda_signers, + payer, + ) +} - let rent = Rent::get()?; +/// Creates a new PDA with an explicit rent-exempt lamport target. +#[inline(always)] +pub(crate) fn create_pda_with_rent_exempt_lamports( + target_account: &AccountView, + owner: &Address, + space: usize, + rent_exempt_lamports: u64, + pda_signers: &[Signer], + payer: &AccountView, +) -> ProgramResult { + // Create the account manually or using the create instruction if target_account.lamports().eq(&0) { // If balance is zero, create account system::CreateAccount { from: payer, to: target_account, - lamports: rent.try_minimum_balance(space)?, + lamports: rent_exempt_lamports, space: space as u64, owner, } @@ -33,9 +51,8 @@ pub(crate) fn create_pda( // Otherwise, if balance is nonzero: // 1) transfer sufficient lamports for rent exemption - let rent_exempt_balance = rent - .try_minimum_balance(space)? - .saturating_sub(target_account.lamports()); + let rent_exempt_balance = + rent_exempt_lamports.saturating_sub(target_account.lamports()); if rent_exempt_balance > 0 { system::Transfer { from: payer, diff --git a/src/v2/processor/fraud_proofs/mod.rs b/src/v2/processor/fraud_proofs/mod.rs index c6d8c16b..bf64af2d 100644 --- a/src/v2/processor/fraud_proofs/mod.rs +++ b/src/v2/processor/fraud_proofs/mod.rs @@ -1 +1,5 @@ //! Processors for v2 fraud-proof instructions. + +mod write_state_buffer; + +pub use write_state_buffer::*; diff --git a/src/v2/processor/fraud_proofs/write_state_buffer.rs b/src/v2/processor/fraud_proofs/write_state_buffer.rs new file mode 100644 index 00000000..6d273fd3 --- /dev/null +++ b/src/v2/processor/fraud_proofs/write_state_buffer.rs @@ -0,0 +1,354 @@ +use dlp_api::{ + error::DlpError, + v2::{ + pda::{PROTOCOL_CONFIG_SEED, STATE_BUFFER_SEED}, + ProtocolConfig, StateBuffer, WriteStateBufferArgs, + }, +}; +use pinocchio::{ + cpi::{Seed, Signer}, + error::ProgramError, + sysvars::{rent::Rent, Sysvar}, + AccountView, ProgramResult, +}; +use wheels::{ + layout::{Decodable, Encodable, MaxLenStorage}, + require_eq, require_eq_keys, require_le, require_n_accounts, require_ne, + require_owned_by, require_signer, +}; + +use crate::{ + processor::fast::utils::pda::create_pda_with_rent_exempt_lamports, + requires::{ + is_uninitialized_account, require_initialized_pda, + require_uninitialized_pda, StandardCtx, + }, +}; + +/// Write payload bytes into a DLP-owned v2 StateBuffer. +/// +/// Accounts: +/// 0: `[signer, writable]` payer for StateBuffer rent +/// 1: `[signer]` authority identity for this buffer +/// 2: `[writable]` StateBuffer PDA +/// 3: `[]` delegated account whose bytes are being uploaded +/// 4: `[]` ProtocolConfig PDA +/// 5: `[]` system program, required by system CPI +#[inline(never)] +pub fn process_write_state_buffer( + accounts: &[AccountView], + data: &[u8], +) -> ProgramResult { + let [ + payer, // force multi-line + authority, + state_buffer, + delegated_account, + protocol_config, + _system_program, + ] = require_n_accounts!(accounts, 6); + + require_signer!(payer); + require_signer!(authority); + require_owned_by!(delegated_account, &crate::fast::ID); + + let args = WriteStateBufferArgs::decode(data)?; + let (offset, write_end) = validate_args(&args)?; + + validate_protocol_config(protocol_config)?; + + let commit_id_bytes = args.commit_id().to_le_bytes(); + let state_buffer_seeds = [ + STATE_BUFFER_SEED, + delegated_account.address().as_ref(), + &commit_id_bytes, + authority.address().as_ref(), + ]; + + if is_uninitialized_account(state_buffer) { + require_eq!(offset, 0, ProgramError::InvalidInstructionData); + + let state_buffer_bump = require_uninitialized_pda( + state_buffer, + &state_buffer_seeds, + &crate::fast::ID, + true, + StandardCtx::new("state buffer"), + )?; + + let initial_payload_capacity = (args.total_len() as usize) + .min(StateBuffer::MAX_INITIAL_PAYLOAD_LEN); + + require_le!( + write_end, + initial_payload_capacity, + ProgramError::InvalidInstructionData + ); + + let initial_len = StateBuffer::data_len_from_payload_capacity( + initial_payload_capacity, + )?; + let final_len = StateBuffer::data_len_from_payload_capacity( + args.total_len() as usize, + )?; + + // Note that rent_exempt_lamports is computed using final_len, not initial_len, + // because later writes can then avoid rent calculation and top-up transfer. + let rent_exempt_lamports = + Rent::get()?.try_minimum_balance(final_len)?; + + create_pda_with_rent_exempt_lamports( + state_buffer, + &crate::fast::ID, + initial_len, + rent_exempt_lamports, + &[Signer::from(&[ + Seed::from(STATE_BUFFER_SEED), + Seed::from(delegated_account.address().as_ref()), + Seed::from(&commit_id_bytes), + Seed::from(authority.address().as_ref()), + Seed::from(&[state_buffer_bump]), + ])], + payer, + )?; + + StateBuffer { + discriminator: StateBuffer::DISCRIMINATOR, + authority: authority.address().clone(), + account_pubkey: delegated_account.address().clone(), + commit_id: args.commit_id(), + data_hash: [0; 32], + total_len: args.total_len(), + finalized: false, + payload: Vec::new(), + } + .encode_to(state_buffer.try_borrow_mut()?.as_mut())?; + } else { + require_initialized_pda( + state_buffer, + &state_buffer_seeds, + &crate::fast::ID, + true, + "state buffer", + )?; + } + + write_chunk( + authority, + delegated_account, + state_buffer, + &args, + offset, + write_end, + ) +} + +fn validate_args( + args: &dlp_api::v2::WriteStateBufferArgsView<'_>, +) -> Result<(usize, usize), ProgramError> { + require_ne!(args.total_len(), 0, ProgramError::InvalidInstructionData); + require_le!( + args.total_len(), + StateBuffer::MAX_TOTAL_PAYLOAD_LEN, + ProgramError::InvalidInstructionData + ); + require_ne!(args.chunk().len(), 0, ProgramError::InvalidInstructionData); + + let offset = args.offset() as usize; + let write_end = offset + .checked_add(args.chunk().len()) + .ok_or(DlpError::Overflow)?; + + require_le!( + write_end, + args.total_len() as usize, + ProgramError::InvalidInstructionData + ); + + Ok((offset, write_end)) +} + +fn validate_protocol_config(protocol_config: &AccountView) -> ProgramResult { + require_initialized_pda( + protocol_config, + &[PROTOCOL_CONFIG_SEED], + &crate::fast::ID, + false, + "protocol config", + )?; + + let protocol_config_data = protocol_config.try_borrow()?; + let protocol_config_state = + ProtocolConfig::decode(protocol_config_data.as_ref())?; + require_eq!( + &protocol_config_state.discriminator(), + &ProtocolConfig::DISCRIMINATOR, + ProgramError::InvalidAccountData + ); + require_eq!( + protocol_config_state.paused(), + false, + ProgramError::InvalidAccountData + ); + + Ok(()) +} + +fn write_chunk( + authority: &AccountView, + delegated_account: &AccountView, + state_buffer_account: &AccountView, + args: &dlp_api::v2::WriteStateBufferArgsView<'_>, + offset: usize, + write_end: usize, +) -> ProgramResult { + let state_fields = validate_state_buffer( + state_buffer_account, + authority, + delegated_account, + args, + )?; + + // Treat exact overlap as an idempotent retry when the prior write landed. + if offset < state_fields.payload_len { + return validate_duplicate_chunk( + state_buffer_account, + args.chunk(), + offset, + write_end, + state_fields.payload_len, + ); + } + + require_eq!( + state_fields.finalized, + false, + ProgramError::InvalidInstructionData + ); + require_eq!( + offset, + state_fields.payload_len, + ProgramError::InvalidInstructionData + ); + + StateBuffer::decode_mut(&MaxLenStorage::new( + state_buffer_account, + StateBuffer::data_len_from_payload_capacity(state_fields.total_len)?, + StateBuffer::MAX_ACCOUNT_DATA_GROWTH_PER_WRITE, + ))? + .payload_mut()? + .extend_from_slice(args.chunk())?; + + let data_hash = if write_end == state_fields.total_len { + let account_data = state_buffer_account.try_borrow()?; + let state = StateBuffer::decode(account_data.as_ref())?; + let payload = state.payload(); + require_eq!( + payload.len(), + state.total_len() as usize, + ProgramError::InvalidAccountData + ); + Some(account_data_hash(payload.as_slice())) + } else { + None + }; + + if let Some(data_hash) = data_hash { + let mut state_mut = StateBuffer::decode_mut(state_buffer_account)?; + *state_mut.data_hash_mut()? = data_hash; + state_mut.finalized_mut()?.set(true)?; + } + + Ok(()) +} + +struct StateBufferFields { + total_len: usize, + payload_len: usize, + finalized: bool, +} + +fn validate_state_buffer( + state_buffer_account: &AccountView, + authority: &AccountView, + delegated_account: &AccountView, + args: &dlp_api::v2::WriteStateBufferArgsView<'_>, +) -> Result { + let account_data = state_buffer_account.try_borrow()?; + let state = StateBuffer::decode(account_data.as_ref())?; + + require_eq!( + &state.discriminator(), + &StateBuffer::DISCRIMINATOR, + ProgramError::InvalidAccountData + ); + require_eq_keys!( + state.authority(), + authority.address(), + DlpError::InvalidAuthority + ); + require_eq_keys!( + state.account_pubkey(), + delegated_account.address(), + ProgramError::InvalidAccountData + ); + require_eq!( + state.commit_id(), + args.commit_id(), + ProgramError::InvalidInstructionData + ); + require_eq!( + state.total_len(), + args.total_len(), + ProgramError::InvalidInstructionData + ); + require_le!( + state.payload().len(), + state.total_len() as usize, + ProgramError::InvalidAccountData + ); + require_le!( + state.payload().capacity(), + state.total_len() as usize, + ProgramError::InvalidAccountData + ); + require_eq!( + account_data.len(), + StateBuffer::data_len_from_payload_capacity( + state.payload().capacity() + )?, + ProgramError::InvalidAccountData + ); + + Ok(StateBufferFields { + total_len: state.total_len() as usize, + payload_len: state.payload().len(), + finalized: state.finalized(), + }) +} + +fn validate_duplicate_chunk( + state_buffer_account: &AccountView, + chunk: &[u8], + offset: usize, + write_end: usize, + payload_len: usize, +) -> ProgramResult { + require_le!(write_end, payload_len, ProgramError::InvalidInstructionData); + + let account_data = state_buffer_account.try_borrow()?; + let state = StateBuffer::decode(account_data.as_ref())?; + let payload = state.payload(); + require_eq!( + &payload.as_slice()[offset..write_end], + chunk, + ProgramError::InvalidInstructionData + ); + + Ok(()) +} + +fn account_data_hash(data: &[u8]) -> [u8; 32] { + solana_sha256_hasher::hashv(&[b"magicblock.account_data.v1", data]) + .to_bytes() +} diff --git a/src/v2/processor/mod.rs b/src/v2/processor/mod.rs index b796cbd8..010d038e 100644 --- a/src/v2/processor/mod.rs +++ b/src/v2/processor/mod.rs @@ -6,6 +6,7 @@ use dlp_api::v2::DlpV2Instruction; use pinocchio::{AccountView, ProgramResult}; pub use bootstrap::*; +pub use fraud_proofs::*; pub fn process_instruction( accounts: &[AccountView], @@ -28,5 +29,8 @@ pub fn process_instruction( DlpV2Instruction::UpdateProtocolConfig => { process_update_protocol_config(accounts, data) } + DlpV2Instruction::WriteStateBuffer => { + process_write_state_buffer(accounts, data) + } } } diff --git a/tests/test_v2_write_state_buffer.rs b/tests/test_v2_write_state_buffer.rs new file mode 100644 index 00000000..90ec358e --- /dev/null +++ b/tests/test_v2_write_state_buffer.rs @@ -0,0 +1,604 @@ +use dlp_api::v2::{ + instruction_builder::write_state_buffer, pda::state_buffer_pda, + StateBuffer, WriteStateBufferArgs, +}; +use solana_program::{hash::Hash, native_token::LAMPORTS_PER_SOL}; +use solana_program_test::{ + BanksClient, BanksClientError, ProgramTestBanksClientExt, +}; +use solana_sdk::{ + pubkey::Pubkey, + rent::Rent, + signature::{Keypair, Signer}, + transaction::Transaction, +}; +use solana_system_interface::instruction as system_instruction; +use wheels::layout::Decodable; + +mod fixtures; + +use crate::fixtures::v2::{ + initialize_protocol_config, setup_program_test_env, + valid_protocol_config_args, +}; + +#[tokio::test] +async fn test_write_state_buffer_unregistered_authority_one_chunk_finalizes() { + let mut env = setup_write_state_buffer_env().await; + let commit_id = 7; + let data = vec![1, 2, 3, 4]; + + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: data.len() as u32, + offset: 0, + chunk: data.clone(), + }, + ) + .await + .unwrap(); + + let state_buffer = state_buffer_pda( + &env.delegated.pubkey(), + commit_id, + &env.writer.pubkey(), + ); + let state_buffer_account = + env.banks.get_account(state_buffer).await.unwrap().unwrap(); + let state = StateBuffer::decode(&state_buffer_account.data).unwrap(); + + assert_eq!(state.discriminator(), StateBuffer::DISCRIMINATOR); + assert_eq!(*state.authority(), env.writer.pubkey()); + assert_eq!(*state.account_pubkey(), env.delegated.pubkey()); + assert_eq!(state.commit_id(), commit_id); + assert_eq!(state.total_len(), data.len() as u32); + assert!(state.finalized()); + assert_eq!(*state.data_hash(), account_data_hash(&data)); + assert_eq!(state.payload().len(), data.len()); + assert_eq!(state.payload().capacity(), data.len()); + assert_eq!(state.payload().as_slice(), data.as_slice()); + assert_eq!( + &state_buffer_account.data[StateBuffer::PAYLOAD_BYTES_OFFSET..], + data.as_slice() + ); +} + +#[tokio::test] +async fn test_write_state_buffer_multiple_chunks_finalize() { + let mut env = setup_write_state_buffer_env().await; + let commit_id = 8; + let total_len = 5; + let state_buffer = state_buffer_pda( + &env.delegated.pubkey(), + commit_id, + &env.writer.pubkey(), + ); + + { + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len, + offset: 0, + chunk: vec![1, 2], + }, + ) + .await + .unwrap(); + + let state_buffer_account = + env.banks.get_account(state_buffer).await.unwrap().unwrap(); + let state = StateBuffer::decode(&state_buffer_account.data).unwrap(); + assert!( + state_buffer_account.lamports + >= Rent::default() + .minimum_balance(state_buffer_account.data.len()) + ); + let payload = state.payload(); + assert_eq!(payload.len(), 2); + assert_eq!(payload.capacity(), total_len as usize); + assert_eq!(payload.as_slice(), &[1, 2]); + + let payload_start = StateBuffer::PAYLOAD_BYTES_OFFSET; + let written_end = payload_start + 2; + let capacity_end = payload_start + total_len as usize; + assert_eq!( + &state_buffer_account.data[payload_start..written_end], + &[1, 2] + ); + assert!(state_buffer_account.data[written_end..capacity_end] + .iter() + .all(|byte| *byte == 0)); + assert!(!state.finalized()); + assert_eq!(*state.data_hash(), [0; 32]); + } + + { + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len, + offset: 2, + chunk: vec![3, 4, 5], + }, + ) + .await + .unwrap(); + + let state_buffer_account = + env.banks.get_account(state_buffer).await.unwrap().unwrap(); + let state = StateBuffer::decode(&state_buffer_account.data).unwrap(); + let payload = state.payload(); + assert_eq!(payload.len(), total_len as usize); + assert_eq!(payload.capacity(), total_len as usize); + assert_eq!(payload.as_slice(), &[1, 2, 3, 4, 5]); + + let payload_start = StateBuffer::PAYLOAD_BYTES_OFFSET; + let payload_end = payload_start + total_len as usize; + assert_eq!( + &state_buffer_account.data[payload_start..payload_end], + &[1, 2, 3, 4, 5] + ); + assert!(state.finalized()); + assert_eq!(*state.data_hash(), account_data_hash(&[1, 2, 3, 4, 5])); + } +} + +#[tokio::test] +async fn test_write_state_buffer_grows_payload_span_past_initial_capacity() { + let mut env = setup_write_state_buffer_env().await; + let commit_id = 15; + let state_buffer = state_buffer_pda( + &env.delegated.pubkey(), + commit_id, + &env.writer.pubkey(), + ); + let initial_payload_capacity = StateBuffer::MAX_INITIAL_PAYLOAD_LEN; + let total_len = initial_payload_capacity + 11; + let final_data_len = + StateBuffer::data_len_from_payload_capacity(total_len).unwrap(); + let first_write_len = initial_payload_capacity - 3; + let mut expected = Vec::with_capacity(total_len); + let mut offset = 0; + + println!("total_len: {total_len}, first_write_len: {first_write_len}"); + + while offset < first_write_len { + let chunk_len = (first_write_len - offset).min(512); + let chunk = (offset..offset + chunk_len) + .map(|value| (value % u8::MAX as usize) as u8) + .collect::>(); + expected.extend_from_slice(&chunk); + + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: total_len as u32, + offset: offset as u32, + chunk, + }, + ) + .await + .unwrap(); + + offset += chunk_len; + } + + assert_eq!(offset, first_write_len); + + let state_buffer_account = + env.banks.get_account(state_buffer).await.unwrap().unwrap(); + let state = StateBuffer::decode(&state_buffer_account.data).unwrap(); + assert!( + state_buffer_account.lamports + >= Rent::default().minimum_balance(final_data_len) + ); + let payload = state.payload(); + assert_eq!(payload.len(), expected.len()); + assert_eq!(payload.capacity(), initial_payload_capacity); + assert_eq!(payload.as_slice(), expected.as_slice()); + + let payload_start = StateBuffer::PAYLOAD_BYTES_OFFSET; + let written_end = payload_start + expected.len(); + let capacity_end = payload_start + initial_payload_capacity; + assert_eq!( + &state_buffer_account.data[payload_start..written_end], + expected.as_slice() + ); + assert!(state_buffer_account.data[written_end..capacity_end] + .iter() + .all(|byte| *byte == 0)); + assert!(!state.finalized()); + + let crossing_chunk = vec![7, 8, 9, 10, 11]; + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: total_len as u32, + offset: first_write_len as u32, + chunk: crossing_chunk.clone(), + }, + ) + .await + .unwrap(); + expected.extend_from_slice(&crossing_chunk); + + let state_buffer_account = + env.banks.get_account(state_buffer).await.unwrap().unwrap(); + let state = StateBuffer::decode(&state_buffer_account.data).unwrap(); + let payload = state.payload(); + assert_eq!(payload.len(), expected.len()); + assert_eq!(payload.capacity(), total_len); + assert_eq!(payload.as_slice(), expected.as_slice()); + + let payload_start = StateBuffer::PAYLOAD_BYTES_OFFSET; + let written_end = payload_start + expected.len(); + let capacity_end = payload_start + total_len; + assert_eq!( + &state_buffer_account.data[payload_start..written_end], + expected.as_slice() + ); + assert!(state_buffer_account.data[written_end..capacity_end] + .iter() + .all(|byte| *byte == 0)); + assert!(!state.finalized()); +} + +#[tokio::test] +async fn test_write_state_buffer_duplicate_retry() { + let mut env = setup_write_state_buffer_env().await; + let commit_id = 9; + let first_chunk = vec![1, 2]; + + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: 4, + offset: 0, + chunk: first_chunk.clone(), + }, + ) + .await + .unwrap(); + + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: 4, + offset: 0, + chunk: first_chunk, + }, + ) + .await + .unwrap(); + + let state_buffer = state_buffer_pda( + &env.delegated.pubkey(), + commit_id, + &env.writer.pubkey(), + ); + let state_buffer_account = + env.banks.get_account(state_buffer).await.unwrap().unwrap(); + let state = StateBuffer::decode(&state_buffer_account.data).unwrap(); + let payload = state.payload(); + assert_eq!(payload.len(), 2); + assert_eq!(payload.capacity(), 4); + assert_eq!(payload.as_slice(), &[1, 2]); + + let payload_start = StateBuffer::PAYLOAD_BYTES_OFFSET; + let written_end = payload_start + 2; + let capacity_end = payload_start + 4; + assert_eq!( + &state_buffer_account.data[payload_start..written_end], + &[1, 2] + ); + assert!(state_buffer_account.data[written_end..capacity_end] + .iter() + .all(|byte| *byte == 0)); + assert!(!state.finalized()); +} + +#[tokio::test] +async fn test_write_state_buffer_rejects_mismatched_duplicate_retry() { + let mut env = setup_write_state_buffer_env().await; + let commit_id = 10; + + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: 4, + offset: 0, + chunk: vec![1, 2], + }, + ) + .await + .unwrap(); + + let result = write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: 4, + offset: 0, + chunk: vec![1, 9], + }, + ) + .await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_write_state_buffer_rejects_wrong_offset() { + let mut env = setup_write_state_buffer_env().await; + + let result = write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id: 11, + total_len: 4, + offset: 1, + chunk: vec![1, 2], + }, + ) + .await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_write_state_buffer_rejects_oversized_total_len() { + let mut env = setup_write_state_buffer_env().await; + + let result = write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id: 12, + total_len: StateBuffer::MAX_TOTAL_PAYLOAD_LEN + 1, + offset: 0, + chunk: vec![1], + }, + ) + .await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_write_state_buffer_rejects_post_finalize_mutation() { + let mut env = setup_write_state_buffer_env().await; + let commit_id = 13; + let data = vec![1, 2, 3]; + + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: data.len() as u32, + offset: 0, + chunk: data.clone(), + }, + ) + .await + .unwrap(); + + write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: data.len() as u32, + offset: 0, + chunk: data, + }, + ) + .await + .unwrap(); + + let result = write_buffer( + &mut env.banks, + &env.payer, + &env.writer, + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: 3, + offset: 3, + chunk: vec![4], + }, + ) + .await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_write_state_buffer_rejects_wrong_authority_for_buffer() { + let mut env = setup_write_state_buffer_env().await; + let commit_id = 14; + let wrong_authority = Keypair::new(); + fund_account(&mut env.banks, &env.payer, &wrong_authority.pubkey()).await; + + let mut ix = write_state_buffer( + env.payer.pubkey(), + env.writer.pubkey(), + env.delegated.pubkey(), + WriteStateBufferArgs { + commit_id, + total_len: 2, + offset: 0, + chunk: vec![1, 2], + }, + ); + + // use wrong_authority as authority + ix.accounts[1].pubkey = wrong_authority.pubkey(); + + let blockhash = fresh_blockhash(&mut env.banks).await; + let tx = Transaction::new_signed_with_payer( + &[ix], + Some(&env.payer.pubkey()), + &[&env.payer, &wrong_authority], + blockhash, + ); + + assert!(env.banks.process_transaction(tx).await.is_err()); +} + +struct WriteStateBufferEnv { + banks: BanksClient, + payer: Keypair, + writer: Keypair, + delegated: Keypair, +} + +async fn setup_write_state_buffer_env() -> WriteStateBufferEnv { + let (mut banks, payer, authority, blockhash) = + setup_program_test_env().await; + let writer = Keypair::new(); + let delegated = Keypair::new(); + let config_args = valid_protocol_config_args(); + + initialize_protocol_config( + &banks, + &payer, + &authority, + blockhash, + config_args, + ) + .await; + fund_account(&mut banks, &payer, &writer.pubkey()).await; + create_delegated_account(&mut banks, &payer, &delegated).await; + + WriteStateBufferEnv { + banks, + payer, + writer, + delegated, + } +} + +async fn create_delegated_account( + banks: &mut BanksClient, + payer: &Keypair, + delegated: &Keypair, +) { + let lamports = Rent::default().minimum_balance(8); + let ix = system_instruction::create_account( + &payer.pubkey(), + &delegated.pubkey(), + lamports, + 8, + &dlp_api::ID, + ); + let blockhash = fresh_blockhash(banks).await; + let tx = Transaction::new_signed_with_payer( + &[ix], + Some(&payer.pubkey()), + &[payer, delegated], + blockhash, + ); + + banks.process_transaction(tx).await.unwrap(); +} + +async fn fund_account( + banks: &mut BanksClient, + payer: &Keypair, + account: &Pubkey, +) { + let ix = system_instruction::transfer( + &payer.pubkey(), + account, + LAMPORTS_PER_SOL, + ); + let blockhash = fresh_blockhash(banks).await; + let tx = Transaction::new_signed_with_payer( + &[ix], + Some(&payer.pubkey()), + &[payer], + blockhash, + ); + + banks.process_transaction(tx).await.unwrap(); +} + +async fn write_buffer( + banks: &mut BanksClient, + payer: &Keypair, + authority: &Keypair, + account: Pubkey, + args: WriteStateBufferArgs, +) -> Result<(), BanksClientError> { + let ix = + write_state_buffer(payer.pubkey(), authority.pubkey(), account, args); + let blockhash = fresh_blockhash(banks).await; + let tx = Transaction::new_signed_with_payer( + &[ix], + Some(&payer.pubkey()), + &[payer, authority], + blockhash, + ); + + banks.process_transaction(tx).await +} + +async fn fresh_blockhash(banks: &mut BanksClient) -> Hash { + let latest_blockhash = banks.get_latest_blockhash().await.unwrap(); + banks + .get_new_latest_blockhash(&latest_blockhash) + .await + .unwrap() +} + +fn account_data_hash(data: &[u8]) -> [u8; 32] { + solana_sha256_hasher::hashv(&[b"magicblock.account_data.v1", data]) + .to_bytes() +}