solana: Add error handling for Bitmap indexing

Author: johnsaigle

solana/programs/example-native-token-transfers/src/bitmap.rs

@@ -1,6 +1,7 @@
+use crate::error::NTTError;
 use anchor_lang::prelude::*;
 use bitmaps::Bitmap as BM;
-
+use std::result::Result as StdResult;
 #[derive(PartialEq, Eq, Clone, Copy, Debug, AnchorDeserialize, AnchorSerialize, InitSpace)]
 pub struct Bitmap {
     map: u128,
@@ -13,6 +14,8 @@ impl Default for Bitmap {
 }
 
 impl Bitmap {
+    pub const BITS: u8 = 128;
+
     pub fn new() -> Self {
         Bitmap { map: 0 }
     }
@@ -21,19 +24,30 @@ impl Bitmap {
         Bitmap { map: value }
     }
 
-    pub fn set(&mut self, index: u8, value: bool) {
+    pub fn set(&mut self, index: u8, value: bool) -> StdResult<(), NTTError> {
+        if index >= Self::BITS {
+            return Err(NTTError::BitmapIndexOutOfBounds);
+        }
         let mut bm = BM::<128>::from_value(self.map);
-        bm.set(index as usize, value);
+        bm.set(usize::from(index), value);
         self.map = *bm.as_value();
+        Ok(())
     }
 
-    pub fn get(&self, index: u8) -> bool {
-        BM::<128>::from_value(self.map).get(index as usize)
+    pub fn get(&self, index: u8) -> StdResult<bool, NTTError> {
+        if index >= Self::BITS {
+            return Err(NTTError::BitmapIndexOutOfBounds);
+        }
+        Ok(BM::<128>::from_value(self.map).get(usize::from(index)))
     }
 
     pub fn count_enabled_votes(&self, enabled: Bitmap) -> u8 {
         let bm = BM::<128>::from_value(self.map) & BM::<128>::from_value(enabled.map);
-        bm.len() as u8
+        // Conversion from usize to u8 is safe here. The Bitmap uses u128, so its maximum length
+        // (number of true bits) is 128.
+        bm.len()
+            .try_into()
+            .expect("Bitmap length must not exceed the bounds of u8")
     }
 }
 
@@ -46,22 +60,40 @@ mod tests {
         let mut enabled = Bitmap::from_value(u128::MAX);
         let mut bm = Bitmap::new();
         assert_eq!(bm.count_enabled_votes(enabled), 0);
-        bm.set(0, true);
+        bm.set(0, true).unwrap();
         assert_eq!(bm.count_enabled_votes(enabled), 1);
-        assert!(bm.get(0));
-        assert!(!bm.get(1));
-        bm.set(1, true);
+        assert!(bm.get(0).unwrap());
+        assert!(!bm.get(1).unwrap());
+        bm.set(1, true).unwrap();
         assert_eq!(bm.count_enabled_votes(enabled), 2);
-        assert!(bm.get(0));
-        assert!(bm.get(1));
-        bm.set(0, false);
+        assert!(bm.get(0).unwrap());
+        assert!(bm.get(1).unwrap());
+        bm.set(0, false).unwrap();
         assert_eq!(bm.count_enabled_votes(enabled), 1);
-        assert!(!bm.get(0));
-        assert!(bm.get(1));
-        bm.set(18, true);
+        assert!(!bm.get(0).unwrap());
+        assert!(bm.get(1).unwrap());
+        bm.set(18, true).unwrap();
         assert_eq!(bm.count_enabled_votes(enabled), 2);
 
-        enabled.set(18, false);
+        enabled.set(18, false).unwrap();
         assert_eq!(bm.count_enabled_votes(enabled), 1);
     }
+
+    #[test]
+    fn test_bitmap_len() {
+        let max_bitmap = Bitmap::from_value(u128::MAX);
+        assert_eq!(128, max_bitmap.count_enabled_votes(max_bitmap));
+    }
+
+    #[test]
+    fn test_bitmap_get_out_of_bounds() {
+        let bm = Bitmap::new();
+        assert_eq!(bm.get(129), Err(NTTError::BitmapIndexOutOfBounds));
+    }
+
+    #[test]
+    fn test_bitmap_set_out_of_bounds() {
+        let mut bm = Bitmap::new();
+        assert_eq!(bm.set(129, false), Err(NTTError::BitmapIndexOutOfBounds));
+    }
 }

solana/programs/example-native-token-transfers/src/config.rs

@@ -31,7 +31,8 @@ pub struct Config {
     /// The number of transceivers that must attest to a transfer before it is
     /// accepted.
     pub threshold: u8,
-    /// Bitmap of enabled transceivers
+    /// Bitmap of enabled transceivers.
+    /// The maximum number of transceivers is equal to [`Bitmap::BITS`].
     pub enabled_transceivers: Bitmap,
     /// Pause the program. This is useful for upgrades and other maintenance.
     pub paused: bool,

solana/programs/example-native-token-transfers/src/error.rs

@@ -3,6 +3,7 @@ use ntt_messages::errors::ScalingError;
 
 #[error_code]
 // TODO(csongor): rename
+#[derive(PartialEq)]
 pub enum NTTError {
     #[msg("CantReleaseYet")]
     CantReleaseYet,
@@ -48,6 +49,8 @@ pub enum NTTError {
     OverflowExponent,
     #[msg("OverflowScaledAmount")]
     OverflowScaledAmount,
+    #[msg("BitmapIndexOutOfBounds")]
+    BitmapIndexOutOfBounds,
 }
 
 impl From<ScalingError> for NTTError {

solana/programs/example-native-token-transfers/src/instructions/admin.rs

@@ -229,7 +229,7 @@ pub fn register_transceiver(ctx: Context<RegisterTransceiver>) -> Result<()> {
             transceiver_address: ctx.accounts.transceiver.key(),
         });
 
-    ctx.accounts.config.enabled_transceivers.set(id, true);
+    ctx.accounts.config.enabled_transceivers.set(id, true)?;
     Ok(())
 }
 

solana/programs/example-native-token-transfers/src/instructions/redeem.rs

@@ -46,7 +46,7 @@ pub struct Redeem<'info> {
     pub transceiver_message: Account<'info, ValidatedTransceiverMessage<NativeTokenTransfer>>,
 
     #[account(
-        constraint = config.enabled_transceivers.get(transceiver.id) @ NTTError::DisabledTransceiver
+        constraint = config.enabled_transceivers.get(transceiver.id)? @ NTTError::DisabledTransceiver
     )]
     pub transceiver: Account<'info, RegisteredTransceiver>,
 
@@ -131,7 +131,7 @@ pub fn redeem(ctx: Context<Redeem>, _args: RedeemArgs) -> Result<()> {
     }
 
     // idempotent
-    accs.inbox_item.votes.set(accs.transceiver.id, true);
+    accs.inbox_item.votes.set(accs.transceiver.id, true)?;
 
     if accs
         .inbox_item

solana/programs/example-native-token-transfers/src/queue/outbox.rs

@@ -31,11 +31,11 @@ impl OutboxItem {
             return Ok(false);
         }
 
-        if self.released.get(transceiver_index) {
+        if self.released.get(transceiver_index)? {
             return Err(NTTError::MessageAlreadySent.into());
         }
 
-        self.released.set(transceiver_index, true);
+        self.released.set(transceiver_index, true)?;
 
         Ok(true)
     }

solana/programs/example-native-token-transfers/src/transceivers/wormhole/instructions/release_outbound.rs

@@ -19,13 +19,13 @@ pub struct ReleaseOutbound<'info> {
 
     #[account(
         mut,
-        constraint = !outbox_item.released.get(transceiver.id) @ NTTError::MessageAlreadySent,
+        constraint = !outbox_item.released.get(transceiver.id)? @ NTTError::MessageAlreadySent,
     )]
     pub outbox_item: Account<'info, OutboxItem>,
 
     #[account(
         constraint = transceiver.transceiver_address == crate::ID,
-        constraint = config.enabled_transceivers.get(transceiver.id) @ NTTError::DisabledTransceiver
+        constraint = config.enabled_transceivers.get(transceiver.id)? @ NTTError::DisabledTransceiver
     )]
     pub transceiver: Account<'info, RegisteredTransceiver>,
 
@@ -65,7 +65,7 @@ pub fn release_outbound(ctx: Context<ReleaseOutbound>, args: ReleaseOutboundArgs
         }
     }
 
-    assert!(accs.outbox_item.released.get(accs.transceiver.id));
+    assert!(accs.outbox_item.released.get(accs.transceiver.id)?);
     let message: TransceiverMessage<WormholeTransceiver, NativeTokenTransfer> =
         TransceiverMessage::new(
             // TODO: should we just put the ntt id here statically?

solana/programs/example-native-token-transfers/tests/transfer.rs

@@ -370,7 +370,7 @@ async fn assert_queued(ctx: &mut ProgramTestContext, outbox_item: Pubkey) {
 
     let clock: Clock = ctx.banks_client.get_sysvar().await.unwrap();
 
-    assert!(!outbox_item_account.released.get(0));
+    assert!(!outbox_item_account.released.get(0).unwrap());
     assert!(outbox_item_account.release_timestamp > clock.unix_timestamp);
 }
 

solana/programs/ntt-quoter/src/processor/request_relay.rs

@@ -19,7 +19,10 @@ fn check_release_constraint_and_fetch_chain_id(
     let mut acc_data: &[_] = &outbox_item.try_borrow_data()?;
     let item = OutboxItem::try_deserialize(&mut acc_data)?;
 
-    if !item.released.get(registered_ntt.wormhole_transceiver_index) {
+    let released = item
+        .released
+        .get(registered_ntt.wormhole_transceiver_index)?;
+    if !released {
         return Err(NttQuoterError::OutboxItemNotReleased.into());
     }