A close macro implementation for AccountInfo

sdk/pinocchio/src/account_info.rs

@@ -1,5 +1,4 @@
 //! Data structures to represent account information.
-
 use core::{marker::PhantomData, mem::ManuallyDrop, ptr::NonNull, slice::from_raw_parts_mut};
 
 use crate::{program_error::ProgramError, pubkey::Pubkey, syscalls::sol_memset_};
@@ -379,6 +378,89 @@ impl AccountInfo {
         Ok(())
     }
 
+    /// Zero out the the account's data_len, lamports and owner fields, effectively
+    /// closing the account.
+    ///
+    /// Note: This doesn't protect against future reinitialization of the account
+    /// since the account_data will need to be zeroed out as well. Or if the attacker
+    /// has access to the keypair of the account that we're trying to close, they can
+    /// just add the lenght, lamports and owner back before the data is wiped out from
+    /// the ledger.
+    ///
+    /// Note: This works because the 48 bytes before the account data are:
+    /// - 8 bytes for the data_len
+    /// - 8 bytes for the lamports
+    /// - 32 bytes for the owner
+    ///
+    /// # Safety
+    ///
+    /// This method makes assumptions about the layout and location of memory
+    /// referenced by `AccountInfo` fields. It should only be called for
+    /// instances of `AccountInfo` that were created by the runtime and received
+    /// in the `process_instruction` entrypoint of a program.
+    ///
+    /// This method is unsafe because it does not check if the account data is already
+    /// borrowed. It should only be called when the account is not being used.
+    #[inline(always)]
+    pub unsafe fn close(&self) {
+        // Zero out the account owner. While the field is a `Pubkey`, it is quicker
+        // to zero the 32 bytes as 4 x `u64`s.
+        *(self.data_ptr().sub(48) as *mut u64) = 0u64;
+        *(self.data_ptr().sub(40) as *mut u64) = 0u64;
+        *(self.data_ptr().sub(32) as *mut u64) = 0u64;
+        *(self.data_ptr().sub(16) as *mut u64) = 0u64;
+        // Zero the account lamports.
+        (*self.raw).lamports = 0;
+        // Zero the account data length.
+        (*self.raw).data_len = 0;
+    }
+
+    /// Zero out the the account's data_len, lamports and owner fields, effectively
+    /// closing the account.
+    ///
+    /// Note: This doesn't protect against future reinitialization of the account
+    /// since the account_data will need to be zeroed out as well. Or if the attacker
+    /// has access to the keypair of the account that we're trying to close, they can
+    /// just add the lenght, lamports and owner back before the data is wiped out from
+    /// the ledger.
+    ///
+    /// Note: This works because the 48 bytes before the account data are:
+    /// - 8 bytes for the data_len
+    /// - 8 bytes for the lamports
+    /// - 32 bytes for the owner
+    ///
+    /// # Safety
+    ///
+    /// This method makes assumptions about the layout and location of memory
+    /// referenced by `AccountInfo` fields. It should only be called for
+    /// instances of `AccountInfo` that were created by the runtime and received
+    /// in the `process_instruction` entrypoint of a program.
+    ///
+    /// This method set the variable as 0 making sure that we're actually zeroing out
+    /// the bytes.
+    #[inline(always)]
+    pub unsafe fn assembly_close(&self) {
+        #[cfg(target_os = "solana")]
+        unsafe {
+            let var = 0u64;
+            core::arch::asm!(
+                // Zero the data lenght.
+                "stxdw [{0}-8], {1}",
+                // Zero the account lamports.
+                "stxdw [{0}-16], {1}",
+                // Zero out the account owner. While the field is a `Pubkey`, it is quicker
+                // to zero the 32 bytes as 4 x `u64`s.
+                "stxdw [{0}-24], {1}",
+                "stxdw [{0}-32], {1}",
+                "stxdw [{0}-40], {1}",
+                "stxdw [{0}-48], {1}",
+                in(reg) self.borrow_mut_data_unchecked().as_mut_ptr(),
+                in(reg) var,
+                options(nostack, preserves_flags)
+            );
+        }
+    }
+
     /// Returns the memory address of the account data.
     fn data_ptr(&self) -> *mut u8 {
         unsafe { (self.raw as *const _ as *mut u8).add(core::mem::size_of::<Account>()) }

sdk/pinocchio/src/lib.rs

@@ -9,6 +9,7 @@
 //! [`solana-program`]: https://docs.rs/solana-program/latest/solana_program/
 
 #![no_std]
+#![cfg_attr(target_os = "solana", feature(asm_experimental_arch, asm_const))]
 
 pub mod account_info;
 pub mod entrypoint;