brain/tests/common/test_utils.rs
Noor 479b9b6fa4 Implement refund on delete app and vm
fix authentication and refund on delete vm
refactor app daemon register to get delete app req
db function for delete app
sample environment variable
extensive tests for delete app and vm
refactor test utilities
2025-06-04 16:01:42 +00:00

101 lines
3.6 KiB
Rust

use anyhow::Result;
use detee_shared::app_proto as sgx_proto;
use detee_shared::general_proto::brain_general_cli_client::BrainGeneralCliClient;
use detee_shared::general_proto::AirdropReq;
use detee_shared::vm_proto as snp_proto;
use ed25519_dalek::{Signer, SigningKey};
use itertools::Itertools;
use std::sync::OnceLock;
use surreal_brain::constants::TOKEN_DECIMAL;
use tonic::metadata::AsciiMetadataValue;
use tonic::transport::Channel;
use tonic::Request;
pub static ADMIN_KEYS: OnceLock<Vec<Key>> = OnceLock::new();
pub fn admin_keys() -> Vec<Key> {
let admin_keys = ADMIN_KEYS.get_or_init(|| {
let admin_keys = vec![Key::new(), Key::new(), Key::new()];
let admin_pub_keys = admin_keys.iter().map(|k| k.pubkey.clone()).join(", ");
std::env::set_var("ADMIN_PUB_KEYS", admin_pub_keys);
admin_keys.clone()
});
admin_keys.clone()
}
#[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 from(private_key: &str) -> Self {
let signing_key: SigningKey = bs58::decode(private_key)
.into_vec()
.expect("Failed to decode private key")
.as_slice()
.try_into()
.expect("Failed to convert to SigningKey");
let pubkey = bs58::encode(signing_key.verifying_key().to_bytes()).into_string();
Key { sg_key: signing_key, 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_vm(
&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 })
}
pub fn sign_stream_auth_app(&self, contracts: Vec<String>) -> Result<sgx_proto::DaemonAuth> {
let pubkey = self.pubkey.clone();
let timestamp = chrono::Utc::now().to_rfc3339();
let signature =
self.try_sign_message(&(timestamp.to_string() + &format!("{contracts:?}")))?;
Ok(sgx_proto::DaemonAuth { timestamp, pubkey, contracts, signature })
}
}
pub async fn airdrop(brain_channel: &Channel, wallet: &str, amount: u64) -> Result<()> {
let mut client = BrainGeneralCliClient::new(brain_channel.clone());
let airdrop_req = AirdropReq { pubkey: wallet.to_string(), tokens: amount * TOKEN_DECIMAL };
let admin_key = admin_keys()[0].clone();
client.airdrop(admin_key.sign_request(airdrop_req.clone())?).await?;
Ok(())
}