Make BEEFY client keystore generic over BEEFY AuthorityId type (#2258)

This is the significant step to make BEEFY client able to handle both ECDSA and (ECDSA, BLS) type signature. The idea is having BEEFY Client generic on crypto types makes migration to new types smoother. This makes the BEEFY Keystore generic over AuthorityId and extends its tests to cover the case when the AuthorityId is of type (ECDSA, BLS12-377) --------- Co-authored-by: Davide Galassi <davxy@datawok.net> Co-authored-by: Robert Hambrock <roberthambrock@gmail.com>
2026-06-18 09:31:02 +00:00 · 2024-02-08 11:08:51 -05:00
parent bc5a758c0c
commit 0a94124d24
22 changed files with 742 additions and 231 deletions
@@ -16,41 +16,45 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.

-use sp_application_crypto::{key_types::BEEFY as BEEFY_KEY_TYPE, RuntimeAppPublic};
+use sp_application_crypto::{key_types::BEEFY as BEEFY_KEY_TYPE, AppCrypto, RuntimeAppPublic};
+use sp_consensus_beefy::{AuthorityIdBound, BeefyAuthorityId, BeefySignatureHasher};
+use sp_core::ecdsa;
+#[cfg(feature = "bls-experimental")]
+use sp_core::ecdsa_bls377;
 use sp_crypto_hashing::keccak_256;
 use sp_keystore::KeystorePtr;

+use codec::Decode;
 use log::warn;
-
-use sp_consensus_beefy::{
-	ecdsa_crypto::{Public, Signature},
-	BeefyAuthorityId,
-};
+use std::marker::PhantomData;

 use crate::{error, LOG_TARGET};

-/// Hasher used for BEEFY signatures.
-pub(crate) type BeefySignatureHasher = sp_runtime::traits::Keccak256;
-
 /// A BEEFY specific keystore implemented as a `Newtype`. This is basically a
 /// wrapper around [`sp_keystore::Keystore`] and allows to customize
 /// common cryptographic functionality.
-pub(crate) struct BeefyKeystore(Option<KeystorePtr>);
+pub(crate) struct BeefyKeystore<AuthorityId: AuthorityIdBound>(
+	Option<KeystorePtr>,
+	PhantomData<fn() -> AuthorityId>,
+);

-impl BeefyKeystore {
+impl<AuthorityId: AuthorityIdBound> BeefyKeystore<AuthorityId> {
 	/// Check if the keystore contains a private key for one of the public keys
 	/// contained in `keys`. A public key with a matching private key is known
 	/// as a local authority id.
 	///
 	/// Return the public key for which we also do have a private key. If no
 	/// matching private key is found, `None` will be returned.
-	pub fn authority_id(&self, keys: &[Public]) -> Option<Public> {
+	pub fn authority_id(&self, keys: &[AuthorityId]) -> Option<AuthorityId> {
 		let store = self.0.clone()?;

 		// we do check for multiple private keys as a key store sanity check.
-		let public: Vec<Public> = keys
+		let public: Vec<AuthorityId> = keys
 			.iter()
-			.filter(|k| store.has_keys(&[(k.to_raw_vec(), BEEFY_KEY_TYPE)]))
+			.filter(|k| {
+				store
+					.has_keys(&[(<AuthorityId as RuntimeAppPublic>::to_raw_vec(k), BEEFY_KEY_TYPE)])
+			})
 			.cloned()
 			.collect();

@@ -71,55 +75,125 @@ impl BeefyKeystore {
 	/// Note that `message` usually will be pre-hashed before being signed.
 	///
 	/// Return the message signature or an error in case of failure.
-	pub fn sign(&self, public: &Public, message: &[u8]) -> Result<Signature, error::Error> {
+	pub fn sign(
+		&self,
+		public: &AuthorityId,
+		message: &[u8],
+	) -> Result<<AuthorityId as RuntimeAppPublic>::Signature, error::Error> {
 		let store = self.0.clone().ok_or_else(|| error::Error::Keystore("no Keystore".into()))?;

-		let msg = keccak_256(message);
-		let public = public.as_ref();
+		// ECDSA should use ecdsa_sign_prehashed since it needs to be hashed by keccak_256 instead
+		// of blake2. As such we need to deal with producing the signatures case-by-case
+		let signature_byte_array: Vec<u8> = match <AuthorityId as AppCrypto>::CRYPTO_ID {
+			ecdsa::CRYPTO_ID => {
+				let msg_hash = keccak_256(message);
+				let public: ecdsa::Public = ecdsa::Public::try_from(public.as_slice()).unwrap();

-		let sig = store
-			.ecdsa_sign_prehashed(BEEFY_KEY_TYPE, public, &msg)
-			.map_err(|e| error::Error::Keystore(e.to_string()))?
-			.ok_or_else(|| error::Error::Signature("ecdsa_sign_prehashed() failed".to_string()))?;
+				let sig = store
+					.ecdsa_sign_prehashed(BEEFY_KEY_TYPE, &public, &msg_hash)
+					.map_err(|e| error::Error::Keystore(e.to_string()))?
+					.ok_or_else(|| {
+						error::Error::Signature("ecdsa_sign_prehashed() failed".to_string())
+					})?;
+				let sig_ref: &[u8] = sig.as_ref();
+				sig_ref.to_vec()
+			},

-		// check that `sig` has the expected result type
-		let sig = sig.clone().try_into().map_err(|_| {
-			error::Error::Signature(format!("invalid signature {:?} for key {:?}", sig, public))
+			#[cfg(feature = "bls-experimental")]
+			ecdsa_bls377::CRYPTO_ID => {
+				let public: ecdsa_bls377::Public =
+					ecdsa_bls377::Public::try_from(public.as_slice()).unwrap();
+				let sig = store
+					.ecdsa_bls377_sign_with_keccak256(BEEFY_KEY_TYPE, &public, &message)
+					.map_err(|e| error::Error::Keystore(e.to_string()))?
+					.ok_or_else(|| error::Error::Signature("bls377_sign()  failed".to_string()))?;
+				let sig_ref: &[u8] = sig.as_ref();
+				sig_ref.to_vec()
+			},
+
+			_ => Err(error::Error::Keystore("key type is not supported by BEEFY Keystore".into()))?,
+		};
+
+		//check that `sig` has the expected result type
+		let signature = <AuthorityId as RuntimeAppPublic>::Signature::decode(
+			&mut signature_byte_array.as_slice(),
+		)
+		.map_err(|_| {
+			error::Error::Signature(format!(
+				"invalid signature {:?} for key {:?}",
+				signature_byte_array, public
+			))
 		})?;

-		Ok(sig)
+		Ok(signature)
 	}

 	/// Returns a vector of [`sp_consensus_beefy::crypto::Public`] keys which are currently
 	/// supported (i.e. found in the keystore).
-	pub fn public_keys(&self) -> Result<Vec<Public>, error::Error> {
+	pub fn public_keys(&self) -> Result<Vec<AuthorityId>, error::Error> {
 		let store = self.0.clone().ok_or_else(|| error::Error::Keystore("no Keystore".into()))?;

-		let pk: Vec<Public> =
-			store.ecdsa_public_keys(BEEFY_KEY_TYPE).drain(..).map(Public::from).collect();
+		let pk = match <AuthorityId as AppCrypto>::CRYPTO_ID {
+			ecdsa::CRYPTO_ID => store
+				.ecdsa_public_keys(BEEFY_KEY_TYPE)
+				.drain(..)
+				.map(|pk| AuthorityId::try_from(pk.as_ref()))
+				.collect::<Result<Vec<_>, _>>()
+				.or_else(|_| {
+					Err(error::Error::Keystore(
+						"unable to convert public key into authority id".into(),
+					))
+				}),

-		Ok(pk)
+			#[cfg(feature = "bls-experimental")]
+			ecdsa_bls377::CRYPTO_ID => store
+				.ecdsa_bls377_public_keys(BEEFY_KEY_TYPE)
+				.drain(..)
+				.map(|pk| AuthorityId::try_from(pk.as_ref()))
+				.collect::<Result<Vec<_>, _>>()
+				.or_else(|_| {
+					Err(error::Error::Keystore(
+						"unable to convert public key into authority id".into(),
+					))
+				}),
+
+			_ => Err(error::Error::Keystore("key type is not supported by BEEFY Keystore".into())),
+		};
+
+		pk
 	}

 	/// Use the `public` key to verify that `sig` is a valid signature for `message`.
 	///
 	/// Return `true` if the signature is authentic, `false` otherwise.
-	pub fn verify(public: &Public, sig: &Signature, message: &[u8]) -> bool {
+	pub fn verify(
+		public: &AuthorityId,
+		sig: &<AuthorityId as RuntimeAppPublic>::Signature,
+		message: &[u8],
+	) -> bool {
 		BeefyAuthorityId::<BeefySignatureHasher>::verify(public, sig, message)
 	}
 }

-impl From<Option<KeystorePtr>> for BeefyKeystore {
-	fn from(store: Option<KeystorePtr>) -> BeefyKeystore {
-		BeefyKeystore(store)
+impl<AuthorityId: AuthorityIdBound> From<Option<KeystorePtr>> for BeefyKeystore<AuthorityId>
+where
+	<AuthorityId as RuntimeAppPublic>::Signature: Send + Sync,
+{
+	fn from(store: Option<KeystorePtr>) -> BeefyKeystore<AuthorityId> {
+		BeefyKeystore(store, PhantomData)
 	}
 }

 #[cfg(test)]
 pub mod tests {
-	use sp_consensus_beefy::{ecdsa_crypto, Keyring};
-	use sp_core::{ecdsa, Pair};
-	use sp_keystore::testing::MemoryKeystore;
+	#[cfg(feature = "bls-experimental")]
+	use sp_consensus_beefy::ecdsa_bls_crypto;
+	use sp_consensus_beefy::{
+		ecdsa_crypto,
+		test_utils::{BeefySignerAuthority, Keyring},
+	};
+	use sp_core::Pair as PairT;
+	use sp_keystore::{testing::MemoryKeystore, Keystore};

 	use super::*;
 	use crate::error::Error;
@@ -128,152 +202,265 @@ pub mod tests {
 		MemoryKeystore::new().into()
 	}

-	#[test]
-	fn verify_should_work() {
-		let msg = keccak_256(b"I am Alice!");
-		let sig = Keyring::Alice.sign(b"I am Alice!");
+	fn pair_verify_should_work<
+		AuthorityId: AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+	>()
+	where
+		<AuthorityId as sp_runtime::RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<sp_runtime::traits::Keccak256>,
+	{
+		let msg = b"I am Alice!";
+		let sig = Keyring::<AuthorityId>::Alice.sign(b"I am Alice!");

-		assert!(ecdsa::Pair::verify_prehashed(
-			&sig.clone().into(),
-			&msg,
-			&Keyring::Alice.public().into(),
+		assert!(<AuthorityId as BeefyAuthorityId<BeefySignatureHasher>>::verify(
+			&Keyring::Alice.public(),
+			&sig,
+			&msg.as_slice(),
 		));

 		// different public key -> fail
-		assert!(!ecdsa::Pair::verify_prehashed(
-			&sig.clone().into(),
-			&msg,
-			&Keyring::Bob.public().into(),
+		assert!(!<AuthorityId as BeefyAuthorityId<BeefySignatureHasher>>::verify(
+			&Keyring::Bob.public(),
+			&sig,
+			&msg.as_slice(),
 		));

-		let msg = keccak_256(b"I am not Alice!");
+		let msg = b"I am not Alice!";

 		// different msg -> fail
-		assert!(
-			!ecdsa::Pair::verify_prehashed(&sig.into(), &msg, &Keyring::Alice.public().into(),)
-		);
+		assert!(!<AuthorityId as BeefyAuthorityId<BeefySignatureHasher>>::verify(
+			&Keyring::Alice.public(),
+			&sig,
+			&msg.as_slice(),
+		));
+	}
+
+	/// Generate key pair in the given store using the provided seed
+	fn generate_in_store<AuthorityId>(
+		store: KeystorePtr,
+		key_type: sp_application_crypto::KeyTypeId,
+		owner: Option<Keyring<AuthorityId>>,
+	) -> AuthorityId
+	where
+		AuthorityId:
+			AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<BeefySignatureHasher>,
+		<AuthorityId as RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+	{
+		let optional_seed: Option<String> = owner.map(|owner| owner.to_seed());
+
+		match <AuthorityId as AppCrypto>::CRYPTO_ID {
+			ecdsa::CRYPTO_ID => {
+				let pk = store.ecdsa_generate_new(key_type, optional_seed.as_deref()).ok().unwrap();
+				AuthorityId::decode(&mut pk.as_ref()).unwrap()
+			},
+			#[cfg(feature = "bls-experimental")]
+			ecdsa_bls377::CRYPTO_ID => {
+				let pk = store
+					.ecdsa_bls377_generate_new(key_type, optional_seed.as_deref())
+					.ok()
+					.unwrap();
+				AuthorityId::decode(&mut pk.as_ref()).unwrap()
+			},
+			_ => panic!("Requested CRYPTO_ID is not supported by the BEEFY Keyring"),
+		}
 	}

 	#[test]
-	fn pair_works() {
-		let want = ecdsa_crypto::Pair::from_string("//Alice", None)
+	fn pair_verify_should_work_ecdsa() {
+		pair_verify_should_work::<ecdsa_crypto::AuthorityId>();
+	}
+
+	#[cfg(feature = "bls-experimental")]
+	#[test]
+	fn pair_verify_should_work_ecdsa_n_bls() {
+		pair_verify_should_work::<ecdsa_bls_crypto::AuthorityId>();
+	}
+
+	fn pair_works<
+		AuthorityId: AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+	>()
+	where
+		<AuthorityId as sp_runtime::RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<sp_runtime::traits::Keccak256>,
+	{
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//Alice", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::Alice.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::Alice.pair().to_raw_vec();
 		assert_eq!(want, got);

-		let want = ecdsa_crypto::Pair::from_string("//Bob", None)
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//Bob", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::Bob.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::Bob.pair().to_raw_vec();
 		assert_eq!(want, got);

-		let want = ecdsa_crypto::Pair::from_string("//Charlie", None)
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//Charlie", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::Charlie.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::Charlie.pair().to_raw_vec();
 		assert_eq!(want, got);

-		let want = ecdsa_crypto::Pair::from_string("//Dave", None)
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//Dave", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::Dave.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::Dave.pair().to_raw_vec();
 		assert_eq!(want, got);

-		let want = ecdsa_crypto::Pair::from_string("//Eve", None)
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//Eve", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::Eve.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::Eve.pair().to_raw_vec();
 		assert_eq!(want, got);

-		let want = ecdsa_crypto::Pair::from_string("//Ferdie", None)
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//Ferdie", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::Ferdie.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::Ferdie.pair().to_raw_vec();
 		assert_eq!(want, got);

-		let want = ecdsa_crypto::Pair::from_string("//One", None)
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//One", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::One.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::One.pair().to_raw_vec();
 		assert_eq!(want, got);

-		let want = ecdsa_crypto::Pair::from_string("//Two", None)
+		let want = <AuthorityId as AppCrypto>::Pair::from_string("//Two", None)
 			.expect("Pair failed")
 			.to_raw_vec();
-		let got = Keyring::Two.pair().to_raw_vec();
+		let got = Keyring::<AuthorityId>::Two.pair().to_raw_vec();
 		assert_eq!(want, got);
 	}

 	#[test]
-	fn authority_id_works() {
+	fn ecdsa_pair_works() {
+		pair_works::<ecdsa_crypto::AuthorityId>();
+	}
+
+	#[cfg(feature = "bls-experimental")]
+	#[test]
+	fn ecdsa_n_bls_pair_works() {
+		pair_works::<ecdsa_bls_crypto::AuthorityId>();
+	}
+
+	fn authority_id_works<
+		AuthorityId: AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+	>()
+	where
+		<AuthorityId as sp_runtime::RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<sp_runtime::traits::Keccak256>,
+	{
 		let store = keystore();

-		let alice: ecdsa_crypto::Public = store
-			.ecdsa_generate_new(BEEFY_KEY_TYPE, Some(&Keyring::Alice.to_seed()))
-			.ok()
-			.unwrap()
-			.into();
+		generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, Some(Keyring::Alice));
+
+		let alice = Keyring::<AuthorityId>::Alice.public();

 		let bob = Keyring::Bob.public();
 		let charlie = Keyring::Charlie.public();

-		let store: BeefyKeystore = Some(store).into();
+		let beefy_store: BeefyKeystore<AuthorityId> = Some(store).into();

 		let mut keys = vec![bob, charlie];

-		let id = store.authority_id(keys.as_slice());
+		let id = beefy_store.authority_id(keys.as_slice());
 		assert!(id.is_none());

 		keys.push(alice.clone());

-		let id = store.authority_id(keys.as_slice()).unwrap();
+		let id = beefy_store.authority_id(keys.as_slice()).unwrap();
 		assert_eq!(id, alice);
 	}

 	#[test]
-	fn sign_works() {
+	fn authority_id_works_for_ecdsa() {
+		authority_id_works::<ecdsa_crypto::AuthorityId>();
+	}
+
+	#[cfg(feature = "bls-experimental")]
+	#[test]
+	fn authority_id_works_for_ecdsa_n_bls() {
+		authority_id_works::<ecdsa_bls_crypto::AuthorityId>();
+	}
+
+	fn sign_works<
+		AuthorityId: AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+	>()
+	where
+		<AuthorityId as sp_runtime::RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<sp_runtime::traits::Keccak256>,
+	{
 		let store = keystore();

-		let alice: ecdsa_crypto::Public = store
-			.ecdsa_generate_new(BEEFY_KEY_TYPE, Some(&Keyring::Alice.to_seed()))
-			.ok()
-			.unwrap()
-			.into();
+		generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, Some(Keyring::Alice));

-		let store: BeefyKeystore = Some(store).into();
+		let alice = Keyring::Alice.public();
+
+		let store: BeefyKeystore<AuthorityId> = Some(store).into();

 		let msg = b"are you involved or commited?";

 		let sig1 = store.sign(&alice, msg).unwrap();
-		let sig2 = Keyring::Alice.sign(msg);
+		let sig2 = Keyring::<AuthorityId>::Alice.sign(msg);

 		assert_eq!(sig1, sig2);
 	}

 	#[test]
-	fn sign_error() {
+	fn sign_works_for_ecdsa() {
+		sign_works::<ecdsa_crypto::AuthorityId>();
+	}
+
+	#[cfg(feature = "bls-experimental")]
+	#[test]
+	fn sign_works_for_ecdsa_n_bls() {
+		sign_works::<ecdsa_bls_crypto::AuthorityId>();
+	}
+
+	fn sign_error<
+		AuthorityId: AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+	>(
+		expected_error_message: &str,
+	) where
+		<AuthorityId as sp_runtime::RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<sp_runtime::traits::Keccak256>,
+	{
 		let store = keystore();

-		store
-			.ecdsa_generate_new(BEEFY_KEY_TYPE, Some(&Keyring::Bob.to_seed()))
-			.ok()
-			.unwrap();
+		generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, Some(Keyring::Bob));

-		let store: BeefyKeystore = Some(store).into();
+		let store: BeefyKeystore<AuthorityId> = Some(store).into();

 		let alice = Keyring::Alice.public();

 		let msg = b"are you involved or commited?";
 		let sig = store.sign(&alice, msg).err().unwrap();
-		let err = Error::Signature("ecdsa_sign_prehashed() failed".to_string());
+		let err = Error::Signature(expected_error_message.to_string());

 		assert_eq!(sig, err);
 	}

+	#[test]
+	fn sign_error_for_ecdsa() {
+		sign_error::<ecdsa_crypto::AuthorityId>("ecdsa_sign_prehashed() failed");
+	}
+
+	#[cfg(feature = "bls-experimental")]
+	#[test]
+	fn sign_error_for_ecdsa_n_bls() {
+		sign_error::<ecdsa_bls_crypto::AuthorityId>("bls377_sign()  failed");
+	}
+
 	#[test]
 	fn sign_no_keystore() {
-		let store: BeefyKeystore = None.into();
+		let store: BeefyKeystore<ecdsa_crypto::Public> = None.into();

 		let alice = Keyring::Alice.public();
 		let msg = b"are you involved or commited";
@@ -283,17 +470,21 @@ pub mod tests {
 		assert_eq!(sig, err);
 	}

-	#[test]
-	fn verify_works() {
+	fn verify_works<
+		AuthorityId: AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+	>()
+	where
+		<AuthorityId as sp_runtime::RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<sp_runtime::traits::Keccak256>,
+	{
 		let store = keystore();

-		let alice: ecdsa_crypto::Public = store
-			.ecdsa_generate_new(BEEFY_KEY_TYPE, Some(&Keyring::Alice.to_seed()))
-			.ok()
-			.unwrap()
-			.into();
+		generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, Some(Keyring::Alice));

-		let store: BeefyKeystore = Some(store).into();
+		let store: BeefyKeystore<AuthorityId> = Some(store).into();
+
+		let alice = Keyring::Alice.public();

 		// `msg` and `sig` match
 		let msg = b"are you involved or commited?";
@@ -305,32 +496,48 @@ pub mod tests {
 		assert!(!BeefyKeystore::verify(&alice, &sig, msg));
 	}

-	// Note that we use keys with and without a seed for this test.
 	#[test]
-	fn public_keys_works() {
+	fn verify_works_for_ecdsa() {
+		verify_works::<ecdsa_crypto::AuthorityId>();
+	}
+
+	#[cfg(feature = "bls-experimental")]
+	#[test]
+
+	fn verify_works_for_ecdsa_n_bls() {
+		verify_works::<ecdsa_bls_crypto::AuthorityId>();
+	}
+
+	// Note that we use keys with and without a seed for this test.
+	fn public_keys_works<
+		AuthorityId: AuthorityIdBound + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Public>,
+	>()
+	where
+		<AuthorityId as sp_runtime::RuntimeAppPublic>::Signature:
+			Send + Sync + From<<<AuthorityId as AppCrypto>::Pair as AppCrypto>::Signature>,
+		<AuthorityId as AppCrypto>::Pair: BeefySignerAuthority<sp_runtime::traits::Keccak256>,
+	{
 		const TEST_TYPE: sp_application_crypto::KeyTypeId =
 			sp_application_crypto::KeyTypeId(*b"test");

 		let store = keystore();

-		let add_key =
-			|key_type, seed: Option<&str>| store.ecdsa_generate_new(key_type, seed).unwrap();
-
 		// test keys
-		let _ = add_key(TEST_TYPE, Some(Keyring::Alice.to_seed().as_str()));
-		let _ = add_key(TEST_TYPE, Some(Keyring::Bob.to_seed().as_str()));
-
-		let _ = add_key(TEST_TYPE, None);
-		let _ = add_key(TEST_TYPE, None);
+		let _ = generate_in_store::<AuthorityId>(store.clone(), TEST_TYPE, Some(Keyring::Alice));
+		let _ = generate_in_store::<AuthorityId>(store.clone(), TEST_TYPE, Some(Keyring::Bob));

 		// BEEFY keys
-		let _ = add_key(BEEFY_KEY_TYPE, Some(Keyring::Dave.to_seed().as_str()));
-		let _ = add_key(BEEFY_KEY_TYPE, Some(Keyring::Eve.to_seed().as_str()));
+		let _ =
+			generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, Some(Keyring::Dave));
+		let _ = generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, Some(Keyring::Eve));

-		let key1: ecdsa_crypto::Public = add_key(BEEFY_KEY_TYPE, None).into();
-		let key2: ecdsa_crypto::Public = add_key(BEEFY_KEY_TYPE, None).into();
+		let _ = generate_in_store::<AuthorityId>(store.clone(), TEST_TYPE, None);
+		let _ = generate_in_store::<AuthorityId>(store.clone(), TEST_TYPE, None);

-		let store: BeefyKeystore = Some(store).into();
+		let key1 = generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, None);
+		let key2 = generate_in_store::<AuthorityId>(store.clone(), BEEFY_KEY_TYPE, None);
+
+		let store: BeefyKeystore<AuthorityId> = Some(store).into();

 		let keys = store.public_keys().ok().unwrap();

@@ -340,4 +547,16 @@ pub mod tests {
 		assert!(keys.contains(&key1));
 		assert!(keys.contains(&key2));
 	}
+
+	#[test]
+	fn public_keys_works_for_ecdsa_keystore() {
+		public_keys_works::<ecdsa_crypto::AuthorityId>();
+	}
+
+	#[cfg(feature = "bls-experimental")]
+	#[test]
+
+	fn public_keys_works_for_ecdsa_n_bls() {
+		public_keys_works::<ecdsa_bls_crypto::AuthorityId>();
+	}
 }