Skip to content

Commit

Permalink
Rewrite quickstart example command line interface (#186)
Browse files Browse the repository at this point in the history
  • Loading branch information
@mssun
mssun committed Dec 20, 2019
1 parent 96984eb commit 708831d
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 54 deletions.
2 changes: 2 additions & 0 deletions examples/quickstart/Cargo.toml
View file
Expand Up @@ -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" }
141 changes: 87 additions & 54 deletions examples/quickstart/src/main.rs
View file
Expand Up @@ -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<path::PathBuf>,
}

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),
}
}

0 comments on commit 708831d

Please sign in to comment.