diff --git a/examples/quickstart/Cargo.toml b/examples/quickstart/Cargo.toml index 7e1ba570c..73b4ce80e 100644 --- a/examples/quickstart/Cargo.toml +++ b/examples/quickstart/Cargo.toml @@ -7,5 +7,7 @@ license = "Apache-2.0" edition = "2018" [dependencies] +structopt = "0.2" +serde = "1.0.92" mesatee_sdk = { path = "../../mesatee_sdk" } lazy_static = { version = "1.1.0" } diff --git a/examples/quickstart/src/main.rs b/examples/quickstart/src/main.rs index 8606cba7c..e9bc40c3b 100644 --- a/examples/quickstart/src/main.rs +++ b/examples/quickstart/src/main.rs @@ -16,93 +16,126 @@ // under the License. use lazy_static::lazy_static; -use mesatee_sdk::{Error, ErrorKind, Mesatee, MesateeEnclaveInfo, Result, TaskStatus}; +use mesatee_sdk::{Mesatee, MesateeEnclaveInfo, TaskStatus}; use std::net::SocketAddr; +use std::path; +use std::path::PathBuf; use std::{thread, time}; +use structopt::StructOpt; lazy_static! { static ref TMS_ADDR: SocketAddr = "127.0.0.1:5554".parse().unwrap(); static ref TDFS_ADDR: SocketAddr = "127.0.0.1:5065".parse().unwrap(); } -static PSI_TEST_DATA_1_PATH: &str = "./test_data/psi1.txt"; -static PSI_TEST_DATA_2_PATH: &str = "./test_data/psi2.txt"; -fn single_party_task(enclave_info: &MesateeEnclaveInfo) -> Result<()> { - println!("[+] This is a single-party task: echo"); +#[derive(Debug, StructOpt)] +struct EchoOpt { + #[structopt(short = "e", required = true)] + enclave_info: PathBuf, - let mesatee = Mesatee::new(enclave_info, "uid1", "token1", *TMS_ADDR, *TDFS_ADDR)?; - let task = mesatee.create_task("echo")?; - let ret = task.invoke_with_payload("haha")?; - println!("{}", ret); + #[structopt(short = "m", required = true)] + message: String, +} + +#[derive(Debug, StructOpt)] +struct PsiOpt { + #[structopt(short = "e", required = true)] + enclave_info: PathBuf, + + #[structopt(short = "f", required = true)] + files: Vec, +} - Ok(()) +#[derive(Debug, StructOpt)] +enum Command { + /// Echo + #[structopt(name = "echo")] + Echo(EchoOpt), + + /// Private set intersection + #[structopt(name = "psi")] + Psi(PsiOpt), +} + +#[derive(Debug, StructOpt)] +/// Quickstart example. +struct Cli { + #[structopt(subcommand)] + command: Command, } -fn multi_party_task(enclave_info: &MesateeEnclaveInfo) -> Result<()> { - println!("[+] This is a multi-party task: psi"); +fn echo(args: EchoOpt) { + println!("[+] Invoke echo function"); + let auditors = vec![]; // legacy + + let enclave_info = + MesateeEnclaveInfo::load(auditors, args.enclave_info.to_str().unwrap()).expect("load"); + + let mesatee = + Mesatee::new(&enclave_info, "uid1", "token1", *TMS_ADDR, *TDFS_ADDR).expect("new"); + let task = mesatee.create_task("echo").expect("create"); + let response = task.invoke_with_payload(&args.message).expect("invoke"); + println!("{}", response); +} + +fn psi(args: PsiOpt) { + println!("[+] Invoke psi function"); + let auditors = vec![]; // legacy + + let enclave_info = + MesateeEnclaveInfo::load(auditors, args.enclave_info.to_str().unwrap()).expect("load"); // Party 1 creates one PSI task - let mesatee1 = Mesatee::new(enclave_info, "uid1", "token1", *TMS_ADDR, *TDFS_ADDR)?; - let file1_id = mesatee1.upload_file(PSI_TEST_DATA_1_PATH)?; - let task = mesatee1.create_task_with_collaborators("psi", &["uid2"], &[file1_id.as_str()])?; + let mesatee1 = + Mesatee::new(&enclave_info, "uid1", "token1", *TMS_ADDR, *TDFS_ADDR).expect("new"); + let file1_id = mesatee1 + .upload_file(args.files[0].to_str().unwrap()) + .expect("upload_file"); + let task = mesatee1 + .create_task_with_collaborators("psi", &["uid2"], &[file1_id.as_str()]) + .expect("create_task"); // Party 2 approves the task and invokes the task let task_id = task.task_id.clone(); - let mesatee2 = Mesatee::new(enclave_info, "uid2", "token2", *TMS_ADDR, *TDFS_ADDR)?; - let file2_id = mesatee2.upload_file(PSI_TEST_DATA_2_PATH)?; - mesatee2.approve_task_with_files(&task_id, &[file2_id.as_str()])?; + let mesatee2 = + Mesatee::new(&enclave_info, "uid2", "token2", *TMS_ADDR, *TDFS_ADDR).expect("new"); + let file2_id = mesatee2 + .upload_file(args.files[1].to_str().unwrap()) + .expect("upload_file"); + mesatee2 + .approve_task_with_files(&task_id, &[file2_id.as_str()]) + .expect("aprrove_task_with_files"); - let _ = task.invoke()?; + let _ = task.invoke().expect("invoke"); // Party 1 waits for PSI results and get results from trusted FS let mut task_info = mesatee1 - .get_task(&task_id)? + .get_task(&task_id) + .expect("get_task") .task_info - .ok_or_else(|| Error::from(ErrorKind::MissingValue))?; + .expect("task_info"); while task_info.status != TaskStatus::Finished { task_info = mesatee1 - .get_task(&task_id)? + .get_task(&task_id) + .expect("get_task") .task_info - .ok_or_else(|| Error::from(ErrorKind::MissingValue))?; + .expect("task_info"); thread::sleep(time::Duration::from_secs(1)); } - let results = mesatee1.get_task_results(&task_id)?; + let results = mesatee1 + .get_task_results(&task_id) + .expect("get_task_results"); println!("{:?}", results); - let content = mesatee1.get_file(&results[0]); + let content = mesatee1.get_file(&results[0]).expect("get_file"); println!("{:?}", content); - - Ok(()) } fn main() { - // Load auditors' public keys and endorsement to TEE enclaves (digital signatures) - let auditors = vec![ - ( - "../services/auditors/godzilla/godzilla.public.der", - "../services/auditors/godzilla/godzilla.sign.sha256", - ), - ( - "../services/auditors/optimus_prime/optimus_prime.public.der", - "../services/auditors/optimus_prime/optimus_prime.sign.sha256", - ), - ( - "../services/auditors/albus_dumbledore/albus_dumbledore.public.der", - "../services/auditors/albus_dumbledore/albus_dumbledore.sign.sha256", - ), - ]; - let enclave_info_file_path = "../services/enclave_info.toml"; - - let mesatee_enclave_info = MesateeEnclaveInfo::load(auditors, enclave_info_file_path).unwrap(); - - match single_party_task(&mesatee_enclave_info) { - Ok(_) => println!("[+] successfully invoke single-party task echo"), - Err(e) => println!("[-] single-party task echo error: {}", e), - } - - match multi_party_task(&mesatee_enclave_info) { - Ok(_) => println!("[+] successfully invoke multi-party task psi"), - Err(e) => println!("[-] multi-party task psi error: {}", e), + let args = Cli::from_args(); + match args.command { + Command::Echo(echo_args) => echo(echo_args), + Command::Psi(psi_args) => psi(psi_args), } }