Skip to content

Commit

Permalink
use trait in session historical
Browse files Browse the repository at this point in the history
  • Loading branch information
shawntabrizi committed Mar 28, 2024
1 parent 6ab4bfb commit 1cfb29f
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 26 deletions.
47 changes: 26 additions & 21 deletions substrate/frame/session/src/historical/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,9 @@ mod shared;

use codec::{Decode, Encode};
use sp_runtime::{
proving_trie::ProvingTrie,
traits::{Convert, OpaqueKeys},
KeyTypeId,
DispatchError, KeyTypeId,
};
use sp_session::{MembershipProof, ValidatorCount};
use sp_staking::SessionIndex;
Expand Down Expand Up @@ -176,7 +177,7 @@ impl<T: Config, I: SessionManager<T::ValidatorId, T::FullIdentification>> NoteHi

if let Some(new_validators) = new_validators_and_id {
let count = new_validators.len() as ValidatorCount;
match ProvingTrie::<T>::generate_for(new_validators) {
match ValidatorProvingTrie::<T>::generate_for(new_validators) {
Ok(trie) => <HistoricalSessions<T>>::insert(new_index, &(trie.root, count)),
Err(reason) => {
print("Failed to generate historical ancestry-inclusion proof.");
Expand Down Expand Up @@ -221,13 +222,15 @@ pub type IdentificationTuple<T> =
(<T as pallet_session::Config>::ValidatorId, <T as Config>::FullIdentification);

/// A trie instance for checking and generating proofs.
pub struct ProvingTrie<T: Config> {
pub struct ValidatorProvingTrie<T: Config> {
db: MemoryDB<T::Hashing>,
root: T::Hash,
}

impl<T: Config> ProvingTrie<T> {
fn generate_for<I>(validators: I) -> Result<Self, &'static str>
impl<T: Config> ProvingTrie<T::Hashing, T::Hash, IdentificationTuple<T>>
for ValidatorProvingTrie<T>
{
fn generate_for<I>(validators: I) -> Result<Self, DispatchError>
where
I: IntoIterator<Item = (T::ValidatorId, T::FullIdentification)>,
{
Expand Down Expand Up @@ -261,22 +264,11 @@ impl<T: Config> ProvingTrie<T> {
}
}

Ok(ProvingTrie { db, root })
}

fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> Self {
use sp_trie::HashDBT;

let mut memory_db = MemoryDB::default();
for node in nodes {
HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
}

ProvingTrie { db: memory_db, root }
Ok(Self { db, root })
}

/// Prove the full verification data for a given key and key ID.
pub fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
let mut recorder = Recorder::<LayoutV0<T::Hashing>>::new();
{
let trie =
Expand All @@ -296,7 +288,7 @@ impl<T: Config> ProvingTrie<T> {
}

/// Access the underlying trie root.
pub fn root(&self) -> &T::Hash {
fn root(&self) -> &T::Hash {
&self.root
}

Expand All @@ -316,6 +308,19 @@ impl<T: Config> ProvingTrie<T> {
}
}

impl<T: Config> ValidatorProvingTrie<T> {
fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> Self {
use sp_trie::HashDBT;

let mut memory_db = MemoryDB::default();
for node in nodes {
HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
}

Self { db: memory_db, root }
}
}

impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T> {
type Proof = MembershipProof;
type IdentificationTuple = IdentificationTuple<T>;
Expand All @@ -332,7 +337,7 @@ impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T

let count = validators.len() as ValidatorCount;

let trie = ProvingTrie::<T>::generate_for(validators).ok()?;
let trie = ValidatorProvingTrie::<T>::generate_for(validators).ok()?;

let (id, data) = key;
trie.prove(id, data.as_ref()).map(|trie_nodes| MembershipProof {
Expand Down Expand Up @@ -364,7 +369,7 @@ impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T
return None
}

let trie = ProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
let trie = ValidatorProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
trie.query(id, data.as_ref())
}
}
Expand Down
11 changes: 6 additions & 5 deletions substrate/frame/session/src/historical/offchain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,18 @@
//! Validator-set extracting an iterator from an off-chain worker stored list containing historical
//! validator-sets. Based on the logic of historical slashing, but the validation is done off-chain.
//! Use [`fn store_current_session_validator_set_to_offchain()`](super::onchain) to store the
//! required data to the offchain validator set. This is used in conjunction with [`ProvingTrie`]
//! and the off-chain indexing API.
//! required data to the offchain validator set. This is used in conjunction with
//! [`ValidatorProvingTrie`] and the off-chain indexing API.
use sp_runtime::{
offchain::storage::{MutateStorageError, StorageRetrievalError, StorageValueRef},
proving_trie::ProvingTrie,
KeyTypeId,
};
use sp_session::MembershipProof;
use sp_std::prelude::*;

use super::{shared, Config, IdentificationTuple, ProvingTrie};
use super::{shared, Config, IdentificationTuple, ValidatorProvingTrie};
use crate::{Pallet as SessionModule, SessionIndex};

/// A set of validators, which was used for a fixed session index.
Expand Down Expand Up @@ -59,7 +60,7 @@ impl<T: Config> ValidatorSet<T> {
}

/// Implement conversion into iterator for usage
/// with [ProvingTrie](super::ProvingTrie::generate_for).
/// with [ValidatorProvingTrie](super::ValidatorProvingTrie::generate_for).
impl<T: Config> sp_std::iter::IntoIterator for ValidatorSet<T> {
type Item = (T::ValidatorId, T::FullIdentification);
type IntoIter = sp_std::vec::IntoIter<Self::Item>;
Expand All @@ -79,7 +80,7 @@ pub fn prove_session_membership<T: Config, D: AsRef<[u8]>>(
) -> Option<MembershipProof> {
let validators = ValidatorSet::<T>::load_from_offchain_db(session_index)?;
let count = validators.len() as u32;
let trie = ProvingTrie::<T>::generate_for(validators.into_iter()).ok()?;
let trie = ValidatorProvingTrie::<T>::generate_for(validators.into_iter()).ok()?;

let (id, data) = session_key;
trie.prove(id, data.as_ref()).map(|trie_nodes| MembershipProof {
Expand Down

0 comments on commit 1cfb29f

Please sign in to comment.