1#[cfg(test)]
2use crate::services::mocks::mock_kms_client_service::MockKmsClientService;
3use crate::services::{
4 aws_kms_client_service::AWSKmsClientService, crypto_service::CryptoService,
5 kms_client_service::KmsClientService,
6};
7use crate::{config::Config, constants::SIGNATURE_RSV_LEN};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10use tracing::error;
11
12#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
13pub struct KeyEntry {
14 pub address: Arc<String>,
15 pub public_key: Arc<Option<String>>,
16 pub public_key_base64: Arc<String>,
17 pub key_id: Arc<String>,
18}
19
20#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
21pub struct SigEntry {
22 pub address: Arc<String>,
23 pub public_key: Arc<String>,
24 pub signature: Arc<String>,
25}
26
27#[async_trait::async_trait]
28pub trait KeysServiceTrait: Send + Sync {
29 async fn create_key(&mut self, config: &Config) -> Result<KeyEntry, String>;
35
36 async fn sign_transaction_hash(
45 &mut self,
46 config: &Config,
47 transaction_hash: &str,
48 key: &str,
49 ) -> Result<SigEntry, String>;
50
51 async fn sign_transaction(
60 &mut self,
61 config: &Config,
62 transaction_str: &str,
63 key: &str,
64 ) -> Result<String, String>;
65
66 async fn verify(
77 &mut self,
78 transaction_hash_hex: &str,
79 signature_hex: &str,
80 key: &str,
81 ) -> Result<bool, String>;
82
83 async fn verify_via_kms(
94 &mut self,
95 transaction_hash_hex: &str,
96 signature_hex: &str,
97 key: &str,
98 ) -> Result<bool, String>;
99
100 async fn delete_key(&mut self, key: &str) -> Result<bool, String>;
108
109 async fn list_keys(&mut self) -> Result<Vec<KeyEntry>, String>;
114}
115
116pub struct KeysService {
117 pub kms_client_service: Arc<dyn KmsClientService>,
118 pub crypto_service: CryptoService,
119}
120
121impl KeysService {
122 pub async fn new(config: Config, crypto_service: CryptoService) -> Result<Self, String> {
131 let kms_client_service: Arc<dyn KmsClientService> = if config.is_aws_mode() {
132 let client = AWSKmsClientService::new(config.get_aws_config().clone())
133 .await
134 .map_err(|e| format!("Failed to initialize AWSKmsClientService: {e}"))?;
135 Arc::new(client)
136 } else if config.is_testing_mode() {
137 Self::get_mock_kms_client_service()
139 } else {
140 return Err("Unsupported KMS mode and not in testing".to_string());
141 };
142
143 Ok(Self {
144 kms_client_service,
145 crypto_service,
146 })
147 }
148
149 #[cfg(test)]
150 fn get_mock_kms_client_service() -> Arc<dyn KmsClientService> {
151 Arc::new(MockKmsClientService)
152 }
153
154 #[cfg(not(test))]
155 fn get_mock_kms_client_service() -> Arc<dyn KmsClientService> {
156 unimplemented!("Mock KMS Client is only available in tests")
157 }
158
159 pub async fn sign(
179 &mut self,
180 transaction_hash_hex: &str,
181 key: &str,
182 prefix: Option<&str>,
183 ) -> Result<String, String> {
184 let signature = self
185 .kms_client_service
186 .sign(transaction_hash_hex, key)
187 .await
188 .map_err(|e| {
189 let msg = format!("Failed to sign transaction with KMS: {e}");
190 error!("{}", msg);
191 msg
192 })?;
193
194 let mut signature = self.crypto_service.convert(&signature).map_err(|e| {
195 let msg = format!("Failed to convert signature: {e}");
196 error!("{}", msg);
197 msg
198 })?;
199
200 if let Some(pref) = prefix {
201 signature = format!("{pref}{signature}");
202 }
203
204 Ok(signature)
205 }
206
207 pub fn verify(
223 &mut self,
224 transaction_hash_hex: &str,
225 signature_hex: &str,
226 public_key: &str,
227 ) -> Result<bool, String> {
228 self.crypto_service
229 .verify(transaction_hash_hex, signature_hex, public_key)
230 .map_err(|e| {
231 let msg = format!("Signature verification failed: {e}");
232 error!("{}", msg);
233 msg
234 })
235 }
236
237 pub fn verify_eip155(
253 &mut self,
254 transaction_hash_hex: &str,
255 signature_hex: &str,
256 public_key: &str,
257 ) -> Result<bool, String> {
258 self.crypto_service
259 .verify_eip155(transaction_hash_hex, signature_hex, public_key)
260 .map_err(|e| {
261 let msg = format!("Signature verification failed: {e}");
262 error!("{}", msg);
263 msg
264 })
265 }
266
267 pub async fn verify_via_kms(
287 &mut self,
288 transaction_hash_hex: &str,
289 signature_hex: &str,
290 public_key: &str,
291 ) -> Result<bool, String> {
292 self.verify_via_kms_internal(transaction_hash_hex, signature_hex, public_key, false)
293 .await
294 }
295
296 pub async fn verify_via_kms_eip155(
317 &mut self,
318 transaction_hash_hex: &str,
319 signature_hex: &str,
320 public_key: &str,
321 ) -> Result<bool, String> {
322 self.verify_via_kms_internal(transaction_hash_hex, signature_hex, public_key, true)
323 .await
324 }
325
326 async fn verify_via_kms_internal(
327 &mut self,
328 transaction_hash_hex: &str,
329 mut signature_hex: &str,
330 public_key: &str,
331 use_eip155: bool,
332 ) -> Result<bool, String> {
333 let is_verified = if use_eip155 {
334 let verify_eip155 = self
335 .verify_eip155(transaction_hash_hex, signature_hex, public_key)
336 .map_err(|e| {
337 let msg = format!("Signature eip155 verification failed: {e}");
338 error!("{}", msg);
339 msg
340 })?;
341
342 if signature_hex.len() == SIGNATURE_RSV_LEN {
344 signature_hex = &signature_hex[..signature_hex.len() - 2];
345 }
346
347 verify_eip155
348 } else {
349 self.verify(transaction_hash_hex, signature_hex, public_key)
350 .map_err(|e| {
351 let msg = format!("Signature verification failed: {e}");
352 error!("{}", msg);
353 msg
354 })?
355 };
356
357 if !is_verified {
358 return Ok(false);
359 }
360
361 let signature = self.crypto_service.unconvert(signature_hex).map_err(|e| {
362 let msg = format!("Signature conversion failed: {e}");
363 error!("{}", msg);
364 msg
365 })?;
366
367 let alias = if use_eip155 {
368 self.crypto_service.address_eth(public_key).map_err(|e| {
369 let msg = format!("Failed to convert public key to address: {e:?}");
370 error!("{}", &msg);
371 msg
372 })?
373 } else {
374 public_key.to_string()
375 };
376
377 self.kms_client_service
378 .verify(transaction_hash_hex, &signature, &alias)
379 .await
380 .map_err(|e| {
381 let msg = format!("Failed to verify signature with KMS: {e}");
382 error!("{}", msg);
383 msg
384 })
385 }
386
387 pub async fn delete_key(&mut self, key: &str) -> Result<bool, String> {
401 self.kms_client_service.delete_key(key).await.map_err(|e| {
402 let msg = format!("Key deletion failed: {e}");
403 error!("{}", msg);
404 msg
405 })
406 }
407
408 pub async fn list_keys(&mut self) -> Result<Vec<KeyEntry>, String> {
418 let entries = self.kms_client_service.list_keys().await.map_err(|e| {
419 let msg = format!("Listing keys failed: {e}");
420 error!("{}", msg);
421 msg
422 })?;
423
424 let mut result = Vec::with_capacity(entries.len());
425
426 for mut entry in entries {
427 if entry.public_key.is_none() {
428 let public_key_base64 = entry.public_key_base64.to_string();
429
430 let public_key =
431 self.crypto_service
432 .public_key(&public_key_base64)
433 .map_err(|e| {
434 let msg = format!("public_key conversion failed: {e:?}");
435 error!("{}", &msg);
436 msg
437 })?;
438
439 entry.public_key = Some(public_key).into();
440 }
441 result.push(entry);
442 }
443
444 Ok(result)
445 }
446}
447
448#[cfg(test)]
449mod tests {
450 use super::*;
451 use crate::{
452 config::ConfigBuilder,
453 constants::{
454 CASPER_PUBLIC_KEY_PREFIXED, CASPER_SECP_PREFIX, ETH_PUBLIC_KEY, ETH_SIGNATURE,
455 ETH_TRANSACTION_HASH, SIGNATURE_RS_LEN, TRANSACTION_HASH, WASM_PATH,
456 },
457 services::crypto_service::CryptoService,
458 wasm_loader::WasmLoader,
459 };
460 use base64::Engine;
461 use base64::engine::general_purpose::STANDARD;
462
463 #[tokio::test]
464 async fn test_sign_successful() {
465 let config = ConfigBuilder::new().with_casper_mode().build();
466
467 let wasm_loader = WasmLoader::new(WASM_PATH)
468 .await
469 .expect("Failed to load WASM module");
470
471 let crypto_service =
472 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
473
474 let mut service = KeysService::new(config.clone(), crypto_service)
476 .await
477 .expect("Failed to create KeysService");
478
479 let public_key = CASPER_PUBLIC_KEY_PREFIXED;
480
481 let result = service
483 .sign(TRANSACTION_HASH, public_key, Some(CASPER_SECP_PREFIX))
484 .await;
485 assert!(result.is_ok(), "sign failed: {result:?}");
486 let signature = result.unwrap();
487 assert!(
488 signature.starts_with(CASPER_SECP_PREFIX),
489 "Signature missing prefix"
490 );
491 assert_eq!(
492 signature.len(),
493 CASPER_SECP_PREFIX.len() + SIGNATURE_RS_LEN,
494 "Signature length invalid"
495 );
496
497 let result = service.sign(TRANSACTION_HASH, public_key, None).await;
499 assert!(result.is_ok(), "sign failed: {result:?}");
500 let signature = result.unwrap();
501
502 assert_eq!(
503 signature.len(),
504 SIGNATURE_RS_LEN,
505 "Signature length invalid"
506 );
507 }
508
509 #[tokio::test]
541 async fn test_verify_successful() {
542 let config = ConfigBuilder::new().with_casper_mode().build();
543
544 let wasm_loader = WasmLoader::new(WASM_PATH)
545 .await
546 .expect("Failed to load WASM module");
547
548 let crypto_service =
549 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
550
551 let mut service = KeysService::new(config.clone(), crypto_service)
552 .await
553 .expect("Failed to create KeysService");
554
555 let public_key = CASPER_PUBLIC_KEY_PREFIXED;
556
557 let signature = service
559 .sign(TRANSACTION_HASH, public_key, None)
560 .await
561 .expect("Signing failed");
562
563 let result = service.verify(TRANSACTION_HASH, &signature, public_key);
565
566 assert!(result.is_ok(), "verify failed: {result:?}");
567 assert!(result.unwrap(), "signature verification returned false");
568 }
569
570 #[tokio::test]
571 async fn test_verify_eip155_successful() {
572 let config = ConfigBuilder::new().with_ethereum_mode().build();
573
574 let wasm_loader = WasmLoader::new(WASM_PATH)
575 .await
576 .expect("Failed to load WASM module");
577
578 let crypto_service =
579 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
580
581 let mut service = KeysService::new(config.clone(), crypto_service)
582 .await
583 .expect("Failed to create KeysService");
584
585 let transaction_hash = ETH_TRANSACTION_HASH;
586 let public_key = ETH_PUBLIC_KEY;
587 let signature = ETH_SIGNATURE;
588
589 let result = service.verify_eip155(transaction_hash, signature, public_key);
591
592 assert!(result.is_ok(), "verify_eip155 failed: {result:?}");
593 assert!(
594 result.unwrap(),
595 "EIP-155 signature verification returned false"
596 );
597 }
598
599 #[tokio::test]
600 async fn test_verify_via_kms_successful() {
601 let config = ConfigBuilder::new().with_casper_mode().build();
602
603 let wasm_loader = WasmLoader::new(WASM_PATH)
604 .await
605 .expect("Failed to load WASM module");
606
607 let crypto_service =
608 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
609
610 let mut service = KeysService::new(config, crypto_service)
611 .await
612 .expect("Failed to create KeysService");
613
614 let signature = service
616 .sign(TRANSACTION_HASH, CASPER_PUBLIC_KEY_PREFIXED, None)
617 .await
618 .expect("Signing failed");
619
620 let result = service
622 .verify_via_kms(TRANSACTION_HASH, &signature, CASPER_PUBLIC_KEY_PREFIXED)
623 .await;
624
625 assert!(result.is_ok(), "verify_via_kms failed: {result:?}");
626 assert!(result.unwrap(), "KMS verification returned false");
627 }
628
629 #[tokio::test]
630 async fn test_verify_via_kms_eip155_successful() {
631 let config = ConfigBuilder::new().with_ethereum_mode().build();
632
633 let wasm_loader = WasmLoader::new(WASM_PATH)
634 .await
635 .expect("Failed to load WASM module");
636
637 let crypto_service =
638 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
639
640 let mut service = KeysService::new(config, crypto_service)
641 .await
642 .expect("Failed to create KeysService");
643
644 let result = service
646 .verify_via_kms_eip155(ETH_TRANSACTION_HASH, ETH_SIGNATURE, ETH_PUBLIC_KEY)
647 .await;
648
649 assert!(result.is_ok(), "verify_via_kms_eip155 failed: {result:?}");
650 assert!(result.unwrap(), "EIP-155 + KMS verification returned false");
651 }
652
653 #[tokio::test]
654 async fn test_delete_key_successful() {
655 let config = ConfigBuilder::new().build();
656
657 let wasm_loader = WasmLoader::new(WASM_PATH)
658 .await
659 .expect("Failed to load WASM module");
660
661 let crypto_service =
662 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
663
664 let mut service = KeysService::new(config, crypto_service)
665 .await
666 .expect("Failed to create KeysService");
667
668 let result = service.delete_key("known_public_key_xyz").await;
670
671 assert!(result.is_ok(), "delete_key failed unexpectedly: {result:?}");
672 assert!(result.unwrap(), "Expected key deletion to succeed");
673 }
674
675 #[tokio::test]
676 async fn test_delete_key_not_found() {
677 let config = ConfigBuilder::new().build();
678
679 let wasm_loader = WasmLoader::new(WASM_PATH)
680 .await
681 .expect("Failed to load WASM module");
682
683 let crypto_service =
684 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
685
686 let mut service = KeysService::new(config, crypto_service)
687 .await
688 .expect("Failed to create KeysService");
689
690 let result = service.delete_key("some_other_key").await;
692
693 assert!(result.is_ok(), "delete_key failed unexpectedly: {result:?}");
694 assert!(
695 !result.unwrap(),
696 "Expected deletion to return false for unknown key"
697 );
698 }
699
700 #[tokio::test]
701 async fn test_list_keys_successful() {
702 let config = ConfigBuilder::new().build();
703 let wasm_loader = WasmLoader::new(WASM_PATH)
704 .await
705 .expect("Failed to load WASM module");
706
707 let crypto_service =
708 CryptoService::new(&wasm_loader).expect("Failed to initialize CryptoService");
709
710 let mut service = KeysService::new(config, crypto_service)
711 .await
712 .expect("Failed to create KeysService");
713
714 let result = service.list_keys().await;
715
716 assert!(result.is_ok(), "list_keys failed: {result:?}");
717
718 let keys = result.unwrap();
719 assert_eq!(keys.len(), 2, "Expected two keys from mocked KMS");
720 assert_eq!(
721 keys[0],
722 KeyEntry {
723 address: "address_1".to_string().into(),
724 public_key_base64: STANDARD.encode("public_key_1_base64").into(),
725 public_key: Some("public_key_1".to_string()).into(),
726 key_id: "key_id_1".to_string().into(),
727 }
728 );
729 assert_eq!(
730 keys[1],
731 KeyEntry {
732 address: "address_2".to_string().into(),
733 public_key_base64: STANDARD.encode("public_key_2_base64").into(),
734 public_key: Some("public_key_2".to_string()).into(),
735 key_id: "key_id_2".to_string().into(),
736 }
737 );
738 }
739}