1use 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#[derive(Default, Clone)]
36pub struct MemoryKeystore {
37 keys: Arc<RwLock<HashMap<KeyTypeId, HashMap<Vec<u8>, String>>>>,
39}
40
41impl MemoryKeystore {
42 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 hash = [0xff; 32];
473
474 let res = store.ecdsa_sign_prehashed(ECDSA, &pair.public(), &hash).unwrap();
476 assert!(res.is_none());
477
478 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 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 assert!(!ecdsa_bls381::Pair::verify(&res.unwrap(), &msg[..], &pair.public()));
508
509 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}