Make keystore return None when a key doesn't exist (#8163)

* Make keystore return `None` when a key doesn't exist

* Fixes

* More fixes

* Update comment

* Update primitives/keystore/src/lib.rs

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

* Update client/keystore/src/local.rs

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

* Address comments

* Update client/keystore/src/local.rs

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

substrate/primitives/finality-grandpa/src/lib.rs

@@ -400,7 +400,7 @@ where
 		AuthorityId::ID,
 		&public.to_public_crypto_pair(),
 		&encoded[..],
-	).ok()?.try_into().ok()?;
+	).ok().flatten()?.try_into().ok()?;
 
 	Some(grandpa::SignedMessage {
 		message,

substrate/primitives/io/src/lib.rs

@@ -474,8 +474,9 @@ pub trait Crypto {
 		let keystore = &***self.extension::<KeystoreExt>()
 			.expect("No `keystore` associated for the current context!");
 		SyncCryptoStore::sign_with(keystore, id, &pub_key.into(), msg)
-			.map(|sig| ed25519::Signature::from_slice(sig.as_slice()))
 			.ok()
+			.flatten()
+			.map(|sig| ed25519::Signature::from_slice(sig.as_slice()))
 	}
 
 	/// Verify `ed25519` signature.
@@ -600,8 +601,9 @@ pub trait Crypto {
 		let keystore = &***self.extension::<KeystoreExt>()
 			.expect("No `keystore` associated for the current context!");
 		SyncCryptoStore::sign_with(keystore, id, &pub_key.into(), msg)
-			.map(|sig| sr25519::Signature::from_slice(sig.as_slice()))
 			.ok()
+			.flatten()
+			.map(|sig| sr25519::Signature::from_slice(sig.as_slice()))
 	}
 
 	/// Verify an `sr25519` signature.
@@ -646,8 +648,9 @@ pub trait Crypto {
 		let keystore = &***self.extension::<KeystoreExt>()
 			.expect("No `keystore` associated for the current context!");
 		SyncCryptoStore::sign_with(keystore, id, &pub_key.into(), msg)
-			.map(|sig| ecdsa::Signature::from_slice(sig.as_slice()))
 			.ok()
+			.flatten()
+			.map(|sig| ecdsa::Signature::from_slice(sig.as_slice()))
 	}
 
 	/// Verify `ecdsa` signature.

substrate/primitives/keystore/src/lib.rs

@@ -34,9 +34,6 @@ pub enum Error {
 	/// Public key type is not supported
 	#[display(fmt="Key not supported: {:?}", _0)]
 	KeyNotSupported(KeyTypeId),
-	/// Pair not found for public key and KeyTypeId
-	#[display(fmt="Pair was not found: {}", _0)]
-	PairNotFound(String),
 	/// Validation error
 	#[display(fmt="Validation error: {}", _0)]
 	ValidationError(String),
@@ -125,37 +122,39 @@ pub trait CryptoStore: Send + Sync {
 	/// Signs a message with the private key that matches
 	/// the public key passed.
 	///
-	/// Returns the SCALE encoded signature if key is found & supported,
-	/// an error otherwise.
+	/// Returns the SCALE encoded signature if key is found and supported, `None` if the key doesn't
+	/// exist or an error when something failed.
 	async fn sign_with(
 		&self,
 		id: KeyTypeId,
 		key: &CryptoTypePublicPair,
 		msg: &[u8],
-	) -> Result<Vec<u8>, Error>;
+	) -> Result<Option<Vec<u8>>, Error>;
 
 	/// Sign with any key
 	///
 	/// Given a list of public keys, find the first supported key and
 	/// sign the provided message with that key.
 	///
-	/// Returns a tuple of the used key and the SCALE encoded signature.
+	/// Returns a tuple of the used key and the SCALE encoded signature or `None` if no key could
+	/// be found to sign.
 	async fn sign_with_any(
 		&self,
 		id: KeyTypeId,
 		keys: Vec<CryptoTypePublicPair>,
 		msg: &[u8]
-	) -> Result<(CryptoTypePublicPair, Vec<u8>), Error> {
+	) -> Result<Option<(CryptoTypePublicPair, Vec<u8>)>, Error> {
 		if keys.len() == 1 {
-			return self.sign_with(id, &keys[0], msg).await.map(|s| (keys[0].clone(), s));
+			return Ok(self.sign_with(id, &keys[0], msg).await?.map(|s| (keys[0].clone(), s)));
 		} else {
 			for k in self.supported_keys(id, keys).await? {
-				if let Ok(sign) = self.sign_with(id, &k, msg).await {
-					return Ok((k, sign));
+				if let Ok(Some(sign)) = self.sign_with(id, &k, msg).await {
+					return Ok(Some((k, sign)));
 				}
 			}
 		}
-		Err(Error::KeyNotSupported(id))
+
+		Ok(None)
 	}
 
 	/// Sign with all keys
@@ -164,13 +163,13 @@ pub trait CryptoStore: Send + Sync {
 	/// each key given that the key is supported.
 	///
 	/// Returns a list of `Result`s each representing the SCALE encoded
-	/// signature of each key or a Error for non-supported keys.
+	/// signature of each key, `None` if the key doesn't exist or a error when something failed.
 	async fn sign_with_all(
 		&self,
 		id: KeyTypeId,
 		keys: Vec<CryptoTypePublicPair>,
 		msg: &[u8],
-	) -> Result<Vec<Result<Vec<u8>, Error>>, ()> {
+	) -> Result<Vec<Result<Option<Vec<u8>>, Error>>, ()> {
 		let futs = keys.iter()
 			.map(|k| self.sign_with(id, k, msg));
 
@@ -187,16 +186,14 @@ pub trait CryptoStore: Send + Sync {
 	/// Namely, VRFOutput and VRFProof which are returned
 	/// inside the `VRFSignature` container struct.
 	///
-	/// This function will return an error in the cases where
-	/// the public key and key type provided do not match a private
-	/// key in the keystore. Or, in the context of remote signing
-	/// an error could be a network one.
+	/// This function will return `None` if the given `key_type` and `public` combination
+	/// doesn't exist in the keystore or an `Err` when something failed.
 	async fn sr25519_vrf_sign(
 		&self,
 		key_type: KeyTypeId,
 		public: &sr25519::Public,
 		transcript_data: VRFTranscriptData,
-	) -> Result<VRFSignature, Error>;
+	) -> Result<Option<VRFSignature>, Error>;
 }
 
 /// Sync version of the CryptoStore
@@ -285,37 +282,41 @@ pub trait SyncCryptoStore: CryptoStore + Send + Sync {
 	/// Signs a message with the private key that matches
 	/// the public key passed.
 	///
-	/// Returns the SCALE encoded signature if key is found & supported,
-	/// an error otherwise.
+	/// Returns the SCALE encoded signature if key is found and supported, `None` if the key doesn't
+	/// exist or an error when something failed.
 	fn sign_with(
 		&self,
 		id: KeyTypeId,
 		key: &CryptoTypePublicPair,
 		msg: &[u8],
-	) -> Result<Vec<u8>, Error>;
+	) -> Result<Option<Vec<u8>>, Error>;
 
 	/// Sign with any key
 	///
 	/// Given a list of public keys, find the first supported key and
 	/// sign the provided message with that key.
 	///
-	/// Returns a tuple of the used key and the SCALE encoded signature.
+	/// Returns a tuple of the used key and the SCALE encoded signature or `None` if no key could
+	/// be found to sign.
 	fn sign_with_any(
 		&self,
 		id: KeyTypeId,
 		keys: Vec<CryptoTypePublicPair>,
 		msg: &[u8]
-	) -> Result<(CryptoTypePublicPair, Vec<u8>), Error> {
+	) -> Result<Option<(CryptoTypePublicPair, Vec<u8>)>, Error> {
 		if keys.len() == 1 {
-			return SyncCryptoStore::sign_with(self, id, &keys[0], msg).map(|s| (keys[0].clone(), s));
+			return Ok(
+				SyncCryptoStore::sign_with(self, id, &keys[0], msg)?.map(|s| (keys[0].clone(), s)),
+			)
 		} else {
 			for k in SyncCryptoStore::supported_keys(self, id, keys)? {
-				if let Ok(sign) = SyncCryptoStore::sign_with(self, id, &k, msg) {
-					return Ok((k, sign));
+				if let Ok(Some(sign)) = SyncCryptoStore::sign_with(self, id, &k, msg) {
+					return Ok(Some((k, sign)));
 				}
 			}
 		}
-		Err(Error::KeyNotSupported(id))
+
+		Ok(None)
 	}
 
 	/// Sign with all keys
@@ -324,13 +325,13 @@ pub trait SyncCryptoStore: CryptoStore + Send + Sync {
 	/// each key given that the key is supported.
 	///
 	/// Returns a list of `Result`s each representing the SCALE encoded
-	/// signature of each key or a Error for non-supported keys.
+	/// signature of each key, `None` if the key doesn't exist or an error when something failed.
 	fn sign_with_all(
 		&self,
 		id: KeyTypeId,
 		keys: Vec<CryptoTypePublicPair>,
 		msg: &[u8],
-	) -> Result<Vec<Result<Vec<u8>, Error>>, ()>{
+	) -> Result<Vec<Result<Option<Vec<u8>>, Error>>, ()> {
 		Ok(keys.iter().map(|k| SyncCryptoStore::sign_with(self, id, k, msg)).collect())
 	}
 
@@ -344,16 +345,14 @@ pub trait SyncCryptoStore: CryptoStore + Send + Sync {
 	/// Namely, VRFOutput and VRFProof which are returned
 	/// inside the `VRFSignature` container struct.
 	///
-	/// This function will return an error in the cases where
-	/// the public key and key type provided do not match a private
-	/// key in the keystore. Or, in the context of remote signing
-	/// an error could be a network one.
+	/// This function will return `None` if the given `key_type` and `public` combination
+	/// doesn't exist in the keystore or an `Err` when something failed.
 	fn sr25519_vrf_sign(
 		&self,
 		key_type: KeyTypeId,
 		public: &sr25519::Public,
 		transcript_data: VRFTranscriptData,
-	) -> Result<VRFSignature, Error>;
+	) -> Result<Option<VRFSignature>, Error>;
 }
 
 /// A pointer to a keystore.

substrate/primitives/keystore/src/testing.rs

@@ -132,7 +132,7 @@ impl CryptoStore for KeyStore {
 		id: KeyTypeId,
 		key: &CryptoTypePublicPair,
 		msg: &[u8],
-	) -> Result<Vec<u8>, Error> {
+	) -> Result<Option<Vec<u8>>, Error> {
 		SyncCryptoStore::sign_with(self, id, key, msg)
 	}
 
@@ -141,7 +141,7 @@ impl CryptoStore for KeyStore {
 		key_type: KeyTypeId,
 		public: &sr25519::Public,
 		transcript_data: VRFTranscriptData,
-	) -> Result<VRFSignature, Error> {
+	) -> Result<Option<VRFSignature>, Error> {
 		SyncCryptoStore::sr25519_vrf_sign(self, key_type, public, transcript_data)
 	}
 }
@@ -280,27 +280,27 @@ impl SyncCryptoStore for KeyStore {
 		id: KeyTypeId,
 		key: &CryptoTypePublicPair,
 		msg: &[u8],
-	) -> Result<Vec<u8>, Error> {
+	) -> Result<Option<Vec<u8>>, Error> {
 		use codec::Encode;
 
 		match key.0 {
 			ed25519::CRYPTO_ID => {
-				let key_pair: ed25519::Pair = self
-					.ed25519_key_pair(id, &ed25519::Public::from_slice(key.1.as_slice()))
-					.ok_or_else(|| Error::PairNotFound("ed25519".to_owned()))?;
-				return Ok(key_pair.sign(msg).encode());
+				let key_pair = self
+					.ed25519_key_pair(id, &ed25519::Public::from_slice(key.1.as_slice()));
+
+				key_pair.map(|k| k.sign(msg).encode()).map(Ok).transpose()
 			}
 			sr25519::CRYPTO_ID => {
-				let key_pair: sr25519::Pair = self
-					.sr25519_key_pair(id, &sr25519::Public::from_slice(key.1.as_slice()))
-					.ok_or_else(|| Error::PairNotFound("sr25519".to_owned()))?;
-				return Ok(key_pair.sign(msg).encode());
+				let key_pair = self
+					.sr25519_key_pair(id, &sr25519::Public::from_slice(key.1.as_slice()));
+
+				key_pair.map(|k| k.sign(msg).encode()).map(Ok).transpose()
 			}
 			ecdsa::CRYPTO_ID => {
-				let key_pair: ecdsa::Pair = self
-					.ecdsa_key_pair(id, &ecdsa::Public::from_slice(key.1.as_slice()))
-					.ok_or_else(|| Error::PairNotFound("ecdsa".to_owned()))?;
-				return Ok(key_pair.sign(msg).encode());
+				let key_pair = self
+					.ecdsa_key_pair(id, &ecdsa::Public::from_slice(key.1.as_slice()));
+
+				key_pair.map(|k| k.sign(msg).encode()).map(Ok).transpose()
 			}
 			_ => Err(Error::KeyNotSupported(id))
 		}
@@ -311,15 +311,19 @@ impl SyncCryptoStore for KeyStore {
 		key_type: KeyTypeId,
 		public: &sr25519::Public,
 		transcript_data: VRFTranscriptData,
-	) -> Result<VRFSignature, Error> {
+	) -> Result<Option<VRFSignature>, Error> {
 		let transcript = make_transcript(transcript_data);
-		let pair = self.sr25519_key_pair(key_type, public)
-			.ok_or_else(|| Error::PairNotFound("Not found".to_owned()))?;
+		let pair = if let Some(k) = self.sr25519_key_pair(key_type, public) {
+			k
+		} else {
+			return Ok(None)
+		};
+
 		let (inout, proof, _) = pair.as_ref().vrf_sign(transcript);
-		Ok(VRFSignature {
+		Ok(Some(VRFSignature {
 			output: inout.to_output(),
 			proof,
-		})
+		}))
 	}
 }
 
@@ -394,7 +398,7 @@ mod tests {
 			&key_pair.public(),
 			transcript_data.clone(),
 		);
-		assert!(result.is_err());
+		assert!(result.unwrap().is_none());
 
 		SyncCryptoStore::insert_unknown(
 			&store,
@@ -410,6 +414,6 @@ mod tests {
 			transcript_data,
 		);
 
-		assert!(result.is_ok());
+		assert!(result.unwrap().is_some());
 	}
 }