solana_ownable/
ownable_processor.rs

1//! Ownable program
2
3use crate::ownable_instruction::OwnableError;
4use bincode::serialize_into;
5use solana_sdk::{
6    account::{ReadableAccount, WritableAccount},
7    instruction::InstructionError,
8    keyed_account::{keyed_account_at_index, KeyedAccount},
9    process_instruction::InvokeContext,
10    program_utils::limited_deserialize,
11    pubkey::Pubkey,
12};
13
14fn set_owner(
15    account_owner_pubkey: &mut Pubkey,
16    new_owner_pubkey: Pubkey,
17    owner_keyed_account: &KeyedAccount,
18) -> Result<(), InstructionError> {
19    match owner_keyed_account.signer_key() {
20        None => return Err(InstructionError::MissingRequiredSignature),
21        Some(signer_key) => {
22            if account_owner_pubkey != signer_key {
23                return Err(OwnableError::IncorrectOwner.into());
24            }
25            *account_owner_pubkey = new_owner_pubkey;
26        }
27    }
28    Ok(())
29}
30
31pub fn process_instruction(
32    _program_id: &Pubkey,
33    data: &[u8],
34    invoke_context: &mut dyn InvokeContext,
35) -> Result<(), InstructionError> {
36    let keyed_accounts = invoke_context.get_keyed_accounts()?;
37
38    let new_owner_pubkey: Pubkey = limited_deserialize(data)?;
39    let account_keyed_account = &mut keyed_account_at_index(keyed_accounts, 0)?;
40    let mut account_owner_pubkey: Pubkey =
41        limited_deserialize(account_keyed_account.try_account_ref()?.data())?;
42
43    if account_owner_pubkey == Pubkey::default() {
44        account_owner_pubkey = new_owner_pubkey;
45    } else {
46        let owner_keyed_account = &mut keyed_account_at_index(keyed_accounts, 1)?;
47        set_owner(
48            &mut account_owner_pubkey,
49            new_owner_pubkey,
50            owner_keyed_account,
51        )?;
52    }
53
54    let mut account = account_keyed_account.try_account_ref_mut()?;
55    serialize_into(account.data_as_mut_slice(), &account_owner_pubkey)
56        .map_err(|_| InstructionError::AccountDataTooSmall)
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use crate::ownable_instruction;
63    use solana_runtime::{bank::Bank, bank_client::BankClient};
64    use solana_sdk::{
65        account::AccountSharedData,
66        client::SyncClient,
67        genesis_config::create_genesis_config,
68        message::Message,
69        signature::{Keypair, Signature, Signer},
70        system_program,
71        transport::Result,
72    };
73
74    fn create_bank(lamports: u64) -> (Bank, Keypair) {
75        let (genesis_config, mint_keypair) = create_genesis_config(lamports);
76        let mut bank = Bank::new(&genesis_config);
77        bank.add_builtin("ownable_program", crate::id(), process_instruction);
78        (bank, mint_keypair)
79    }
80
81    fn create_bank_client(lamports: u64) -> (BankClient, Keypair) {
82        let (bank, mint_keypair) = create_bank(lamports);
83        (BankClient::new(bank), mint_keypair)
84    }
85
86    fn create_ownable_account(
87        bank_client: &BankClient,
88        payer_keypair: &Keypair,
89        account_keypair: &Keypair,
90        owner_pubkey: &Pubkey,
91        lamports: u64,
92    ) -> Result<Signature> {
93        let instructions = ownable_instruction::create_account(
94            &payer_keypair.pubkey(),
95            &account_keypair.pubkey(),
96            owner_pubkey,
97            lamports,
98        );
99        let message = Message::new(&instructions, Some(&payer_keypair.pubkey()));
100        bank_client.send_and_confirm_message(&[payer_keypair, account_keypair], message)
101    }
102
103    fn send_set_owner(
104        bank_client: &BankClient,
105        payer_keypair: &Keypair,
106        account_pubkey: &Pubkey,
107        old_owner_keypair: &Keypair,
108        new_owner_pubkey: &Pubkey,
109    ) -> Result<Signature> {
110        let instruction = ownable_instruction::set_owner(
111            account_pubkey,
112            &old_owner_keypair.pubkey(),
113            new_owner_pubkey,
114        );
115        let message = Message::new(&[instruction], Some(&payer_keypair.pubkey()));
116        bank_client.send_and_confirm_message(&[payer_keypair, old_owner_keypair], message)
117    }
118
119    #[test]
120    fn test_ownable_set_owner() {
121        let (bank_client, payer_keypair) = create_bank_client(2);
122        let account_keypair = Keypair::new();
123        let account_pubkey = account_keypair.pubkey();
124        let owner_keypair = Keypair::new();
125        let owner_pubkey = owner_keypair.pubkey();
126
127        create_ownable_account(
128            &bank_client,
129            &payer_keypair,
130            &account_keypair,
131            &owner_pubkey,
132            1,
133        )
134        .unwrap();
135
136        let new_owner_keypair = Keypair::new();
137        let new_owner_pubkey = new_owner_keypair.pubkey();
138        send_set_owner(
139            &bank_client,
140            &payer_keypair,
141            &account_pubkey,
142            &owner_keypair,
143            &new_owner_pubkey,
144        )
145        .unwrap();
146
147        let account_data = bank_client
148            .get_account_data(&account_pubkey)
149            .unwrap()
150            .unwrap();
151        let account_owner_pubkey: Pubkey = limited_deserialize(&account_data).unwrap();
152        assert_eq!(account_owner_pubkey, new_owner_pubkey);
153    }
154
155    #[test]
156    fn test_ownable_missing_owner_signature() {
157        let mut account_owner_pubkey = solana_sdk::pubkey::new_rand();
158        let owner_pubkey = account_owner_pubkey;
159        let new_owner_pubkey = solana_sdk::pubkey::new_rand();
160        let account = AccountSharedData::new_ref(1, 0, &system_program::id());
161        let owner_keyed_account = KeyedAccount::new(&owner_pubkey, false, &account); // <-- Attack! Setting owner without the original owner's signature.
162        let err = set_owner(
163            &mut account_owner_pubkey,
164            new_owner_pubkey,
165            &owner_keyed_account,
166        )
167        .unwrap_err();
168        assert_eq!(err, InstructionError::MissingRequiredSignature);
169    }
170
171    #[test]
172    fn test_ownable_incorrect_owner() {
173        let mut account_owner_pubkey = solana_sdk::pubkey::new_rand();
174        let new_owner_pubkey = solana_sdk::pubkey::new_rand();
175        let account = AccountSharedData::new_ref(1, 0, &system_program::id());
176        let mallory_pubkey = solana_sdk::pubkey::new_rand(); // <-- Attack! Signing with wrong pubkey
177        let owner_keyed_account = KeyedAccount::new(&mallory_pubkey, true, &account);
178        let err = set_owner(
179            &mut account_owner_pubkey,
180            new_owner_pubkey,
181            &owner_keyed_account,
182        )
183        .unwrap_err();
184        assert_eq!(err, OwnableError::IncorrectOwner.into());
185    }
186}
solana_ownable/ownable_processor.rs

solana_ownable/
ownable_processor.rs
