This is an automated email from the ASF dual-hosted git repository.
mssun pushed a commit to branch develop
in repository https://gitbox.apache.org/repos/asf/incubator-teaclave.git
The following commit(s) were added to refs/heads/develop by this push:
new 0d5f478 [types] Remove ser/de for worker structs with polished tests
(#217)
0d5f478 is described below
commit 0d5f478900c76736a6c99c92c01263932fcf4299
Author: Zhaofeng Chen <[email protected]>
AuthorDate: Tue Feb 4 22:09:54 2020 -0800
[types] Remove ser/de for worker structs with polished tests (#217)
---
services/proto/src/teaclave_common.rs | 51 ++------
services/proto/src/teaclave_execution_service.rs | 11 +-
.../enclave/src/teaclave_execution_service.rs | 103 ++++++---------
.../enclave/src/teaclave_worker.rs | 121 +++++++----------
types/Cargo.toml | 1 +
types/src/crypto.rs | 145 ++++++++++++++++-----
types/src/worker.rs | 101 +++++++++++---
worker/src/function/context.rs | 2 +-
worker/src/function/gbdt_training.rs | 95 ++++++--------
worker/src/function/mesapy.rs | 29 +----
10 files changed, 350 insertions(+), 309 deletions(-)
diff --git a/services/proto/src/teaclave_common.rs
b/services/proto/src/teaclave_common.rs
index 9d140c9..aa62d80 100644
--- a/services/proto/src/teaclave_common.rs
+++ b/services/proto/src/teaclave_common.rs
@@ -1,13 +1,9 @@
#[cfg(feature = "mesalock_sgx")]
use std::prelude::v1::*;
-use anyhow::{bail, ensure, Error, Result};
-use std::format;
-use teaclave_types::{
- AesGcm128CryptoInfo, AesGcm256CryptoInfo, TeaclaveFileCryptoInfo,
TeaclaveFileRootKey128,
-};
-
use crate::teaclave_common_proto as proto;
+use anyhow::{Error, Result};
+use teaclave_types::TeaclaveFileCryptoInfo;
#[derive(Debug)]
pub struct UserCredential {
@@ -49,48 +45,17 @@ impl From<UserCredential> for proto::UserCredential {
impl std::convert::TryFrom<proto::FileCryptoInfo> for TeaclaveFileCryptoInfo {
type Error = Error;
fn try_from(proto: proto::FileCryptoInfo) -> Result<Self> {
- let info = match proto.schema.as_str() {
- "aes_gcm_128" => {
- let info = AesGcm128CryptoInfo::new(&proto.key, &proto.iv)?;
- TeaclaveFileCryptoInfo::AesGcm128(info)
- }
- "aes_gcm_256" => {
- let info = AesGcm256CryptoInfo::new(&proto.key, &proto.iv)?;
- TeaclaveFileCryptoInfo::AesGcm256(info)
- }
- "teaclave_file_root_key_128" => {
- ensure!(
- proto.iv.is_empty(),
- "IV is not empty for teaclave_file_root_key_128"
- );
- let info = TeaclaveFileRootKey128::new(&proto.key)?;
- TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(info)
- }
- _ => bail!("Invalid crypto schema: {}", proto.schema.as_str()),
- };
-
- Ok(info)
+ TeaclaveFileCryptoInfo::new(&proto.schema, &proto.key, &proto.iv)
}
}
impl std::convert::From<TeaclaveFileCryptoInfo> for proto::FileCryptoInfo {
fn from(crypto: TeaclaveFileCryptoInfo) -> Self {
- match crypto {
- TeaclaveFileCryptoInfo::AesGcm128(info) => proto::FileCryptoInfo {
- schema: "aes_gcm_128".to_string(),
- key: info.key.to_vec(),
- iv: info.iv.to_vec(),
- },
- TeaclaveFileCryptoInfo::AesGcm256(info) => proto::FileCryptoInfo {
- schema: "aes_gcm_256".to_string(),
- key: info.key.to_vec(),
- iv: info.iv.to_vec(),
- },
- TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(info) =>
proto::FileCryptoInfo {
- schema: "teaclave_file_root_key_128".to_string(),
- key: info.key.to_vec(),
- iv: Vec::new(),
- },
+ let (key, iv) = crypto.key_iv();
+ proto::FileCryptoInfo {
+ schema: crypto.schema(),
+ key,
+ iv,
}
}
}
diff --git a/services/proto/src/teaclave_execution_service.rs
b/services/proto/src/teaclave_execution_service.rs
index 90b2af0..7d1b3f2 100644
--- a/services/proto/src/teaclave_execution_service.rs
+++ b/services/proto/src/teaclave_execution_service.rs
@@ -1,5 +1,3 @@
-use serde::{Deserialize, Serialize};
-#[cfg(feature = "mesalock_sgx")]
use std::prelude::v1::*;
use anyhow::{anyhow, Error, Result};
@@ -12,13 +10,12 @@ pub use proto::TeaclaveExecutionClient;
pub use proto::TeaclaveExecutionRequest;
pub use proto::TeaclaveExecutionResponse;
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Debug)]
pub struct StagedFunctionExecuteRequest {
- #[serde(flatten)]
- invocation: WorkerInvocation,
+ pub invocation: WorkerInvocation,
}
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Debug)]
pub struct StagedFunctionExecuteResponse {
pub summary: std::string::String,
}
@@ -43,7 +40,7 @@ impl
std::convert::TryFrom<proto::StagedFunctionExecuteRequest> for StagedFuncti
let ret = Self {
invocation: WorkerInvocation {
runtime_name: proto.runtime_name,
- executor_type: proto.executor_type.try_into()?,
+ executor_type: proto.executor_type.as_str().try_into()?,
function_name: proto.function_name,
function_payload: proto.function_payload,
function_args: TeaclaveFunctionArguments {
diff --git a/tests/functional_tests/enclave/src/teaclave_execution_service.rs
b/tests/functional_tests/enclave/src/teaclave_execution_service.rs
index 6733cb3..b7059c7 100644
--- a/tests/functional_tests/enclave/src/teaclave_execution_service.rs
+++ b/tests/functional_tests/enclave/src/teaclave_execution_service.rs
@@ -1,12 +1,16 @@
-use anyhow;
-use serde_json;
-use std::io::Write;
use std::prelude::v1::*;
-use std::untrusted::fs;
+use std::convert::TryInto;
use teaclave_proto::teaclave_execution_service::*;
use teaclave_rpc::endpoint::Endpoint;
-use teaclave_types::AesGcm128CryptoInfo;
+
+use teaclave_types::convert_plaintext_file;
+use teaclave_types::hashmap;
+use teaclave_types::TeaclaveFileCryptoInfo;
+use teaclave_types::TeaclaveFunctionArguments;
+use teaclave_types::TeaclaveWorkerFileInfo;
+use teaclave_types::TeaclaveWorkerFileRegistry;
+use teaclave_types::WorkerInvocation;
pub fn run_tests() -> bool {
use teaclave_test_utils::*;
@@ -18,66 +22,45 @@ fn setup_client() -> TeaclaveExecutionClient {
TeaclaveExecutionClient::new(channel).unwrap()
}
-fn enc_input_file(input_crypto: &str, plain_input: &str, enc_input: &str) ->
anyhow::Result<()> {
- let crypto_info: AesGcm128CryptoInfo = serde_json::from_str(input_crypto)?;
- let mut bytes = fs::read_to_string(plain_input)?.into_bytes();
- crypto_info.encrypt(&mut bytes)?;
-
- let mut file = fs::File::create(enc_input)?;
- file.write_all(&bytes)?;
- Ok(())
-}
-
fn test_invoke_success() {
- let request_payload = r#"{
- "runtime_name": "default",
- "executor_type": "native",
- "function_name": "gbdt_training",
- "function_payload": "",
- "function_args": {
- "feature_size": "4",
- "max_depth": "4",
- "iterations": "100",
- "shrinkage": "0.1",
- "feature_sample_ratio": "1.0",
- "data_sample_ratio": "1.0",
- "min_leaf_size": "1",
- "loss": "LAD",
- "training_optimization_level": "2"
- },
- "input_files": {
- "training_data": {
- "path": "test_cases/gbdt_training/train.enc",
- "crypto_info": {
- "aes_gcm128": {
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
14, 15],
- "iv": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
- }
- }
- }
- },
- "output_files": {
- "trained_model": {
- "path": "test_cases/gbdt_training/model.enc.out",
- "crypto_info": {
- "teaclave_file_root_key128": {
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
14, 15]
- }
- }
- }
- }
- }"#;
+ let mut client = setup_client();
+
+ let function_args = TeaclaveFunctionArguments::new(&hashmap!(
+ "feature_size" => "4",
+ "max_depth" => "4",
+ "iterations" => "100",
+ "shrinkage" => "0.1",
+ "feature_sample_ratio" => "1.0",
+ "data_sample_ratio" => "1.0",
+ "min_leaf_size" => "1",
+ "loss" => "LAD",
+ "training_optimization_level" => "2"
+ ));
- let input_crypto = r#"{
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
- "iv": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
- }"#;
let enc_input = "test_cases/gbdt_training/train.enc";
let plain_input = "test_cases/gbdt_training/train.txt";
+ let enc_output = "test_cases/gbdt_training/model.enc.out";
+
+ let input_info = convert_plaintext_file(plain_input, enc_input).unwrap();
+ let input_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ "training_data".to_string() => input_info));
+
+ let output_info = TeaclaveWorkerFileInfo::new(enc_output,
TeaclaveFileCryptoInfo::default());
+ let output_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ "trained_model".to_string() => output_info));
+
+ let request = StagedFunctionExecuteRequest {
+ invocation: WorkerInvocation {
+ runtime_name: "default".to_string(),
+ executor_type: "native".try_into().unwrap(),
+ function_name: "gbdt_training".to_string(),
+ function_payload: String::new(),
+ function_args,
+ input_files,
+ output_files,
+ },
+ };
- enc_input_file(input_crypto, plain_input, enc_input).unwrap();
- let request: StagedFunctionExecuteRequest =
serde_json::from_str(request_payload).unwrap();
- let mut client = setup_client();
let response_result = client.invoke_function(request);
assert!(response_result.is_ok());
}
diff --git a/tests/integration_tests/enclave/src/teaclave_worker.rs
b/tests/integration_tests/enclave/src/teaclave_worker.rs
index 7951929..5a967f7 100644
--- a/tests/integration_tests/enclave/src/teaclave_worker.rs
+++ b/tests/integration_tests/enclave/src/teaclave_worker.rs
@@ -1,95 +1,66 @@
-use protected_fs::ProtectedFile;
-use std::io::{Read, Write};
+use std::convert::TryInto;
+use std::io::Read;
use std::prelude::v1::*;
-use std::untrusted::fs;
-use serde_json;
-
-use anyhow;
-use teaclave_types::AesGcm128CryptoInfo;
-use teaclave_types::TeaclaveFileRootKey128;
+use teaclave_types::convert_plaintext_file;
+use teaclave_types::hashmap;
+use teaclave_types::read_all_bytes;
+use teaclave_types::TeaclaveFileCryptoInfo;
+use teaclave_types::TeaclaveFunctionArguments;
+use teaclave_types::TeaclaveWorkerFileInfo;
+use teaclave_types::TeaclaveWorkerFileRegistry;
use teaclave_types::WorkerInvocation;
-use teaclave_worker::Worker;
-fn enc_input_file(input_crypto: &str, plain_input: &str, enc_input: &str) ->
anyhow::Result<()> {
- let crypto_info: AesGcm128CryptoInfo = serde_json::from_str(input_crypto)?;
- let mut bytes = fs::read_to_string(plain_input)?.into_bytes();
- crypto_info.encrypt(&mut bytes)?;
-
- let mut file = fs::File::create(enc_input)?;
- file.write_all(&bytes)?;
- Ok(())
-}
-
-fn dec_output_file(output_crypto: &str, enc_output: &str) ->
anyhow::Result<String> {
- let crypto: TeaclaveFileRootKey128 = serde_json::from_str(output_crypto)?;
- let mut file = ProtectedFile::open_ex(enc_output, &crypto.key)?;
- let mut result = String::new();
- file.read_to_string(&mut result)?;
- Ok(result)
-}
+use teaclave_worker::Worker;
fn test_start_worker() {
- let request_payload = r#"{
- "runtime_name": "default",
- "executor_type": "native",
- "function_name": "gbdt_training",
- "function_payload": "",
- "function_args": {
- "feature_size": "4",
- "max_depth": "4",
- "iterations": "100",
- "shrinkage": "0.1",
- "feature_sample_ratio": "1.0",
- "data_sample_ratio": "1.0",
- "min_leaf_size": "1",
- "loss": "LAD",
- "training_optimization_level": "2"
- },
- "input_files": {
- "training_data": {
- "path": "test_cases/gbdt_training/train.enc",
- "crypto_info": {
- "aes_gcm128": {
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
14, 15],
- "iv": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
- }
- }
- }
- },
- "output_files": {
- "trained_model": {
- "path": "test_cases/gbdt_training/model.enc.out",
- "crypto_info": {
- "teaclave_file_root_key128": {
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
14, 15]
- }
- }
- }
- }
- }"#;
+ let function_args = TeaclaveFunctionArguments::new(&hashmap!(
+ "feature_size" => "4",
+ "max_depth" => "4",
+ "iterations" => "100",
+ "shrinkage" => "0.1",
+ "feature_sample_ratio" => "1.0",
+ "data_sample_ratio" => "1.0",
+ "min_leaf_size" => "1",
+ "loss" => "LAD",
+ "training_optimization_level" => "2"
+ ));
- let input_crypto = r#"{
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
- "iv": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
- }"#;
let enc_input = "test_cases/gbdt_training/train.enc";
let plain_input = "test_cases/gbdt_training/train.txt";
- let output_crypto = r#"{
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
- }"#;
let enc_output = "test_cases/gbdt_training/model.enc.out";
let expected_output = "test_cases/gbdt_training/expected_model.txt";
- enc_input_file(input_crypto, plain_input, enc_input).unwrap();
- let request: WorkerInvocation =
serde_json::from_str(request_payload).unwrap();
+ let input_info = convert_plaintext_file(plain_input, enc_input).unwrap();
+
+ let input_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ "training_data".to_string() => input_info));
+
+ let output_info = TeaclaveWorkerFileInfo::new(enc_output,
TeaclaveFileCryptoInfo::default());
+
+ let output_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ "trained_model".to_string() => output_info.clone()));
+
+ let request = WorkerInvocation {
+ runtime_name: "default".to_string(),
+ executor_type: "native".try_into().unwrap(),
+ function_name: "gbdt_training".to_string(),
+ function_payload: String::new(),
+ function_args,
+ input_files,
+ output_files,
+ };
+
let worker = Worker::default();
let summary = worker.invoke_function(request).unwrap();
assert_eq!(summary, "Trained 120 lines of data.");
- let result = dec_output_file(output_crypto, enc_output).unwrap();
- let expected = fs::read_to_string(&expected_output).unwrap();
+ let mut f = output_info.get_readable_io().unwrap();
+ let mut result = Vec::new();
+ f.read_to_end(&mut result).unwrap();
+
+ let expected = read_all_bytes(expected_output).unwrap();
assert_eq!(&result[..], &expected[..]);
}
diff --git a/types/Cargo.toml b/types/Cargo.toml
index b9b60ec..0fb3d53 100644
--- a/types/Cargo.toml
+++ b/types/Cargo.toml
@@ -21,6 +21,7 @@ protected_fs_rs = { path = "../common/protected_fs_rs",
default-features = fals
anyhow = { version = "1.0.26" }
sgx_types = { version = "1.1.0" }
+rand = { version = "0.7.0" }
hex = { version = "0.4.0" }
serde = { version = "1.0.92", features = ["derive"] }
toml = { version = "0.5.3" }
diff --git a/types/src/crypto.rs b/types/src/crypto.rs
index 6399bdb..134163e 100644
--- a/types/src/crypto.rs
+++ b/types/src/crypto.rs
@@ -2,35 +2,36 @@
use std::prelude::v1::*;
use anyhow;
+use rand::prelude::RngCore;
use ring;
-use serde::{Deserialize, Serialize};
use std::format;
const AES_GCM_256_KEY_LENGTH: usize = 32;
const AES_GCM_256_IV_LENGTH: usize = 12;
-#[derive(Default, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug)]
pub struct AesGcm256CryptoInfo {
pub key: [u8; AES_GCM_256_KEY_LENGTH],
pub iv: [u8; AES_GCM_256_IV_LENGTH],
}
impl AesGcm256CryptoInfo {
- pub fn new(key: &[u8], iv: &[u8]) -> anyhow::Result<Self> {
+ pub fn new(in_key: &[u8], in_iv: &[u8]) -> anyhow::Result<Self> {
anyhow::ensure!(
- key.len() == AES_GCM_256_KEY_LENGTH,
+ in_key.len() == AES_GCM_256_KEY_LENGTH,
"Invalid key length for AesGcm256: {}",
- key.len()
+ in_key.len()
);
anyhow::ensure!(
- iv.len() == AES_GCM_256_IV_LENGTH,
+ in_iv.len() == AES_GCM_256_IV_LENGTH,
"Invalid iv length for AesGcm256: {}",
- iv.len()
+ in_iv.len()
);
- let mut info = AesGcm256CryptoInfo::default();
- info.key.copy_from_slice(key);
- info.iv.copy_from_slice(iv);
- Ok(info)
+ let mut key = [0u8; AES_GCM_256_KEY_LENGTH];
+ let mut iv = [0u8; AES_GCM_256_IV_LENGTH];
+ key.copy_from_slice(in_key);
+ iv.copy_from_slice(in_iv);
+ Ok(AesGcm256CryptoInfo { key, iv })
}
pub fn decrypt(&self, in_out: &mut Vec<u8>) -> anyhow::Result<()> {
@@ -45,33 +46,45 @@ impl AesGcm256CryptoInfo {
}
}
+impl Default for AesGcm256CryptoInfo {
+ fn default() -> Self {
+ let mut key = [0u8; AES_GCM_256_KEY_LENGTH];
+ let mut iv = [0u8; AES_GCM_256_IV_LENGTH];
+ let mut rng = rand::thread_rng();
+ rng.fill_bytes(&mut key);
+ rng.fill_bytes(&mut iv);
+ AesGcm256CryptoInfo { key, iv }
+ }
+}
+
const AES_GCM_128_KEY_LENGTH: usize = 16;
const AES_GCM_128_IV_LENGTH: usize = 12;
-#[derive(Default, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug)]
pub struct AesGcm128CryptoInfo {
pub key: [u8; AES_GCM_128_KEY_LENGTH],
pub iv: [u8; AES_GCM_128_IV_LENGTH],
}
impl AesGcm128CryptoInfo {
- pub fn new(key: &[u8], iv: &[u8]) -> anyhow::Result<Self> {
+ pub fn new(in_key: &[u8], in_iv: &[u8]) -> anyhow::Result<Self> {
anyhow::ensure!(
- key.len() == AES_GCM_128_KEY_LENGTH,
+ in_key.len() == AES_GCM_128_KEY_LENGTH,
"Invalid key length for AesGcm128: {}",
- key.len()
+ in_key.len()
);
anyhow::ensure!(
- iv.len() == AES_GCM_128_IV_LENGTH,
+ in_iv.len() == AES_GCM_128_IV_LENGTH,
"Invalid iv length for AesGcm128: {}",
- iv.len()
+ in_iv.len()
);
- let mut info = AesGcm128CryptoInfo::default();
- info.key.copy_from_slice(key);
- info.iv.copy_from_slice(iv);
- Ok(info)
+ let mut key = [0u8; AES_GCM_128_KEY_LENGTH];
+ let mut iv = [0u8; AES_GCM_128_IV_LENGTH];
+ key.copy_from_slice(in_key);
+ iv.copy_from_slice(in_iv);
+ Ok(AesGcm128CryptoInfo { key, iv })
}
pub fn decrypt(&self, in_out: &mut Vec<u8>) -> anyhow::Result<()> {
@@ -86,34 +99,104 @@ impl AesGcm128CryptoInfo {
}
}
+impl Default for AesGcm128CryptoInfo {
+ fn default() -> Self {
+ let mut key = [0u8; AES_GCM_128_KEY_LENGTH];
+ let mut iv = [0u8; AES_GCM_128_IV_LENGTH];
+ let mut rng = rand::thread_rng();
+ rng.fill_bytes(&mut key);
+ rng.fill_bytes(&mut iv);
+ AesGcm128CryptoInfo { key, iv }
+ }
+}
+
const TEACLAVE_FILE_ROOT_KEY_128_LENGTH: usize = 16;
-#[derive(Default, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug)]
pub struct TeaclaveFileRootKey128 {
- pub key: [u8; 16],
+ pub key: [u8; TEACLAVE_FILE_ROOT_KEY_128_LENGTH],
}
impl TeaclaveFileRootKey128 {
- pub fn new(key: &[u8]) -> anyhow::Result<Self> {
+ pub fn new(in_key: &[u8]) -> anyhow::Result<Self> {
anyhow::ensure!(
- key.len() == TEACLAVE_FILE_ROOT_KEY_128_LENGTH,
+ in_key.len() == TEACLAVE_FILE_ROOT_KEY_128_LENGTH,
"Invalid key length for teaclave_file_root_key_128: {}",
- key.len()
+ in_key.len()
);
- let mut info = TeaclaveFileRootKey128::default();
- info.key.copy_from_slice(key);
- Ok(info)
+ let mut key = [0u8; TEACLAVE_FILE_ROOT_KEY_128_LENGTH];
+ key.copy_from_slice(in_key);
+ Ok(TeaclaveFileRootKey128 { key })
+ }
+}
+
+impl Default for TeaclaveFileRootKey128 {
+ fn default() -> Self {
+ let mut key = [0u8; TEACLAVE_FILE_ROOT_KEY_128_LENGTH];
+ let mut rng = rand::thread_rng();
+ rng.fill_bytes(&mut key);
+ TeaclaveFileRootKey128 { key }
}
}
-#[derive(Debug, Serialize, Deserialize)]
-#[serde(rename_all(deserialize = "snake_case"))]
+#[derive(Clone, Debug)]
pub enum TeaclaveFileCryptoInfo {
AesGcm128(AesGcm128CryptoInfo),
AesGcm256(AesGcm256CryptoInfo),
TeaclaveFileRootKey128(TeaclaveFileRootKey128),
}
+impl TeaclaveFileCryptoInfo {
+ pub fn new(schema: &str, key: &[u8], iv: &[u8]) -> anyhow::Result<Self> {
+ let info = match schema {
+ "aes_gcm_128" => {
+ let info = AesGcm128CryptoInfo::new(key, iv)?;
+ TeaclaveFileCryptoInfo::AesGcm128(info)
+ }
+ "aes_gcm_256" => {
+ let info = AesGcm256CryptoInfo::new(key, iv)?;
+ TeaclaveFileCryptoInfo::AesGcm256(info)
+ }
+ "teaclave_file_root_key_128" => {
+ anyhow::ensure!(
+ iv.is_empty(),
+ "IV is not empty for teaclave_file_root_key_128"
+ );
+ let info = TeaclaveFileRootKey128::new(key)?;
+ TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(info)
+ }
+ _ => anyhow::bail!("Invalid crypto schema: {}", schema),
+ };
+ Ok(info)
+ }
+
+ pub fn schema(&self) -> String {
+ match self {
+ TeaclaveFileCryptoInfo::AesGcm128(_) => "aes_gcm_128".to_string(),
+ TeaclaveFileCryptoInfo::AesGcm256(_) => "aes_gcm_256".to_string(),
+ TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(_) => {
+ "teaclave_file_root_key_128".to_string()
+ }
+ }
+ }
+
+ pub fn key_iv(&self) -> (Vec<u8>, Vec<u8>) {
+ match self {
+ TeaclaveFileCryptoInfo::AesGcm128(crypto) => (crypto.key.to_vec(),
crypto.iv.to_vec()),
+ TeaclaveFileCryptoInfo::AesGcm256(crypto) => (crypto.key.to_vec(),
crypto.iv.to_vec()),
+ TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto) => {
+ (crypto.key.to_vec(), Vec::new())
+ }
+ }
+ }
+}
+
+impl Default for TeaclaveFileCryptoInfo {
+ fn default() -> Self {
+
TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(TeaclaveFileRootKey128::default())
+ }
+}
+
fn make_teaclave_aad() -> ring::aead::Aad<[u8; 8]> {
let bytes = [0u8; 8];
ring::aead::Aad::from(bytes)
diff --git a/types/src/worker.rs b/types/src/worker.rs
index 5dcb10c..4051b41 100644
--- a/types/src/worker.rs
+++ b/types/src/worker.rs
@@ -1,12 +1,7 @@
-#[cfg(feature = "mesalock_sgx")]
-use std::prelude::v1::*;
-
-use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::format;
use std::io::{self, Read};
-
-use protected_fs::ProtectedFile;
+use std::prelude::v1::*;
#[cfg(feature = "mesalock_sgx")]
use std::untrusted::fs::File;
@@ -17,19 +12,29 @@ use std::fs::File;
use anyhow;
use crate::TeaclaveFileCryptoInfo;
+use protected_fs::ProtectedFile;
+use serde::{Deserialize, Serialize};
-#[derive(Serialize, Deserialize, Debug)]
-#[serde(rename_all(deserialize = "snake_case"))]
+#[macro_export]
+macro_rules! hashmap {
+ ($( $key: expr => $val: expr ),*) => {{
+ let mut map = ::std::collections::HashMap::new();
+ $( map.insert($key, $val); )*
+ map
+ }}
+}
+
+#[derive(Debug)]
pub enum TeaclaveExecutorSelector {
Native,
Python,
}
-impl std::convert::TryFrom<String> for TeaclaveExecutorSelector {
+impl std::convert::TryFrom<&str> for TeaclaveExecutorSelector {
type Error = anyhow::Error;
- fn try_from(selector: String) -> anyhow::Result<Self> {
- let sel = match selector.as_ref() {
+ fn try_from(selector: &str) -> anyhow::Result<Self> {
+ let sel = match selector {
"python" => TeaclaveExecutorSelector::Python,
"native" => TeaclaveExecutorSelector::Native,
_ => anyhow::bail!("Invalid executor selector: {}", selector),
@@ -69,19 +74,64 @@ impl io::Read for ReadBuffer {
}
}
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Clone, Debug)]
pub struct TeaclaveWorkerFileInfo {
pub path: std::path::PathBuf,
pub crypto_info: TeaclaveFileCryptoInfo,
}
-fn read_all_bytes<P: AsRef<std::path::Path>>(path: P) ->
anyhow::Result<Vec<u8>> {
+impl TeaclaveWorkerFileInfo {
+ pub fn new(
+ path: impl std::convert::Into<std::path::PathBuf>,
+ crypto_info: TeaclaveFileCryptoInfo,
+ ) -> Self {
+ TeaclaveWorkerFileInfo {
+ path: path.into(),
+ crypto_info,
+ }
+ }
+}
+
+pub fn read_all_bytes(path: impl AsRef<std::path::Path>) ->
anyhow::Result<Vec<u8>> {
let mut content = Vec::new();
let mut file = File::open(path)?;
file.read_to_end(&mut content)?;
Ok(content)
}
+fn teaclave_file_with_bytes(path: &str, bytes: &[u8]) ->
anyhow::Result<TeaclaveWorkerFileInfo> {
+ let crypto_info = TeaclaveFileCryptoInfo::default();
+ let file_info = TeaclaveWorkerFileInfo::new(path, crypto_info);
+ let mut f = file_info.get_writable_io()?;
+ f.write_all(bytes)?;
+ Ok(file_info)
+}
+
+pub fn convert_plaintext_file(src: &str, dst: &str) ->
anyhow::Result<TeaclaveWorkerFileInfo> {
+ let bytes = read_all_bytes(src)?;
+ teaclave_file_with_bytes(dst, &bytes)
+}
+
+pub fn convert_encrypted_file(
+ src: TeaclaveWorkerFileInfo,
+ dst: &str,
+) -> anyhow::Result<TeaclaveWorkerFileInfo> {
+ let plain_text = match &src.crypto_info {
+ TeaclaveFileCryptoInfo::AesGcm128(crypto) => {
+ let mut bytes = read_all_bytes(src.path)?;
+ crypto.decrypt(&mut bytes)?;
+ bytes
+ }
+ TeaclaveFileCryptoInfo::AesGcm256(crypto) => {
+ let mut bytes = read_all_bytes(src.path)?;
+ crypto.decrypt(&mut bytes)?;
+ bytes
+ }
+ TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(_) => return Ok(src),
+ };
+ teaclave_file_with_bytes(dst, &plain_text)
+}
+
impl TeaclaveWorkerFileInfo {
pub fn get_readable_io(&self) -> anyhow::Result<Box<dyn io::Read>> {
let readable: Box<dyn io::Read> = match &self.crypto_info {
@@ -115,12 +165,17 @@ impl TeaclaveWorkerFileInfo {
}
}
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Debug)]
pub struct TeaclaveWorkerFileRegistry {
- #[serde(flatten)]
pub entries: HashMap<String, TeaclaveWorkerFileInfo>,
}
+impl TeaclaveWorkerFileRegistry {
+ pub fn new(entries: HashMap<String, TeaclaveWorkerFileInfo>) -> Self {
+ TeaclaveWorkerFileRegistry { entries }
+ }
+}
+
impl<T> std::convert::TryFrom<HashMap<String, T>> for
TeaclaveWorkerFileRegistry
where
T: std::convert::TryInto<TeaclaveWorkerFileInfo, Error = anyhow::Error>,
@@ -156,11 +211,23 @@ where
#[derive(Serialize, Deserialize, Debug)]
pub struct TeaclaveFunctionArguments {
- #[serde(flatten)]
pub args: HashMap<String, String>,
}
impl TeaclaveFunctionArguments {
+ pub fn new<K, V>(input: &HashMap<K, V>) -> Self
+ where
+ K: std::string::ToString,
+ V: std::string::ToString,
+ {
+ let args = input.iter().fold(HashMap::new(), |mut acc, (k, v)| {
+ acc.insert(k.to_string(), v.to_string());
+ acc
+ });
+
+ TeaclaveFunctionArguments { args }
+ }
+
pub fn try_get<T: std::str::FromStr>(&self, key: &str) ->
anyhow::Result<T> {
self.args
.get(key)
@@ -181,7 +248,7 @@ impl TeaclaveFunctionArguments {
}
}
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Debug)]
pub struct WorkerInvocation {
pub runtime_name: String,
pub executor_type: TeaclaveExecutorSelector, // "native" | "python"
diff --git a/worker/src/function/context.rs b/worker/src/function/context.rs
index 791ca04..f0ac184 100644
--- a/worker/src/function/context.rs
+++ b/worker/src/function/context.rs
@@ -219,11 +219,11 @@ pub fn rtc_close_handle(f: FileHandle) ->
anyhow::Result<()> {
#[cfg(feature = "enclave_unit_test")]
pub mod tests {
use super::*;
- use crate::hashmap;
use crate::runtime::RawIoRuntime;
use std::path::PathBuf;
use std::str::FromStr;
use teaclave_test_utils::*;
+ use teaclave_types::hashmap;
use teaclave_types::AesGcm256CryptoInfo;
use teaclave_types::TeaclaveFileCryptoInfo;
use teaclave_types::TeaclaveWorkerFileInfo;
diff --git a/worker/src/function/gbdt_training.rs
b/worker/src/function/gbdt_training.rs
index 4406380..6b67fe6 100644
--- a/worker/src/function/gbdt_training.rs
+++ b/worker/src/function/gbdt_training.rs
@@ -35,6 +35,9 @@ use gbdt::gradient_boost::GBDT;
#[derive(Default)]
pub struct GbdtTraining;
+static IN_DATA: &str = "training_data";
+static OUT_MODEL: &str = "trained_model";
+
impl TeaclaveFunction for GbdtTraining {
fn execute(
&self,
@@ -52,7 +55,7 @@ impl TeaclaveFunction for GbdtTraining {
let training_optimization_level: u8 =
args.try_get("training_optimization_level")?;
// read input
- let training_file = runtime.open_input("training_data")?;
+ let training_file = runtime.open_input(IN_DATA)?;
let mut train_dv = parse_training_data(training_file, feature_size)?;
let data_size = train_dv.len();
@@ -75,7 +78,7 @@ impl TeaclaveFunction for GbdtTraining {
let model_json = serde_json::to_string(&gbdt_train_mod)?;
// save the model to output
- let mut model_file = runtime.create_output("trained_model")?;
+ let mut model_file = runtime.create_output(OUT_MODEL)?;
model_file.write_all(model_json.as_bytes())?;
let summary = format!("Trained {} lines of data.", data_size);
@@ -127,69 +130,54 @@ pub mod tests {
use super::*;
use teaclave_test_utils::*;
+ use std::untrusted::fs;
+
+ use teaclave_types::hashmap;
+ use teaclave_types::TeaclaveFileCryptoInfo;
+ use teaclave_types::TeaclaveFunctionArguments;
+ use teaclave_types::TeaclaveWorkerFileInfo;
+ use teaclave_types::TeaclaveWorkerFileRegistry;
+
use crate::function::TeaclaveFunction;
use crate::runtime::RawIoRuntime;
- use std::untrusted::fs;
- use teaclave_types::{TeaclaveFunctionArguments,
TeaclaveWorkerFileRegistry};
pub fn run_tests() -> bool {
run_tests!(test_gbdt_training, test_gbdt_parse_training_data,)
}
fn test_gbdt_training() {
- let args_json = r#"
- {
- "feature_size": "4",
- "max_depth": "4",
- "iterations": "100",
- "shrinkage": "0.1",
- "feature_sample_ratio": "1.0",
- "data_sample_ratio": "1.0",
- "min_leaf_size": "1",
- "loss": "LAD",
- "training_optimization_level": "2"
- }
- "#;
-
- let input_json = r#"
- {
- "training_data": {
- "path": "test_cases/gbdt_training/train.txt",
- "crypto_info": {
- "aes_gcm128": {
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
13, 14, 15],
- "iv": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
- }
- }
- }
- }
- "#;
-
- let output_json = r#"
- {
- "trained_model": {
- "path": "test_cases/gbdt_training/training_model.txt.out",
- "crypto_info": {
- "aes_gcm128": {
- "key": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
13, 14, 15],
- "iv": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
- }
- }
- }
- }
- "#;
-
- let func_args: TeaclaveFunctionArguments =
serde_json::from_str(&args_json).unwrap();
- let input_files: TeaclaveWorkerFileRegistry =
serde_json::from_str(&input_json).unwrap();
- let output_files: TeaclaveWorkerFileRegistry =
serde_json::from_str(&output_json).unwrap();
+ let func_args = TeaclaveFunctionArguments::new(&hashmap!(
+ "feature_size" => "4",
+ "max_depth" => "4",
+ "iterations" => "100",
+ "shrinkage" => "0.1",
+ "feature_sample_ratio" => "1.0",
+ "data_sample_ratio" => "1.0",
+ "min_leaf_size" => "1",
+ "loss" => "LAD",
+ "training_optimization_level" => "2"
+ ));
+
+ let plain_input = "test_cases/gbdt_training/train.txt";
+ let plain_output = "test_cases/gbdt_training/training_model.txt.out";
+ let expected_output = "test_cases/gbdt_training/expected_model.txt";
+
+ let input_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ IN_DATA.to_string() =>
+ TeaclaveWorkerFileInfo::new(plain_input,
TeaclaveFileCryptoInfo::default())
+ ));
+
+ let output_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ OUT_MODEL.to_string() =>
+ TeaclaveWorkerFileInfo::new(plain_output,
TeaclaveFileCryptoInfo::default())
+ ));
+
let runtime = Box::new(RawIoRuntime::new(input_files, output_files));
let function = GbdtTraining;
let summary = function.execute(runtime, func_args).unwrap();
assert_eq!(summary, "Trained 120 lines of data.");
- let plain_output = "test_cases/gbdt_training/training_model.txt.out";
- let expected_output = "test_cases/gbdt_training/expected_model.txt";
let result = fs::read_to_string(&plain_output).unwrap();
let expected = fs::read_to_string(&expected_output).unwrap();
assert_eq!(&result[..], &expected[..]);
@@ -198,6 +186,9 @@ pub mod tests {
fn test_gbdt_parse_training_data() {
let line = "4.8,3.0,1.4,0.3,3.0";
let result = parse_data_line(&line, 4);
- assert_eq!(result.is_ok(), true);
+ assert!(result.is_ok());
+
+ let result = parse_data_line(&line, 3);
+ assert!(result.is_err());
}
}
diff --git a/worker/src/function/mesapy.rs b/worker/src/function/mesapy.rs
index 7cfe721..b386c21 100644
--- a/worker/src/function/mesapy.rs
+++ b/worker/src/function/mesapy.rs
@@ -100,15 +100,6 @@ impl TeaclaveFunction for Mesapy {
}
}
-#[macro_export]
-macro_rules! hashmap {
- ($( $key: expr => $val: expr ),*) => {{
- let mut map = ::std::collections::HashMap::new();
- $( map.insert($key, $val); )*
- map
- }}
-}
-
#[cfg(feature = "enclave_unit_test")]
pub mod tests {
use super::*;
@@ -116,9 +107,7 @@ pub mod tests {
use crate::function::TeaclaveFunction;
use crate::runtime::RawIoRuntime;
- use std::path::PathBuf;
- use std::str::FromStr;
- use teaclave_types::AesGcm256CryptoInfo;
+ use teaclave_types::hashmap;
use teaclave_types::TeaclaveFileCryptoInfo;
use teaclave_types::TeaclaveFunctionArguments;
use teaclave_types::TeaclaveWorkerFileInfo;
@@ -129,7 +118,7 @@ pub mod tests {
}
fn test_mesapy() {
- let py_args = hashmap!("--name" => "Teaclave");
+ let py_args = TeaclaveFunctionArguments::new(&hashmap!("--name" =>
"Teaclave"));
let py_payload = "
import sys
def entrypoint(argv):
@@ -137,18 +126,12 @@ def entrypoint(argv):
print argv[1]
";
- let input = PathBuf::from_str("test_cases/mesapy/input.txt").unwrap();
- let output =
PathBuf::from_str("test_cases/mesapy/output.txt").unwrap();
+ let input = "test_cases/mesapy/input.txt";
+ let output = "test_cases/mesapy/output.txt";
- let input_info = TeaclaveWorkerFileInfo {
- path: input,
- crypto_info:
TeaclaveFileCryptoInfo::AesGcm256(AesGcm256CryptoInfo::default()),
- };
+ let input_info = TeaclaveWorkerFileInfo::new(input,
TeaclaveFileCryptoInfo::default());
- let output_info = TeaclaveWorkerFileInfo {
- path: output,
- crypto_info:
TeaclaveFileCryptoInfo::AesGcm256(AesGcm256CryptoInfo::default()),
- };
+ let output_info = TeaclaveWorkerFileInfo::new(output,
TeaclaveFileCryptoInfo::default());
let input_files = TeaclaveWorkerFileRegistry {
entries: hashmap!("in_f1".to_string() => input_info),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
