balius_runtime/sign/
mod.rs

1pub mod in_memory;
2
3use std::sync::Arc;
4
5use tokio::sync::Mutex;
6
7use crate::{metrics::Metrics, wit::balius::app::sign as wit};
8
9#[derive(Clone)]
10pub enum Signer {
11    InMemory(in_memory::Signer),
12    Custom(Arc<Mutex<dyn SignerProvider + Send + Sync>>),
13}
14
15impl From<in_memory::Signer> for Signer {
16    fn from(signer: in_memory::Signer) -> Self {
17        Signer::InMemory(signer)
18    }
19}
20
21pub struct SignerHost {
22    worker_id: String,
23    provider: Signer,
24    metrics: Arc<Metrics>,
25}
26impl SignerHost {
27    pub fn new(worker_id: &str, provider: &Signer, metrics: &Arc<Metrics>) -> Self {
28        Self {
29            worker_id: worker_id.to_string(),
30            provider: provider.clone(),
31            metrics: metrics.clone(),
32        }
33    }
34
35    pub async fn add_key(&mut self, key_name: String, algorithm: String) -> Vec<u8> {
36        match &mut self.provider {
37            Signer::InMemory(signer) => signer.add_key(&self.worker_id, key_name, algorithm).await,
38            Signer::Custom(signer) => {
39                let mut lock = signer.lock().await;
40                lock.add_key(&self.worker_id, key_name, algorithm).await
41            }
42        }
43    }
44}
45
46#[async_trait::async_trait]
47pub trait SignerProvider {
48    async fn add_key(&mut self, worker_id: &str, key_name: String, algorithm: String) -> Vec<u8>;
49    async fn sign_payload(
50        &mut self,
51        worker_id: &str,
52        key_name: String,
53        payload: wit::Payload,
54    ) -> Result<wit::Signature, wit::SignError>;
55}
56
57impl wit::Host for SignerHost {
58    async fn sign_payload(
59        &mut self,
60        key_name: String,
61        payload: wit::Payload,
62    ) -> Result<wit::Signature, wit::SignError> {
63        self.metrics.signer_sign_payload(&self.worker_id);
64        match &mut self.provider {
65            Signer::InMemory(signer) => {
66                signer
67                    .sign_payload(&self.worker_id, key_name, payload)
68                    .await
69            }
70            Signer::Custom(signer) => {
71                let mut lock = signer.lock().await;
72                lock.sign_payload(&self.worker_id, key_name, payload).await
73            }
74        }
75    }
76}