Repository: oasisprotocol/oasis-sdk (rust)

Issue:
Invariant checker
Each Runtime SDK module should be able to define an invariant checker (e.g., by introducing a new `InvariantHandler` trait) that can optionally be invoked on the state checking for any invariant violations. To invoke the invariant checker one would use a special query defined by the `core` module.

Examples of invariants for the accounts module include:
* All account balances should sum up to the total supply.

Examples of invariants for the consensus layer interaction module (#11) include:
* Total supply of the designated consensus layer token denomination should be less than or equal to the balance of the general runtime's account in the consensus layer.

In E2E tests we should then invoke the invariant checker query.

Pull request:
runtime-sdk: Add invariant checker
Closes #18.

Patch:
diff --git a/runtime-sdk/src/context.rs b/runtime-sdk/src/context.rs
index 5641187075..99a7db3629 100644
--- a/runtime-sdk/src/context.rs
+++ b/runtime-sdk/src/context.rs
@@ -11,9 +11,10 @@ use io_context::Context as IoContext;
 use slog::{self, o};
 
 use oasis_core_runtime::{
-    common::logger::get_logger,
+    common::{logger::get_logger, namespace::Namespace},
     consensus,
     consensus::roothash,
+    protocol::HostInfo,
     storage::mkvs,
     transaction::{context::Context as RuntimeContext, tags::Tags},
 };
@@ -76,6 +77,14 @@ pub trait Context {
         self.mode() == Mode::SimulateTx
     }
 
+    /// Information about the host environment.
+    fn host_info(&self) -> &HostInfo;
+
+    /// Runtime ID.
+    fn runtime_id(&self) -> &Namespace {
+        &self.host_info().runtime_id
+    }
+
     /// Last runtime block header.
     fn runtime_header(&self) -> &roothash::Header;
 
@@ -159,6 +168,7 @@ pub trait TxContext: Context {
 pub struct RuntimeBatchContext<'a, R: runtime::Runtime, S: storage::Store> {
     mode: Mode,
 
+    host_info: &'a HostInfo,
     runtime_header: &'a roothash::Header,
     runtime_round_results: &'a roothash::RoundResults,
     runtime_storage: S,
@@ -186,6 +196,7 @@ impl<'a, R: runtime::Runtime, S: storage::Store> RuntimeBatchContext<'a, R, S> {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
         mode: Mode,
+        host_info: &'a HostInfo,
         runtime_header: &'a roothash::Header,
         runtime_round_results: &'a roothash::RoundResults,
         runtime_storage: S,
@@ -196,6 +207,7 @@ impl<'a, R: runtime::Runtime, S: storage::Store> RuntimeBatchContext<'a, R, S> {
     ) -> Self {
         Self {
             mode,
+            host_info,
             runtime_header,
             runtime_round_results,
             runtime_storage,
@@ -216,6 +228,7 @@ impl<'a, R: runtime::Runtime, S: storage::Store> RuntimeBatchContext<'a, R, S> {
     pub(crate) fn from_runtime(
         ctx: &'a RuntimeContext<'_>,
         mkvs: &'a mut dyn mkvs::MKVS,
+        host_info: &'a HostInfo,
     ) -> RuntimeBatchContext<'a, R, storage::MKVSStore<&'a mut dyn mkvs::MKVS>> {
         let mode = if ctx.check_only {
             Mode::CheckTx
@@ -224,6 +237,7 @@ impl<'a, R: runtime::Runtime, S: storage::Store> RuntimeBatchContext<'a, R, S> {
         };
         RuntimeBatchContext {
             mode,
+            host_info,
             runtime_header: ctx.header,
             runtime_round_results: ctx.round_results,
             runtime_storage: storage::MKVSStore::new(ctx.io_ctx.clone(), mkvs),
@@ -253,6 +267,10 @@ impl<'a, R: runtime::Runtime, S: storage::Store> Context for RuntimeBatchContext
         self.mode
     }
 
+    fn host_info(&self) -> &HostInfo {
+        &self.host_info
+    }
+
     fn runtime_header(&self) -> &roothash::Header {
         &self.runtime_header
     }
@@ -300,6 +318,7 @@ impl<'a, R: runtime::Runtime, S: storage::Store> Context for RuntimeBatchContext
 
         let sim_ctx = RuntimeBatchContext {
             mode: Mode::SimulateTx,
+            host_info: self.host_info,
             runtime_header: self.runtime_header,
             runtime_round_results: self.runtime_round_results,
             runtime_storage: store,
@@ -333,6 +352,7 @@ impl<'a, R: runtime::Runtime, S: storage::Store> BatchContext for RuntimeBatchCo
 
         let tx_ctx = RuntimeTxContext {
             mode: self.mode,
+            host_info: self.host_info,
             runtime_header: self.runtime_header,
             runtime_round_results: self.runtime_round_results,
             consensus_state: self.consensus_state,
@@ -372,6 +392,7 @@ impl<'a, R: runtime::Runtime, S: storage::Store> BatchContext for RuntimeBatchCo
 pub struct RuntimeTxContext<'round, 'store, R: runtime::Runtime, S: storage::Store> {
     mode: Mode,
 
+    host_info: &'round HostInfo,
     runtime_header: &'round roothash::Header,
     runtime_round_results: &'round roothash::RoundResults,
     consensus_state: &'round consensus::state::ConsensusState,
@@ -415,6 +436,10 @@ impl<'round, 'store, R: runtime::Runtime, S: storage::Store> Context
         self.mode
     }
 
+    fn host_info(&self) -> &HostInfo {
+        &self.host_info
+    }
+
     fn runtime_header(&self) -> &roothash::Header {
         self.runtime_header
     }
@@ -463,6 +488,7 @@ impl<'round, 'store, R: runtime::Runtime, S: storage::Store> Context
 
         let sim_ctx = RuntimeBatchContext {
             mode: Mode::SimulateTx,
+            host_info: self.host_info,
             runtime_header: self.runtime_header,
             runtime_round_results: self.runtime_round_results,
             runtime_storage: store,
diff --git a/runtime-sdk/src/dispatcher.rs b/runtime-sdk/src/dispatcher.rs
index f7b1070573..632277c86c 100644
--- a/runtime-sdk/src/dispatcher.rs
+++ b/runtime-sdk/src/dispatcher.rs
@@ -11,6 +11,7 @@ use thiserror::Error;
 use oasis_core_runtime::{
     self,
     common::cbor,
+    protocol::HostInfo,
     storage::{context::StorageContext, mkvs},
     transaction::{
         self,
@@ -73,6 +74,7 @@ impl From<types::transaction::CallResult> for DispatchResult {
 
 /// The runtime dispatcher.
 pub struct Dispatcher<R: Runtime> {
+    host_info: HostInfo,
     _runtime: PhantomData<R>,
 }
 
@@ -81,8 +83,9 @@ impl<R: Runtime> Dispatcher<R> {
     ///
     /// Note that the dispatcher is fully static and the constructor is only needed so that the
     /// instance can be used directly with the dispatcher system provided by Oasis Core.
-    pub(super) fn new() -> Self {
+    pub(super) fn new(host_info: HostInfo) -> Self {
         Self {
+            host_info,
             _runtime: PhantomData,
         }
     }
@@ -307,7 +310,9 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
             // Prepare dispatch context.
             let mut ctx =
                 RuntimeBatchContext::<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>>::from_runtime(
-                    &rt_ctx, mkvs,
+                    &rt_ctx,
+                    mkvs,
+                    &self.host_info,
                 );
             // Perform state migrations if required.
             R::migrate(&mut ctx);
@@ -356,7 +361,9 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
             // Prepare dispatch context.
             let mut ctx =
                 RuntimeBatchContext::<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>>::from_runtime(
-                    &ctx, mkvs,
+                    &ctx,
+                    mkvs,
+                    &self.host_info,
                 );
             // Perform state migrations if required.
             R::migrate(&mut ctx);
@@ -386,7 +393,9 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
             // Prepare dispatch context.
             let mut ctx =
                 RuntimeBatchContext::<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>>::from_runtime(
-                    &ctx, mkvs,
+                    &ctx,
+                    mkvs,
+                    &self.host_info,
                 );
             // Perform state migrations if required.
             R::migrate(&mut ctx);
diff --git a/runtime-sdk/src/module.rs b/runtime-sdk/src/module.rs
index 98a115e874..e7e6dda874 100644
--- a/runtime-sdk/src/module.rs
+++ b/runtime-sdk/src/module.rs
@@ -256,6 +256,24 @@ impl BlockHandler for Tuple {
     }
 }
 
+/// Invariant handler.
+pub trait InvariantHandler {
+    /// Check invariants.
+    fn check_invariants<C: Context>(_ctx: &mut C) -> Result<(), modules::core::Error> {
+        // Default implementation doesn't do anything.
+        Ok(())
+    }
+}
+
+#[impl_for_tuples(30)]
+impl InvariantHandler for Tuple {
+    /// Check the invariants in all modules in the tuple.
+    fn check_invariants<C: Context>(ctx: &mut C) -> Result<(), modules::core::Error> {
+        for_tuples!( #( Tuple::check_invariants(ctx)?; )* );
+        Ok(())
+    }
+}
+
 /// A runtime module.
 pub trait Module {
     /// Module name.
diff --git a/runtime-sdk/src/modules/accounts/mod.rs b/runtime-sdk/src/modules/accounts/mod.rs
index cd54216d21..f2e6a25e65 100644
--- a/runtime-sdk/src/modules/accounts/mod.rs
+++ b/runtime-sdk/src/modules/accounts/mod.rs
@@ -15,7 +15,7 @@ use crate::{
     module,
     module::Module as _,
     modules,
-    modules::core::{Module as Core, API as _},
+    modules::core::{Error as CoreError, Module as Core, API as _},
     storage,
     types::{
         address::Address,
@@ -139,6 +139,11 @@ pub trait API {
         state: S,
         address: Address,
     ) -> Result<types::AccountBalances, Error>;
+
+    /// Fetch total supplies.
+    fn get_total_supplies<S: storage::Store>(
+        state: S,
+    ) -> Result<BTreeMap<token::Denomination, token::Quantity>, Error>;
 }
 
 /// State schema constants.
@@ -160,6 +165,31 @@ pub static ADDRESS_COMMON_POOL: Lazy<Address> =
 pub static ADDRESS_FEE_ACCUMULATOR: Lazy<Address> =
     Lazy::new(|| Address::from_module(MODULE_NAME, "fee-accumulator"));
 
+/// This is needed to properly iterate over the BALANCES map.
+#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
+struct AddressWithDenomination(Address, token::Denomination);
+
+#[derive(Error, Debug)]
+enum AWDError {
+    #[error("malformed address")]
+    MalformedAddress,
+
+    #[error("malformed denomination")]
+    MalformedDenomination,
+}
+
+impl std::convert::TryFrom<&[u8]> for AddressWithDenomination {
+    type Error = AWDError;
+
+    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+        let address =
+            Address::try_from(&bytes[..Address::SIZE]).map_err(|_| AWDError::MalformedAddress)?;
+        let denomination = token::Denomination::try_from(&bytes[Address::SIZE..])
+            .map_err(|_| AWDError::MalformedDenomination)?;
+        Ok(AddressWithDenomination(address, denomination))
+    }
+}
+
 impl Module {
     /// Add given amount of tokens to the specified account's balance.
     fn add_amount<S: storage::Store>(
@@ -228,6 +258,39 @@ impl Module {
         total_supplies.insert(amount.denomination(), &total_supply);
         Ok(())
     }
+
+    /// Get all balances.
+    fn get_all_balances<S: storage::Store>(
+        state: S,
+    ) -> Result<BTreeMap<Address, BTreeMap<token::Denomination, token::Quantity>>, Error> {
+        let store = storage::PrefixStore::new(state, &MODULE_NAME);
+        let balances = storage::TypedStore::new(storage::PrefixStore::new(store, &state::BALANCES));
+
+        // Unfortunately, we can't just return balances.iter().collect() here,
+        // because the stored format doesn't match -- we need this workaround
+        // instead.
+
+        let balmap: BTreeMap<AddressWithDenomination, token::Quantity> = balances.iter().collect();
+
+        let mut b: BTreeMap<Address, BTreeMap<token::Denomination, token::Quantity>> =
+            BTreeMap::new();
+
+        for (addrden, amt) in &balmap {
+            let addr = &addrden.0;
+            let den = &addrden.1;
+
+            // Fetch existing account's balances or insert blank ones.
+            let addr_bals = b.entry(*addr).or_insert_with(BTreeMap::new);
+
+            // Add to given denomination's balance or insert it if new.
+            addr_bals
+                .entry(den.clone())
+                .and_modify(|a| *a += amt)
+                .or_insert_with(|| amt.clone());
+        }
+
+        Ok(b)
+    }
 }
 
 impl API for Module {
@@ -300,6 +363,15 @@ impl API for Module {
             balances: account.iter().collect(),
         })
     }
+
+    fn get_total_supplies<S: storage::Store>(
+        state: S,
+    ) -> Result<BTreeMap<token::Denomination, token::Quantity>, Error> {
+        let store = storage::PrefixStore::new(state, &MODULE_NAME);
+        let ts = storage::TypedStore::new(storage::PrefixStore::new(store, &state::TOTAL_SUPPLY));
+
+        Ok(ts.iter().collect())
+    }
 }
 
 impl Module {
@@ -602,3 +674,63 @@ impl module::BlockHandler for Module {
         }
     }
 }
+
+impl module::InvariantHandler for Module {
+    /// Check invariants.
+    fn check_invariants<C: Context>(ctx: &mut C) -> Result<(), CoreError> {
+        // All account balances should sum up to the total supply for their
+        // corresponding denominations.
+
+        #[allow(clippy::or_fun_call)]
+        let balances = Self::get_all_balances(ctx.runtime_state()).or(Err(
+            CoreError::InvariantViolation("unable to get balances of all accounts".to_string()),
+        ))?;
+        #[allow(clippy::or_fun_call)]
+        let total_supplies = Self::get_total_supplies(ctx.runtime_state()).or(Err(
+            CoreError::InvariantViolation("unable to get total supplies".to_string()),
+        ))?;
+
+        // First, compute total supplies based on account balances.
+        let mut computed_ts: BTreeMap<token::Denomination, token::Quantity> = BTreeMap::new();
+
+        for bals in balances.values() {
+            for (den, amt) in bals {
+                computed_ts
+                    .entry(den.clone())
+                    .and_modify(|a| *a += amt)
+                    .or_insert_with(|| amt.clone());
+            }
+        }
+
+        // Now check if the computed and given total supplies match.
+        for (den, ts) in &total_supplies {
+            // Return error if total supplies have a denomination that we
+            // didn't encounter when computing total supplies based on account
+            // balances.
+            #[allow(clippy::or_fun_call)]
+            let computed = computed_ts
+                .remove(&den)
+                .ok_or(CoreError::InvariantViolation(
+                    "unexpected denomination".to_string(),
+                ))?;
+
+            if &computed != ts {
+                // Computed and actual total supplies don't match.
+                return Err(CoreError::InvariantViolation(
+                    "computed and actual total supplies don't match".to_string(),
+                ));
+            }
+        }
+
+        // There should be no remaining denominations in the computed supplies,
+        // because that would mean that accounts have denominations that don't
+        // appear in the total supplies table, which would obviously be wrong.
+        if computed_ts.is_empty() {
+            Ok(())
+        } else {
+            Err(CoreError::InvariantViolation(
+                "encountered denomination that isn't present in total supplies table".to_string(),
+            ))
+        }
+    }
+}
diff --git a/runtime-sdk/src/modules/consensus/mod.rs b/runtime-sdk/src/modules/consensus/mod.rs
index 692fc19977..5edd508083 100644
--- a/runtime-sdk/src/modules/consensus/mod.rs
+++ b/runtime-sdk/src/modules/consensus/mod.rs
@@ -317,3 +317,5 @@ impl module::MigrationHandler for Module {
 impl module::AuthHandler for Module {}
 
 impl module::BlockHandler for Module {}
+
+impl module::InvariantHandler for Module {}
diff --git a/runtime-sdk/src/modules/consensus_accounts/mod.rs b/runtime-sdk/src/modules/consensus_accounts/mod.rs
index 183f56c146..f0383685fb 100644
--- a/runtime-sdk/src/modules/consensus_accounts/mod.rs
+++ b/runtime-sdk/src/modules/consensus_accounts/mod.rs
@@ -13,7 +13,7 @@ use crate::{
     module,
     module::Module as _,
     modules,
-    modules::core::{Module as Core, API as _},
+    modules::core::{Error as CoreError, Module as Core, API as _},
     types::{
         address::Address,
         message::{MessageEvent, MessageEventHookInvocation, MessageResult},
@@ -396,3 +396,38 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API> modul
     for Module<Accounts, Consensus>
 {
 }
+
+impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API> module::InvariantHandler
+    for Module<Accounts, Consensus>
+{
+    /// Check invariants.
+    fn check_invariants<C: Context>(ctx: &mut C) -> Result<(), CoreError> {
+        // Total supply of the designated consensus layer token denomination
+        // should be less than or equal to the balance of the runtime's general
+        // account in the consensus layer.
+
+        let den = Consensus::consensus_denomination(ctx).unwrap();
+        #[allow(clippy::or_fun_call)]
+        let ts = Accounts::get_total_supplies(ctx.runtime_state()).or(Err(
+            CoreError::InvariantViolation("unable to get total supplies".to_string()),
+        ))?;
+
+        let rt_addr = Address::from_runtime_id(ctx.runtime_id());
+        let rt_acct = Consensus::account(ctx, rt_addr).unwrap_or_default();
+        let rt_ga_balance = rt_acct.general.balance;
+
+        match ts.get(&den) {
+            Some(total_supply) => {
+                if total_supply <= &rt_ga_balance {
+                    Ok(())
+                } else {
+                    Err(CoreError::InvariantViolation(
+                        "total supply is greater than runtime's general account balance"
+                            .to_string(),
+                    ))
+                }
+            }
+            None => Ok(()), // Having no total supply also satisfies above invariant.
+        }
+    }
+}
diff --git a/runtime-sdk/src/modules/core/mod.rs b/runtime-sdk/src/modules/core/mod.rs
index 18882bffc2..7981ba6824 100644
--- a/runtime-sdk/src/modules/core/mod.rs
+++ b/runtime-sdk/src/modules/core/mod.rs
@@ -8,10 +8,11 @@ use crate::{
     context::{BatchContext, Context, TxContext},
     core::common::cbor,
     dispatcher, error,
-    module::{self, Module as _},
+    module::{self, InvariantHandler as _, Module as _},
     types::transaction::{
         self, AddressSpec, AuthProof, Call, TransactionWeight, UnverifiedTransaction,
     },
+    Runtime,
 };
 
 #[cfg(test)]
@@ -83,6 +84,10 @@ pub enum Error {
     #[error("multisig too many signers")]
     #[sdk_error(code = 16)]
     MultisigTooManySigners,
+
+    #[error("invariant violation: {0}")]
+    #[sdk_error(code = 17)]
+    InvariantViolation(String),
 }
 
 /// Gas costs.
@@ -273,6 +278,11 @@ impl Module {
             })
         })
     }
+
+    /// Check invariants of all modules in the runtime.
+    fn query_check_invariants<C: Context>(ctx: &mut C) -> Result<(), Error> {
+        <C::Runtime as Runtime>::Modules::check_invariants(ctx)
+    }
 }
 
 impl module::Module for Module {
@@ -362,6 +372,10 @@ impl module::MethodHandler for Module {
                 let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
                 Ok(cbor::to_value(&Self::query_estimate_gas(ctx, args)?))
             })()),
+            "core.CheckInvariants" => module::DispatchResult::Handled((|| {
+                let _ = Self::query_check_invariants(ctx)?;
+                Ok(cbor::to_value(true))
+            })()),
             _ => module::DispatchResult::Unhandled(args),
         }
     }
@@ -377,3 +391,5 @@ impl module::BlockHandler for Module {
         res
     }
 }
+
+impl module::InvariantHandler for Module {}
diff --git a/runtime-sdk/src/modules/rewards/mod.rs b/runtime-sdk/src/modules/rewards/mod.rs
index d37e99f1ff..adb18c4b55 100644
--- a/runtime-sdk/src/modules/rewards/mod.rs
+++ b/runtime-sdk/src/modules/rewards/mod.rs
@@ -244,6 +244,8 @@ impl<Accounts: modules::accounts::API> module::BlockHandler for Module<Accounts>
     }
 }
 
+impl<Accounts: modules::accounts::API> module::InvariantHandler for Module<Accounts> {}
+
 /// A trait that exists solely to convert `beacon::EpochTime` to bytes for use as a storage key.
 trait ToStorageKey {
     fn to_storage_key(&self) -> [u8; 8];
diff --git a/runtime-sdk/src/runtime.rs b/runtime-sdk/src/runtime.rs
index 40ae7e408b..22dd6ca19a 100644
--- a/runtime-sdk/src/runtime.rs
+++ b/runtime-sdk/src/runtime.rs
@@ -8,7 +8,7 @@ use oasis_core_runtime::{
 use crate::{
     context::Context,
     crypto, dispatcher,
-    module::{AuthHandler, BlockHandler, MethodHandler, MigrationHandler},
+    module::{AuthHandler, BlockHandler, InvariantHandler, MethodHandler, MigrationHandler},
     modules, storage,
 };
 
@@ -17,7 +17,7 @@ pub trait Runtime {
     /// Runtime version.
     const VERSION: version::Version;
 
-    type Modules: AuthHandler + MigrationHandler + MethodHandler + BlockHandler;
+    type Modules: AuthHandler + MigrationHandler + MethodHandler + BlockHandler + InvariantHandler;
 
     /// Genesis state for the runtime.
     fn genesis_state() -> <Self::Modules as MigrationHandler>::Genesis;
@@ -65,7 +65,7 @@ pub trait Runtime {
             );
 
             // Register runtime's methods.
-            let dispatcher = dispatcher::Dispatcher::<Self>::new();
+            let dispatcher = dispatcher::Dispatcher::<Self>::new(hi);
             Some(Box::new(dispatcher))
         };
 
diff --git a/runtime-sdk/src/types/address.rs b/runtime-sdk/src/types/address.rs
index 24e31cdf66..5891352fbc 100644
--- a/runtime-sdk/src/types/address.rs
+++ b/runtime-sdk/src/types/address.rs
@@ -5,7 +5,8 @@ use bech32::{self, FromBase32, ToBase32, Variant};
 use thiserror::Error;
 
 use oasis_core_runtime::{
-    common::crypto::hash::Hash, consensus::address::Address as ConsensusAddress,
+    common::{crypto::hash::Hash, namespace::Namespace},
+    consensus::address::Address as ConsensusAddress,
 };
 
 use crate::crypto::{multisig, signature::PublicKey};
@@ -23,6 +24,10 @@ const ADDRESS_V0_SECP256K1_CONTEXT: &[u8] = b"oasis-runtime-sdk/address: secp256
 
 const ADDRESS_V0_MODULE_CONTEXT: &[u8] = b"oasis-runtime-sdk/address: module";
 
+// V0 runtime address.
+const ADDRESS_RUNTIME_V0_CONTEXT: &[u8] = b"oasis-core/address: runtime";
+const ADDRESS_RUNTIME_V0_VERSION: u8 = 0;
+
 /// V0 multisig address context.
 const ADDRESS_V0_MULTISIG_CONTEXT: &[u8] = b"oasis-runtime-sdk/address: multisig";
 
@@ -40,6 +45,9 @@ pub enum Error {
 pub struct Address([u8; ADDRESS_SIZE]);
 
 impl Address {
+    /// Size of an address in bytes.
+    pub const SIZE: usize = ADDRESS_SIZE;
+
     /// Creates a new address from a context, version and data.
     pub fn new(ctx: &'static [u8], version: u8, data: &[u8]) -> Self {
         let h = Hash::digest_bytes_list(&[ctx, &[version], data]);
@@ -72,6 +80,15 @@ impl Address {
         )
     }
 
+    /// Creates a new runtime address.
+    pub fn from_runtime_id(id: &Namespace) -> Self {
+        Address::new(
+            ADDRESS_RUNTIME_V0_CONTEXT,
+            ADDRESS_RUNTIME_V0_VERSION,
+            id.as_ref(),
+        )
+    }
+
     /// Creates a new address from a public key.
     pub fn from_pk(pk: &PublicKey) -> Self {
         match pk {
@@ -293,4 +310,15 @@ mod test {
         let consensus_addr: ConsensusAddress = addr.into();
         assert_eq!(addr.to_bech32(), consensus_addr.to_bech32())
     }
+
+    #[test]
+    fn test_address_from_runtime_id() {
+        let runtime_id =
+            Namespace::from("80000000000000002aff7f6dfb62720cfd735f2b037b81572fad1b7937d826b3");
+        let addr = Address::from_runtime_id(&runtime_id);
+        assert_eq!(
+            addr.to_bech32(),
+            "oasis1qpllh99nhwzrd56px4txvl26atzgg4f3a58jzzad"
+        );
+    }
 }


Test patch:
diff --git a/runtime-sdk/src/modules/accounts/test.rs b/runtime-sdk/src/modules/accounts/test.rs
index 9918b684ed..ba77fbab40 100644
--- a/runtime-sdk/src/modules/accounts/test.rs
+++ b/runtime-sdk/src/modules/accounts/test.rs
@@ -5,7 +5,7 @@ use oasis_core_runtime::common::cbor;
 
 use crate::{
     context::{BatchContext, Context},
-    module::{AuthHandler, BlockHandler},
+    module::{AuthHandler, BlockHandler, InvariantHandler},
     modules::core,
     testing::{keys, mock},
     types::{
@@ -487,3 +487,326 @@ fn test_fee_disbursement() {
         "remainder should be disbursed to the common pool"
     );
 }
+
+#[test]
+fn test_get_all_balances_and_total_supplies_basic() {
+    let mut mock = mock::Mock::default();
+    let mut ctx = mock.create_ctx();
+
+    let alice = keys::alice::address();
+    let bob = keys::bob::address();
+
+    let gen = Genesis {
+        balances: {
+            let mut balances = BTreeMap::new();
+            // Alice.
+            balances.insert(alice, {
+                let mut denominations = BTreeMap::new();
+                denominations.insert(Denomination::NATIVE, 1_000_000.into());
+                denominations
+            });
+            // Bob.
+            balances.insert(bob, {
+                let mut denominations = BTreeMap::new();
+                denominations.insert(Denomination::NATIVE, 2_000_000.into());
+                denominations
+            });
+            balances
+        },
+        total_supplies: {
+            let mut total_supplies = BTreeMap::new();
+            total_supplies.insert(Denomination::NATIVE, 3_000_000.into());
+            total_supplies
+        },
+        ..Default::default()
+    };
+
+    Accounts::init(&mut ctx, &gen);
+
+    let all_bals =
+        Accounts::get_all_balances(ctx.runtime_state()).expect("get_all_balances should succeed");
+    for (addr, bals) in &all_bals {
+        assert_eq!(bals.len(), 1, "exactly one denomination should be present");
+        assert!(
+            bals.contains_key(&Denomination::NATIVE),
+            "only native denomination should be present"
+        );
+        if addr == &alice {
+            assert_eq!(
+                bals[&Denomination::NATIVE],
+                1_000_000.into(),
+                "Alice's balance should be 1000000"
+            );
+        } else if addr == &bob {
+            assert_eq!(
+                bals[&Denomination::NATIVE],
+                2_000_000.into(),
+                "Bob's balance should be 2000000"
+            );
+        } else {
+            panic!("invalid address");
+        }
+    }
+
+    let ts = Accounts::get_total_supplies(ctx.runtime_state())
+        .expect("get_total_supplies should succeed");
+    assert_eq!(
+        ts.len(),
+        1,
+        "exactly one denomination should be present in total supplies"
+    );
+    assert!(
+        ts.contains_key(&Denomination::NATIVE),
+        "only native denomination should be present in total supplies"
+    );
+    assert_eq!(
+        ts[&Denomination::NATIVE],
+        3_000_000.into(),
+        "total supply should be 3000000"
+    );
+}
+
+#[test]
+fn test_get_all_balances_and_total_supplies_more() {
+    let mut mock = mock::Mock::default();
+    let mut ctx = mock.create_ctx();
+
+    let dn = Denomination::NATIVE;
+    let d1: Denomination = "den1".parse().unwrap();
+    let d2: Denomination = "den2".parse().unwrap();
+    let d3: Denomination = "den3".parse().unwrap();
+
+    let alice = keys::alice::address();
+    let bob = keys::bob::address();
+
+    let gen = Genesis {
+        balances: {
+            let mut balances = BTreeMap::new();
+            // Alice.
+            balances.insert(alice, {
+                let mut denominations = BTreeMap::new();
+                denominations.insert(dn.clone(), 1_000_000.into());
+                denominations.insert(d1.clone(), 1_000.into());
+                denominations.insert(d2.clone(), 100.into());
+                denominations
+            });
+            // Bob.
+            balances.insert(bob, {
+                let mut denominations = BTreeMap::new();
+                denominations.insert(d1.clone(), 2_000.into());
+                denominations.insert(d3.clone(), 200.into());
+                denominations
+            });
+            balances
+        },
+        total_supplies: {
+            let mut total_supplies = BTreeMap::new();
+            total_supplies.insert(dn.clone(), 1_000_000.into());
+            total_supplies.insert(d1.clone(), 3_000.into());
+            total_supplies.insert(d2.clone(), 100.into());
+            total_supplies.insert(d3.clone(), 200.into());
+            total_supplies
+        },
+        ..Default::default()
+    };
+
+    Accounts::init(&mut ctx, &gen);
+
+    let all_bals =
+        Accounts::get_all_balances(ctx.runtime_state()).expect("get_all_balances should succeed");
+    for (addr, bals) in &all_bals {
+        if addr == &alice {
+            assert_eq!(bals.len(), 3, "Alice should have exactly 3 denominations");
+            assert_eq!(
+                bals[&dn],
+                1_000_000.into(),
+                "Alice's native balance should be 1000000"
+            );
+            assert_eq!(
+                bals[&d1],
+                1_000.into(),
+                "Alice's den1 balance should be 1000"
+            );
+            assert_eq!(bals[&d2], 100.into(), "Alice's den2 balance should be 100");
+        } else if addr == &bob {
+            assert_eq!(bals.len(), 2, "Bob should have exactly 2 denominations");
+            assert_eq!(bals[&d1], 2_000.into(), "Bob's den1 balance should be 2000");
+            assert_eq!(bals[&d3], 200.into(), "Bob's den3 balance should be 200");
+        } else {
+            panic!("invalid address");
+        }
+    }
+
+    let ts = Accounts::get_total_supplies(ctx.runtime_state())
+        .expect("get_total_supplies should succeed");
+    assert_eq!(
+        ts.len(),
+        4,
+        "exactly 4 denominations should be present in total supplies"
+    );
+    assert!(
+        ts.contains_key(&dn),
+        "native denomination should be present in total supplies"
+    );
+    assert!(
+        ts.contains_key(&d1),
+        "den1 denomination should be present in total supplies"
+    );
+    assert!(
+        ts.contains_key(&d2),
+        "den2 denomination should be present in total supplies"
+    );
+    assert!(
+        ts.contains_key(&d3),
+        "den3 denomination should be present in total supplies"
+    );
+    assert_eq!(
+        ts[&dn],
+        1_000_000.into(),
+        "native total supply should be 1000000"
+    );
+    assert_eq!(ts[&d1], 3_000.into(), "den1 total supply should be 3000");
+    assert_eq!(ts[&d2], 100.into(), "den2 total supply should be 100");
+    assert_eq!(ts[&d3], 200.into(), "den3 total supply should be 200");
+}
+
+#[test]
+fn test_check_invariants_basic() {
+    let mut mock = mock::Mock::default();
+    let mut ctx = mock.create_ctx();
+
+    init_accounts(&mut ctx);
+
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_ok(),
+        "invariants check should succeed"
+    );
+}
+
+#[test]
+fn test_check_invariants_more() {
+    let mut mock = mock::Mock::default();
+    let mut ctx = mock.create_ctx();
+
+    let dn = Denomination::NATIVE;
+    let d1: Denomination = "den1".parse().unwrap();
+    let d2: Denomination = "den2".parse().unwrap();
+    let d3: Denomination = "den3".parse().unwrap();
+
+    let alice = keys::alice::address();
+    let bob = keys::bob::address();
+    let charlie = keys::charlie::address();
+
+    let gen = Genesis {
+        balances: {
+            let mut balances = BTreeMap::new();
+            // Alice.
+            balances.insert(alice, {
+                let mut denominations = BTreeMap::new();
+                denominations.insert(dn.clone(), 1_000_000.into());
+                denominations.insert(d1.clone(), 1_000.into());
+                denominations.insert(d2.clone(), 100.into());
+                denominations
+            });
+            // Bob.
+            balances.insert(bob, {
+                let mut denominations = BTreeMap::new();
+                denominations.insert(d1.clone(), 2_000.into());
+                denominations.insert(d3.clone(), 200.into());
+                denominations
+            });
+            balances
+        },
+        total_supplies: {
+            let mut total_supplies = BTreeMap::new();
+            total_supplies.insert(dn.clone(), 1_000_000.into());
+            total_supplies.insert(d1.clone(), 3_000.into());
+            total_supplies.insert(d2.clone(), 100.into());
+            total_supplies.insert(d3.clone(), 200.into());
+            total_supplies
+        },
+        ..Default::default()
+    };
+
+    Accounts::init(&mut ctx, &gen);
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_ok(),
+        "initial inv chk should succeed"
+    );
+
+    assert!(
+        Accounts::add_amount(
+            ctx.runtime_state(),
+            charlie,
+            &BaseUnits::new(100.into(), d1.clone())
+        )
+        .is_ok(),
+        "giving Charlie money should succeed"
+    );
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_err(),
+        "inv chk 1 should fail"
+    );
+
+    assert!(
+        Accounts::inc_total_supply(ctx.runtime_state(), &BaseUnits::new(100.into(), d1.clone()))
+            .is_ok(),
+        "increasing total supply should succeed"
+    );
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_ok(),
+        "inv chk 2 should succeed"
+    );
+
+    let d4: Denomination = "den4".parse().unwrap();
+
+    assert!(
+        Accounts::add_amount(
+            ctx.runtime_state(),
+            charlie,
+            &BaseUnits::new(300.into(), d4.clone())
+        )
+        .is_ok(),
+        "giving Charlie more money should succeed"
+    );
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_err(),
+        "inv chk 3 should fail"
+    );
+
+    assert!(
+        Accounts::inc_total_supply(ctx.runtime_state(), &BaseUnits::new(300.into(), d4.clone()))
+            .is_ok(),
+        "increasing total supply should succeed"
+    );
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_ok(),
+        "inv chk 4 should succeed"
+    );
+
+    let d5: Denomination = "den5".parse().unwrap();
+
+    assert!(
+        Accounts::inc_total_supply(ctx.runtime_state(), &BaseUnits::new(123.into(), d5.clone()))
+            .is_ok(),
+        "increasing total supply should succeed"
+    );
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_err(),
+        "inv chk 5 should fail"
+    );
+
+    assert!(
+        Accounts::add_amount(
+            ctx.runtime_state(),
+            charlie,
+            &BaseUnits::new(123.into(), d5.clone())
+        )
+        .is_ok(),
+        "giving Charlie more money should succeed"
+    );
+    assert!(
+        Accounts::check_invariants(&mut ctx).is_ok(),
+        "inv chk 6 should succeed"
+    );
+}
diff --git a/runtime-sdk/src/modules/core/test.rs b/runtime-sdk/src/modules/core/test.rs
index 271670fd09..816162685c 100644
--- a/runtime-sdk/src/modules/core/test.rs
+++ b/runtime-sdk/src/modules/core/test.rs
@@ -104,6 +104,7 @@ impl module::AuthHandler for GasWasterModule {}
 impl module::MigrationHandler for GasWasterModule {
     type Genesis = ();
 }
+impl module::InvariantHandler for GasWasterModule {}
 
 // Runtime that knows how to waste gas.
 struct GasWasterRuntime;
diff --git a/runtime-sdk/src/testing/mock.rs b/runtime-sdk/src/testing/mock.rs
index 47fb6a5ae4..1b96c98c95 100644
--- a/runtime-sdk/src/testing/mock.rs
+++ b/runtime-sdk/src/testing/mock.rs
@@ -1,9 +1,12 @@
 //! Mock dispatch context for use in tests.
+use std::collections::BTreeMap;
+
 use io_context::Context as IoContext;
 
 use oasis_core_runtime::{
-    common::{cbor, version::Version},
+    common::{cbor, namespace::Namespace, version::Version},
     consensus::{beacon, roothash, state::ConsensusState},
+    protocol::HostInfo,
     storage::mkvs,
 };
 
@@ -31,6 +34,7 @@ impl Runtime for EmptyRuntime {
 
 /// Mock dispatch context factory.
 pub struct Mock {
+    pub host_info: HostInfo,
     pub runtime_header: roothash::Header,
     pub runtime_round_results: roothash::RoundResults,
     pub mkvs: Box<dyn mkvs::MKVS>,
@@ -55,6 +59,7 @@ impl Mock {
     ) -> RuntimeBatchContext<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>> {
         RuntimeBatchContext::new(
             mode,
+            &self.host_info,
             &self.runtime_header,
             &self.runtime_round_results,
             storage::MKVSStore::new(IoContext::background().freeze(), self.mkvs.as_mut()),
@@ -78,6 +83,13 @@ impl Default for Mock {
             .new(Box::new(mkvs::sync::NoopReadSyncer));
 
         Self {
+            host_info: HostInfo {
+                runtime_id: Namespace::default(),
+                consensus_backend: "mock".to_string(),
+                consensus_protocol_version: Version::default(),
+                consensus_chain_context: "test".to_string(),
+                local_config: BTreeMap::new(),
+            },
             runtime_header: roothash::Header::default(),
             runtime_round_results: roothash::RoundResults::default(),
             mkvs: Box::new(mkvs),
diff --git a/tests/e2e/runtime.go b/tests/e2e/runtime.go
index eb71d9c846..00d400726b 100644
--- a/tests/e2e/runtime.go
+++ b/tests/e2e/runtime.go
@@ -27,6 +27,7 @@ import (
 	"github.com/oasisprotocol/oasis-sdk/client-sdk/go/testing"
 
 	"github.com/oasisprotocol/oasis-sdk/client-sdk/go/client"
+	"github.com/oasisprotocol/oasis-sdk/tests/e2e/txgen"
 )
 
 const (
@@ -284,6 +285,8 @@ func (sc *RuntimeScenario) waitNodesSynced() error {
 }
 
 func (sc *RuntimeScenario) Run(childEnv *env.Env) error {
+	ctx := context.Background()
+
 	// Start the test network.
 	if err := sc.Net.Start(); err != nil {
 		return err
@@ -306,6 +309,12 @@ func (sc *RuntimeScenario) Run(childEnv *env.Env) error {
 	}
 	rtc := client.New(conn, runtimeID)
 
+	// Do an initial invariants check.
+	if err = txgen.CheckInvariants(ctx, rtc); err != nil {
+		sc.Logger.Error("initial invariants check failed", "err", err)
+		return err
+	}
+
 	// Run the given tests for this runtime.
 	for _, test := range sc.RunTest {
 		testName := runtime.FuncForPC(reflect.ValueOf(test).Pointer()).Name()
@@ -319,6 +328,14 @@ func (sc *RuntimeScenario) Run(childEnv *env.Env) error {
 			return testErr
 		}
 		sc.Logger.Info("test passed", "test", testName)
+
+		// Do an invariants check after each test.
+		if err = txgen.CheckInvariants(ctx, rtc); err != nil {
+			sc.Logger.Error("invariants check failed after test",
+				"test", testName,
+				"err", err)
+			return err
+		}
 	}
 
 	return sc.Net.CheckLogWatchers()
diff --git a/tests/e2e/txgen/txgen.go b/tests/e2e/txgen/txgen.go
index af6d6e7b45..52accb45ea 100644
--- a/tests/e2e/txgen/txgen.go
+++ b/tests/e2e/txgen/txgen.go
@@ -75,6 +75,19 @@ func EstimateGas(ctx context.Context, rtc client.RuntimeClient, tx types.Transac
 	return tx
 }
 
+// CheckInvariants issues a check of invariants in all modules in the runtime.
+func CheckInvariants(ctx context.Context, rtc client.RuntimeClient) error {
+	var ok bool
+	err := rtc.Query(ctx, client.RoundLatest, "core.CheckInvariants", nil, &ok)
+	if err != nil {
+		return err
+	}
+	if !ok {
+		return fmt.Errorf("invariants check failed")
+	}
+	return nil
+}
+
 // SignAndSubmitTx signs and submits the given transaction.
 // Gas estimation is done automatically.
 func SignAndSubmitTx(ctx context.Context, rtc client.RuntimeClient, signer signature.Signer, tx types.Transaction) error {
@@ -166,10 +179,14 @@ func Generate(ctx context.Context, rtc client.RuntimeClient, rng *rand.Rand, acc
 		okCount     uint64
 	)
 
+	errCh := make(chan error)
+
 	for {
 		select {
 		case <-ctx.Done():
 			return genErrCount, subErrCount, okCount, nil
+		case err := <-errCh:
+			return genErrCount, subErrCount, okCount, err
 		case <-ticker.C:
 			// Choose random account and txn generator.
 			acct := accounts[rng.Intn(len(accounts))]
@@ -195,6 +212,17 @@ func Generate(ctx context.Context, rtc client.RuntimeClient, rng *rand.Rand, acc
 						atomic.AddUint64(&okCount, 1)
 					}
 				}
+
+				// Perform an invariants check.
+				if err := CheckInvariants(ctx, rtc); err != nil {
+					// This is wrapped in a select block to make sure that if
+					// multiple goroutines report an error, they don't get
+					// blocked here.
+					select {
+					case errCh <- err:
+					default:
+					}
+				}
 			}(acct, gen)
 		}
 	}
diff --git a/tests/runtimes/simple-keyvalue/src/keyvalue.rs b/tests/runtimes/simple-keyvalue/src/keyvalue.rs
index 71860bf907..413f1f9051 100644
--- a/tests/runtimes/simple-keyvalue/src/keyvalue.rs
+++ b/tests/runtimes/simple-keyvalue/src/keyvalue.rs
@@ -89,6 +89,7 @@ impl sdk::module::Module for Module {
 
 impl sdk::module::AuthHandler for Module {}
 impl sdk::module::BlockHandler for Module {}
+impl sdk::module::InvariantHandler for Module {}
 
 impl sdk::module::MethodHandler for Module {
     fn dispatch_call<C: TxContext>(