Merge pull request #4 from jito-foundation/ck/fee-hotfix

Precision Fee

core/src/fees.rs

@@ -1,7 +1,8 @@
 use bytemuck::{Pod, Zeroable};
 use jito_bytemuck::types::PodU64;
 use shank::ShankType;
-use solana_program::pubkey::Pubkey;
+use solana_program::{msg, pubkey::Pubkey};
+use spl_math::precise_number::PreciseNumber;
 
 use crate::{error::TipRouterError, MAX_FEE_BPS};
 
@@ -52,10 +53,28 @@ impl Fees {
         }
     }
 
+    pub fn check_fees_okay(&self, current_epoch: u64) -> Result<(), TipRouterError> {
+        let _ = self.precise_block_engine_fee(current_epoch)?;
+        let _ = self.precise_dao_fee(current_epoch)?;
+        let _ = self.precise_ncn_fee(current_epoch)?;
+
+        Ok(())
+    }
+
     pub fn block_engine_fee(&self, current_epoch: u64) -> u64 {
         self.current_fee(current_epoch).block_engine_fee_bps()
     }
 
+    pub fn precise_block_engine_fee(
+        &self,
+        current_epoch: u64,
+    ) -> Result<PreciseNumber, TipRouterError> {
+        let fee = self.current_fee(current_epoch);
+
+        PreciseNumber::new(fee.block_engine_fee_bps() as u128)
+            .ok_or(TipRouterError::NewPreciseNumberError)
+    }
+
     /// Calculate fee as a portion of remaining BPS after block engine fee
     /// new_fee = dao_fee_bps / ((10000 - block_engine_fee_bps) / 10000)
     /// = dao_fee_bps * 10000 / (10000 - block_engine_fee_bps)
@@ -67,21 +86,68 @@ impl Fees {
         fee.dao_share_bps()
             .checked_mul(MAX_FEE_BPS)
             .and_then(|x| x.checked_div(remaining_bps))
-            .ok_or(TipRouterError::ArithmeticOverflow)
+            .ok_or(TipRouterError::DenominatorIsZero)
+    }
+
+    pub fn precise_dao_fee(&self, current_epoch: u64) -> Result<PreciseNumber, TipRouterError> {
+        let fee = self.current_fee(current_epoch);
+
+        let remaining_bps = MAX_FEE_BPS
+            .checked_sub(fee.block_engine_fee_bps())
+            .ok_or(TipRouterError::ArithmeticOverflow)?;
+
+        let precise_remaining_bps = PreciseNumber::new(remaining_bps as u128)
+            .ok_or(TipRouterError::NewPreciseNumberError)?;
+
+        let dao_fee = fee
+            .ncn_share_bps()
+            .checked_mul(MAX_FEE_BPS)
+            .ok_or(TipRouterError::ArithmeticOverflow)?;
+
+        let precise_dao_fee =
+            PreciseNumber::new(dao_fee as u128).ok_or(TipRouterError::NewPreciseNumberError)?;
+
+        precise_dao_fee
+            .checked_div(&precise_remaining_bps)
+            .ok_or(TipRouterError::DenominatorIsZero)
     }
 
     /// Calculate fee as a portion of remaining BPS after block engine fee
     /// new_fee = ncn_fee_bps / ((10000 - block_engine_fee_bps) / 10000)
     /// = ncn_fee_bps * 10000 / (10000 - block_engine_fee_bps)
     pub fn ncn_fee(&self, current_epoch: u64) -> Result<u64, TipRouterError> {
         let fee = self.current_fee(current_epoch);
+
         let remaining_bps = MAX_FEE_BPS
             .checked_sub(fee.block_engine_fee_bps())
             .ok_or(TipRouterError::ArithmeticOverflow)?;
         fee.ncn_share_bps()
             .checked_mul(MAX_FEE_BPS)
             .and_then(|x| x.checked_div(remaining_bps))
-            .ok_or(TipRouterError::ArithmeticOverflow)
+            .ok_or(TipRouterError::DenominatorIsZero)
+    }
+
+    pub fn precise_ncn_fee(&self, current_epoch: u64) -> Result<PreciseNumber, TipRouterError> {
+        let fee = self.current_fee(current_epoch);
+
+        let remaining_bps = MAX_FEE_BPS
+            .checked_sub(fee.block_engine_fee_bps())
+            .ok_or(TipRouterError::ArithmeticOverflow)?;
+
+        let precise_remaining_bps = PreciseNumber::new(remaining_bps as u128)
+            .ok_or(TipRouterError::NewPreciseNumberError)?;
+
+        let ncn_fee = fee
+            .ncn_share_bps()
+            .checked_mul(MAX_FEE_BPS)
+            .ok_or(TipRouterError::ArithmeticOverflow)?;
+
+        let precise_ncn_fee =
+            PreciseNumber::new(ncn_fee as u128).ok_or(TipRouterError::NewPreciseNumberError)?;
+
+        precise_ncn_fee
+            .checked_div(&precise_remaining_bps)
+            .ok_or(TipRouterError::DenominatorIsZero)
     }
 
     pub fn fee_wallet(&self, current_epoch: u64) -> Pubkey {
@@ -130,19 +196,27 @@ impl Fees {
             new_fees.set_ncn_share_bps(new_ncn_fee_bps);
         }
         if let Some(new_block_engine_fee_bps) = new_block_engine_fee_bps {
-            if new_block_engine_fee_bps > MAX_FEE_BPS {
+            // Block engine fee must be less than MAX_FEE_BPS,
+            // otherwise we'll divide by zero when calculating
+            // the other fees
+            if new_block_engine_fee_bps >= MAX_FEE_BPS {
+                msg!("Block engine fee cannot equal or exceed MAX_FEE_BPS");
                 return Err(TipRouterError::FeeCapExceeded);
             }
             new_fees.set_block_engine_fee_bps(new_block_engine_fee_bps);
         }
         if let Some(new_wallet) = new_wallet {
             new_fees.wallet = new_wallet;
         }
-        new_fees.set_activation_epoch(
-            current_epoch
-                .checked_add(1)
-                .ok_or(TipRouterError::ArithmeticOverflow)?,
-        );
+
+        let next_epoch = current_epoch
+            .checked_add(1)
+            .ok_or(TipRouterError::ArithmeticOverflow)?;
+
+        new_fees.set_activation_epoch(next_epoch);
+
+        self.check_fees_okay(next_epoch)?;
+
         Ok(())
     }
 }
@@ -225,6 +299,18 @@ mod tests {
         assert_eq!(fees.fee_1.activation_epoch(), 11);
     }
 
+    #[test]
+    fn test_update_all_fees() {
+        let mut fees = Fees::new(Pubkey::new_unique(), 0, 0, 0, 5);
+
+        fees.set_new_fees(Some(100), Some(200), Some(300), None, 10)
+            .unwrap();
+        assert_eq!(fees.fee_1.dao_share_bps(), 100);
+        assert_eq!(fees.fee_1.ncn_share_bps(), 200);
+        assert_eq!(fees.fee_1.block_engine_fee_bps(), 300);
+        assert_eq!(fees.fee_1.activation_epoch(), 11);
+    }
+
     #[test]
     fn test_update_fees_no_changes() {
         let original = Fee::new(Pubkey::new_unique(), 100, 200, 300, 5);
@@ -246,17 +332,38 @@ mod tests {
     fn test_update_fees_errors() {
         let mut fees = Fees::new(Pubkey::new_unique(), 100, 200, 300, 5);
 
-        assert!(matches!(
+        assert_eq!(
             fees.set_new_fees(Some(10001), None, None, None, 10),
             Err(TipRouterError::FeeCapExceeded)
-        ));
+        );
 
         let mut fees = Fees::new(Pubkey::new_unique(), 100, 200, 300, 5);
 
-        assert!(matches!(
+        assert_eq!(
             fees.set_new_fees(None, None, None, None, u64::MAX),
             Err(TipRouterError::ArithmeticOverflow)
-        ));
+        );
+
+        let mut fees = Fees::new(Pubkey::new_unique(), 100, 200, 300, 5);
+
+        assert_eq!(
+            fees.set_new_fees(None, None, Some(MAX_FEE_BPS), None, 10),
+            Err(TipRouterError::FeeCapExceeded)
+        );
+    }
+
+    #[test]
+    fn test_check_fees_okay() {
+        let fees = Fees::new(Pubkey::new_unique(), 0, 0, 0, 5);
+
+        fees.check_fees_okay(5).unwrap();
+
+        let fees = Fees::new(Pubkey::new_unique(), 0, 0, MAX_FEE_BPS, 5);
+
+        assert_eq!(
+            fees.check_fees_okay(5),
+            Err(TipRouterError::DenominatorIsZero)
+        );
     }
 
     #[test]

program/src/initialize_ncn_config.rs

@@ -57,13 +57,13 @@ pub fn process_initialize_ncn_config(
         return Err(ProgramError::InvalidSeeds);
     }
 
-    if dao_fee_bps > MAX_FEE_BPS {
+    if block_engine_fee_bps >= MAX_FEE_BPS {
         return Err(TipRouterError::FeeCapExceeded.into());
     }
-    if ncn_fee_bps > MAX_FEE_BPS {
+    if dao_fee_bps > MAX_FEE_BPS {
         return Err(TipRouterError::FeeCapExceeded.into());
     }
-    if block_engine_fee_bps > MAX_FEE_BPS {
+    if ncn_fee_bps > MAX_FEE_BPS {
         return Err(TipRouterError::FeeCapExceeded.into());
     }
 
@@ -96,5 +96,7 @@ pub fn process_initialize_ncn_config(
     );
     config.bump = config_bump;
 
+    config.fees.check_fees_okay(epoch)?;
+
     Ok(())
 }