Merge pull request #1103 from gregdhill/fix/supply-unwrap

fix: unsafe usage of unwrap in supply

crates/supply/src/benchmarking.rs

@@ -5,7 +5,7 @@ use super::*;
 use frame_benchmarking::v2::{benchmarks, impl_benchmark_test_suite};
 use frame_support::traits::OnInitialize;
 use frame_system::{self, RawOrigin as SystemOrigin};
-use sp_runtime::traits::One;
+use sp_runtime::traits::{One, Zero};
 use sp_std::prelude::*;
 
 #[benchmarks]

crates/supply/src/lib.rs

@@ -16,6 +16,7 @@ mod tests;
 
 use codec::{Decode, Encode, EncodeLike};
 use frame_support::{
+    pallet_prelude::DispatchResult,
     traits::{Currency, Get, ReservableCurrency},
     transactional,
     weights::Weight,
@@ -24,10 +25,8 @@ use frame_support::{
 use frame_system::ensure_root;
 use primitives::TruncateFixedPointToInt;
 use scale_info::TypeInfo;
-use sp_runtime::{
-    traits::{AccountIdConversion, Saturating, Zero},
-    FixedPointNumber,
-};
+use sp_arithmetic::ArithmeticError;
+use sp_runtime::{traits::AccountIdConversion, FixedPointNumber};
 
 mod default_weights;
 pub use default_weights::WeightInfo;
@@ -90,7 +89,9 @@ pub mod pallet {
     #[pallet::hooks]
     impl<T: Config> Hooks<T::BlockNumber> for Pallet<T> {
         fn on_initialize(n: T::BlockNumber) -> Weight {
-            Self::begin_block(n);
+            if let Err(e) = Self::begin_block(n) {
+                sp_runtime::print(e);
+            }
             T::WeightInfo::on_initialize()
         }
     }
@@ -163,25 +164,25 @@ impl<T: Config> Pallet<T> {
         T::SupplyPalletId::get().into_account_truncating()
     }
 
-    pub(crate) fn begin_block(height: T::BlockNumber) {
+    pub(crate) fn begin_block(height: T::BlockNumber) -> DispatchResult {
         // ignore if uninitialized or not start height
         if let Some(start_height) = <StartHeight<T>>::get().filter(|&start_height| height == start_height) {
             let end_height = start_height + T::InflationPeriod::get();
             <StartHeight<T>>::put(end_height);
 
             let total_supply = T::Currency::total_issuance();
-            let total_supply_as_fixed = T::UnsignedFixedPoint::checked_from_integer(total_supply).unwrap();
-            let total_inflation = total_supply_as_fixed
-                .saturating_mul(<Inflation<T>>::get())
-                .truncate_to_inner()
-                .unwrap_or(Zero::zero());
+            let total_inflation = <Inflation<T>>::get()
+                .checked_mul_int(total_supply)
+                .ok_or(ArithmeticError::Overflow)?;
 
             <LastEmission<T>>::put(total_inflation);
             let supply_account_id = Self::account_id();
             T::Currency::deposit_creating(&supply_account_id, total_inflation);
             T::OnInflation::on_inflation(&supply_account_id, total_inflation);
             Self::deposit_event(Event::<T>::Inflation { total_inflation });
         }
+
+        Ok(())
     }
 }
 

crates/supply/src/mock.rs

@@ -4,13 +4,12 @@ use frame_support::{
     traits::{Everything, GenesisBuild},
     PalletId,
 };
-pub use primitives::CurrencyId;
-use primitives::UnsignedFixedPoint;
+pub use primitives::{CurrencyId, UnsignedFixedPoint};
 use sp_core::H256;
+pub use sp_runtime::FixedPointNumber;
 use sp_runtime::{
     testing::Header,
     traits::{BlakeTwo256, IdentityLookup},
-    FixedPointNumber,
 };
 
 type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<Test>;

crates/supply/src/tests.rs

@@ -1,19 +1,40 @@
 /// Tests for Supply
 use crate::mock::*;
+use frame_support::{assert_err, assert_ok, traits::Currency};
+use sp_arithmetic::ArithmeticError;
 
 #[test]
 fn should_inflate_supply_from_start_height() {
     run_test(|| {
-        Supply::begin_block(0);
+        assert_ok!(Supply::begin_block(0));
         let mut start_height = 100;
         assert_eq!(Supply::start_height(), Some(start_height));
         assert_eq!(Supply::last_emission(), 0);
 
         for emission in [200_000, 204_000] {
-            Supply::begin_block(start_height);
+            assert_ok!(Supply::begin_block(start_height));
             start_height += YEARS;
             assert_eq!(Supply::start_height(), Some(start_height));
             assert_eq!(Supply::last_emission(), emission);
         }
     })
 }
+
+#[test]
+fn should_not_inflate_total_supply() {
+    run_test(|| {
+        Balances::make_free_balance_be(&Supply::account_id(), u128::MAX);
+
+        let start_height = 100;
+        assert_ok!(Supply::set_start_height_and_inflation(
+            RuntimeOrigin::root(),
+            start_height,
+            UnsignedFixedPoint::checked_from_rational(110, 100).unwrap()
+        ));
+        assert_eq!(Supply::start_height(), Some(start_height));
+        assert_eq!(Supply::last_emission(), 0);
+
+        assert_err!(Supply::begin_block(start_height), ArithmeticError::Overflow);
+        assert_eq!(Balances::total_issuance(), u128::MAX);
+    })
+}