Custom RPC for Merkle Mountain Range pallet (#8137)

* Add MMR custom RPC.

* Change RuntimeApi to avoid hardcoding leaf type.

* Properly implement the new RuntimeAPI and wire up RPC.

* Extract Offchain DB as separate execution extension.

* Enable offchain DB access for offchain calls.

* Fix offchain_election tests.

* Skip block initialisation for proof generation.

* Fix integration test setup.

* Fix offchain tests. Not sure how I missed them earlier 🤷.

* Fix long line.

* One more test missing.

* Update mock for multi-phase.

* Address review grumbbles.

* Address review grumbles.

* Fix line width of a comment

substrate/frame/election-provider-multi-phase/src/mock.rs

@@ -27,7 +27,7 @@ use parking_lot::RwLock;
 use sp_core::{
 	offchain::{
 		testing::{PoolState, TestOffchainExt, TestTransactionPoolExt},
-		OffchainExt, TransactionPoolExt,
+		OffchainDbExt, OffchainWorkerExt, TransactionPoolExt,
 	},
 	H256,
 };
@@ -369,7 +369,8 @@ impl ExtBuilder {
 		seed[0..4].copy_from_slice(&iters.to_le_bytes());
 		offchain_state.write().seed = seed;
 
-		ext.register_extension(OffchainExt::new(offchain));
+		ext.register_extension(OffchainDbExt::new(offchain.clone()));
+		ext.register_extension(OffchainWorkerExt::new(offchain));
 		ext.register_extension(TransactionPoolExt::new(pool));
 
 		(ext, pool_state)

substrate/frame/example-offchain-worker/src/tests.rs

@@ -22,7 +22,7 @@ use codec::Decode;
 use frame_support::{assert_ok, parameter_types};
 use sp_core::{
 	H256,
-	offchain::{OffchainExt, TransactionPoolExt, testing},
+	offchain::{OffchainWorkerExt, TransactionPoolExt, testing},
 	sr25519::Signature,
 };
 
@@ -144,7 +144,7 @@ fn it_aggregates_the_price() {
 fn should_make_http_call_and_parse_result() {
 	let (offchain, state) = testing::TestOffchainExt::new();
 	let mut t = sp_io::TestExternalities::default();
-	t.register_extension(OffchainExt::new(offchain));
+	t.register_extension(OffchainWorkerExt::new(offchain));
 
 	price_oracle_response(&mut state.write());
 
@@ -160,7 +160,7 @@ fn should_make_http_call_and_parse_result() {
 fn knows_how_to_mock_several_http_calls() {
 	let (offchain, state) = testing::TestOffchainExt::new();
 	let mut t = sp_io::TestExternalities::default();
-	t.register_extension(OffchainExt::new(offchain));
+	t.register_extension(OffchainWorkerExt::new(offchain));
 
 	{
 		let mut state = state.write();
@@ -217,7 +217,7 @@ fn should_submit_signed_transaction_on_chain() {
 
 
 	let mut t = sp_io::TestExternalities::default();
-	t.register_extension(OffchainExt::new(offchain));
+	t.register_extension(OffchainWorkerExt::new(offchain));
 	t.register_extension(TransactionPoolExt::new(pool));
 	t.register_extension(KeystoreExt(Arc::new(keystore)));
 
@@ -255,7 +255,7 @@ fn should_submit_unsigned_transaction_on_chain_for_any_account() {
 		.clone();
 
 	let mut t = sp_io::TestExternalities::default();
-	t.register_extension(OffchainExt::new(offchain));
+	t.register_extension(OffchainWorkerExt::new(offchain));
 	t.register_extension(TransactionPoolExt::new(pool));
 	t.register_extension(KeystoreExt(Arc::new(keystore)));
 
@@ -308,7 +308,7 @@ fn should_submit_unsigned_transaction_on_chain_for_all_accounts() {
 		.clone();
 
 	let mut t = sp_io::TestExternalities::default();
-	t.register_extension(OffchainExt::new(offchain));
+	t.register_extension(OffchainWorkerExt::new(offchain));
 	t.register_extension(TransactionPoolExt::new(pool));
 	t.register_extension(KeystoreExt(Arc::new(keystore)));
 
@@ -349,7 +349,7 @@ fn should_submit_raw_unsigned_transaction_on_chain() {
 	let keystore = KeyStore::new();
 
 	let mut t = sp_io::TestExternalities::default();
-	t.register_extension(OffchainExt::new(offchain));
+	t.register_extension(OffchainWorkerExt::new(offchain));
 	t.register_extension(TransactionPoolExt::new(pool));
 	t.register_extension(KeystoreExt(Arc::new(keystore)));
 

substrate/frame/im-online/src/tests.rs

@@ -23,7 +23,8 @@ use super::*;
 use crate::mock::*;
 use sp_core::OpaquePeerId;
 use sp_core::offchain::{
-	OffchainExt,
+	OffchainDbExt,
+	OffchainWorkerExt,
 	TransactionPoolExt,
 	testing::{TestOffchainExt, TestTransactionPoolExt},
 };
@@ -205,7 +206,8 @@ fn should_generate_heartbeats() {
 	let mut ext = new_test_ext();
 	let (offchain, _state) = TestOffchainExt::new();
 	let (pool, state) = TestTransactionPoolExt::new();
-	ext.register_extension(OffchainExt::new(offchain));
+	ext.register_extension(OffchainDbExt::new(offchain.clone()));
+	ext.register_extension(OffchainWorkerExt::new(offchain));
 	ext.register_extension(TransactionPoolExt::new(pool));
 
 	ext.execute_with(|| {
@@ -310,7 +312,8 @@ fn should_not_send_a_report_if_already_online() {
 	let mut ext = new_test_ext();
 	let (offchain, _state) = TestOffchainExt::new();
 	let (pool, pool_state) = TestTransactionPoolExt::new();
-	ext.register_extension(OffchainExt::new(offchain));
+	ext.register_extension(OffchainDbExt::new(offchain.clone()));
+	ext.register_extension(OffchainWorkerExt::new(offchain));
 	ext.register_extension(TransactionPoolExt::new(pool));
 
 	ext.execute_with(|| {

substrate/frame/merkle-mountain-range/primitives/src/lib.rs

@@ -360,6 +360,11 @@ impl OpaqueLeaf {
 	pub fn from_encoded_leaf(encoded_leaf: Vec<u8>) -> Self {
 		OpaqueLeaf(encoded_leaf)
 	}
+
+	/// Attempt to decode the leaf into expected concrete type.
+	pub fn try_decode<T: codec::Decode>(&self) -> Option<T> {
+		codec::Decode::decode(&mut &*self.0).ok()
+	}
 }
 
 impl FullLeaf for OpaqueLeaf {
@@ -368,18 +373,49 @@ impl FullLeaf for OpaqueLeaf {
 	}
 }
 
+/// A type-safe wrapper for the concrete leaf type.
+///
+/// This structure serves merely to avoid passing raw `Vec<u8>` around.
+/// It must be `Vec<u8>`-encoding compatible.
+///
+/// It is different from [`OpaqueLeaf`], because it does implement `Codec`
+/// and the encoding has to match raw `Vec<u8>` encoding.
+#[derive(codec::Encode, codec::Decode, RuntimeDebug, PartialEq, Eq)]
+pub struct EncodableOpaqueLeaf(pub Vec<u8>);
+
+impl EncodableOpaqueLeaf {
+	/// Convert a concrete leaf into encodable opaque version.
+	pub fn from_leaf<T: FullLeaf>(leaf: &T) -> Self {
+		let opaque = OpaqueLeaf::from_leaf(leaf);
+		Self::from_opaque_leaf(opaque)
+	}
+
+	/// Given an opaque leaf, make it encodable.
+	pub fn from_opaque_leaf(opaque: OpaqueLeaf) -> Self {
+		Self(opaque.0)
+	}
+
+	/// Try to convert into a [OpaqueLeaf].
+	pub fn into_opaque_leaf(self) -> OpaqueLeaf {
+		// wrap into `OpaqueLeaf` type
+		OpaqueLeaf::from_encoded_leaf(self.0)
+	}
+}
+
 sp_api::decl_runtime_apis! {
 	/// API to interact with MMR pallet.
-	pub trait MmrApi<Leaf: codec::Codec, Hash: codec::Codec> {
+	pub trait MmrApi<Hash: codec::Codec> {
 		/// Generate MMR proof for a leaf under given index.
-		fn generate_proof(leaf_index: u64) -> Result<(Leaf, Proof<Hash>), Error>;
+		#[skip_initialize_block]
+		fn generate_proof(leaf_index: u64) -> Result<(EncodableOpaqueLeaf, Proof<Hash>), Error>;
 
 		/// Verify MMR proof against on-chain MMR.
 		///
 		/// Note this function will use on-chain MMR root hash and check if the proof
 		/// matches the hash.
 		/// See [Self::verify_proof_stateless] for a stateless verifier.
-		fn verify_proof(leaf: Leaf, proof: Proof<Hash>) -> Result<(), Error>;
+		#[skip_initialize_block]
+		fn verify_proof(leaf: EncodableOpaqueLeaf, proof: Proof<Hash>) -> Result<(), Error>;
 
 		/// Verify MMR proof against given root hash.
 		///
@@ -387,7 +423,8 @@ sp_api::decl_runtime_apis! {
 		/// proof is verified against given MMR root hash.
 		///
 		/// The leaf data is expected to be encoded in it's compact form.
-		fn verify_proof_stateless(root: Hash, leaf: Vec<u8>, proof: Proof<Hash>)
+		#[skip_initialize_block]
+		fn verify_proof_stateless(root: Hash, leaf: EncodableOpaqueLeaf, proof: Proof<Hash>)
 			-> Result<(), Error>;
 	}
 }
@@ -535,7 +572,7 @@ mod tests {
 	}
 
 	#[test]
-	fn opaque_leaves_should_be_scale_compatible_with_concrete_ones() {
+	fn opaque_leaves_should_be_full_leaf_compatible() {
 		// given
 		let a = Test::Data("Hello World!".into());
 		let b = Test::Data("".into());
@@ -564,4 +601,36 @@ mod tests {
 			opaque,
 		);
 	}
+
+	#[test]
+	fn encode_opaque_leaf_should_be_scale_compatible() {
+		use codec::Encode;
+
+		// given
+		let a = Test::Data("Hello World!".into());
+		let case1 = EncodableOpaqueLeaf::from_leaf(&a);
+		let case2 = EncodableOpaqueLeaf::from_opaque_leaf(OpaqueLeaf(a.encode()));
+		let case3 = a.encode().encode();
+
+		// when
+		let encoded = vec![&case1, &case2]
+			.into_iter()
+			.map(|x| x.encode())
+			.collect::<Vec<_>>();
+		let decoded = vec![&*encoded[0], &*encoded[1], &*case3]
+			.into_iter()
+			.map(|x| EncodableOpaqueLeaf::decode(&mut &*x))
+			.collect::<Vec<_>>();
+
+		// then
+		assert_eq!(case1, case2);
+		assert_eq!(encoded[0], encoded[1]);
+		// then encoding should also match double-encoded leaf.
+		assert_eq!(encoded[0], case3);
+
+		assert_eq!(decoded[0], decoded[1]);
+		assert_eq!(decoded[1], decoded[2]);
+		assert_eq!(decoded[0], Ok(case2));
+		assert_eq!(decoded[1], Ok(case1));
+	}
 }

substrate/frame/merkle-mountain-range/rpc/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "pallet-mmr-rpc"
+version = "3.0.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2018"
+license = "Apache-2.0"
+homepage = "https://substrate.dev"
+repository = "https://github.com/paritytech/substrate/"
+description = "Node-specific RPC methods for interaction with Merkle Mountain Range pallet."
+publish = false
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]
+
+[dependencies]
+codec = { package = "parity-scale-codec", version = "2.0.0" }
+jsonrpc-core = "15.1.0"
+jsonrpc-core-client = "15.1.0"
+jsonrpc-derive = "15.1.0"
+pallet-mmr-primitives = { version = "3.0.0", path = "../primitives" }
+serde = { version = "1.0.101", features = ["derive"] }
+sp-api = { version = "3.0.0", path = "../../../primitives/api" }
+sp-blockchain = { version = "3.0.0", path = "../../../primitives/blockchain" }
+sp-core = { version = "3.0.0", path = "../../../primitives/core" }
+sp-rpc = { version = "3.0.0", path = "../../../primitives/rpc" }
+sp-runtime = { version = "3.0.0", path = "../../../primitives/runtime" }
+
+[dev-dependencies]
+serde_json = "1.0.41"

substrate/frame/merkle-mountain-range/rpc/src/lib.rs

@@ -0,0 +1,222 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#![warn(missing_docs)]
+
+//! Node-specific RPC methods for interaction with Merkle Mountain Range pallet.
+
+use std::sync::Arc;
+
+use codec::{Codec, Encode};
+use jsonrpc_core::{Error, ErrorCode, Result};
+use jsonrpc_derive::rpc;
+use serde::{Deserialize, Serialize};
+use sp_api::ProvideRuntimeApi;
+use sp_blockchain::HeaderBackend;
+use sp_core::Bytes;
+use sp_runtime::{
+	generic::BlockId,
+	traits::{Block as BlockT},
+};
+use pallet_mmr_primitives::{Error as MmrError, Proof};
+
+pub use pallet_mmr_primitives::MmrApi as MmrRuntimeApi;
+
+/// Retrieved MMR leaf and its proof.
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
+#[serde(rename_all = "camelCase")]
+pub struct LeafProof<BlockHash> {
+	/// Block hash the proof was generated for.
+	pub block_hash: BlockHash,
+	/// SCALE-encoded leaf data.
+	pub leaf: Bytes,
+	/// SCALE-encoded proof data. See [pallet_mmr_primitives::Proof].
+	pub proof: Bytes,
+}
+
+impl<BlockHash> LeafProof<BlockHash> {
+	/// Create new `LeafProof` from given concrete `leaf` and `proof`.
+	pub fn new<Leaf, MmrHash>(
+		block_hash: BlockHash,
+		leaf: Leaf,
+		proof: Proof<MmrHash>,
+	) -> Self where
+		Leaf: Encode,
+		MmrHash: Encode,
+	{
+		Self {
+			block_hash,
+			leaf: Bytes(leaf.encode()),
+			proof: Bytes(proof.encode()),
+		}
+	}
+}
+
+/// MMR RPC methods.
+#[rpc]
+pub trait MmrApi<BlockHash> {
+	/// Generate MMR proof for given leaf index.
+	///
+	/// This method calls into a runtime with MMR pallet included and attempts to generate
+	/// MMR proof for leaf at given `leaf_index`.
+	/// Optionally, a block hash at which the runtime should be queried can be specified.
+	///
+	/// Returns the (full) leaf itself and a proof for this leaf (compact encoding, i.e. hash of
+	/// the leaf). Both parameters are SCALE-encoded.
+	#[rpc(name = "mmr_generateProof")]
+	fn generate_proof(
+		&self,
+		leaf_index: u64,
+		at: Option<BlockHash>,
+	) -> Result<LeafProof<BlockHash>>;
+}
+
+/// An implementation of MMR specific RPC methods.
+pub struct Mmr<C, B> {
+	client: Arc<C>,
+	_marker: std::marker::PhantomData<B>,
+}
+
+impl<C, B> Mmr<C, B> {
+	/// Create new `Mmr` with the given reference to the client.
+	pub fn new(client: Arc<C>) -> Self {
+		Self {
+			client,
+			_marker: Default::default(),
+		}
+	}
+}
+
+impl<C, Block, MmrHash> MmrApi<<Block as BlockT>::Hash,> for Mmr<C, (Block, MmrHash)>
+where
+	Block: BlockT,
+	C: Send + Sync + 'static + ProvideRuntimeApi<Block> + HeaderBackend<Block>,
+	C::Api: MmrRuntimeApi<
+		Block,
+		MmrHash,
+	>,
+	MmrHash: Codec + Send + Sync + 'static,
+{
+	fn generate_proof(
+		&self,
+		leaf_index: u64,
+		at: Option<<Block as BlockT>::Hash>,
+	) -> Result<LeafProof<<Block as BlockT>::Hash>> {
+		let api = self.client.runtime_api();
+		let block_hash = at.unwrap_or_else(||
+			// If the block hash is not supplied assume the best block.
+			self.client.info().best_hash
+		);
+
+		let (leaf, proof) = api
+			.generate_proof_with_context(
+				&BlockId::hash(block_hash),
+				sp_core::ExecutionContext::OffchainCall(None),
+				leaf_index,
+			)
+			.map_err(runtime_error_into_rpc_error)?
+			.map_err(mmr_error_into_rpc_error)?;
+
+		Ok(LeafProof::new(block_hash, leaf, proof))
+	}
+}
+
+const RUNTIME_ERROR: i64 = 8000;
+const MMR_ERROR: i64 = 8010;
+
+/// Converts a mmr-specific error into an RPC error.
+fn mmr_error_into_rpc_error(err: MmrError) -> Error {
+	match err {
+		MmrError::LeafNotFound => Error {
+			code: ErrorCode::ServerError(MMR_ERROR + 1),
+			message: "Leaf was not found".into(),
+			data: Some(format!("{:?}", err).into()),
+		},
+		MmrError::GenerateProof => Error {
+			code: ErrorCode::ServerError(MMR_ERROR + 2),
+			message: "Error while generating the proof".into(),
+			data: Some(format!("{:?}", err).into()),
+		},
+		_ => Error {
+			code: ErrorCode::ServerError(MMR_ERROR),
+			message: "Unexpected MMR error".into(),
+			data: Some(format!("{:?}", err).into()),
+		},
+	}
+}
+
+/// Converts a runtime trap into an RPC error.
+fn runtime_error_into_rpc_error(err: impl std::fmt::Debug) -> Error {
+	Error {
+		code: ErrorCode::ServerError(RUNTIME_ERROR),
+		message: "Runtime trapped".into(),
+		data: Some(format!("{:?}", err).into()),
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use sp_core::H256;
+
+	#[test]
+	fn should_serialize_leaf_proof() {
+		// given
+		let leaf = vec![1_u8, 2, 3, 4];
+		let proof = Proof {
+			leaf_index: 1,
+			leaf_count: 9,
+			items: vec![H256::repeat_byte(1), H256::repeat_byte(2)],
+		};
+
+		let leaf_proof = LeafProof::new(H256::repeat_byte(0), leaf, proof);
+
+		// when
+		let actual = serde_json::to_string(&leaf_proof).unwrap();
+
+		// then
+		assert_eq!(
+			actual,
+			r#"{"blockHash":"0x0000000000000000000000000000000000000000000000000000000000000000","leaf":"0x1001020304","proof":"0x010000000000000009000000000000000801010101010101010101010101010101010101010101010101010101010101010202020202020202020202020202020202020202020202020202020202020202"}"#
+		);
+	}
+
+	#[test]
+	fn should_deserialize_leaf_proof() {
+		// given
+		let expected = LeafProof {
+			block_hash: H256::repeat_byte(0),
+			leaf: Bytes(vec![1_u8, 2, 3, 4].encode()),
+			proof: Bytes(Proof {
+				leaf_index: 1,
+				leaf_count: 9,
+				items: vec![H256::repeat_byte(1), H256::repeat_byte(2)],
+			}.encode()),
+		};
+
+		// when
+		let actual: LeafProof<H256> = serde_json::from_str(r#"{
+			"blockHash":"0x0000000000000000000000000000000000000000000000000000000000000000",
+			"leaf":"0x1001020304",
+			"proof":"0x010000000000000009000000000000000801010101010101010101010101010101010101010101010101010101010101010202020202020202020202020202020202020202020202020202020202020202"
+		}"#).unwrap();
+
+		// then
+		assert_eq!(actual, expected);
+
+	}
+}

substrate/frame/merkle-mountain-range/src/mmr/storage.rs

@@ -62,10 +62,9 @@ impl<T, I, L> mmr_lib::MMRStore<NodeOf<T, I, L>> for Storage<OffchainStorage, T,
 	fn get_elem(&self, pos: u64) -> mmr_lib::Result<Option<NodeOf<T, I, L>>> {
 		let key = Module::<T, I>::offchain_key(pos);
 		// Retrieve the element from Off-chain DB.
-		Ok(
-			sp_io::offchain ::local_storage_get(sp_core::offchain::StorageKind::PERSISTENT, &key)
-				.and_then(|v| codec::Decode::decode(&mut &*v).ok())
-		)
+		Ok(sp_io::offchain
+			::local_storage_get(sp_core::offchain::StorageKind::PERSISTENT, &key)
+			.and_then(|v| codec::Decode::decode(&mut &*v).ok()))
 	}
 
 	fn append(&mut self, _: u64, _: Vec<NodeOf<T, I, L>>) -> mmr_lib::Result<()> {
@@ -95,9 +94,8 @@ impl<T, I, L> mmr_lib::MMRStore<NodeOf<T, I, L>> for Storage<RuntimeStorage, T,
 			// on-chain we only store the hash (even if it's a leaf)
 			<Nodes<T, I>>::insert(size, elem.hash());
 			// Indexing API is used to store the full leaf content.
-			elem.using_encoded(|elem| {
-				sp_io::offchain_index::set(&Module::<T, I>::offchain_key(size), elem)
-			});
+			let key = Module::<T, I>::offchain_key(size);
+			elem.using_encoded(|elem| sp_io::offchain_index::set(&key, elem));
 			size += 1;
 
 			if let Node::Data(..) = elem {

substrate/frame/merkle-mountain-range/src/tests.rs

@@ -23,7 +23,7 @@ use sp_core::{
 	H256,
 	offchain::{
 		testing::TestOffchainExt,
-		OffchainExt,
+		OffchainWorkerExt, OffchainDbExt,
 	},
 };
 use pallet_mmr_primitives::{Proof, Compact};
@@ -34,7 +34,8 @@ pub(crate) fn new_test_ext() -> sp_io::TestExternalities {
 
 fn register_offchain_ext(ext: &mut sp_io::TestExternalities) {
 	let (offchain, _offchain_state) = TestOffchainExt::with_offchain_db(ext.offchain_db());
-	ext.register_extension(OffchainExt::new(offchain));
+	ext.register_extension(OffchainDbExt::new(offchain.clone()));
+	ext.register_extension(OffchainWorkerExt::new(offchain));
 }
 
 fn new_block() -> u64 {

substrate/frame/session/src/historical/offchain.rs

@@ -147,7 +147,8 @@ mod tests {
 	use sp_core::crypto::key_types::DUMMY;
 	use sp_core::offchain::{
 		testing::TestOffchainExt,
-		OffchainExt,
+		OffchainDbExt,
+		OffchainWorkerExt,
 		StorageKind,
 	};
 
@@ -181,7 +182,8 @@ mod tests {
 		seed[0..4].copy_from_slice(&ITERATIONS.to_le_bytes());
 		offchain_state.write().seed = seed;
 
-		ext.register_extension(OffchainExt::new(offchain));
+		ext.register_extension(OffchainDbExt::new(offchain.clone()));
+		ext.register_extension(OffchainWorkerExt::new(offchain));
 		ext
 	}
 

substrate/frame/staking/fuzzer/src/submit_solution.rs

@@ -26,7 +26,7 @@ use pallet_staking::testing_utils::*;
 use frame_support::{assert_ok, storage::StorageValue, traits::UnfilteredDispatchable};
 use frame_system::RawOrigin;
 use sp_runtime::DispatchError;
-use sp_core::offchain::{testing::TestOffchainExt, OffchainExt};
+use sp_core::offchain::{testing::TestOffchainExt, OffchainWorkerExt, OffchainDbExt};
 use pallet_staking::{EraElectionStatus, ElectionStatus, Module as Staking, Call as StakingCall};
 
 mod mock;
@@ -55,7 +55,8 @@ pub fn new_test_ext(iterations: u32) -> sp_io::TestExternalities {
 	seed[0..4].copy_from_slice(&iterations.to_le_bytes());
 	offchain_state.write().seed = seed;
 
-	ext.register_extension(OffchainExt::new(offchain));
+	ext.register_extension(OffchainDbExt::new(offchain.clone()));
+	ext.register_extension(OffchainWorkerExt::new(offchain));
 
 	ext
 }

substrate/frame/staking/src/tests.rs

@@ -2917,7 +2917,7 @@ mod offchain_election {
 	use parking_lot::RwLock;
 	use sp_core::offchain::{
 		testing::{PoolState, TestOffchainExt, TestTransactionPoolExt},
-		OffchainExt, TransactionPoolExt,
+		OffchainWorkerExt, TransactionPoolExt, OffchainDbExt,
 	};
 	use sp_io::TestExternalities;
 	use sp_npos_elections::StakedAssignment;
@@ -2960,7 +2960,8 @@ mod offchain_election {
 		seed[0..4].copy_from_slice(&iterations.to_le_bytes());
 		offchain_state.write().seed = seed;
 
-		ext.register_extension(OffchainExt::new(offchain));
+		ext.register_extension(OffchainDbExt::new(offchain.clone()));
+		ext.register_extension(OffchainWorkerExt::new(offchain));
 		ext.register_extension(TransactionPoolExt::new(pool));
 
 		pool_state