sp_keystore/
testing.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! Types that should only be used for testing!
19
20use crate::{Error, Keystore, KeystorePtr};
21
22#[cfg(feature = "bandersnatch-experimental")]
23use sp_core::bandersnatch;
24#[cfg(feature = "bls-experimental")]
25use sp_core::{bls381, ecdsa_bls381, KeccakHasher};
26use sp_core::{
27	crypto::{ByteArray, KeyTypeId, Pair, VrfSecret},
28	ecdsa, ed25519, sr25519,
29};
30
31use parking_lot::RwLock;
32use std::{collections::HashMap, sync::Arc};
33
34/// A keystore implementation usable in tests.
35#[derive(Default, Clone)]
36pub struct MemoryKeystore {
37	/// `KeyTypeId` maps to public keys and public keys map to private keys.
38	keys: Arc<RwLock<HashMap<KeyTypeId, HashMap<Vec<u8>, String>>>>,
39}
40
41impl MemoryKeystore {
42	/// Creates a new instance of `Self`.
43	pub fn new() -> Self {
44		Self::default()
45	}
46
47	fn pair<T: Pair>(&self, key_type: KeyTypeId, public: &T::Public) -> Option<T> {
48		self.keys.read().get(&key_type).and_then(|inner| {
49			inner
50				.get(public.as_slice())
51				.map(|s| T::from_string(s, None).expect("seed slice is valid"))
52		})
53	}
54
55	fn public_keys<T: Pair>(&self, key_type: KeyTypeId) -> Vec<T::Public> {
56		self.keys
57			.read()
58			.get(&key_type)
59			.map(|keys| {
60				keys.values()
61					.map(|s| T::from_string(s, None).expect("seed slice is valid"))
62					.map(|p| p.public())
63					.collect()
64			})
65			.unwrap_or_default()
66	}
67
68	fn generate_new<T: Pair>(
69		&self,
70		key_type: KeyTypeId,
71		seed: Option<&str>,
72	) -> Result<T::Public, Error> {
73		match seed {
74			Some(seed) => {
75				let pair = T::from_string(seed, None)
76					.map_err(|_| Error::ValidationError("Generates a pair.".to_owned()))?;
77				self.keys
78					.write()
79					.entry(key_type)
80					.or_default()
81					.insert(pair.public().to_raw_vec(), seed.into());
82				Ok(pair.public())
83			},
84			None => {
85				let (pair, phrase, _) = T::generate_with_phrase(None);
86				self.keys
87					.write()
88					.entry(key_type)
89					.or_default()
90					.insert(pair.public().to_raw_vec(), phrase);
91				Ok(pair.public())
92			},
93		}
94	}
95
96	fn sign<T: Pair>(
97		&self,
98		key_type: KeyTypeId,
99		public: &T::Public,
100		msg: &[u8],
101	) -> Result<Option<T::Signature>, Error> {
102		let sig = self.pair::<T>(key_type, public).map(|pair| pair.sign(msg));
103		Ok(sig)
104	}
105
106	fn vrf_sign<T: Pair + VrfSecret>(
107		&self,
108		key_type: KeyTypeId,
109		public: &T::Public,
110		data: &T::VrfSignData,
111	) -> Result<Option<T::VrfSignature>, Error> {
112		let sig = self.pair::<T>(key_type, public).map(|pair| pair.vrf_sign(data));
113		Ok(sig)
114	}
115
116	fn vrf_pre_output<T: Pair + VrfSecret>(
117		&self,
118		key_type: KeyTypeId,
119		public: &T::Public,
120		input: &T::VrfInput,
121	) -> Result<Option<T::VrfPreOutput>, Error> {
122		let pre_output = self.pair::<T>(key_type, public).map(|pair| pair.vrf_pre_output(input));
123		Ok(pre_output)
124	}
125}
126
127impl Keystore for MemoryKeystore {
128	fn sr25519_public_keys(&self, key_type: KeyTypeId) -> Vec<sr25519::Public> {
129		self.public_keys::<sr25519::Pair>(key_type)
130	}
131
132	fn sr25519_generate_new(
133		&self,
134		key_type: KeyTypeId,
135		seed: Option<&str>,
136	) -> Result<sr25519::Public, Error> {
137		self.generate_new::<sr25519::Pair>(key_type, seed)
138	}
139
140	fn sr25519_sign(
141		&self,
142		key_type: KeyTypeId,
143		public: &sr25519::Public,
144		msg: &[u8],
145	) -> Result<Option<sr25519::Signature>, Error> {
146		self.sign::<sr25519::Pair>(key_type, public, msg)
147	}
148
149	fn sr25519_vrf_sign(
150		&self,
151		key_type: KeyTypeId,
152		public: &sr25519::Public,
153		data: &sr25519::vrf::VrfSignData,
154	) -> Result<Option<sr25519::vrf::VrfSignature>, Error> {
155		self.vrf_sign::<sr25519::Pair>(key_type, public, data)
156	}
157
158	fn sr25519_vrf_pre_output(
159		&self,
160		key_type: KeyTypeId,
161		public: &sr25519::Public,
162		input: &sr25519::vrf::VrfInput,
163	) -> Result<Option<sr25519::vrf::VrfPreOutput>, Error> {
164		self.vrf_pre_output::<sr25519::Pair>(key_type, public, input)
165	}
166
167	fn ed25519_public_keys(&self, key_type: KeyTypeId) -> Vec<ed25519::Public> {
168		self.public_keys::<ed25519::Pair>(key_type)
169	}
170
171	fn ed25519_generate_new(
172		&self,
173		key_type: KeyTypeId,
174		seed: Option<&str>,
175	) -> Result<ed25519::Public, Error> {
176		self.generate_new::<ed25519::Pair>(key_type, seed)
177	}
178
179	fn ed25519_sign(
180		&self,
181		key_type: KeyTypeId,
182		public: &ed25519::Public,
183		msg: &[u8],
184	) -> Result<Option<ed25519::Signature>, Error> {
185		self.sign::<ed25519::Pair>(key_type, public, msg)
186	}
187
188	fn ecdsa_public_keys(&self, key_type: KeyTypeId) -> Vec<ecdsa::Public> {
189		self.public_keys::<ecdsa::Pair>(key_type)
190	}
191
192	fn ecdsa_generate_new(
193		&self,
194		key_type: KeyTypeId,
195		seed: Option<&str>,
196	) -> Result<ecdsa::Public, Error> {
197		self.generate_new::<ecdsa::Pair>(key_type, seed)
198	}
199
200	fn ecdsa_sign(
201		&self,
202		key_type: KeyTypeId,
203		public: &ecdsa::Public,
204		msg: &[u8],
205	) -> Result<Option<ecdsa::Signature>, Error> {
206		self.sign::<ecdsa::Pair>(key_type, public, msg)
207	}
208
209	fn ecdsa_sign_prehashed(
210		&self,
211		key_type: KeyTypeId,
212		public: &ecdsa::Public,
213		msg: &[u8; 32],
214	) -> Result<Option<ecdsa::Signature>, Error> {
215		let sig = self.pair::<ecdsa::Pair>(key_type, public).map(|pair| pair.sign_prehashed(msg));
216		Ok(sig)
217	}
218
219	#[cfg(feature = "bandersnatch-experimental")]
220	fn bandersnatch_public_keys(&self, key_type: KeyTypeId) -> Vec<bandersnatch::Public> {
221		self.public_keys::<bandersnatch::Pair>(key_type)
222	}
223
224	#[cfg(feature = "bandersnatch-experimental")]
225	fn bandersnatch_generate_new(
226		&self,
227		key_type: KeyTypeId,
228		seed: Option<&str>,
229	) -> Result<bandersnatch::Public, Error> {
230		self.generate_new::<bandersnatch::Pair>(key_type, seed)
231	}
232
233	#[cfg(feature = "bandersnatch-experimental")]
234	fn bandersnatch_sign(
235		&self,
236		key_type: KeyTypeId,
237		public: &bandersnatch::Public,
238		msg: &[u8],
239	) -> Result<Option<bandersnatch::Signature>, Error> {
240		self.sign::<bandersnatch::Pair>(key_type, public, msg)
241	}
242
243	#[cfg(feature = "bandersnatch-experimental")]
244	fn bandersnatch_vrf_sign(
245		&self,
246		key_type: KeyTypeId,
247		public: &bandersnatch::Public,
248		data: &bandersnatch::vrf::VrfSignData,
249	) -> Result<Option<bandersnatch::vrf::VrfSignature>, Error> {
250		self.vrf_sign::<bandersnatch::Pair>(key_type, public, data)
251	}
252
253	#[cfg(feature = "bandersnatch-experimental")]
254	fn bandersnatch_ring_vrf_sign(
255		&self,
256		key_type: KeyTypeId,
257		public: &bandersnatch::Public,
258		data: &bandersnatch::vrf::VrfSignData,
259		prover: &bandersnatch::ring_vrf::RingProver,
260	) -> Result<Option<bandersnatch::ring_vrf::RingVrfSignature>, Error> {
261		let sig = self
262			.pair::<bandersnatch::Pair>(key_type, public)
263			.map(|pair| pair.ring_vrf_sign(data, prover));
264		Ok(sig)
265	}
266
267	#[cfg(feature = "bandersnatch-experimental")]
268	fn bandersnatch_vrf_pre_output(
269		&self,
270		key_type: KeyTypeId,
271		public: &bandersnatch::Public,
272		input: &bandersnatch::vrf::VrfInput,
273	) -> Result<Option<bandersnatch::vrf::VrfPreOutput>, Error> {
274		self.vrf_pre_output::<bandersnatch::Pair>(key_type, public, input)
275	}
276
277	#[cfg(feature = "bls-experimental")]
278	fn bls381_public_keys(&self, key_type: KeyTypeId) -> Vec<bls381::Public> {
279		self.public_keys::<bls381::Pair>(key_type)
280	}
281
282	#[cfg(feature = "bls-experimental")]
283	fn bls381_generate_new(
284		&self,
285		key_type: KeyTypeId,
286		seed: Option<&str>,
287	) -> Result<bls381::Public, Error> {
288		self.generate_new::<bls381::Pair>(key_type, seed)
289	}
290
291	#[cfg(feature = "bls-experimental")]
292	fn bls381_sign(
293		&self,
294		key_type: KeyTypeId,
295		public: &bls381::Public,
296		msg: &[u8],
297	) -> Result<Option<bls381::Signature>, Error> {
298		self.sign::<bls381::Pair>(key_type, public, msg)
299	}
300
301	#[cfg(feature = "bls-experimental")]
302	fn ecdsa_bls381_public_keys(&self, key_type: KeyTypeId) -> Vec<ecdsa_bls381::Public> {
303		self.public_keys::<ecdsa_bls381::Pair>(key_type)
304	}
305
306	#[cfg(feature = "bls-experimental")]
307	fn ecdsa_bls381_generate_new(
308		&self,
309		key_type: KeyTypeId,
310		seed: Option<&str>,
311	) -> Result<ecdsa_bls381::Public, Error> {
312		self.generate_new::<ecdsa_bls381::Pair>(key_type, seed)
313	}
314
315	#[cfg(feature = "bls-experimental")]
316	fn ecdsa_bls381_sign(
317		&self,
318		key_type: KeyTypeId,
319		public: &ecdsa_bls381::Public,
320		msg: &[u8],
321	) -> Result<Option<ecdsa_bls381::Signature>, Error> {
322		self.sign::<ecdsa_bls381::Pair>(key_type, public, msg)
323	}
324
325	#[cfg(feature = "bls-experimental")]
326	fn ecdsa_bls381_sign_with_keccak256(
327		&self,
328		key_type: KeyTypeId,
329		public: &ecdsa_bls381::Public,
330		msg: &[u8],
331	) -> Result<Option<ecdsa_bls381::Signature>, Error> {
332		let sig = self
333			.pair::<ecdsa_bls381::Pair>(key_type, public)
334			.map(|pair| pair.sign_with_hasher::<KeccakHasher>(msg));
335		Ok(sig)
336	}
337
338	fn insert(&self, key_type: KeyTypeId, suri: &str, public: &[u8]) -> Result<(), ()> {
339		self.keys
340			.write()
341			.entry(key_type)
342			.or_default()
343			.insert(public.to_owned(), suri.to_string());
344		Ok(())
345	}
346
347	fn keys(&self, key_type: KeyTypeId) -> Result<Vec<Vec<u8>>, Error> {
348		let keys = self
349			.keys
350			.read()
351			.get(&key_type)
352			.map(|map| map.keys().cloned().collect())
353			.unwrap_or_default();
354		Ok(keys)
355	}
356
357	fn has_keys(&self, public_keys: &[(Vec<u8>, KeyTypeId)]) -> bool {
358		public_keys
359			.iter()
360			.all(|(k, t)| self.keys.read().get(t).and_then(|s| s.get(k)).is_some())
361	}
362}
363
364impl Into<KeystorePtr> for MemoryKeystore {
365	fn into(self) -> KeystorePtr {
366		Arc::new(self)
367	}
368}
369
370#[cfg(test)]
371mod tests {
372	use super::*;
373	use sp_core::{
374		sr25519,
375		testing::{ECDSA, ED25519, SR25519},
376	};
377
378	#[test]
379	fn store_key_and_extract() {
380		let store = MemoryKeystore::new();
381
382		let public = store.ed25519_generate_new(ED25519, None).expect("Generates key");
383
384		let public_keys = store.ed25519_public_keys(ED25519);
385
386		assert!(public_keys.contains(&public.into()));
387	}
388
389	#[test]
390	fn store_unknown_and_extract_it() {
391		let store = MemoryKeystore::new();
392
393		let secret_uri = "//Alice";
394		let key_pair = sr25519::Pair::from_string(secret_uri, None).expect("Generates key pair");
395
396		store
397			.insert(SR25519, secret_uri, key_pair.public().as_ref())
398			.expect("Inserts unknown key");
399
400		let public_keys = store.sr25519_public_keys(SR25519);
401
402		assert!(public_keys.contains(&key_pair.public().into()));
403	}
404
405	#[test]
406	fn sr25519_vrf_sign() {
407		let store = MemoryKeystore::new();
408
409		let secret_uri = "//Alice";
410		let key_pair = sr25519::Pair::from_string(secret_uri, None).expect("Generates key pair");
411
412		let data = sr25519::vrf::VrfInput::new(
413			b"Test",
414			&[
415				(b"one", &1_u64.to_le_bytes()),
416				(b"two", &2_u64.to_le_bytes()),
417				(b"three", "test".as_bytes()),
418			],
419		)
420		.into_sign_data();
421
422		let result = store.sr25519_vrf_sign(SR25519, &key_pair.public(), &data);
423		assert!(result.unwrap().is_none());
424
425		store
426			.insert(SR25519, secret_uri, key_pair.public().as_ref())
427			.expect("Inserts unknown key");
428
429		let result = store.sr25519_vrf_sign(SR25519, &key_pair.public(), &data);
430
431		assert!(result.unwrap().is_some());
432	}
433
434	#[test]
435	fn sr25519_vrf_pre_output() {
436		let store = MemoryKeystore::new();
437
438		let secret_uri = "//Alice";
439		let pair = sr25519::Pair::from_string(secret_uri, None).expect("Generates key pair");
440
441		let input = sr25519::vrf::VrfInput::new(
442			b"Test",
443			&[
444				(b"one", &1_u64.to_le_bytes()),
445				(b"two", &2_u64.to_le_bytes()),
446				(b"three", "test".as_bytes()),
447			],
448		);
449
450		let result = store.sr25519_vrf_pre_output(SR25519, &pair.public(), &input);
451		assert!(result.unwrap().is_none());
452
453		store
454			.insert(SR25519, secret_uri, pair.public().as_ref())
455			.expect("Inserts unknown key");
456
457		let pre_output =
458			store.sr25519_vrf_pre_output(SR25519, &pair.public(), &input).unwrap().unwrap();
459
460		let result = pre_output.make_bytes::<32>(b"rand", &input, &pair.public());
461		assert!(result.is_ok());
462	}
463
464	#[test]
465	fn ecdsa_sign_prehashed_works() {
466		let store = MemoryKeystore::new();
467
468		let suri = "//Alice";
469		let pair = ecdsa::Pair::from_string(suri, None).unwrap();
470
471		// Let's pretend this to be the hash output as content doesn't really matter here.
472		let hash = [0xff; 32];
473
474		// no key in key store
475		let res = store.ecdsa_sign_prehashed(ECDSA, &pair.public(), &hash).unwrap();
476		assert!(res.is_none());
477
478		// insert key, sign again
479		store.insert(ECDSA, suri, pair.public().as_ref()).unwrap();
480
481		let res = store.ecdsa_sign_prehashed(ECDSA, &pair.public(), &hash).unwrap();
482		assert!(res.is_some());
483	}
484
485	#[test]
486	#[cfg(feature = "bls-experimental")]
487	fn ecdsa_bls381_sign_with_keccak_works() {
488		use sp_core::testing::ECDSA_BLS377;
489
490		let store = MemoryKeystore::new();
491
492		let suri = "//Alice";
493		let pair = ecdsa_bls381::Pair::from_string(suri, None).unwrap();
494
495		let msg = b"this should be a normal unhashed message not a hash of a message because bls scheme comes with its own hashing";
496
497		// insert key, sign again
498		store.insert(ECDSA_BLS377, suri, pair.public().as_ref()).unwrap();
499
500		let res = store
501			.ecdsa_bls381_sign_with_keccak256(ECDSA_BLS377, &pair.public(), &msg[..])
502			.unwrap();
503
504		assert!(res.is_some());
505
506		// does not verify with default out-of-the-box verification
507		assert!(!ecdsa_bls381::Pair::verify(&res.unwrap(), &msg[..], &pair.public()));
508
509		// should verify using keccak256 as hasher
510		assert!(ecdsa_bls381::Pair::verify_with_hasher::<KeccakHasher>(
511			&res.unwrap(),
512			msg,
513			&pair.public()
514		));
515	}
516
517	#[test]
518	#[cfg(feature = "bandersnatch-experimental")]
519	fn bandersnatch_vrf_sign() {
520		use sp_core::testing::BANDERSNATCH;
521
522		let store = MemoryKeystore::new();
523
524		let secret_uri = "//Alice";
525		let key_pair =
526			bandersnatch::Pair::from_string(secret_uri, None).expect("Generates key pair");
527
528		let in1 = bandersnatch::vrf::VrfInput::new("in", "foo");
529		let sign_data =
530			bandersnatch::vrf::VrfSignData::new_unchecked(b"Test", Some("m1"), Some(in1));
531
532		let result = store.bandersnatch_vrf_sign(BANDERSNATCH, &key_pair.public(), &sign_data);
533		assert!(result.unwrap().is_none());
534
535		store
536			.insert(BANDERSNATCH, secret_uri, key_pair.public().as_ref())
537			.expect("Inserts unknown key");
538
539		let result = store.bandersnatch_vrf_sign(BANDERSNATCH, &key_pair.public(), &sign_data);
540
541		assert!(result.unwrap().is_some());
542	}
543
544	#[test]
545	#[cfg(feature = "bandersnatch-experimental")]
546	fn bandersnatch_ring_vrf_sign() {
547		use sp_core::testing::BANDERSNATCH;
548
549		let store = MemoryKeystore::new();
550
551		let ring_ctx = bandersnatch::ring_vrf::RingContext::<1024>::new_testing();
552
553		let mut pks: Vec<_> = (0..16)
554			.map(|i| bandersnatch::Pair::from_seed(&[i as u8; 32]).public())
555			.collect();
556
557		let prover_idx = 3;
558		let prover = ring_ctx.prover(&pks, prover_idx).unwrap();
559
560		let secret_uri = "//Alice";
561		let pair = bandersnatch::Pair::from_string(secret_uri, None).expect("Generates key pair");
562		pks[prover_idx] = pair.public();
563
564		let in1 = bandersnatch::vrf::VrfInput::new("in1", "foo");
565		let sign_data =
566			bandersnatch::vrf::VrfSignData::new_unchecked(b"Test", &["m1", "m2"], [in1]);
567
568		let result =
569			store.bandersnatch_ring_vrf_sign(BANDERSNATCH, &pair.public(), &sign_data, &prover);
570		assert!(result.unwrap().is_none());
571
572		store
573			.insert(BANDERSNATCH, secret_uri, pair.public().as_ref())
574			.expect("Inserts unknown key");
575
576		let result =
577			store.bandersnatch_ring_vrf_sign(BANDERSNATCH, &pair.public(), &sign_data, &prover);
578
579		assert!(result.unwrap().is_some());
580	}
581}