diff --git a/program/rust/src/rust_oracle.rs b/program/rust/src/rust_oracle.rs index 2ff23916b..e91870d96 100644 --- a/program/rust/src/rust_oracle.rs +++ b/program/rust/src/rust_oracle.rs @@ -75,17 +75,19 @@ pub fn init_mapping( accounts: &[AccountInfo], instruction_data: &[u8], ) -> OracleResult { - let [_funding_account, fresh_mapping_account] = match accounts { - [x, y] - if valid_funding_account(x) - && valid_signable_account(program_id, y, size_of::()) - && valid_fresh_account(y) => - { - Ok([x, y]) - } + let [funding_account, fresh_mapping_account] = match accounts { + [x, y] => Ok([x, y]), _ => Err(ProgramError::InvalidArgument), }?; + check_valid_funding_account(funding_account)?; + check_valid_signable_account( + program_id, + fresh_mapping_account, + size_of::(), + )?; + check_valid_fresh_account(fresh_mapping_account)?; + // Initialize by setting to zero again (just in case) and populating the account header let hdr = load::(instruction_data)?; initialize_mapping_account(fresh_mapping_account, hdr.ver_)?; @@ -98,18 +100,16 @@ pub fn add_mapping( accounts: &[AccountInfo], instruction_data: &[u8], ) -> OracleResult { - let [_funding_account, cur_mapping, next_mapping] = match accounts { - [x, y, z] - if valid_funding_account(x) - && valid_signable_account(program_id, y, size_of::()) - && valid_signable_account(program_id, z, size_of::()) - && valid_fresh_account(z) => - { - Ok([x, y, z]) - } + let [funding_account, cur_mapping, next_mapping] = match accounts { + [x, y, z] => Ok([x, y, z]), _ => Err(ProgramError::InvalidArgument), }?; + check_valid_funding_account(funding_account)?; + check_valid_signable_account(program_id, cur_mapping, size_of::())?; + check_valid_signable_account(program_id, next_mapping, size_of::())?; + check_valid_fresh_account(next_mapping)?; + let hdr = load::(instruction_data)?; let mut cur_mapping = load_mapping_account_mut(cur_mapping, hdr.ver_)?; pyth_assert( @@ -142,18 +142,16 @@ pub fn add_price( return Err(ProgramError::InvalidArgument); } - let [_funding_account, product_account, price_account] = match accounts { - [x, y, z] - if valid_funding_account(x) - && valid_signable_account(program_id, y, PC_PROD_ACC_SIZE as usize) - && valid_signable_account(program_id, z, size_of::()) - && valid_fresh_account(z) => - { - Ok([x, y, z]) - } + let [funding_account, product_account, price_account] = match accounts { + [x, y, z] => Ok([x, y, z]), _ => Err(ProgramError::InvalidArgument), }?; + check_valid_funding_account(funding_account)?; + check_valid_signable_account(program_id, product_account, PC_PROD_ACC_SIZE as usize)?; + check_valid_signable_account(program_id, price_account, size_of::())?; + check_valid_fresh_account(price_account)?; + let mut product_data = load_product_account_mut(product_account, cmd_args.ver_)?; clear_account(price_account)?; @@ -177,18 +175,20 @@ pub fn add_product( accounts: &[AccountInfo], instruction_data: &[u8], ) -> OracleResult { - let [_funding_account, tail_mapping_account, new_product_account] = match accounts { - [x, y, z] - if valid_funding_account(x) - && valid_signable_account(program_id, y, size_of::()) - && valid_signable_account(program_id, z, PC_PROD_ACC_SIZE as usize) - && valid_fresh_account(z) => - { - Ok([x, y, z]) - } + let [funding_account, tail_mapping_account, new_product_account] = match accounts { + [x, y, z] => Ok([x, y, z]), _ => Err(ProgramError::InvalidArgument), }?; + check_valid_funding_account(funding_account)?; + check_valid_signable_account( + program_id, + tail_mapping_account, + size_of::(), + )?; + check_valid_signable_account(program_id, new_product_account, PC_PROD_ACC_SIZE as usize)?; + check_valid_fresh_account(new_product_account)?; + let hdr = load::(instruction_data)?; let mut mapping_data = load_mapping_account_mut(tail_mapping_account, hdr.ver_)?; // The mapping account must have free space to add the product account @@ -217,6 +217,13 @@ fn valid_funding_account(account: &AccountInfo) -> bool { account.is_signer && account.is_writable } +fn check_valid_funding_account(account: &AccountInfo) -> Result<(), ProgramError> { + pyth_assert( + valid_funding_account(account), + ProgramError::InvalidArgument, + ) +} + fn valid_signable_account(program_id: &Pubkey, account: &AccountInfo, minimum_size: usize) -> bool { account.is_signer && account.is_writable @@ -225,6 +232,17 @@ fn valid_signable_account(program_id: &Pubkey, account: &AccountInfo, minimum_si && Rent::default().is_exempt(account.lamports(), account.data_len()) } +fn check_valid_signable_account( + program_id: &Pubkey, + account: &AccountInfo, + minimum_size: usize, +) -> Result<(), ProgramError> { + pyth_assert( + valid_signable_account(program_id, account, minimum_size), + ProgramError::InvalidArgument, + ) +} + /// Returns `true` if the `account` is fresh, i.e., its data can be overwritten. /// Use this check to prevent accidentally overwriting accounts whose data is already populated. fn valid_fresh_account(account: &AccountInfo) -> bool { @@ -235,6 +253,10 @@ fn valid_fresh_account(account: &AccountInfo) -> bool { } } +fn check_valid_fresh_account(account: &AccountInfo) -> Result<(), ProgramError> { + pyth_assert(valid_fresh_account(account), ProgramError::InvalidArgument) +} + /// Sets the data of account to all-zero pub fn clear_account(account: &AccountInfo) -> Result<(), ProgramError> { let mut data = account @@ -275,7 +297,7 @@ pub fn initialize_mapping_account(account: &AccountInfo, version: u32) -> Result mapping_account.ver_ = version; mapping_account.type_ = PC_ACCTYPE_MAPPING; mapping_account.size_ = - (size_of::() - size_of_val(&mapping_account.prod_)) as u32; + try_convert(size_of::() - size_of_val(&mapping_account.prod_))?; Ok(()) }