Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rewrite quickstart example command line interface #186

Merged
merged 2 commits into from Dec 20, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 2 additions & 0 deletions examples/quickstart/Cargo.toml
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
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),
}
}