[message-buffer 14/x] - Message buffer/zero copy fixes (#811)

* test(message-buffer): add rust integration tests

add rust-toolchain.toml to pin rust version, add integration tests, update cpi caller auth seeds

* refactor(message-buffer): remove unused test ix

* chore(message-buffer): clean up

* refactor(message-buffer): simple refactor

* test(message-buffer): refactor integration test structure

* refactor(message-buffer): rename

* fix(message-buffer): fix min size check when shrinking msg buffer

* chore(message-buffer): cleanup

* fix(message-buffer): resize borrow bug fix

* test(message-buffer): refactor test util methods into MessageBufferTestContext for less duplication

* test(message-buffer): resolve merge conflicts from repo restructure

* chore(message-buffer): delete commented out code

* feat(message-buffer): use AccountLoader for zero-copy for resize/delete & put_all

* chore(message-buffer): clean up

* style(message-buffer): clean up

* fix(message-buffer): address PR feedback

revert to old put_all impl of using bytemuck for deserializing header, update tests add back old
check

* chore(message-buffer): clean up

Author: swimricky

pythnet/message_buffer/programs/message_buffer/src/instructions/create_buffer.rs

@@ -76,7 +76,6 @@ pub fn create_buffer<'info>(
         }
         loader.exit(&crate::ID)?;
     } else {
-        // FIXME: change this to be emit!(Event)
         msg!("Buffer account already initialized");
     }
 

pythnet/message_buffer/programs/message_buffer/src/instructions/delete_buffer.rs

@@ -1,51 +1,21 @@
 use {
-    crate::{
-        state::*,
-        MessageBufferError,
-        MESSAGE,
-    },
+    crate::state::*,
     anchor_lang::prelude::*,
 };
 
 pub fn delete_buffer<'info>(
     ctx: Context<'_, '_, '_, 'info, DeleteBuffer<'info>>,
     allowed_program_auth: Pubkey,
-    base_account_key: Pubkey,
-    bump: u8,
+    _base_account_key: Pubkey,
 ) -> Result<()> {
-    let message_buffer_account_info = ctx
-        .remaining_accounts
-        .first()
-        .ok_or(MessageBufferError::MessageBufferNotProvided)?;
-
     ctx.accounts
         .whitelist
         .is_allowed_program_auth(&allowed_program_auth)?;
-
-    MessageBuffer::check_discriminator(message_buffer_account_info)?;
-
-    let expected_key = Pubkey::create_program_address(
-        &[
-            allowed_program_auth.as_ref(),
-            MESSAGE.as_bytes(),
-            base_account_key.as_ref(),
-            &[bump],
-        ],
-        &crate::ID,
-    )
-    .map_err(|_| MessageBufferError::InvalidPDA)?;
-
-    require_keys_eq!(
-        message_buffer_account_info.key(),
-        expected_key,
-        MessageBufferError::InvalidPDA
-    );
-    let loader = AccountLoader::<MessageBuffer>::try_from(message_buffer_account_info)?;
-    loader.close(ctx.accounts.admin.to_account_info())?;
     Ok(())
 }
 
 #[derive(Accounts)]
+#[instruction(allowed_program_auth: Pubkey, base_account_key: Pubkey)]
 pub struct DeleteBuffer<'info> {
     #[account(
         seeds = [b"message".as_ref(), b"whitelist".as_ref()],
@@ -57,5 +27,12 @@ pub struct DeleteBuffer<'info> {
     // Also the recipient of the lamports from closing the buffer account
     #[account(mut)]
     pub admin: Signer<'info>,
-    // remaining_account:  - [AccumulatorInput PDA]
+
+    #[account(
+    mut,
+    close = admin,
+    seeds = [allowed_program_auth.as_ref(), b"message".as_ref(), base_account_key.as_ref()],
+    bump = message_buffer.load()?.bump,
+    )]
+    pub message_buffer: AccountLoader<'info, MessageBuffer>,
 }

pythnet/message_buffer/programs/message_buffer/src/instructions/put_all.rs

@@ -1,53 +1,40 @@
 use {
-    crate::{
-        state::*,
-        MessageBufferError,
-    },
+    crate::state::*,
     anchor_lang::prelude::*,
-    std::mem,
 };
 
 
 pub fn put_all<'info>(
     ctx: Context<'_, '_, '_, 'info, PutAll<'info>>,
-    base_account_key: Pubkey,
+    _base_account_key: Pubkey,
     messages: Vec<Vec<u8>>,
 ) -> Result<()> {
-    let cpi_caller_auth = ctx.accounts.whitelist_verifier.is_allowed()?;
-    let message_buffer_account_info = ctx
-        .remaining_accounts
-        .first()
-        .ok_or(MessageBufferError::MessageBufferNotProvided)?;
+    ctx.accounts.whitelist_verifier.is_allowed()?;
 
-    MessageBuffer::check_discriminator(message_buffer_account_info)?;
-
-    let account_data = &mut message_buffer_account_info.try_borrow_mut_data()?;
-    let header_end_index = mem::size_of::<MessageBuffer>() + 8;
+    let msg_buffer_ai = ctx.accounts.message_buffer.to_account_info();
+    let account_data = &mut msg_buffer_ai.try_borrow_mut_data()?;
+    let header_end_index = MessageBuffer::HEADER_LEN as usize;
 
     let (header_bytes, body_bytes) = account_data.split_at_mut(header_end_index);
 
     let message_buffer: &mut MessageBuffer = bytemuck::from_bytes_mut(&mut header_bytes[8..]);
 
-    message_buffer.validate(
-        message_buffer_account_info.key(),
-        cpi_caller_auth,
-        base_account_key,
-    )?;
-
     message_buffer.refresh_header();
-
     let (num_msgs, num_bytes) = message_buffer.put_all_in_buffer(body_bytes, &messages);
-
     if num_msgs != messages.len() {
         msg!("unable to fit all messages in MessageBuffer account. Wrote {}/{} messages and {} bytes", num_msgs, messages.len(), num_bytes);
     }
-
     Ok(())
 }
 
 #[derive(Accounts)]
-#[instruction( base_account_key: Pubkey)]
+#[instruction(base_account_key: Pubkey)]
 pub struct PutAll<'info> {
     pub whitelist_verifier: WhitelistVerifier<'info>,
-    // remaining_accounts:  - [AccumulatorInput PDA]
+    #[account(
+        mut,
+        seeds = [whitelist_verifier.cpi_caller_auth.key().as_ref(), b"message".as_ref(), base_account_key.as_ref()],
+        bump = message_buffer.load()?.bump,
+    )]
+    pub message_buffer:     AccountLoader<'info, MessageBuffer>,
 }

pythnet/message_buffer/programs/message_buffer/src/instructions/resize_buffer.rs

@@ -2,112 +2,35 @@ use {
     crate::{
         state::*,
         MessageBufferError,
-        MESSAGE,
-    },
-    anchor_lang::{
-        prelude::*,
-        solana_program::entrypoint::MAX_PERMITTED_DATA_INCREASE,
-        system_program::{
-            self,
-            Transfer,
-        },
     },
+    anchor_lang::prelude::*,
 };
 
+
 pub fn resize_buffer<'info>(
     ctx: Context<'_, '_, '_, 'info, ResizeBuffer<'info>>,
     allowed_program_auth: Pubkey,
-    base_account_key: Pubkey,
-    buffer_bump: u8,
+    _base_account_key: Pubkey,
     target_size: u32,
 ) -> Result<()> {
-    let message_buffer_account_info = ctx
-        .remaining_accounts
-        .first()
-        .ok_or(MessageBufferError::MessageBufferNotProvided)?;
-
     ctx.accounts
         .whitelist
         .is_allowed_program_auth(&allowed_program_auth)?;
-    MessageBuffer::check_discriminator(message_buffer_account_info)?;
-
 
-    let target_size = target_size as usize;
-
-    let current_account_size = message_buffer_account_info.data_len();
-    let target_size_delta = target_size.saturating_sub(current_account_size);
+    let message_buffer = &ctx.accounts.message_buffer.load()?;
+    let max_end_offset = message_buffer.end_offsets.iter().max().unwrap();
+    let minimum_size = max_end_offset + message_buffer.header_len;
     require_gte!(
-        MAX_PERMITTED_DATA_INCREASE,
-        target_size_delta,
-        MessageBufferError::TargetSizeDeltaExceeded
-    );
-
-    let expected_key = Pubkey::create_program_address(
-        &[
-            allowed_program_auth.as_ref(),
-            MESSAGE.as_bytes(),
-            base_account_key.as_ref(),
-            &[buffer_bump],
-        ],
-        &crate::ID,
-    )
-    .map_err(|_| MessageBufferError::InvalidPDA)?;
-
-    require_keys_eq!(
-        message_buffer_account_info.key(),
-        expected_key,
-        MessageBufferError::InvalidPDA
+        target_size as usize,
+        minimum_size as usize,
+        MessageBufferError::MessageBufferTooSmall
     );
-
-    // allow for target_size == account_size in case Rent requirements have changed
-    // and additional lamports need to be transferred.
-    // the realloc step will be a no-op in this case.
-    if target_size >= current_account_size {
-        let target_rent = Rent::get()?.minimum_balance(target_size);
-        if message_buffer_account_info.lamports() < target_rent {
-            system_program::transfer(
-                CpiContext::new(
-                    ctx.accounts.system_program.to_account_info(),
-                    Transfer {
-                        from: ctx.accounts.admin.to_account_info(),
-                        to:   message_buffer_account_info.to_account_info(),
-                    },
-                ),
-                target_rent - message_buffer_account_info.lamports(),
-            )?;
-        }
-        message_buffer_account_info
-            .realloc(target_size, false)
-            .map_err(|_| MessageBufferError::ReallocFailed)?;
-    } else {
-        // Check that account doesn't get resized to smaller than the amount of
-        // data it is currently holding (if any)
-        {
-            let account_data = &message_buffer_account_info.try_borrow_data()?;
-            let header_end_index = std::mem::size_of::<MessageBuffer>() + 8;
-            let (header_bytes, _) = account_data.split_at(header_end_index);
-            let message_buffer: &MessageBuffer = bytemuck::from_bytes(&header_bytes[8..]);
-            let max_end_offset = message_buffer.end_offsets.iter().max().unwrap();
-            let minimum_size = max_end_offset + message_buffer.header_len;
-            require_gte!(
-                target_size,
-                minimum_size as usize,
-                MessageBufferError::MessageBufferTooSmall
-            );
-        }
-
-
-        // Not transferring excess lamports back to admin.
-        // Account will retain more lamports than necessary.
-        message_buffer_account_info.realloc(target_size, false)?;
-    }
     Ok(())
 }
 
 #[derive(Accounts)]
 #[instruction(
-    allowed_program_auth: Pubkey, base_account_key: Pubkey,
-    buffer_bump: u8, target_size: u32
+    allowed_program_auth: Pubkey, base_account_key: Pubkey, target_size: u32
 )]
 pub struct ResizeBuffer<'info> {
     #[account(
@@ -122,5 +45,18 @@ pub struct ResizeBuffer<'info> {
     pub admin: Signer<'info>,
 
     pub system_program: Program<'info, System>,
-    // remaining_accounts:  - [AccumulatorInput PDA]
+
+    /// If decreasing, Anchor will automatically check
+    /// if target_size is too small and if so,then load() will fail.
+    /// If increasing, Anchor also automatically checks if target_size delta
+    /// exceeds MAX_PERMITTED_DATA_INCREASE
+    #[account(
+        mut,
+        realloc = target_size as usize,
+        realloc::zero = false,
+        realloc::payer = admin,
+        seeds = [allowed_program_auth.as_ref(), b"message".as_ref(), base_account_key.as_ref()],
+        bump = message_buffer.load()?.bump,
+    )]
+    pub message_buffer: AccountLoader<'info, MessageBuffer>,
 }

pythnet/message_buffer/programs/message_buffer/src/lib.rs

@@ -112,21 +112,13 @@ pub mod message_buffer {
     /// *`target_size`          -  Size to re-allocate for the
     ///                           `MessageBuffer` PDA. If increasing the size,
     ///                           max delta of current_size & target_size is 10240
-    /// *`buffer_bump`          -  Bump seed for the `MessageBuffer` PDA
     pub fn resize_buffer<'info>(
         ctx: Context<'_, '_, '_, 'info, ResizeBuffer<'info>>,
         allowed_program_auth: Pubkey,
         base_account_key: Pubkey,
-        buffer_bump: u8,
         target_size: u32,
     ) -> Result<()> {
-        instructions::resize_buffer(
-            ctx,
-            allowed_program_auth,
-            base_account_key,
-            buffer_bump,
-            target_size,
-        )
+        instructions::resize_buffer(ctx, allowed_program_auth, base_account_key, target_size)
     }
 
     /// Closes the buffer account and transfers the remaining lamports to the
@@ -138,14 +130,12 @@ pub mod message_buffer {
     /// * `base_account_key`    - Pubkey of the original account the
     ///                           `MessageBuffer` is derived from
     ///                           (e.g. pyth price account)
-    /// *`buffer_bump`          -  Bump seed for the `MessageBuffer` PDA
     pub fn delete_buffer<'info>(
         ctx: Context<'_, '_, '_, 'info, DeleteBuffer<'info>>,
         allowed_program_auth: Pubkey,
         base_account_key: Pubkey,
-        buffer_bump: u8,
     ) -> Result<()> {
-        instructions::delete_buffer(ctx, allowed_program_auth, base_account_key, buffer_bump)
+        instructions::delete_buffer(ctx, allowed_program_auth, base_account_key)
     }
 }