diff --git a/Cargo.lock b/Cargo.lock index 200e33bb..721fdd4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6539,6 +6539,7 @@ dependencies = [ "reqwest 0.12.12", "schnorrkel", "secp256k1", + "seismic-enclave-derive", "serde", "serde_json", "sha2", @@ -6547,6 +6548,15 @@ dependencies = [ "tracing", ] +[[package]] +name = "seismic-enclave-derive" +version = "0.1.0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.99", +] + [[package]] name = "seismic-enclave-server" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 87b59071..bc7c9292 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,9 +2,104 @@ resolver = "2" members = [ "crates/enclave", + "crates/enclave/derive", "crates/enclave-server", ] +[workspace.lints] +rust.missing_debug_implementations = "warn" +rust.missing_docs = "warn" +rust.rust_2018_idioms = { level = "deny", priority = -1 } +rust.unreachable_pub = "warn" +rust.unused_must_use = "deny" + +[workspace.lints.clippy] +# These are some of clippy's nursery (i.e., experimental) lints that we like. +# By default, nursery lints are allowed. Some of the lints below have made good +# suggestions which we fixed. The others didn't have any findings, so we can +# assume they don't have that many false positives. Let's enable them to +# prevent future problems. +borrow_as_ptr = "warn" +branches_sharing_code = "warn" +clear_with_drain = "warn" +cloned_instead_of_copied = "warn" +collection_is_never_read = "warn" +dbg_macro = "warn" +derive_partial_eq_without_eq = "warn" +doc_markdown = "warn" +empty_line_after_doc_comments = "warn" +empty_line_after_outer_attr = "warn" +enum_glob_use = "warn" +equatable_if_let = "warn" +explicit_into_iter_loop = "warn" +explicit_iter_loop = "warn" +flat_map_option = "warn" +from_iter_instead_of_collect = "warn" +if_not_else = "warn" +if_then_some_else_none = "warn" +implicit_clone = "warn" +imprecise_flops = "warn" +iter_on_empty_collections = "warn" +iter_on_single_items = "warn" +iter_with_drain = "warn" +iter_without_into_iter = "warn" +large_stack_frames = "warn" +manual_assert = "warn" +manual_clamp = "warn" +manual_is_variant_and = "warn" +manual_string_new = "warn" +match_same_arms = "warn" +missing_const_for_fn = "warn" +mutex_integer = "warn" +naive_bytecount = "warn" +needless_bitwise_bool = "warn" +needless_continue = "warn" +needless_for_each = "warn" +needless_pass_by_ref_mut = "warn" +nonstandard_macro_braces = "warn" +option_as_ref_cloned = "warn" +or_fun_call = "warn" +path_buf_push_overwrite = "warn" +read_zero_byte_vec = "warn" +redundant_clone = "warn" +redundant_else = "warn" +single_char_pattern = "warn" +string_lit_as_bytes = "warn" +string_lit_chars_any = "warn" +suboptimal_flops = "warn" +suspicious_operation_groupings = "warn" +trailing_empty_array = "warn" +trait_duplication_in_bounds = "warn" +transmute_undefined_repr = "warn" +trivial_regex = "warn" +tuple_array_conversions = "warn" +type_repetition_in_bounds = "warn" +uninhabited_references = "warn" +unnecessary_self_imports = "warn" +unnecessary_struct_initialization = "warn" +unnested_or_patterns = "warn" +unused_peekable = "warn" +unused_rounding = "warn" +use_self = "warn" +useless_let_if_seq = "warn" +while_float = "warn" +zero_sized_map_values = "warn" + +# These are nursery lints which have findings. Allow them for now. Some are not +# quite mature enough for use in our codebase and some we don't really want. +# Explicitly listing should make it easier to fix in the future. +as_ptr_cast_mut = "allow" +cognitive_complexity = "allow" +debug_assert_with_mut_call = "allow" +fallible_impl_from = "allow" +future_not_send = "allow" +needless_collect = "allow" +non_send_fields_in_send_ty = "allow" +redundant_pub_crate = "allow" +significant_drop_in_scrutinee = "allow" +significant_drop_tightening = "allow" +too_long_first_doc_paragraph = "allow" + [workspace.package] edition = "2021" version = "0.1.0" @@ -15,6 +110,9 @@ license = "MIT" readme = "README.md" [workspace.dependencies] +seismic-enclave-derive = { path = "crates/enclave/derive" } +seismic-enclave = { path = "crates/enclave" } + aes-gcm = "0.10" anyhow = "1.0" az-tdx-vtpm = "0.7.1" @@ -40,5 +138,8 @@ serde_json = "1.0" sha2 = "0.10" strum = { version = "0.26", features = ["derive"] } tokio = { version = "1.44", features = ["full"] } -tracing = { version = "0.1" } +tracing = { version = "0.1"} tracing-subscriber = { version = "0.3", default-features = false, features = ["env-filter", "fmt", "ansi", "json"] } +proc-macro2 = "1.0" +quote = "1.0" +syn = "2.0" diff --git a/crates/enclave-server/Cargo.toml b/crates/enclave-server/Cargo.toml index 48c5952b..5ace3ee6 100644 --- a/crates/enclave-server/Cargo.toml +++ b/crates/enclave-server/Cargo.toml @@ -14,10 +14,11 @@ name = "seismic_enclave_server" path = "src/lib.rs" [dependencies] +seismic-enclave.workspace = true + # attestation-service depends on attestation-agent, ensure versions are compatible when updating attestation-service = { git = "https://github.com/confidential-containers/trustee", features = ["all-verifier"], rev="1fdd67d"} attestation-agent = { git = "https://github.com/confidential-containers/guest-components", features = ["az-tdx-vtpm-attester"], rev="e6999a3"} -seismic-enclave = { path = "../enclave" } aes-gcm.workspace = true anyhow.workspace = true diff --git a/crates/enclave/Cargo.toml b/crates/enclave/Cargo.toml index a657f032..a9f9520c 100644 --- a/crates/enclave/Cargo.toml +++ b/crates/enclave/Cargo.toml @@ -14,6 +14,8 @@ name = "seismic_enclave" path = "src/lib.rs" [dependencies] +seismic-enclave-derive.workspace = true + aes-gcm.workspace = true anyhow.workspace = true bincode.workspace = true diff --git a/crates/enclave/derive/Cargo.toml b/crates/enclave/derive/Cargo.toml new file mode 100644 index 00000000..ab643b26 --- /dev/null +++ b/crates/enclave/derive/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "seismic-enclave-derive" +version.workspace = true +edition.workspace = true +repository.workspace = true +homepage.workspace = true +authors.workspace = true +license.workspace = true +readme.workspace = true +description = "Tools for building and interfacing with the Seismic enclave" + +[lib] +proc-macro = true + +[dependencies] +proc-macro2.workspace = true +quote.workspace = true +syn = { version = "2.0", features = ["full"] } diff --git a/crates/enclave/derive/src/lib.rs b/crates/enclave/derive/src/lib.rs new file mode 100644 index 00000000..0bff69e6 --- /dev/null +++ b/crates/enclave/derive/src/lib.rs @@ -0,0 +1,53 @@ +use proc_macro::TokenStream; +use quote::quote; +use syn::{parse_macro_input, ItemTrait, ReturnType, TraitItem}; + +/// Derive a sync client from an async rpc trait. +#[proc_macro_attribute] +pub fn derive_sync_client_trait(_attr: TokenStream, item: TokenStream) -> TokenStream { + let input = parse_macro_input!(item as ItemTrait); + let trait_name = &input.ident; + let sync_trait_name = syn::Ident::new(&format!("Sync{}Client", trait_name), trait_name.span()); + + let methods = input.items.iter().filter_map(|item| { + if let TraitItem::Fn(m) = item { + let sig = &m.sig; + let method_name = &sig.ident; + let inputs = &sig.inputs; + match &sig.output { + + ReturnType::Type(_, ty) =>{ + if let syn::Type::Path(type_path) = &**ty { + if let Some(segment) = type_path.path.segments.last() { + if segment.ident == "RpcResult" { + if let syn::PathArguments::AngleBracketed(args) = &segment.arguments { + if let Some(syn::GenericArgument::Type(inner_ty)) = args.args.first() { + return Some(quote! { + fn #method_name(#inputs) -> Result<#inner_ty, jsonrpsee::core::client::Error>; + }) + } + } + } + } + } + panic!("Method {} has an invalid return type. Return type must be RpcResult", method_name); + } + _ => { + panic!("Method {} has an invalid return type. Return type must be RpcResult", method_name); + } + } + } else { + None + } + }); + + let expanded = quote! { + pub trait #sync_trait_name { + #(#methods)* + } + + #input // Keep the original async trait unchanged + }; + + TokenStream::from(expanded) +} diff --git a/crates/enclave/src/client/client.rs b/crates/enclave/src/client/client.rs index 1af7f471..0abe688d 100644 --- a/crates/enclave/src/client/client.rs +++ b/crates/enclave/src/client/client.rs @@ -1,17 +1,42 @@ -use jsonrpsee::http_client::HttpClient; +use jsonrpsee::{core::ClientError, http_client::HttpClient}; use std::{ + future::Future, net::{IpAddr, Ipv4Addr}, ops::Deref, + sync::OnceLock, }; +use tokio::runtime::{Handle, Runtime}; + +use crate::{ + coco_aa::{AttestationGetEvidenceRequest, AttestationGetEvidenceResponse}, + coco_as::{AttestationEvalEvidenceRequest, AttestationEvalEvidenceResponse}, + genesis::GenesisDataResponse, + signing::{ + Secp256k1SignRequest, Secp256k1SignResponse, Secp256k1VerifyRequest, + Secp256k1VerifyResponse, + }, + snapshot::{ + PrepareEncryptedSnapshotRequest, PrepareEncryptedSnapshotResponse, + RestoreFromEncryptedSnapshotRequest, RestoreFromEncryptedSnapshotResponse, + }, + snapsync::{SnapSyncRequest, SnapSyncResponse}, + tx_io::{IoDecryptionRequest, IoDecryptionResponse, IoEncryptionRequest, IoEncryptionResponse}, +}; + +use super::rpc::{EnclaveApiClient, SyncEnclaveApiClient}; pub const ENCLAVE_DEFAULT_ENDPOINT_ADDR: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED); pub const ENCLAVE_DEFAULT_ENDPOINT_PORT: u16 = 7878; +static ENCLAVE_CLIENT_RUNTIME: OnceLock = OnceLock::new(); + /// A client for the enclave API. #[derive(Debug, Clone)] pub struct EnclaveClient { /// The inner HTTP client. - inner: HttpClient, + async_client: HttpClient, + /// The runtime for the client. + handle: Handle, } impl Default for EnclaveClient { @@ -27,7 +52,7 @@ impl Deref for EnclaveClient { type Target = HttpClient; fn deref(&self) -> &Self::Target { - &self.inner + &self.async_client } } @@ -37,11 +62,143 @@ impl EnclaveClient { let inner = jsonrpsee::http_client::HttpClientBuilder::default() .build(url) .unwrap(); - Self { inner } + let handle = Handle::try_current().unwrap_or_else(|_| { + let runtime = ENCLAVE_CLIENT_RUNTIME.get_or_init(|| Runtime::new().unwrap()); + runtime.handle().clone() + }); + Self { + async_client: inner, + handle, + } } - /// + /// Create a new enclave client from an address and port. pub fn new_from_addr_port(addr: impl Into, port: u16) -> Self { Self::new(format!("http://{}:{}", addr.into(), port)) } + + /// Block on a future with the runtime. + pub fn block_on_with_runtime(&self, future: F) -> T + where + F: Future, + { + tokio::task::block_in_place(|| self.handle.block_on(future)) + } +} + +macro_rules! impl_sync_client_trait { + ($(fn $method_name:ident(&self $(, $param:ident: $param_ty:ty)*) -> $return_ty:ty),* $(,)?) => { + impl SyncEnclaveApiClient for EnclaveClient { + $( + fn $method_name(&self, $($param: $param_ty),*) -> $return_ty { + self.block_on_with_runtime(self.async_client.$method_name($($param),*)) + } + )+ + } + }; +} + +impl_sync_client_trait!( + fn health_check(&self) -> Result, + fn get_public_key(&self) -> Result, + fn get_genesis_data(&self) -> Result, + fn get_snapsync_backup(&self, _req: SnapSyncRequest) -> Result, + fn sign(&self, _req: Secp256k1SignRequest) -> Result, + fn encrypt(&self, req: IoEncryptionRequest) -> Result, + fn decrypt(&self, req: IoDecryptionRequest) -> Result, + fn get_eph_rng_keypair(&self) -> Result, + fn verify(&self, _req: Secp256k1VerifyRequest) -> Result, + fn get_attestation_evidence(&self, _req: AttestationGetEvidenceRequest) -> Result, + fn eval_attestation_evidence(&self, _req: AttestationEvalEvidenceRequest) -> Result, + fn prepare_encrypted_snapshot(&self, _req: PrepareEncryptedSnapshotRequest) -> Result, + fn restore_from_encrypted_snapshot(&self, _req: RestoreFromEncryptedSnapshotRequest) -> Result, +); + +#[cfg(test)] +pub mod tests { + use crate::{get_unsecure_sample_secp256k1_pk, rpc::BuildableServer, MockEnclaveServer}; + + use super::*; + use secp256k1::{rand, Secp256k1}; + use std::{ + net::{SocketAddr, TcpListener}, + time::Duration, + }; + use tokio::time::sleep; + + #[test] + fn test_client_sync_context() { + // testing if sync client can be created in a sync runtime + let port = 1888; + let addr = SocketAddr::from((ENCLAVE_DEFAULT_ENDPOINT_ADDR, port)); + let _ = EnclaveClient::new(format!("http://{}:{}", addr.ip(), addr.port())); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_sync_client() { + // spawn a seperate thread for the server, otherwise the test will hang + let port = get_random_port(); + let addr = SocketAddr::from((ENCLAVE_DEFAULT_ENDPOINT_ADDR, port)); + println!("addr: {:?}", addr); + let _server_handle = MockEnclaveServer::new(addr).start().await.unwrap(); + let _ = sleep(Duration::from_secs(2)); + + let client = EnclaveClient::new(format!("http://{}:{}", addr.ip(), addr.port())); + sync_test_health_check(&client); + sync_test_get_public_key(&client); + sync_test_get_eph_rng_keypair(&client); + sync_test_tx_io_encrypt_decrypt(&client); + } + + pub fn get_random_port() -> u16 { + TcpListener::bind("127.0.0.1:0") // 0 means OS assigns a free port + .expect("Failed to bind to a port") + .local_addr() + .unwrap() + .port() + } + + pub fn sync_test_tx_io_encrypt_decrypt(client: &C) { + // make the request struct + let secp = Secp256k1::new(); + let (_secret_key, public_key) = secp.generate_keypair(&mut rand::thread_rng()); + let data_to_encrypt = vec![72, 101, 108, 108, 111]; + let mut nonce = vec![0u8; 4]; // 4 leading zeros + nonce.extend_from_slice(&(12345678u64).to_be_bytes()); // Append the 8-byte u64 + let encryption_request = IoEncryptionRequest { + key: public_key, + data: data_to_encrypt.clone(), + nonce: nonce.clone().into(), + }; + + // make the http request + let encryption_response = client.encrypt(encryption_request).unwrap(); + + // check the response + assert!(!encryption_response.encrypted_data.is_empty()); + + let decryption_request = IoDecryptionRequest { + key: public_key, + data: encryption_response.encrypted_data, + nonce: nonce.into(), + }; + + let decryption_response = client.decrypt(decryption_request).unwrap(); + assert_eq!(decryption_response.decrypted_data, data_to_encrypt); + } + + pub fn sync_test_health_check(client: &C) { + let resposne = client.health_check().unwrap(); + assert_eq!(resposne, "OK"); + } + + pub fn sync_test_get_public_key(client: &C) { + let res = client.get_public_key().unwrap(); + assert_eq!(res, get_unsecure_sample_secp256k1_pk()); + } + + pub fn sync_test_get_eph_rng_keypair(client: &C) { + let res = client.get_eph_rng_keypair().unwrap(); + println!("eph_rng_keypair: {:?}", res); + } } diff --git a/crates/enclave/src/client/mock.rs b/crates/enclave/src/client/mock.rs new file mode 100644 index 00000000..ea3f9864 --- /dev/null +++ b/crates/enclave/src/client/mock.rs @@ -0,0 +1,356 @@ +use std::{ + net::{IpAddr, SocketAddr}, + str::FromStr, +}; + +use anyhow::Result; +use jsonrpsee::{ + core::{async_trait, ClientError, RpcResult}, + server::ServerHandle, + Methods, +}; + +use crate::{ + coco_aa::{AttestationGetEvidenceRequest, AttestationGetEvidenceResponse}, + coco_as::{AttestationEvalEvidenceRequest, AttestationEvalEvidenceResponse}, + ecdh_decrypt, ecdh_encrypt, + genesis::GenesisDataResponse, + get_unsecure_sample_schnorrkel_keypair, get_unsecure_sample_secp256k1_pk, + get_unsecure_sample_secp256k1_sk, + signing::{ + Secp256k1SignRequest, Secp256k1SignResponse, Secp256k1VerifyRequest, + Secp256k1VerifyResponse, + }, + snapshot::{ + PrepareEncryptedSnapshotRequest, PrepareEncryptedSnapshotResponse, + RestoreFromEncryptedSnapshotRequest, RestoreFromEncryptedSnapshotResponse, + }, + snapsync::{SnapSyncRequest, SnapSyncResponse}, + tx_io::{IoDecryptionRequest, IoDecryptionResponse, IoEncryptionRequest, IoEncryptionResponse}, +}; + +use super::{ + rpc::{BuildableServer, EnclaveApiServer, SyncEnclaveApiClient}, + ENCLAVE_DEFAULT_ENDPOINT_ADDR, ENCLAVE_DEFAULT_ENDPOINT_PORT, +}; + +pub struct MockEnclaveServer { + addr: SocketAddr, +} + +impl MockEnclaveServer { + pub fn new(addr: impl Into) -> Self { + Self { addr: addr.into() } + } + + pub fn new_from_addr_port(addr: String, port: u16) -> Self { + Self::new((IpAddr::from_str(&addr).unwrap(), port)) + } + + /// Mock implementation of the health check method. + pub fn health_check() -> String { + "OK".to_string() + } + + /// Mock implementation of the get_eph_rng_keypair method. + pub fn get_eph_rng_keypair() -> schnorrkel::keys::Keypair { + // Return a sample Schnorrkel keypair for testing + get_unsecure_sample_schnorrkel_keypair() + } + + /// Mock implementation of the encrypt method. + pub fn encrypt(req: IoEncryptionRequest) -> IoEncryptionResponse { + // Use the sample secret key for encryption + let encrypted_data = ecdh_encrypt( + &req.key, + &get_unsecure_sample_secp256k1_sk(), + &req.data, + req.nonce, + ) + .unwrap(); + + IoEncryptionResponse { encrypted_data } + } + + /// Mock implementation of the decrypt method. + pub fn decrypt(req: IoDecryptionRequest) -> IoDecryptionResponse { + // Use the sample secret key for decryption + let decrypted_data = ecdh_decrypt( + &req.key, + &get_unsecure_sample_secp256k1_sk(), + &req.data, + req.nonce, + ) + .unwrap(); + + IoDecryptionResponse { decrypted_data } + } + + /// Mock implementation of the get_public_key method. + pub fn get_public_key() -> secp256k1::PublicKey { + get_unsecure_sample_secp256k1_pk() + } + + /// Mock implementation of the sign method. + pub fn sign(_req: Secp256k1SignRequest) -> Secp256k1SignResponse { + unimplemented!("sign not implemented for mock server") + } + + /// Mock implementation of the verify method. + pub fn verify(_req: Secp256k1VerifyRequest) -> Secp256k1VerifyResponse { + unimplemented!("verify not implemented for mock server") + } + + /// Mock implementation of the get_genesis_data method. + pub fn get_genesis_data() -> GenesisDataResponse { + unimplemented!("get_genesis_data not implemented for mock server") + } + + /// Mock implementation of the get_snapsync_backup method. + pub fn get_snapsync_backup(_req: SnapSyncRequest) -> SnapSyncResponse { + unimplemented!("get_snapsync_backup not implemented for mock server") + } + + /// Mock implementation of the get_attestation_evidence method. + pub fn get_attestation_evidence( + _req: AttestationGetEvidenceRequest, + ) -> AttestationGetEvidenceResponse { + unimplemented!("get_attestation_evidence not implemented for mock server") + } + + /// Mock implementation of the eval_attestation_evidence method. + pub fn eval_attestation_evidence( + _req: AttestationEvalEvidenceRequest, + ) -> AttestationEvalEvidenceResponse { + unimplemented!("eval_attestation_evidence not implemented for mock server") + } + + /// Mock implementation of the prepare_encrypted_snapshot method. + pub fn prepare_encrypted_snapshot( + _req: PrepareEncryptedSnapshotRequest, + ) -> PrepareEncryptedSnapshotResponse { + unimplemented!("prepare_encrypted_snapshot not implemented for mock server") + } + + /// Mock implementation of the restore_from_encrypted_snapshot method. + pub fn restore_from_encrypted_snapshot( + _req: RestoreFromEncryptedSnapshotRequest, + ) -> RestoreFromEncryptedSnapshotResponse { + unimplemented!("restore_from_encrypted_snapshot not implemented for mock server") + } +} + +impl Default for MockEnclaveServer { + fn default() -> Self { + Self::new((ENCLAVE_DEFAULT_ENDPOINT_ADDR, ENCLAVE_DEFAULT_ENDPOINT_PORT)) + } +} + +impl BuildableServer for MockEnclaveServer { + fn addr(&self) -> SocketAddr { + self.addr + } + + fn methods(self) -> Methods { + self.into_rpc().into() + } + + async fn start(self) -> Result { + BuildableServer::start_rpc_server(self).await + } +} + +#[async_trait] +impl EnclaveApiServer for MockEnclaveServer { + /// Handler for: `getPublicKey` + async fn get_public_key(&self) -> RpcResult { + Ok(MockEnclaveServer::get_public_key()) + } + + /// Handler for: `healthCheck` + async fn health_check(&self) -> RpcResult { + Ok(MockEnclaveServer::health_check()) + } + + /// Handler for: `getGenesisData` + async fn get_genesis_data(&self) -> RpcResult { + Ok(MockEnclaveServer::get_genesis_data()) + } + + /// Handler for: `getSnapsyncBackup` + async fn get_snapsync_backup(&self, request: SnapSyncRequest) -> RpcResult { + Ok(MockEnclaveServer::get_snapsync_backup(request)) + } + + /// Handler for: `encrypt` + async fn encrypt(&self, req: IoEncryptionRequest) -> RpcResult { + Ok(MockEnclaveServer::encrypt(req)) + } + + /// Handler for: `decrypt` + async fn decrypt(&self, req: IoDecryptionRequest) -> RpcResult { + Ok(MockEnclaveServer::decrypt(req)) + } + + /// Handler for: `getAttestationEvidence` + async fn get_attestation_evidence( + &self, + req: AttestationGetEvidenceRequest, + ) -> RpcResult { + Ok(MockEnclaveServer::get_attestation_evidence(req)) + } + + /// Handler for: `evalAttestationEvidence` + async fn eval_attestation_evidence( + &self, + req: AttestationEvalEvidenceRequest, + ) -> RpcResult { + Ok(MockEnclaveServer::eval_attestation_evidence(req)) + } + + /// Handler for: `sign` + async fn sign(&self, req: Secp256k1SignRequest) -> RpcResult { + Ok(MockEnclaveServer::sign(req)) + } + + /// Handler for: `verify` + async fn verify(&self, req: Secp256k1VerifyRequest) -> RpcResult { + Ok(MockEnclaveServer::verify(req)) + } + + /// Handler for: 'eph_rng.get_keypair' + async fn get_eph_rng_keypair(&self) -> RpcResult { + Ok(MockEnclaveServer::get_eph_rng_keypair()) + } + + /// Handler for: 'snapshot.prepare_encrypted_snapshot' + async fn prepare_encrypted_snapshot( + &self, + req: PrepareEncryptedSnapshotRequest, + ) -> RpcResult { + Ok(MockEnclaveServer::prepare_encrypted_snapshot(req)) + } + + /// Handler for: 'snapshot.restore_from_encrypted_snapshot' + async fn restore_from_encrypted_snapshot( + &self, + req: RestoreFromEncryptedSnapshotRequest, + ) -> RpcResult { + Ok(MockEnclaveServer::restore_from_encrypted_snapshot(req)) + } +} + +pub struct MockEnclaveClient; +impl MockEnclaveClient { + pub fn new() -> Self { + Self {} + } +} + +macro_rules! impl_mock_sync_client_trait { + ($(fn $method_name:ident(&self $(, $param:ident: $param_ty:ty)*) -> $return_ty:ty),* $(,)?) => { + impl SyncEnclaveApiClient for MockEnclaveClient { + $( + fn $method_name(&self, $($param: $param_ty),*) -> $return_ty { + Ok(MockEnclaveServer::$method_name($($param),*)) + } + )+ + } + }; +} + +impl_mock_sync_client_trait!( + fn health_check(&self) -> Result, + fn get_public_key(&self) -> Result, + fn get_genesis_data(&self) -> Result, + fn get_snapsync_backup(&self, _req: SnapSyncRequest) -> Result, + fn sign(&self, _req: Secp256k1SignRequest) -> Result, + fn encrypt(&self, req: IoEncryptionRequest) -> Result, + fn decrypt(&self, req: IoDecryptionRequest) -> Result, + fn get_eph_rng_keypair(&self) -> Result, + fn verify(&self, _req: Secp256k1VerifyRequest) -> Result, + fn get_attestation_evidence(&self, _req: AttestationGetEvidenceRequest) -> Result, + fn eval_attestation_evidence(&self, _req: AttestationEvalEvidenceRequest) -> Result, + fn prepare_encrypted_snapshot(&self, _req: PrepareEncryptedSnapshotRequest) -> Result, + fn restore_from_encrypted_snapshot(&self, _req: RestoreFromEncryptedSnapshotRequest) -> Result, +); + +#[cfg(test)] +mod tests { + use std::{ops::Deref, time::Duration}; + + use secp256k1::{rand, Secp256k1}; + use tokio::time::sleep; + + use super::*; + use crate::{client::tests::*, rpc::EnclaveApiClient, EnclaveClient}; + + #[test] + fn test_mock_client() { + let client = MockEnclaveClient {}; + sync_test_health_check(&client); + sync_test_get_public_key(&client); + sync_test_get_eph_rng_keypair(&client); + sync_test_tx_io_encrypt_decrypt(&client); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_mock_server_and_sync_client() { + // spawn a seperate thread for the server, otherwise the test will hang + let port = get_random_port(); + let addr = SocketAddr::from((ENCLAVE_DEFAULT_ENDPOINT_ADDR, port)); + println!("addr: {:?}", addr); + let _server_handle = MockEnclaveServer::new(addr).start().await.unwrap(); + let _ = sleep(Duration::from_secs(2)); + + let client = EnclaveClient::new(format!("http://{}:{}", addr.ip(), addr.port())); + async_test_health_check(&client).await; + async_test_get_public_key(&client).await; + async_test_get_eph_rng_keypair(&client).await; + async_test_tx_io_encrypt_decrypt(&client).await; + } + + async fn async_test_tx_io_encrypt_decrypt(client: &EnclaveClient) { + // make the request struct + let secp = Secp256k1::new(); + let (_secret_key, public_key) = secp.generate_keypair(&mut rand::thread_rng()); + let data_to_encrypt = vec![72, 101, 108, 108, 111]; + let mut nonce = vec![0u8; 4]; // 4 leading zeros + nonce.extend_from_slice(&(12345678u64).to_be_bytes()); // Append the 8-byte u64 + let encryption_request = IoEncryptionRequest { + key: public_key, + data: data_to_encrypt.clone(), + nonce: nonce.clone().into(), + }; + + // make the http request + let encryption_response = client.deref().encrypt(encryption_request).await.unwrap(); + + // check the response + assert!(!encryption_response.encrypted_data.is_empty()); + + let decryption_request = IoDecryptionRequest { + key: public_key, + data: encryption_response.encrypted_data, + nonce: nonce.into(), + }; + + let decryption_response = client.decrypt(decryption_request).unwrap(); + assert_eq!(decryption_response.decrypted_data, data_to_encrypt); + } + + async fn async_test_health_check(client: &EnclaveClient) { + let resposne = client.deref().health_check().await.unwrap(); + assert_eq!(resposne, "OK"); + } + + async fn async_test_get_public_key(client: &EnclaveClient) { + let res = client.deref().get_public_key().await.unwrap(); + assert_eq!(res, get_unsecure_sample_secp256k1_pk()); + } + + async fn async_test_get_eph_rng_keypair(client: &EnclaveClient) { + let res = client.deref().get_eph_rng_keypair().await.unwrap(); + println!("eph_rng_keypair: {:?}", res); + } +} diff --git a/crates/enclave/src/client/mock_server.rs b/crates/enclave/src/client/mock_server.rs deleted file mode 100644 index bf3c2d18..00000000 --- a/crates/enclave/src/client/mock_server.rs +++ /dev/null @@ -1,154 +0,0 @@ -use std::{ - net::{IpAddr, SocketAddr}, - str::FromStr, -}; - -use anyhow::Result; -use jsonrpsee::{ - core::{async_trait, RpcResult}, - server::ServerHandle, - Methods, -}; - -use crate::{ - coco_aa::{AttestationGetEvidenceRequest, AttestationGetEvidenceResponse}, - coco_as::{AttestationEvalEvidenceRequest, AttestationEvalEvidenceResponse}, - ecdh_decrypt, ecdh_encrypt, - genesis::GenesisDataResponse, - get_unsecure_sample_schnorrkel_keypair, get_unsecure_sample_secp256k1_pk, - get_unsecure_sample_secp256k1_sk, rpc_invalid_ciphertext_error, - signing::{ - Secp256k1SignRequest, Secp256k1SignResponse, Secp256k1VerifyRequest, - Secp256k1VerifyResponse, - }, - snapshot::{ - PrepareEncryptedSnapshotRequest, PrepareEncryptedSnapshotResponse, - RestoreFromEncryptedSnapshotRequest, RestoreFromEncryptedSnapshotResponse, - }, - snapsync::{SnapSyncRequest, SnapSyncResponse}, - tx_io::{IoDecryptionRequest, IoDecryptionResponse, IoEncryptionRequest, IoEncryptionResponse}, -}; - -use super::{ - rpc::{BuildableServer, EnclaveApiServer}, - ENCLAVE_DEFAULT_ENDPOINT_ADDR, ENCLAVE_DEFAULT_ENDPOINT_PORT, -}; - -pub struct MockEnclaveServer { - addr: SocketAddr, -} - -impl MockEnclaveServer { - pub fn new(addr: impl Into) -> Self { - Self { addr: addr.into() } - } - - pub fn new_from_addr_port(addr: String, port: u16) -> Self { - Self::new((IpAddr::from_str(&addr).unwrap(), port)) - } -} - -impl Default for MockEnclaveServer { - fn default() -> Self { - Self::new((ENCLAVE_DEFAULT_ENDPOINT_ADDR, ENCLAVE_DEFAULT_ENDPOINT_PORT)) - } -} - -impl BuildableServer for MockEnclaveServer { - fn addr(&self) -> SocketAddr { - self.addr - } - - fn methods(self) -> Methods { - self.into_rpc().into() - } - - async fn start(self) -> Result { - BuildableServer::start_rpc_server(self).await - } -} - -#[async_trait] -impl EnclaveApiServer for MockEnclaveServer { - async fn get_public_key(&self) -> RpcResult { - Ok(get_unsecure_sample_secp256k1_pk()) - } - - async fn health_check(&self) -> RpcResult { - Ok("OK".to_string()) - } - - async fn get_genesis_data(&self) -> RpcResult { - unimplemented!("genesis_get_data not implemented for mock server") - } - - async fn get_snapsync_backup(&self, _request: SnapSyncRequest) -> RpcResult { - unimplemented!("provide_snapsync_backup not implemented for mock server") - } - - async fn sign(&self, _req: Secp256k1SignRequest) -> RpcResult { - unimplemented!("secp256k1_sign not implemented for mock server") - } - - async fn verify(&self, _req: Secp256k1VerifyRequest) -> RpcResult { - unimplemented!("secp256k1_verify not implemented for mock server") - } - - async fn get_attestation_evidence( - &self, - _req: AttestationGetEvidenceRequest, - ) -> RpcResult { - unimplemented!("attestation_get_evidence not implemented for mock server") - } - - async fn eval_attestation_evidence( - &self, - _req: AttestationEvalEvidenceRequest, - ) -> RpcResult { - unimplemented!("attestation_eval_evidence not implemented for mock server") - } - - async fn encrypt(&self, request: IoEncryptionRequest) -> RpcResult { - // load key and encrypt data - let encrypted_data = ecdh_encrypt( - &request.key, - &get_unsecure_sample_secp256k1_sk(), - &request.data, - request.nonce, - ) - .unwrap(); - - Ok(IoEncryptionResponse { encrypted_data }) - } - - async fn decrypt(&self, request: IoDecryptionRequest) -> RpcResult { - // load key and decrypt data - let decrypted_data = ecdh_decrypt( - &request.key, - &get_unsecure_sample_secp256k1_sk(), - &request.data, - request.nonce, - ) - .map_err(|e| rpc_invalid_ciphertext_error(e))?; - - Ok(IoDecryptionResponse { decrypted_data }) - } - - async fn get_eph_rng_keypair(&self) -> RpcResult { - Ok(get_unsecure_sample_schnorrkel_keypair()) - } - - async fn prepare_encrypted_snapshot( - &self, - _request: PrepareEncryptedSnapshotRequest, - ) -> RpcResult { - unimplemented!("prepare_encrypted_snapshot not implemented for mock server") - } - - async fn restore_from_encrypted_snapshot( - &self, - _request: RestoreFromEncryptedSnapshotRequest, - ) -> RpcResult { - unimplemented!("restore_from_encrypted_snapshot not implemented for mock server") - } -} diff --git a/crates/enclave/src/client/mod.rs b/crates/enclave/src/client/mod.rs index cce1bf07..fed4b215 100644 --- a/crates/enclave/src/client/mod.rs +++ b/crates/enclave/src/client/mod.rs @@ -5,8 +5,8 @@ //! traits define the API and implementation for the TEE client. #![allow(async_fn_in_trait)] pub mod client; -pub mod mock_server; +pub mod mock; pub mod rpc; pub use client::*; -pub use mock_server::*; +pub use mock::*; diff --git a/crates/enclave/src/client/rpc.rs b/crates/enclave/src/client/rpc.rs index 4f015f8c..94cfc966 100644 --- a/crates/enclave/src/client/rpc.rs +++ b/crates/enclave/src/client/rpc.rs @@ -6,6 +6,7 @@ use jsonrpsee::core::RpcResult; use jsonrpsee::proc_macros::rpc; use jsonrpsee::server::{ServerBuilder, ServerHandle}; use jsonrpsee::Methods; +use seismic_enclave_derive::derive_sync_client_trait; use crate::coco_aa::{AttestationGetEvidenceRequest, AttestationGetEvidenceResponse}; use crate::coco_as::{AttestationEvalEvidenceRequest, AttestationEvalEvidenceResponse}; @@ -40,7 +41,8 @@ pub trait BuildableServer { } } -#[rpc(client, server)] +#[derive_sync_client_trait] // get SyncEnclaveApi trait +#[rpc(client, server)] // get EnclaveApiClient EnclaveApiServer trait pub trait EnclaveApi { /// Health check endpoint that returns "OK" if service is running #[method(name = "healthCheck")] @@ -56,37 +58,37 @@ pub trait EnclaveApi { /// Provides backup data for snapshot synchronization #[method(name = "getSnapsyncBackup")] - async fn get_snapsync_backup(&self, request: SnapSyncRequest) -> RpcResult; + async fn get_snapsync_backup(&self, _req: SnapSyncRequest) -> RpcResult; /// Signs a message using secp256k1 private key #[method(name = "sign")] - async fn sign(&self, req: Secp256k1SignRequest) -> RpcResult; + async fn sign(&self, _req: Secp256k1SignRequest) -> RpcResult; /// Verifies a secp256k1 signature against a message #[method(name = "verify")] - async fn verify(&self, req: Secp256k1VerifyRequest) -> RpcResult; + async fn verify(&self, _req: Secp256k1VerifyRequest) -> RpcResult; /// Generates attestation evidence from the attestation authority #[method(name = "getAttestationEvidence")] async fn get_attestation_evidence( &self, - req: AttestationGetEvidenceRequest, + _req: AttestationGetEvidenceRequest, ) -> RpcResult; /// Evaluates provided attestation evidence #[method(name = "evalAttestationEvidence")] async fn eval_attestation_evidence( &self, - req: AttestationEvalEvidenceRequest, + _req: AttestationEvalEvidenceRequest, ) -> RpcResult; /// Encrypts transaction data using ECDH and AES #[method(name = "encrypt")] - async fn encrypt(&self, req: IoEncryptionRequest) -> RpcResult; + async fn encrypt(&self, _req: IoEncryptionRequest) -> RpcResult; /// Decrypts transaction data using ECDH and AES #[method(name = "decrypt")] - async fn decrypt(&self, req: IoDecryptionRequest) -> RpcResult; + async fn decrypt(&self, _req: IoDecryptionRequest) -> RpcResult; /// Generates an ephemeral keypair #[method(name = "eph_rng.get_keypair")] diff --git a/crates/enclave/tests/lib.rs b/crates/enclave/tests/lib.rs index 56e6286c..8b137891 100644 --- a/crates/enclave/tests/lib.rs +++ b/crates/enclave/tests/lib.rs @@ -1 +1 @@ -mod mock_server; + diff --git a/crates/enclave/tests/mock_server.rs b/crates/enclave/tests/mock_server.rs deleted file mode 100644 index 2f89d743..00000000 --- a/crates/enclave/tests/mock_server.rs +++ /dev/null @@ -1,90 +0,0 @@ -#[cfg(test)] -use seismic_enclave::client::rpc::BuildableServer; -use seismic_enclave::client::EnclaveClient; -use seismic_enclave::client::ENCLAVE_DEFAULT_ENDPOINT_ADDR; -use seismic_enclave::get_unsecure_sample_secp256k1_pk; -use seismic_enclave::request_types::tx_io::*; -use seismic_enclave::rpc::EnclaveApiClient; -use seismic_enclave::MockEnclaveServer; -use std::net::SocketAddr; -use std::net::TcpListener; -use std::thread::sleep; -use std::time::Duration; - -fn get_random_port() -> u16 { - TcpListener::bind("127.0.0.1:0") // 0 means OS assigns a free port - .expect("Failed to bind to a port") - .local_addr() - .unwrap() - .port() -} - -async fn test_tx_io_encrypt_decrypt(client: &EnclaveClient) { - // make the request struct - let data_to_encrypt = vec![72, 101, 108, 108, 111]; - let mut nonce = vec![0u8; 4]; // 4 leading zeros - nonce.extend_from_slice(&(12345678u64).to_be_bytes()); // Append the 8-byte u64 - let encryption_request = IoEncryptionRequest { - key: get_unsecure_sample_secp256k1_pk(), - data: data_to_encrypt.clone(), - nonce: nonce.clone().into(), - }; - - // make the http request - let encryption_response = client.encrypt(encryption_request).await.unwrap(); - - // check the response - assert!(!encryption_response.encrypted_data.is_empty()); - - let decryption_request = IoDecryptionRequest { - key: get_unsecure_sample_secp256k1_pk(), - data: encryption_response.encrypted_data, - nonce: nonce.into(), - }; - - let decryption_response = client.decrypt(decryption_request).await.unwrap(); - assert_eq!(decryption_response.decrypted_data, data_to_encrypt); -} - -async fn test_health_check(client: &EnclaveClient) { - let resposne = client.health_check().await.unwrap(); - assert_eq!(resposne, "OK"); -} - -async fn test_get_public_key(client: &EnclaveClient) { - let res = client.get_public_key().await.unwrap(); - assert_eq!(res, get_unsecure_sample_secp256k1_pk()); -} - -async fn test_get_eph_rng_keypair(client: &EnclaveClient) { - let res = client.get_eph_rng_keypair().await.unwrap(); - println!("eph_rng_keypair: {:?}", res); -} - -#[tokio::test] -async fn test_mock_server() { - // spawn a seperate thread for the server, otherwise the test will hang - let port = get_random_port(); - let addr = SocketAddr::from((ENCLAVE_DEFAULT_ENDPOINT_ADDR, port)); - println!("addr: {:?}", addr); - let _server_handle = MockEnclaveServer::new(addr).start().await.unwrap(); - sleep(Duration::from_secs(4)); - let client = EnclaveClient::new(format!("http://{}:{}", addr.ip(), addr.port())); - println!("client: {:?}", client); - - test_health_check(&client).await; - test_tx_io_encrypt_decrypt(&client).await; - test_get_public_key(&client).await; - test_get_eph_rng_keypair(&client).await; - - let client = EnclaveClient::new(format!("http://{}:{}", addr.ip(), addr.port())); - - let handle = tokio::spawn(async move { - println!("client 2: {:?}", client); - test_health_check(&client).await; - test_tx_io_encrypt_decrypt(&client).await; - test_get_public_key(&client).await; - test_get_eph_rng_keypair(&client).await; - }); - handle.await.unwrap(); -}