Implement add_mapping instruction (#221)

* add_mapping instruction

* cleanup

Co-authored-by: Jayant Krishnamurthy <[email protected]>

Author: jayantk

program/rust/src/lib.rs

@@ -11,6 +11,7 @@ mod log;
 mod processor;
 mod rust_oracle;
 mod time_machine_types;
+mod utils;
 
 use crate::c_oracle_header::SUCCESSFULLY_UPDATED_AGGREGATE;
 use crate::error::{

program/rust/src/processor.rs

@@ -8,6 +8,7 @@ use solana_program::sysvar::slot_history::AccountInfo;
 use crate::c_entrypoint_wrapper;
 use crate::c_oracle_header::{
     cmd_hdr,
+    command_t_e_cmd_add_mapping,
     command_t_e_cmd_agg_price,
     command_t_e_cmd_init_mapping,
     command_t_e_cmd_upd_account_version,
@@ -20,6 +21,7 @@ use crate::error::{
     OracleResult,
 };
 use crate::rust_oracle::{
+    add_mapping,
     init_mapping,
     update_price,
     update_version,
@@ -57,6 +59,7 @@ pub fn process_instruction(
             update_version(program_id, accounts, instruction_data)
         }
         command_t_e_cmd_init_mapping => init_mapping(program_id, accounts, instruction_data),
+        command_t_e_cmd_add_mapping => add_mapping(program_id, accounts, instruction_data),
         _ => c_entrypoint_wrapper(input),
     }
 }

program/rust/src/rust_oracle.rs

@@ -26,9 +26,12 @@ use crate::c_oracle_header::{
     pc_map_table_t,
     PC_ACCTYPE_MAPPING,
     PC_MAGIC,
+    PC_MAP_TABLE_SIZE,
 };
 use crate::error::OracleResult;
 
+use crate::utils::pyth_assert;
+
 use super::c_entrypoint_wrapper;
 
 ///Calls the c oracle update_price, and updates the Time Machine if needed
@@ -75,14 +78,44 @@ pub fn init_mapping(
     }?;
 
     // Initialize by setting to zero again (just in case) and populating the account header
-    clear_account(fresh_mapping_account)?;
+    let hdr = load::<cmd_hdr_t>(instruction_data)?;
+    initialize_mapping_account(fresh_mapping_account, hdr.ver_)?;
+
+    Ok(SUCCESS)
+}
+
+pub fn add_mapping(
+    program_id: &Pubkey,
+    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::<pc_map_table_t>())
+                && valid_signable_account(program_id, z, size_of::<pc_map_table_t>())
+                && valid_fresh_account(z) =>
+        {
+            Ok([x, y, z])
+        }
+        _ => Err(ProgramError::InvalidArgument),
+    }?;
 
     let hdr = load::<cmd_hdr_t>(instruction_data)?;
-    let mut mapping_data = load_account_as_mut::<pc_map_table_t>(fresh_mapping_account)?;
-    mapping_data.magic_ = PC_MAGIC;
-    mapping_data.ver_ = hdr.ver_;
-    mapping_data.type_ = PC_ACCTYPE_MAPPING;
-    mapping_data.size_ = (size_of::<pc_map_table_t>() - size_of_val(&mapping_data.prod_)) as u32;
+    let mut cur_mapping = load_mapping_account_mut(cur_mapping, hdr.ver_)?;
+    pyth_assert(
+        cur_mapping.num_ == PC_MAP_TABLE_SIZE
+            && unsafe { cur_mapping.next_.k8_.iter().all(|x| *x == 0) },
+        ProgramError::InvalidArgument,
+    )?;
+
+    initialize_mapping_account(next_mapping, hdr.ver_)?;
+    unsafe {
+        cur_mapping
+            .next_
+            .k1_
+            .copy_from_slice(&next_mapping.key.to_bytes());
+    }
 
     Ok(SUCCESS)
 }
@@ -150,3 +183,38 @@ fn load_account_as_mut<'a, T: Pod>(
         bytemuck::from_bytes_mut(&mut data[0..size_of::<T>()])
     }))
 }
+
+/// Mutably borrow the data in `account` as a mapping account, validating that the account
+/// is properly formatted. Any mutations to the returned value will be reflected in the
+/// account data. Use this to read already-initialized accounts.
+fn load_mapping_account_mut<'a>(
+    account: &'a AccountInfo,
+    expected_version: u32,
+) -> Result<RefMut<'a, pc_map_table_t>, ProgramError> {
+    let mapping_account_ref = load_account_as_mut::<pc_map_table_t>(account)?;
+    let mapping_account = *mapping_account_ref;
+
+    pyth_assert(
+        mapping_account.magic_ == PC_MAGIC
+            && mapping_account.ver_ == expected_version
+            && mapping_account.type_ == PC_ACCTYPE_MAPPING,
+        ProgramError::InvalidArgument,
+    )?;
+
+    Ok(mapping_account_ref)
+}
+
+/// Initialize account as a new mapping account. This function will zero out any existing data in
+/// the account.
+fn initialize_mapping_account(account: &AccountInfo, version: u32) -> Result<(), ProgramError> {
+    clear_account(account)?;
+
+    let mut mapping_account = load_account_as_mut::<pc_map_table_t>(account)?;
+    mapping_account.magic_ = PC_MAGIC;
+    mapping_account.ver_ = version;
+    mapping_account.type_ = PC_ACCTYPE_MAPPING;
+    mapping_account.size_ =
+        (size_of::<pc_map_table_t>() - size_of_val(&mapping_account.prod_)) as u32;
+
+    Ok(())
+}

program/rust/src/utils.rs

@@ -0,0 +1,9 @@
+use solana_program::program_error::ProgramError;
+
+pub fn pyth_assert(condition: bool, error_code: ProgramError) -> Result<(), ProgramError> {
+    if !condition {
+        Result::Err(error_code)
+    } else {
+        Result::Ok(())
+    }
+}