balius_runtime/sign/
in_memory.rs

1use std::collections::HashMap;
2
3use pallas::crypto::key::ed25519;
4use rand::rngs::OsRng;
5
6use crate::wit::balius::app::sign as wit;
7
8use super::SignerProvider;
9
10#[derive(Default, Clone)]
11pub struct Signer {
12    map: HashMap<String, HashMap<String, SignerKey>>,
13}
14
15impl Signer {
16    pub fn new() -> Self {
17        Self::default()
18    }
19}
20
21impl From<HashMap<String, HashMap<String, SignerKey>>> for Signer {
22    fn from(value: HashMap<String, HashMap<String, SignerKey>>) -> Self {
23        Self { map: value }
24    }
25}
26
27#[async_trait::async_trait]
28impl SignerProvider for Signer {
29    async fn add_key(&mut self, worker_id: &str, key_name: String, algorithm: String) -> Vec<u8> {
30        if algorithm != "ed25519" {
31            panic!("Unsupported algorithm")
32        }
33        let keys = self.map.entry(worker_id.to_string()).or_default();
34        let secret_key = keys
35            .entry(key_name)
36            .or_insert(ed25519::SecretKey::new(OsRng).into());
37        secret_key.public_key()
38    }
39
40    async fn sign_payload(
41        &mut self,
42        worker_id: &str,
43        key_name: String,
44        payload: wit::Payload,
45    ) -> Result<wit::Signature, wit::SignError> {
46        let Some(key) = self
47            .map
48            .entry(worker_id.to_string())
49            .or_default()
50            .get(&key_name)
51        else {
52            return Err(wit::SignError::KeyNotFound(key_name.to_string()));
53        };
54        key.sign_payload(payload)
55    }
56}
57
58#[derive(Clone, Debug)]
59pub enum SignerKey {
60    Ed25519(Ed25519Key),
61}
62
63impl From<ed25519::SecretKey> for SignerKey {
64    fn from(value: ed25519::SecretKey) -> Self {
65        Self::Ed25519(Ed25519Key::SecretKey(value))
66    }
67}
68
69impl From<ed25519::SecretKeyExtended> for SignerKey {
70    fn from(value: ed25519::SecretKeyExtended) -> Self {
71        Self::Ed25519(Ed25519Key::SecretKeyExtended(value))
72    }
73}
74
75impl SignerKey {
76    pub fn sign_payload(&self, payload: wit::Payload) -> Result<wit::Signature, wit::SignError> {
77        match self {
78            Self::Ed25519(key) => Ok(key.sign_payload(payload)),
79        }
80    }
81
82    pub fn public_key(&self) -> Vec<u8> {
83        match self {
84            Self::Ed25519(key) => key.public_key(),
85        }
86    }
87}
88
89#[derive(Clone, Debug)]
90pub enum Ed25519Key {
91    SecretKey(ed25519::SecretKey),
92    SecretKeyExtended(ed25519::SecretKeyExtended),
93}
94impl Ed25519Key {
95    fn sign_payload(&self, payload: wit::Payload) -> wit::Signature {
96        let signature = match self {
97            Self::SecretKey(key) => key.sign(payload),
98            Self::SecretKeyExtended(key) => key.sign(payload),
99        };
100        signature.as_ref().to_vec()
101    }
102    fn public_key(&self) -> Vec<u8> {
103        match self {
104            Self::SecretKey(key) => key.public_key().as_ref().to_vec(),
105            Self::SecretKeyExtended(key) => key.public_key().as_ref().to_vec(),
106        }
107    }
108}