brain/tests/common/test_utils.rs
2025-05-06 00:55:58 +03:00

48 lines
1.7 KiB
Rust

use anyhow::Result;
use detee_shared::vm_proto as snp_proto;
use ed25519_dalek::{Signer, SigningKey};
use tonic::metadata::AsciiMetadataValue;
use tonic::Request;
#[derive(Debug, Clone)]
pub struct Key {
pub sg_key: SigningKey,
pub pubkey: String,
}
impl Key {
pub fn new() -> Self {
let sk = SigningKey::generate(&mut rand::rngs::OsRng);
let pubkey = bs58::encode(sk.verifying_key().to_bytes()).into_string();
Key { sg_key: sk, pubkey }
}
pub fn sign_request<T: std::fmt::Debug>(&self, req: T) -> Result<Request<T>> {
let pubkey = self.pubkey.clone();
let timestamp = chrono::Utc::now().to_rfc3339();
let signature = self.try_sign_message(&format!("{timestamp}{req:?}"))?;
let timestamp: AsciiMetadataValue = timestamp.parse()?;
let pubkey: AsciiMetadataValue = pubkey.parse()?;
let signature: AsciiMetadataValue = signature.parse()?;
let mut req = Request::new(req);
req.metadata_mut().insert("timestamp", timestamp);
req.metadata_mut().insert("pubkey", pubkey);
req.metadata_mut().insert("request-signature", signature);
Ok(req)
}
pub fn try_sign_message(&self, message: &str) -> Result<String> {
let key = self.sg_key.clone();
Ok(bs58::encode(key.sign(message.as_bytes()).to_bytes()).into_string())
}
pub fn sign_stream_auth(&self, contracts: Vec<String>) -> Result<snp_proto::DaemonStreamAuth> {
let pubkey = self.pubkey.clone();
let timestamp = chrono::Utc::now().to_rfc3339();
let signature =
self.try_sign_message(&(timestamp.to_string() + &format!("{contracts:?}")))?;
Ok(snp_proto::DaemonStreamAuth { timestamp, pubkey, contracts, signature })
}
}