Skip to content

Commit

Permalink
add_benchmark_with_dilithium2_verification
Browse files Browse the repository at this point in the history
  • Loading branch information
kostiask committed Nov 12, 2023
1 parent d6ad3e2 commit 7292ffe
Show file tree
Hide file tree
Showing 12 changed files with 341 additions and 65 deletions.
53 changes: 52 additions & 1 deletion client/keystore/src/local.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
use async_trait::async_trait;
use parking_lot::RwLock;
use sp_application_crypto::{ecdsa, ed25519, sr25519, AppKey, AppPair, IsWrappedBy};
use sp_application_crypto::{ecdsa, ed25519, sr25519, dilithium2, AppKey, AppPair, IsWrappedBy};
use sp_core::{
crypto::{
ByteArray, CryptoTypePublicPair, ExposeSecret, KeyTypeId, Pair as PairT, SecretString,
Expand Down Expand Up @@ -102,6 +102,18 @@ impl CryptoStore for LocalKeystore {
SyncCryptoStore::ed25519_generate_new(self, id, seed)
}

async fn dilithium2_public_keys(&self, id: KeyTypeId) -> Vec<dilithium2::Public> {
SyncCryptoStore::dilithium2_public_keys(self, id)
}

async fn dilithium2_generate_new(
&self,
id: KeyTypeId,
seed: Option<&str>,
) -> std::result::Result<dilithium2::Public, TraitError> {
SyncCryptoStore::dilithium2_generate_new(self, id, seed)
}

async fn ecdsa_public_keys(&self, id: KeyTypeId) -> Vec<ecdsa::Public> {
SyncCryptoStore::ecdsa_public_keys(self, id)
}
Expand Down Expand Up @@ -169,6 +181,7 @@ impl SyncCryptoStore for LocalKeystore {
Ok(raw_keys.into_iter().fold(Vec::new(), |mut v, k| {
v.push(CryptoTypePublicPair(sr25519::CRYPTO_ID, k.clone()));
v.push(CryptoTypePublicPair(ed25519::CRYPTO_ID, k.clone()));
v.push(CryptoTypePublicPair(dilithium2::CRYPTO_ID, k.clone()));
v.push(CryptoTypePublicPair(ecdsa::CRYPTO_ID, k));
v
}))
Expand Down Expand Up @@ -201,6 +214,17 @@ impl SyncCryptoStore for LocalKeystore {
.map_err(TraitError::from)?;
key_pair.map(|k| k.sign(msg).encode()).map(Ok).transpose()
},
dilithium2::CRYPTO_ID => {
let pub_key = dilithium2::Public::from_slice(key.1.as_slice()).map_err(|()| {
TraitError::Other("Corrupted public key - Invalid size".into())
})?;
let key_pair = self
.0
.read()
.key_pair_by_type::<dilithium2::Pair>(&pub_key, id)
.map_err(TraitError::from)?;
key_pair.map(|k| k.sign(msg).encode()).map(Ok).transpose()
},
sr25519::CRYPTO_ID => {
let pub_key = sr25519::Public::from_slice(key.1.as_slice()).map_err(|()| {
TraitError::Other("Corrupted public key - Invalid size".into())
Expand Down Expand Up @@ -281,6 +305,33 @@ impl SyncCryptoStore for LocalKeystore {
Ok(pair.public())
}

fn dilithium2_public_keys(&self, key_type: KeyTypeId) -> Vec<dilithium2::Public> {
self.0
.read()
.raw_public_keys(key_type)
.map(|v| {
v.into_iter()
.filter_map(|k| dilithium2::Public::from_slice(k.as_slice()).ok())
.collect()
})
.unwrap_or_default()
}

fn dilithium2_generate_new(
&self,
id: KeyTypeId,
seed: Option<&str>,
) -> std::result::Result<dilithium2::Public, TraitError> {
let pair = match seed {
Some(seed) =>
self.0.write().insert_ephemeral_from_seed_by_type::<dilithium2::Pair>(seed, id),
None => self.0.write().generate_by_type::<dilithium2::Pair>(id),
}
.map_err(|e| -> TraitError { e.into() })?;

Ok(pair.public())
}

fn ecdsa_public_keys(&self, key_type: KeyTypeId) -> Vec<ecdsa::Public> {
self.0
.read()
Expand Down
25 changes: 22 additions & 3 deletions frame/benchmarking/src/baseline.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,13 +32,19 @@ use sp_std::prelude::*;

pub const TEST_KEY_TYPE_ID: KeyTypeId = KeyTypeId(*b"test");

mod app_sr25519 {
use super::TEST_KEY_TYPE_ID;
use sp_application_crypto::{app_crypto, sr25519};
app_crypto!(sr25519, TEST_KEY_TYPE_ID);
}
mod app_dilithium2 {
use super::TEST_KEY_TYPE_ID;
use sp_application_crypto::{app_crypto, dilithium2};
app_crypto!(dilithium2, TEST_KEY_TYPE_ID);
}

type SignerId = app_dilithium2::Public;
type SignerId = app_sr25519::Public;
type SignerId2 = app_dilithium2::Public;

pub struct Pallet<T: Config>(System<T>);
pub trait Config: frame_system::Config {}
Expand Down Expand Up @@ -89,11 +95,24 @@ benchmarks! {
assert!(hash != T::Hash::default());
}

dilithium2_verification {
sr25519_verification {
let i in 1 .. 100;

let public = SignerId::generate_pair(None);
let sigs_count: u8 = i.try_into().unwrap();
let msg_and_sigs: Vec<_> = (0..sigs_count).map(|j| {
let msg = vec![j, j];
(msg.clone(), public.sign(&msg).unwrap())
})
.collect();
}: {
msg_and_sigs.iter().for_each(|(msg, sig)| {
assert!(sig.verify(&msg[..], &public));
});
}

dilithium2_verification {
let i in 1 .. 100;
let public = SignerId2::generate_pair(None);
let sigs_count: u8 = i.try_into().unwrap();
let msg_and_sigs: Vec<_> = (0..sigs_count).map(|j| {
let msg = vec![j, j];
Expand Down
5 changes: 3 additions & 2 deletions frame/benchmarking/src/weights.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ pub trait WeightInfo {
fn division(i: u32, ) -> Weight;
fn hashing(i: u32, ) -> Weight;
fn dilithium2_verification(i: u32, ) -> Weight;
fn sr25519_verification(i: u32, ) -> Weight;
fn storage_read(i: u32, ) -> Weight;
fn storage_write(i: u32, ) -> Weight;
}
Expand All @@ -72,7 +73,7 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
fn hashing(_i: u32, ) -> Weight {
(20_865_902_000 as Weight)
}
fn dilithium2_verification(i: u32, ) -> Weight {
fn sr25519_verification(i: u32, ) -> Weight {
(319_000 as Weight)
// Standard Error: 8_000
.saturating_add((47_171_000 as Weight).saturating_mul(i as Weight))
Expand Down Expand Up @@ -110,7 +111,7 @@ impl WeightInfo for () {
fn hashing(_i: u32, ) -> Weight {
(20_865_902_000 as Weight)
}
fn dilithium2_verification(i: u32, ) -> Weight {
fn sr25519_verification(i: u32, ) -> Weight {
(319_000 as Weight)
// Standard Error: 8_000
.saturating_add((47_171_000 as Weight).saturating_mul(i as Weight))
Expand Down
41 changes: 41 additions & 0 deletions primitives/application-crypto/test/src/dilithium2.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// This file is part of Substrate.

// Copyright (C) 2019-2022 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.

//! Integration tests for dilithium2
use sp_api::ProvideRuntimeApi;
use sp_application_crypto::dilithium2::{AppPair, AppPublic};
use sp_core::{crypto::Pair, testing::DILITHIUM2};
use sp_keystore::{testing::KeyStore, SyncCryptoStore};
use sp_runtime::generic::BlockId;
use std::sync::Arc;
use substrate_test_runtime_client::{
runtime::TestAPI, DefaultTestClientBuilderExt, TestClientBuilder, TestClientBuilderExt,
};

#[test]
fn dilithium2_works_in_runtime() {
let keystore = Arc::new(KeyStore::new());
let test_client = TestClientBuilder::new().set_keystore(keystore.clone()).build();
let (signature, public) = test_client
.runtime_api()
.test_dilithium2_crypto(&BlockId::Number(0))
.expect("Tests `dilithium2` crypto.");
let supported_keys = SyncCryptoStore::keys(&*keystore, DILITHIUM2).unwrap();
assert!(supported_keys.contains(&public.clone().into()));
assert!(AppPair::verify(&signature, "dilithium2", &AppPublic::from(public)));
}
3 changes: 3 additions & 0 deletions primitives/application-crypto/test/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,3 +23,6 @@ mod ecdsa;
mod ed25519;
#[cfg(test)]
mod sr25519;

#[cfg(test)]
mod dilithium2;
8 changes: 6 additions & 2 deletions primitives/core/src/dilithium2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -437,11 +437,14 @@ impl TraitPair for Pair {
fn derive<Iter: Iterator<Item=DeriveJunction>>(
&self,
path: Iter,
_: Option<Seed>,
_seed: Option<Seed>,
) -> Result<(Self, Option<Seed>), Self::DeriveError> {
let acc = self.secret.0;
let mut seed = [0u8; 32];
seed.copy_from_slice(&acc[0..32]);
match _seed {
Some(s) => seed.copy_from_slice(&s[0..32]),
None => seed.copy_from_slice(&acc[0..32])
};

for j in path {
match j {
Expand All @@ -452,6 +455,7 @@ impl TraitPair for Pair {

Ok((Self::from_seed(&seed), Some(seed)))
}

fn from_seed(seed: &Self::Seed) -> Self {
Self::from_seed_slice(&seed[..]).expect("seed has valid length; qed")
}
Expand Down
2 changes: 1 addition & 1 deletion primitives/core/src/testing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::crypto::KeyTypeId;
pub const ED25519: KeyTypeId = KeyTypeId(*b"ed25");
/// Key type for generic Sr 25519 key.
pub const SR25519: KeyTypeId = KeyTypeId(*b"sr25");
/// Key type for generic Sr 25519 key.
/// Key type for generic Sr Dilithium2 key.
pub const DILITHIUM2: KeyTypeId = KeyTypeId(*b"dth2");
/// Key type for generic ECDSA key.
pub const ECDSA: KeyTypeId = KeyTypeId(*b"ecds");
Expand Down
14 changes: 13 additions & 1 deletion primitives/io/src/batch_verifier.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
//! Batch/parallel verification.
use futures::{channel::oneshot, future::FutureExt};
use sp_core::{crypto::Pair, ecdsa, ed25519, sr25519, traits::SpawnNamed};
use sp_core::{crypto::Pair, ecdsa, ed25519, sr25519, dilithium2, traits::SpawnNamed};
use std::sync::{
atomic::{AtomicBool, Ordering as AtomicOrdering},
Arc,
Expand Down Expand Up @@ -107,6 +107,18 @@ impl BatchVerifier {
)
}

pub fn push_dilithium2(
&mut self,
signature: dilithium2::Signature,
pub_key: dilithium2::Public,
message: Vec<u8>,
) -> bool {
self.spawn_verification_task(
move || dilithium2::Pair::verify(&signature, &message, &pub_key),
"substrate_dilithium2_verify",
)
}

/// Push sr25519 signature to verify.
///
/// Returns false if some of the pushed signatures before already failed the check.
Expand Down
Loading

0 comments on commit 7292ffe

Please sign in to comment.