diff --git a/kalypso-cli/Cargo.toml b/kalypso-cli/Cargo.toml index b3e1742..1e67b6a 100644 --- a/kalypso-cli/Cargo.toml +++ b/kalypso-cli/Cargo.toml @@ -14,6 +14,7 @@ env_logger = "0.11" flate2 = "1.0.28" futures = "0.3" generator_client = {path = "../generator_client", package = "generator_client"} +generator = {path = "../generator", package = "generator"} hex = "0.4.3" kalypso_helper = {path = "../helper", package = "helper"} log = "0.4" diff --git a/kalypso-cli/src/common_deps.rs b/kalypso-cli/src/common_deps.rs index c20c4a7..bdb1525 100644 --- a/kalypso-cli/src/common_deps.rs +++ b/kalypso-cli/src/common_deps.rs @@ -1266,3 +1266,35 @@ impl CommonDeps { }) } } + +pub struct TestEnclaveConnectionInfo { + pub enclave_client_url: String, +} + +impl CommonDeps { + pub fn test_connection_info( + config: &std::collections::HashMap, + ) -> Result { + get_config_ref!(config, "enclave_client_url", enclave_client_url); + + Ok(TestEnclaveConnectionInfo { + enclave_client_url: enclave_client_url.to_string(), + }) + } +} + +pub struct BenchmarkInfo { + pub benchmark_url: String, +} + +impl CommonDeps { + pub fn benchmark_info( + config: &std::collections::HashMap, + ) -> Result { + get_config_ref!(config, "benchmark_url", benchmark_url); + + Ok(BenchmarkInfo { + benchmark_url: benchmark_url.to_string(), + }) + } +} diff --git a/kalypso-cli/src/config.json b/kalypso-cli/src/config.json index 060dcb5..1b85ace 100644 --- a/kalypso-cli/src/config.json +++ b/kalypso-cli/src/config.json @@ -276,6 +276,12 @@ "prompt": "name of the program to start in the enclave", "secret": false, "env_var": "PROVER_PROGRAM_NAME" + }, + { + "field": "benchmark_url", + "prompt": "Enter url of the server that does benchmarking and returns benchmark result", + "secret": false, + "env_var": "BENCHMARK_URL" } ], "operations": [ @@ -538,6 +544,20 @@ "generator_client_url", "prover_program_name" ] + }, + { + "name": "Test Enclave Connection", + "description": "Test Enclave Connection", + "required_prompts": [ + "enclave_client_url" + ] + }, + { + "name": "Benchmark Prover", + "description": "Get Benchmark results of prover", + "required_prompts": [ + "benchmark_url" + ] } ] } diff --git a/kalypso-cli/src/operations/mod.rs b/kalypso-cli/src/operations/mod.rs index cb7dc46..e01fb8f 100644 --- a/kalypso-cli/src/operations/mod.rs +++ b/kalypso-cli/src/operations/mod.rs @@ -79,6 +79,8 @@ pub fn get_operation(name: &str) -> Option> { "Load Generator Config" => Some(Box::new(generator_config::GeneratorConfig)), "Start Enclave Program" => Some(Box::new(programs::StartProgam)), "Stop Enclave Program" => Some(Box::new(programs::StopProgram)), + "Test Enclave Connection" => Some(Box::new(programs::TestConnection)), + "Benchmark Prover" => Some(Box::new(programs::Benchmark)), _ => unimplemented!(), } } diff --git a/kalypso-cli/src/operations/programs.rs b/kalypso-cli/src/operations/programs.rs index 838250a..69ba1ba 100644 --- a/kalypso-cli/src/operations/programs.rs +++ b/kalypso-cli/src/operations/programs.rs @@ -6,6 +6,36 @@ use crate::common_deps::CommonDeps; use super::Operation; +pub struct Benchmark; + +#[async_trait] +impl Operation for Benchmark { + async fn execute(&self, config: HashMap) -> Result<(), String> { + let benchmark_info = CommonDeps::benchmark_info(&config)?; + + benchmark(benchmark_info.benchmark_url) + .await + .map_err(|e| format!("Failed to benchmark prover: {}", e))?; + + Ok(()) + } +} + +pub struct TestConnection; + +#[async_trait] +impl Operation for TestConnection { + async fn execute(&self, config: HashMap) -> Result<(), String> { + let test_connection_info = CommonDeps::test_connection_info(&config)?; + + test_connection(test_connection_info.enclave_client_url) + .await + .map_err(|e| format!("Failed to connect enclave client: {}", e))?; + + Ok(()) + } +} + pub struct StopProgram; #[async_trait] @@ -157,3 +187,69 @@ async fn stop_program( .into()) } } + +async fn test_connection(enclave_client_url: String) -> Result<(), Box> { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(10)) + .build()?; + let full_url = format!("{}/api/test", enclave_client_url); + + // Prepare the headers + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(reqwest::header::CONTENT_TYPE, "application/json".parse()?); + + // Send the POST request + let response = client.get(&full_url).headers(headers).send().await?; + + // Capture the HTTP status before moving the response + let status = response.status(); + + // Check the HTTP status and the JSON message + if status.is_success() { + Ok(()) + } else { + // Return error with message + Err(format!("Failed to test connection. HTTP Status: {}", status).into()) + } +} + +async fn benchmark(benchmark_url: String) -> Result<(), Box> { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(10)) + .build()?; + let full_url = format!("{}/api/benchmark", benchmark_url); + + // Prepare the headers + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(reqwest::header::CONTENT_TYPE, "application/json".parse()?); + + // Send the POST request + let response = client.get(&full_url).headers(headers).send().await?; + + // Capture the HTTP status before moving the response + let status = response.status(); + + // Read the response body as text + let response_text = response.text().await?; + + let json_response: generator::models::BenchmarkResponse = + match serde_json::from_str(&response_text) { + Ok(resp) => resp, + Err(e) => { + return Err(format!( + "Failed to parse JSON response. HTTP Status: {}, Body: {}, Error: {}", + status, response_text, e + ) + .into()) + } + }; + + // Check the HTTP status and the JSON message + if status.is_success() { + println!("{:?}", json_response); + Ok(()) + } else { + // Return error with message + Err(format!("Failed to benchmark prover. HTTP Status: {}", status).into()) + } +} diff --git a/kalypso-cli/src/prompts.rs b/kalypso-cli/src/prompts.rs index 3d0804f..5aee1a3 100644 --- a/kalypso-cli/src/prompts.rs +++ b/kalypso-cli/src/prompts.rs @@ -84,6 +84,7 @@ impl<'a> Prompter<'a> { validators.insert("internal_prover_port".to_string(), validate_dec_str_id); validators.insert("input_verification_url".to_string(), validate_rpc_url); validators.insert("prover_program_name".to_string(), validate_string); + validators.insert("benchmark_url".to_string(), validate_rpc_url); validators.insert( "confirmation".to_string(),