Chain extension: unit tests (#710)

Author: pmikolajczyk41

bin/runtime/src/chain_extension/environment.rs

@@ -0,0 +1,54 @@
+use frame_support::weights::Weight;
+use pallet_contracts::{
+    chain_extension::{BufInBufOutState, Environment as SubstrateEnvironment, Ext, SysConfig},
+    ChargedAmount,
+};
+use sp_core::crypto::UncheckedFrom;
+use sp_runtime::DispatchError;
+use sp_std::vec::Vec;
+
+use crate::chain_extension::ByteCount;
+
+/// Abstraction around `pallet_contracts::chain_extension::Environment`. Makes testing easier.
+///
+/// Gathers all the methods that are used by `SnarcosChainExtension`. For now, all operations are
+/// performed in the `BufInBufOut` mode, so we don't have to take care of other modes.
+///
+/// Each method is already documented in `pallet_contracts::chain_extension`.
+pub(super) trait Environment: Sized {
+    /// A type returned by `charge_weight` and passed to `adjust_weight`.
+    ///
+    /// The original type `ChargedAmount` has only a private constructor and thus we have to
+    /// abstract it as well to make testing it possible.
+    type ChargedAmount;
+
+    fn in_len(&self) -> ByteCount;
+    fn read(&self, max_len: u32) -> Result<Vec<u8>, DispatchError>;
+
+    fn charge_weight(&mut self, amount: Weight) -> Result<Self::ChargedAmount, DispatchError>;
+    fn adjust_weight(&mut self, charged: Self::ChargedAmount, actual_weight: Weight);
+}
+
+/// Transparent delegation.
+impl<E: Ext> Environment for SubstrateEnvironment<'_, '_, E, BufInBufOutState>
+where
+    <E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
+{
+    type ChargedAmount = ChargedAmount;
+
+    fn in_len(&self) -> ByteCount {
+        self.in_len()
+    }
+
+    fn read(&self, max_len: u32) -> Result<Vec<u8>, DispatchError> {
+        self.read(max_len)
+    }
+
+    fn charge_weight(&mut self, amount: Weight) -> Result<ChargedAmount, DispatchError> {
+        self.charge_weight(amount)
+    }
+
+    fn adjust_weight(&mut self, charged: ChargedAmount, actual_weight: Weight) {
+        self.adjust_weight(charged, actual_weight)
+    }
+}

bin/runtime/src/chain_extension/executor.rs

@@ -0,0 +1,49 @@
+use pallet_snarcos::{Error, Pallet as Snarcos, ProvingSystem, VerificationKeyIdentifier};
+use sp_std::vec::Vec;
+
+use crate::Runtime;
+
+/// Abstraction around `Runtime`. Makes testing easier.
+///
+/// Gathers all the methods that are used by `SnarcosChainExtension`.
+///
+/// Each method is already documented in `pallet_snarcos`.
+pub(super) trait Executor: Sized {
+    /// The error returned from dispatchables is generic. For most purposes however, it doesn't
+    /// matter what type will be passed there. Normally, `Runtime` will be the generic argument,
+    /// but in testing it will be sufficient to instantiate it with `()`.
+    type ErrorGenericType;
+
+    fn store_key(
+        identifier: VerificationKeyIdentifier,
+        key: Vec<u8>,
+    ) -> Result<(), Error<Self::ErrorGenericType>>;
+
+    fn verify(
+        verification_key_identifier: VerificationKeyIdentifier,
+        proof: Vec<u8>,
+        public_input: Vec<u8>,
+        system: ProvingSystem,
+    ) -> Result<(), Error<Self::ErrorGenericType>>;
+}
+
+/// Transparent delegation.
+impl Executor for Runtime {
+    type ErrorGenericType = Runtime;
+
+    fn store_key(
+        identifier: VerificationKeyIdentifier,
+        key: Vec<u8>,
+    ) -> Result<(), Error<Runtime>> {
+        Snarcos::<Runtime>::bare_store_key(identifier, key)
+    }
+
+    fn verify(
+        verification_key_identifier: VerificationKeyIdentifier,
+        proof: Vec<u8>,
+        public_input: Vec<u8>,
+        system: ProvingSystem,
+    ) -> Result<(), Error<Runtime>> {
+        Snarcos::<Runtime>::bare_verify(verification_key_identifier, proof, public_input, system)
+    }
+}

bin/runtime/src/chain_extension/mod.rs

@@ -1,25 +1,29 @@
-use codec::Decode;
-use frame_support::{log::error, pallet_prelude::Weight};
+use codec::{Decode, Encode};
+use environment::Environment;
+use executor::Executor;
+use frame_support::{log::error, weights::Weight};
 use pallet_contracts::chain_extension::{
-    ChainExtension, Environment, Ext, InitState, RetVal, SysConfig,
-};
-use pallet_snarcos::{
-    Config, Error, Pallet as Snarcos, ProvingSystem, VerificationKeyIdentifier, WeightInfo,
+    ChainExtension, Environment as SubstrateEnvironment, Ext, InitState, RetVal, SysConfig,
 };
+use pallet_snarcos::{Config, Error, ProvingSystem, VerificationKeyIdentifier, WeightInfo};
 use sp_core::crypto::UncheckedFrom;
 use sp_runtime::DispatchError;
 use sp_std::{mem::size_of, vec::Vec};
 use Error::*;
 
 use crate::{MaximumVerificationKeyLength, Runtime};
+mod environment;
+mod executor;
+#[cfg(test)]
+mod tests;
 
 pub const SNARCOS_STORE_KEY_FUNC_ID: u32 = 41;
 pub const SNARCOS_VERIFY_FUNC_ID: u32 = 42;
 
 // Return codes for `SNARCOS_STORE_KEY_FUNC_ID`.
 pub const SNARCOS_STORE_KEY_OK: u32 = 10_000;
 pub const SNARCOS_STORE_KEY_TOO_LONG_KEY: u32 = 10_001;
-pub const SNARCOS_STORE_KEY_IN_USE: u32 = 10_002;
+pub const SNARCOS_STORE_KEY_IDENTIFIER_IN_USE: u32 = 10_002;
 pub const SNARCOS_STORE_KEY_ERROR_UNKNOWN: u32 = 10_003;
 
 // Return codes for `SNARCOS_VERIFY_FUNC_ID`.
@@ -35,13 +39,18 @@ pub const SNARCOS_VERIFY_ERROR_UNKNOWN: u32 = 11_007;
 pub struct SnarcosChainExtension;
 
 impl ChainExtension<Runtime> for SnarcosChainExtension {
-    fn call<E: Ext>(func_id: u32, env: Environment<E, InitState>) -> Result<RetVal, DispatchError>
+    fn call<E: Ext>(
+        func_id: u32,
+        env: SubstrateEnvironment<E, InitState>,
+    ) -> Result<RetVal, DispatchError>
     where
         <E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
     {
         match func_id {
-            SNARCOS_STORE_KEY_FUNC_ID => Self::snarcos_store_key(env),
-            SNARCOS_VERIFY_FUNC_ID => Self::snarcos_verify(env),
+            SNARCOS_STORE_KEY_FUNC_ID => {
+                Self::snarcos_store_key::<_, Runtime>(env.buf_in_buf_out())
+            }
+            SNARCOS_VERIFY_FUNC_ID => Self::snarcos_verify::<_, Runtime>(env.buf_in_buf_out()),
             _ => {
                 error!("Called an unregistered `func_id`: {}", func_id);
                 Err(DispatchError::Other("Unimplemented func_id"))
@@ -58,7 +67,7 @@ pub type ByteCount = u32;
 /// the order of values is important.
 ///
 /// It cannot be `MaxEncodedLen` due to `Vec<_>` and thus `Environment::read_as` cannot be used.
-#[derive(Decode)]
+#[derive(Decode, Encode)]
 struct StoreKeyArgs {
     pub identifier: VerificationKeyIdentifier,
     pub key: Vec<u8>,
@@ -70,35 +79,53 @@ struct StoreKeyArgs {
 /// the order of values is important.
 ///
 /// It cannot be `MaxEncodedLen` due to `Vec<_>` and thus `Environment::read_as` cannot be used.
-#[derive(Decode)]
+#[derive(Decode, Encode)]
 struct VerifyArgs {
     pub identifier: VerificationKeyIdentifier,
     pub proof: Vec<u8>,
     pub input: Vec<u8>,
     pub system: ProvingSystem,
 }
 
-impl SnarcosChainExtension {
-    fn snarcos_store_key<E: Ext>(env: Environment<E, InitState>) -> Result<RetVal, DispatchError>
-    where
-        <E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
-    {
-        // We need to read input as plain bytes (encoded args).
-        let mut env = env.buf_in_buf_out();
+/// Provides a weight of `store_key` dispatchable.
+fn weight_of_store_key(key_length: ByteCount) -> Weight {
+    <<Runtime as Config>::WeightInfo as WeightInfo>::store_key(key_length)
+}
+
+/// Provides a weight of `verify` dispatchable depending on the `ProvingSystem`. In case no system
+/// is passed, we return maximal amongst all the systems.
+fn weight_of_verify(system: Option<ProvingSystem>) -> Weight {
+    match system {
+        Some(ProvingSystem::Groth16) => {
+            <<Runtime as Config>::WeightInfo as WeightInfo>::verify_groth16()
+        }
+        Some(ProvingSystem::Gm17) => <<Runtime as Config>::WeightInfo as WeightInfo>::verify_gm17(),
+        Some(ProvingSystem::Marlin) => {
+            <<Runtime as Config>::WeightInfo as WeightInfo>::verify_marlin()
+        }
+        None => weight_of_verify(Some(ProvingSystem::Groth16))
+            .max(weight_of_verify(Some(ProvingSystem::Gm17)))
+            .max(weight_of_verify(Some(ProvingSystem::Marlin))),
+    }
+}
 
-        // Check if it makes sense to read and decode data.
-        let key_length = env
+impl SnarcosChainExtension {
+    fn snarcos_store_key<Env: Environment, Exc: Executor>(
+        mut env: Env,
+    ) -> Result<RetVal, DispatchError> {
+        // Check if it makes sense to read and decode data. This is only an upperbound for the key
+        // length, because this bytes suffix contains (possibly compressed) info about actual key
+        // length (needed for decoding).
+        let approx_key_length = env
             .in_len()
             .saturating_sub(size_of::<VerificationKeyIdentifier>() as ByteCount);
-        if key_length > MaximumVerificationKeyLength::get() {
+        if approx_key_length > MaximumVerificationKeyLength::get() {
             return Ok(RetVal::Converging(SNARCOS_STORE_KEY_TOO_LONG_KEY));
         }
 
         // We charge now - even if decoding fails and we shouldn't touch storage, we have to incur
         // fee for reading memory.
-        env.charge_weight(<<Runtime as Config>::WeightInfo as WeightInfo>::store_key(
-            key_length,
-        ))?;
+        let pre_charged = env.charge_weight(weight_of_store_key(approx_key_length))?;
 
         // Parsing will have to be done here. This is due to the fact that methods
         // `Environment<_,_,_,S: BufIn>::read*` don't move starting pointer and thus we can make
@@ -111,42 +138,27 @@ impl SnarcosChainExtension {
         let args = StoreKeyArgs::decode(&mut &*bytes)
             .map_err(|_| DispatchError::Other("Failed to decode arguments"))?;
 
-        let return_status = match Snarcos::<Runtime>::bare_store_key(args.identifier, args.key) {
+        // Now we know the exact key length.
+        env.adjust_weight(
+            pre_charged,
+            weight_of_store_key(args.key.len() as ByteCount),
+        );
+
+        let return_status = match Exc::store_key(args.identifier, args.key) {
             Ok(_) => SNARCOS_STORE_KEY_OK,
             // In case `DispatchResultWithPostInfo` was returned (or some simpler equivalent for
-            // `bare_store_key`), we could adjust weight. However, for the storing key action it
-            // doesn't make sense.
+            // `bare_store_key`), we could have adjusted weight. However, for the storing key action
+            // it doesn't make much sense.
             Err(VerificationKeyTooLong) => SNARCOS_STORE_KEY_TOO_LONG_KEY,
-            Err(IdentifierAlreadyInUse) => SNARCOS_STORE_KEY_IN_USE,
+            Err(IdentifierAlreadyInUse) => SNARCOS_STORE_KEY_IDENTIFIER_IN_USE,
             _ => SNARCOS_STORE_KEY_ERROR_UNKNOWN,
         };
         Ok(RetVal::Converging(return_status))
     }
 
-    fn weight_of_verify(system: Option<ProvingSystem>) -> Weight {
-        match system {
-            Some(ProvingSystem::Groth16) => {
-                <<Runtime as Config>::WeightInfo as WeightInfo>::verify_groth16()
-            }
-            Some(ProvingSystem::Gm17) => {
-                <<Runtime as Config>::WeightInfo as WeightInfo>::verify_gm17()
-            }
-            Some(ProvingSystem::Marlin) => {
-                <<Runtime as Config>::WeightInfo as WeightInfo>::verify_marlin()
-            }
-            None => Self::weight_of_verify(Some(ProvingSystem::Groth16))
-                .max(Self::weight_of_verify(Some(ProvingSystem::Gm17)))
-                .max(Self::weight_of_verify(Some(ProvingSystem::Marlin))),
-        }
-    }
-
-    fn snarcos_verify<E: Ext>(env: Environment<E, InitState>) -> Result<RetVal, DispatchError>
-    where
-        <E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
-    {
-        // We need to read input as plain bytes (encoded args).
-        let mut env = env.buf_in_buf_out();
-
+    fn snarcos_verify<Env: Environment, Exc: Executor>(
+        mut env: Env,
+    ) -> Result<RetVal, DispatchError> {
         // We charge optimistically, i.e. assuming that decoding succeeds and the verification
         // key is present. However, since we don't know the system yet, we have to charge maximal
         // possible fee. We will adjust it as soon as possible.
@@ -156,23 +168,23 @@ impl SnarcosChainExtension {
         // `pallet_snarcos::WeightInfo::verify_decoding_failure`, we can both charge less here
         // (with further `env.adjust_weight`) and in the pallet itself (returning
         // `DispatchErrorWithPostInfo` reducing actual fee and the block weight).
-        let pre_charge = env.charge_weight(Self::weight_of_verify(None))?;
+        let pre_charge = env.charge_weight(weight_of_verify(None))?;
 
         // Parsing is done here for similar reasons as in `Self::snarcos_store_key`.
         let bytes = env.read(env.in_len())?;
 
         let args: VerifyArgs = VerifyArgs::decode(&mut &*bytes)
             .map_err(|_| DispatchError::Other("Failed to decode arguments"))?;
 
-        env.adjust_weight(pre_charge, Self::weight_of_verify(Some(args.system)));
+        // Now we know the proving system and we can charge appropriate amount of gas.
+        env.adjust_weight(pre_charge, weight_of_verify(Some(args.system)));
 
-        let result =
-            Snarcos::<Runtime>::bare_verify(args.identifier, args.proof, args.input, args.system);
+        let result = Exc::verify(args.identifier, args.proof, args.input, args.system);
 
         let return_status = match result {
             Ok(_) => SNARCOS_VERIFY_OK,
             // In case `DispatchResultWithPostInfo` was returned (or some simpler equivalent for
-            // `bare_store_key`), we could adjust weight. However, we don't support it yet.
+            // `bare_verify`), we could adjust weight. However, we don't support it yet.
             Err(DeserializingProofFailed) => SNARCOS_VERIFY_DESERIALIZING_PROOF_FAIL,
             Err(DeserializingPublicInputFailed) => SNARCOS_VERIFY_DESERIALIZING_INPUT_FAIL,
             Err(UnknownVerificationKeyIdentifier) => SNARCOS_VERIFY_UNKNOWN_IDENTIFIER,