From c13b1d6ec339adbab9f54a5f0bbb920465f492d0 Mon Sep 17 00:00:00 2001
From: febo <febo@anza.xyz>
Date: Tue, 12 May 2026 13:39:55 +0100
Subject: [PATCH] Explicit ownership checks

---
 program/src/processor/set_data.rs | 14 +++++++++++---
 program/src/processor/write.rs    | 15 ++++++++++++++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/program/src/processor/set_data.rs b/program/src/processor/set_data.rs
index c103d60..4474973 100644
--- a/program/src/processor/set_data.rs
+++ b/program/src/processor/set_data.rs
@@ -57,11 +57,19 @@ pub fn set_data(accounts: &mut [AccountView], instruction_data: &[u8]) -> Progra
         validate_authority(header, authority, program, program_data)?;
     }
 
-    let has_buffer = buffer.address() != &crate::ID;
-    // Get data from buffer or remaining instruction data, if any.
-    let data = match (optional_data, has_buffer) {
+    // buffer (if `remaining_data` is `None`)
+    // - must be initialized
+    // - must be owned by the program
+    //
+    // Get data from buffer or remaining instruction data (if any).
+    let data = match (optional_data, buffer.address() != &crate::ID) {
         (Some((data_source, Some(remaining_data))), false) => Some((data_source, remaining_data)),
         (Some((data_source, None)), true) => {
+            // Since we are not writing to the `buffer` account, validate
+            // the ownership of the account.
+            if !buffer.owned_by(&crate::ID) {
+                return Err(ProgramError::InvalidAccountOwner);
+            }
             // SAFETY: single immutable borrow of `buffer` account data.
             let buffer_data = unsafe { buffer.borrow_unchecked() };
 
diff --git a/program/src/processor/write.rs b/program/src/processor/write.rs
index ddaf0eb..60efd3a 100644
--- a/program/src/processor/write.rs
+++ b/program/src/processor/write.rs
@@ -37,6 +37,11 @@ pub fn write(accounts: &mut [AccountView], instruction_data: &[u8]) -> ProgramRe
     // - must be initialized
     // - must be rent exempt (pre-funded account) since we are reallocating the buffer
     //   account
+    //
+    // source_buffer (if `args.data()` is empty)
+    // - must be initialized
+    // - must be owned by the program
+    // - must not be the same account as `target_buffer`
 
     let (required_length, source_data) = {
         // SAFETY: scoped immutable borrow of `buffer` account data. There
@@ -62,6 +67,11 @@ pub fn write(accounts: &mut [AccountView], instruction_data: &[u8]) -> ProgramRe
         };
 
         let source_buffer_data = if source_buffer.address() != &crate::ID {
+            // Since we are not writing to the `source_buffer` account, validate
+            // the ownership of the account.
+            if !source_buffer.owned_by(&crate::ID) {
+                return Err(ProgramError::InvalidAccountOwner);
+            }
             // SAFETY: singe immutable borrow of `source_buffer` account data.
             Some(unsafe { source_buffer.borrow_unchecked() })
         } else {
@@ -71,7 +81,10 @@ pub fn write(accounts: &mut [AccountView], instruction_data: &[u8]) -> ProgramRe
         let source_data = match (instruction_data, source_buffer_data) {
             (Some(instruction_data), None) => instruction_data,
             (None, Some(buffer_data)) => match AccountDiscriminator::try_from_bytes(buffer_data)? {
-                Some(AccountDiscriminator::Buffer) => &buffer_data[Header::LEN..],
+                // `source_buffer` and `target_buffer` must not be the same account.
+                Some(AccountDiscriminator::Buffer) if source_buffer != target_buffer => {
+                    &buffer_data[Header::LEN..]
+                }
                 _ => return Err(ProgramError::InvalidAccountData),
             },
             _ => return Err(ProgramError::InvalidInstructionData),