diff --git a/program/rust/src/accounts/price.rs b/program/rust/src/accounts/price.rs index 421cc827..82f2d62c 100644 --- a/program/rust/src/accounts/price.rs +++ b/program/rust/src/accounts/price.rs @@ -1,4 +1,6 @@ pub use price_pythnet::*; +#[cfg(test)] +use quickcheck::Arbitrary; use { super::{ AccountHeader, @@ -188,6 +190,7 @@ mod price_pythnet { } #[repr(C)] +#[cfg_attr(test, derive(Debug, PartialEq))] #[derive(Copy, Clone, Pod, Zeroable)] pub struct PriceComponent { pub pub_: Pubkey, @@ -195,7 +198,21 @@ pub struct PriceComponent { pub latest_: PriceInfo, } +#[cfg(test)] +impl Arbitrary for PriceComponent { + fn arbitrary(g: &mut quickcheck::Gen) -> Self { + let mut key = [0u8; 32]; + key.iter_mut().for_each(|item| *item = u8::arbitrary(g)); + PriceComponent { + pub_: Pubkey::new_from_array(key), + agg_: PriceInfo::arbitrary(g), + latest_: PriceInfo::arbitrary(g), + } + } +} + #[repr(C)] +#[cfg_attr(test, derive(Debug, PartialEq))] #[derive(Copy, Clone, Pod, Zeroable)] pub struct PriceInfo { pub price_: i64, @@ -205,6 +222,19 @@ pub struct PriceInfo { pub pub_slot_: u64, } +#[cfg(test)] +impl Arbitrary for PriceInfo { + fn arbitrary(g: &mut quickcheck::Gen) -> Self { + PriceInfo { + price_: i64::arbitrary(g), + conf_: u64::arbitrary(g), + status_: u32::arbitrary(g), + corp_act_status_: u32::arbitrary(g), + pub_slot_: u64::arbitrary(g), + } + } +} + #[repr(C)] #[derive(Copy, Clone, Pod, Zeroable)] pub struct PriceEma { diff --git a/program/rust/src/processor/add_publisher.rs b/program/rust/src/processor/add_publisher.rs index 2e6a081e..841de3b9 100644 --- a/program/rust/src/processor/add_publisher.rs +++ b/program/rust/src/processor/add_publisher.rs @@ -24,7 +24,10 @@ use { account_info::AccountInfo, entrypoint::ProgramResult, program_error::ProgramError, - program_memory::sol_memset, + program_memory::{ + sol_memcmp, + sol_memset, + }, pubkey::Pubkey, }, std::mem::size_of, @@ -41,8 +44,7 @@ pub fn add_publisher( let cmd_args = load::(instruction_data)?; pyth_assert( - instruction_data.len() == size_of::() - && cmd_args.publisher != Pubkey::default(), + instruction_data.len() == size_of::(), ProgramError::InvalidArgument, )?; @@ -63,6 +65,14 @@ pub fn add_publisher( let mut price_data = load_checked::(price_account, cmd_args.header.version)?; + // Use the call with the default pubkey (000..) as a trigger to sort the publishers as a + // migration step from unsorted list to sorted list. + if cmd_args.publisher == Pubkey::default() { + let num_comps = try_convert::(price_data.num_)?; + sort_price_comps(&mut price_data.comp_, num_comps)?; + return Ok(()); + } + if price_data.num_ >= PC_NUM_COMP { return Err(ProgramError::InvalidArgument); } @@ -81,6 +91,141 @@ pub fn add_publisher( ); price_data.comp_[current_index].pub_ = cmd_args.publisher; price_data.num_ += 1; + + // Sort the publishers in the list + { + let num_comps = try_convert::(price_data.num_)?; + sort_price_comps(&mut price_data.comp_, num_comps)?; + } + price_data.header.size = try_convert::<_, u32>(PriceAccount::INITIAL_SIZE)?; Ok(()) } + +/// A copy of rust slice/sort.rs heapsort implementation which is small and fast. We couldn't use +/// the sort directly because it was only accessible behind a unstable feature flag at the time of +/// writing this code. +fn heapsort(v: &mut [(Pubkey, usize)]) { + // This binary heap respects the invariant `parent >= child`. + let sift_down = |v: &mut [(Pubkey, usize)], mut node: usize| { + loop { + // Children of `node`. + let mut child = 2 * node + 1; + if child >= v.len() { + break; + } + + // Choose the greater child. + if child + 1 < v.len() + && sol_memcmp(v[child].0.as_ref(), v[child + 1].0.as_ref(), 32) < 0 + { + child += 1; + } + + // Stop if the invariant holds at `node`. + if sol_memcmp(v[node].0.as_ref(), v[child].0.as_ref(), 32) >= 0 { + break; + } + + // Swap `node` with the greater child, move one step down, and continue sifting. + v.swap(node, child); + node = child; + } + }; + + // Build the heap in linear time. + for i in (0..v.len() / 2).rev() { + sift_down(v, i); + } + + // Pop maximal elements from the heap. + for i in (1..v.len()).rev() { + v.swap(0, i); + sift_down(&mut v[..i], 0); + } +} + +/// Sort the publishers price component list in place by performing minimal swaps. +/// This code is inspired by the sort_by_cached_key implementation in the Rust stdlib. +/// The rust stdlib implementation is not used because it uses a fast sort variant that has +/// a large code size. +/// +/// num_publishers is the number of publishers in the list that should be sorted. It is explicitly +/// passed to avoid callers mistake of passing the full slice which may contain uninitialized values. +fn sort_price_comps(comps: &mut [PriceComponent], num_comps: usize) -> Result<(), ProgramError> { + let comps = comps + .get_mut(..num_comps) + .ok_or(ProgramError::InvalidArgument)?; + + // Publishers are likely sorted in ascending order but + // heapsorts creates a max-heap so we reverse the order + // of the keys to make the heapify step faster. + let mut keys = comps + .iter() + .enumerate() + .map(|(i, x)| (x.pub_, i)) + .rev() + .collect::>(); + + heapsort(&mut keys); + + for i in 0..num_comps { + // We know that the publisher with key[i].0 should be at index i in the sorted array and + // want to swap them in-place in O(n). Normally, the publisher at key[i].0 should be at + // key[i].1 but if it is swapped, we need to find the correct index by following the chain + // of swaps. + let mut index = keys[i].1; + + while index < i { + index = keys[index].1; + } + // Setting the final index here is important to make the code linear as we won't + // loop over from i to index again when we reach i again. + keys[i].1 = index; + comps.swap(i, index); + } + + Ok(()) +} + +#[cfg(test)] +mod test { + use { + super::*, + quickcheck_macros::quickcheck, + }; + + #[quickcheck] + pub fn test_sort_price_comps(mut comps: Vec) { + let mut rust_std_sorted_comps = comps.clone(); + rust_std_sorted_comps.sort_by_key(|x| x.pub_); + + let num_comps = comps.len(); + assert_eq!( + sort_price_comps(&mut comps, num_comps + 1), + Err(ProgramError::InvalidArgument) + ); + + assert_eq!(sort_price_comps(&mut comps, num_comps), Ok(())); + assert_eq!(comps, rust_std_sorted_comps); + } + + #[quickcheck] + pub fn test_sort_price_comps_smaller_slice( + mut comps: Vec, + mut num_comps: usize, + ) { + num_comps = if comps.is_empty() { + 0 + } else { + num_comps % comps.len() + }; + + let mut rust_std_sorted_comps = comps.get(..num_comps).unwrap().to_vec(); + rust_std_sorted_comps.sort_by_key(|x| x.pub_); + + + assert_eq!(sort_price_comps(&mut comps, num_comps), Ok(())); + assert_eq!(comps.get(..num_comps).unwrap(), rust_std_sorted_comps); + } +} diff --git a/program/rust/src/processor/upd_price.rs b/program/rust/src/processor/upd_price.rs index 5e62f241..278673b4 100644 --- a/program/rust/src/processor/upd_price.rs +++ b/program/rust/src/processor/upd_price.rs @@ -2,6 +2,7 @@ use { crate::{ accounts::{ PriceAccount, + PriceComponent, PriceInfo, PythOracleSerialize, UPD_PRICE_WRITE_SEED, @@ -31,6 +32,7 @@ use { }, program::invoke_signed, program_error::ProgramError, + program_memory::sol_memcmp, pubkey::Pubkey, sysvar::Sysvar, }, @@ -127,7 +129,7 @@ pub fn upd_price( // Check clock let clock = Clock::from_account_info(clock_account)?; - let mut publisher_index: usize = 0; + let publisher_index: usize; let latest_aggregate_price: PriceInfo; // The price_data borrow happens in a scope because it must be @@ -137,17 +139,15 @@ pub fn upd_price( // Verify that symbol account is initialized let price_data = load_checked::(price_account, cmd_args.header.version)?; - // Verify that publisher is authorized - while publisher_index < try_convert::(price_data.num_)? { - if price_data.comp_[publisher_index].pub_ == *funding_account.key { - break; + publisher_index = match find_publisher_index( + &price_data.comp_[..try_convert::(price_data.num_)?], + funding_account.key, + ) { + Some(index) => index, + None => { + return Err(OracleError::PermissionViolation.into()); } - publisher_index += 1; - } - pyth_assert( - publisher_index < try_convert::(price_data.num_)?, - OracleError::PermissionViolation.into(), - )?; + }; latest_aggregate_price = price_data.agg_; let latest_publisher_price = price_data.comp_[publisher_index].latest_; @@ -281,6 +281,64 @@ pub fn upd_price( Ok(()) } +/// Find the index of the publisher in the list of components. +/// +/// This method first tries to binary search for the publisher's key in the list of components +/// to get the result faster if the list is sorted. If the list is not sorted, it falls back to +/// a linear search. +#[inline(always)] +fn find_publisher_index(comps: &[PriceComponent], key: &Pubkey) -> Option { + // Verify that publisher is authorized by initially binary searching + // for the publisher's component in the price account. The binary + // search might not work if the publisher list is not sorted; therefore + // we fall back to a linear search. + let mut binary_search_result = None; + + { + // Binary search to find the publisher key. Rust std binary search is not used because + // they guarantee valid outcome only if the array is sorted whereas we want to rely on + // a Equal match if it is a result on an unsorted array. Currently the rust + // implementation behaves the same but we do not want to rely on api internals. + let mut left = 0; + let mut right = comps.len(); + while left < right { + let mid = left + (right - left) / 2; + // sol_memcmp is much faster than rust default comparison of Pubkey. It costs + // 10CU whereas rust default comparison costs a few times more. + match sol_memcmp(comps[mid].pub_.as_ref(), key.as_ref(), 32) { + i if i < 0 => { + left = mid + 1; + } + i if i > 0 => { + right = mid; + } + _ => { + binary_search_result = Some(mid); + break; + } + } + } + } + + match binary_search_result { + Some(index) => Some(index), + None => { + let mut index = 0; + while index < comps.len() { + if sol_memcmp(comps[index].pub_.as_ref(), key.as_ref(), 32) == 0 { + break; + } + index += 1; + } + if index == comps.len() { + None + } else { + Some(index) + } + } + } +} + #[allow(dead_code)] // Wrapper struct for the accounts required to add data to the accumulator program. struct MessageBufferAccounts<'a, 'b: 'a> { @@ -289,3 +347,45 @@ struct MessageBufferAccounts<'a, 'b: 'a> { oracle_auth_pda: &'a AccountInfo<'b>, message_buffer_data: &'a AccountInfo<'b>, } + +#[cfg(test)] +mod test { + use { + super::*, + crate::accounts::PriceComponent, + quickcheck_macros::quickcheck, + solana_program::pubkey::Pubkey, + }; + + /// Test the find_publisher_index method works with an unordered list of components. + #[quickcheck] + pub fn test_find_publisher_index_unordered_comp(comps: Vec) { + comps.iter().enumerate().for_each(|(idx, comp)| { + assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx)); + }); + + let mut key_not_in_list = Pubkey::new_unique(); + while comps.iter().any(|comp| comp.pub_ == key_not_in_list) { + key_not_in_list = Pubkey::new_unique(); + } + + assert_eq!(find_publisher_index(&comps, &key_not_in_list), None); + } + + /// Test the find_publisher_index method works with a sorted list of components. + #[quickcheck] + pub fn test_find_publisher_index_ordered_comp(mut comps: Vec) { + comps.sort_by_key(|comp| comp.pub_); + + comps.iter().enumerate().for_each(|(idx, comp)| { + assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx)); + }); + + let mut key_not_in_list = Pubkey::new_unique(); + while comps.iter().any(|comp| comp.pub_ == key_not_in_list) { + key_not_in_list = Pubkey::new_unique(); + } + + assert_eq!(find_publisher_index(&comps, &key_not_in_list), None); + } +} diff --git a/program/rust/src/tests/test_add_publisher.rs b/program/rust/src/tests/test_add_publisher.rs index cc7b3c90..45d6172c 100644 --- a/program/rust/src/tests/test_add_publisher.rs +++ b/program/rust/src/tests/test_add_publisher.rs @@ -144,7 +144,7 @@ fn test_add_publisher() { { let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); assert_eq!(price_data.num_, i + 1); - assert!(price_data.comp_[i as usize].pub_ == cmd.publisher); + assert!(price_data.comp_.iter().any(|c| c.pub_ == cmd.publisher)); assert_eq!(price_data.header.size, PriceAccount::INITIAL_SIZE); } } @@ -163,4 +163,45 @@ fn test_add_publisher() { ), Err(ProgramError::InvalidArgument) ); + + // Make sure that publishers are sorted + { + let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + for i in 1..PC_NUM_COMP { + assert!(price_data.comp_[i as usize].pub_ > price_data.comp_[(i - 1) as usize].pub_); + } + } + + // Test sorting by reordering the publishers to be in reverse order + // and then adding the default (000...) publisher to trigger the sorting. + { + let mut price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + price_data + .comp_ + .get_mut(..PC_NUM_COMP as usize) + .unwrap() + .reverse(); + } + + cmd.publisher = Pubkey::default(); + instruction_data = bytes_of::(&cmd); + assert!(process_instruction( + &program_id, + &[ + funding_account.clone(), + price_account.clone(), + permissions_account.clone(), + ], + instruction_data + ) + .is_ok()); + + // Make sure that publishers get sorted after adding the default publisher + { + let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + assert!(price_data.num_ == PC_NUM_COMP); + for i in 1..PC_NUM_COMP { + assert!(price_data.comp_[i as usize].pub_ > price_data.comp_[(i - 1) as usize].pub_); + } + } } diff --git a/program/rust/src/tests/test_upd_price_v2.rs b/program/rust/src/tests/test_upd_price_v2.rs index 0e0c1f4b..63785819 100644 --- a/program/rust/src/tests/test_upd_price_v2.rs +++ b/program/rust/src/tests/test_upd_price_v2.rs @@ -444,6 +444,87 @@ fn test_upd_price_v2() -> Result<(), Box> { Ok(()) } +#[test] +fn test_upd_works_with_unordered_publisher_set() -> Result<(), Box> { + let mut instruction_data = [0u8; size_of::()]; + + let program_id = Pubkey::new_unique(); + + let mut price_setup = AccountSetup::new::(&program_id); + let mut price_account = price_setup.as_account_info(); + price_account.is_signer = false; + PriceAccount::initialize(&price_account, PC_VERSION).unwrap(); + + let mut publishers_setup: Vec<_> = (0..20).map(|_| AccountSetup::new_funding()).collect(); + let mut publishers: Vec<_> = publishers_setup + .iter_mut() + .map(|s| s.as_account_info()) + .collect(); + + publishers.sort_by_key(|x| x.key); + + { + let mut price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + price_data.num_ = 20; + // Store the publishers in reverse order + publishers + .iter() + .rev() + .enumerate() + .for_each(|(i, account)| { + price_data.comp_[i].pub_ = *account.key; + }); + } + + let mut clock_setup = AccountSetup::new_clock(); + let mut clock_account = clock_setup.as_account_info(); + clock_account.is_signer = false; + clock_account.is_writable = false; + + update_clock_slot(&mut clock_account, 1); + + for (i, publisher) in publishers.iter().enumerate() { + populate_instruction(&mut instruction_data, (i + 100) as i64, 10, 1); + process_instruction( + &program_id, + &[ + publisher.clone(), + price_account.clone(), + clock_account.clone(), + ], + &instruction_data, + )?; + } + + update_clock_slot(&mut clock_account, 2); + + // Trigger the aggregate calculation by sending another price + // update + populate_instruction(&mut instruction_data, 100, 10, 2); + process_instruction( + &program_id, + &[ + publishers[0].clone(), + price_account.clone(), + clock_account.clone(), + ], + &instruction_data, + )?; + + + let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + + // The result will be the following only if all the + // publishers prices are included in the aggregate. + assert_eq!(price_data.valid_slot_, 1); + assert_eq!(price_data.agg_.pub_slot_, 2); + assert_eq!(price_data.agg_.price_, 109); + assert_eq!(price_data.agg_.conf_, 8); + assert_eq!(price_data.agg_.status_, PC_STATUS_TRADING); + + Ok(()) +} + // Create an upd_price instruction with the provided parameters fn populate_instruction(instruction_data: &mut [u8], price: i64, conf: u64, pub_slot: u64) { let mut cmd = load_mut::(instruction_data).unwrap();