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
commit 770b9b2f7f31b601c1302b61eef14cacd5095590 Author: Mingshen Sun <[email protected]> AuthorDate: Sat Mar 21 20:34:35 2020 -0700 [types] Common types refactoring --- services/execution/enclave/src/service.rs | 34 ++- services/management/enclave/src/service.rs | 18 +- tests/integration/enclave/src/teaclave_worker.rs | 32 ++- types/src/lib.rs | 4 + types/src/staged_file.rs | 179 +++++++++++++ types/src/staged_function.rs | 176 +++++++++++++ types/src/staged_task.rs | 20 +- types/src/worker.rs | 305 +---------------------- worker/src/function/context.rs | 15 +- worker/src/function/echo.rs | 6 +- worker/src/function/gbdt_prediction.rs | 16 +- worker/src/function/gbdt_training.rs | 14 +- worker/src/function/mesapy.rs | 15 +- worker/src/runtime/default.rs | 14 +- worker/src/runtime/raw_io.rs | 14 +- worker/src/worker.rs | 12 +- 16 files changed, 466 insertions(+), 408 deletions(-) diff --git a/services/execution/enclave/src/service.rs b/services/execution/enclave/src/service.rs index c9fe674..50bf2e8 100644 --- a/services/execution/enclave/src/service.rs +++ b/services/execution/enclave/src/service.rs @@ -191,7 +191,7 @@ fn prepare_task(task: &StagedTask) -> StagedFunction { let request = FileAgentRequest::new(HandleFileCommand::Download, file_request_info); handle_file_request(request).unwrap(); - let mut converted_input_file_map: HashMap<String, TeaclaveWorkerInputFileInfo> = HashMap::new(); + let mut converted_input_file_map: HashMap<String, StagedInputFile> = HashMap::new(); for (key, value) in input_file_map.iter() { let (from, crypto_info) = value; let mut dest = from.clone(); @@ -201,9 +201,9 @@ fn prepare_task(task: &StagedTask) -> StagedFunction { let input_file_info = convert_encrypted_input_file(from, *crypto_info, &dest).unwrap(); converted_input_file_map.insert(key.to_string(), input_file_info); } - let input_files = TeaclaveWorkerFileRegistry::new(converted_input_file_map); + let input_files = StagedFiles::new(converted_input_file_map); - let mut output_file_map: HashMap<String, TeaclaveWorkerOutputFileInfo> = HashMap::new(); + let mut output_file_map: HashMap<String, StagedOutputFile> = HashMap::new(); for (key, value) in task.output_data.iter() { let mut dest = agent_dir_path.to_path_buf(); dest.push(&format!("{}.out", key)); @@ -211,21 +211,19 @@ fn prepare_task(task: &StagedTask) -> StagedFunction { TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto) => crypto, _ => unimplemented!(), }; - let output_info = - TeaclaveWorkerOutputFileInfo::new(dest.to_string_lossy().to_string(), crypto); + let output_info = StagedOutputFile::new(dest.to_string_lossy().to_string(), crypto); output_file_map.insert(key.to_string(), output_info); } - let output_files = TeaclaveWorkerFileRegistry::new(output_file_map); - - StagedFunction { - name: function_name, - payload: function_payload, - arguments: function_arguments, - input_files, - output_files, - runtime_name, - executor_type, - } + let output_files = StagedFiles::new(output_file_map); + + StagedFunction::new() + .name(function_name) + .payload(function_payload) + .arguments(function_arguments) + .input_files(input_files) + .output_files(output_files) + .runtime_name(runtime_name) + .executor_type(executor_type) } #[cfg(feature = "enclave_unit_test")] @@ -287,8 +285,8 @@ pub mod tests { let crypto = TeaclaveFileRootKey128::new(&[0; 16]).unwrap(); let crypto_info = TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto); - let training_input_data = InputDataValue::new(input_url, "", crypto_info); - let model_output_data = OutputDataValue::new(output_url, crypto_info); + let training_input_data = FunctionInputFile::new(input_url, "", crypto_info); + let model_output_data = FunctionOutputFile::new(output_url, crypto_info); let input_data = hashmap!("training_data".to_string() => training_input_data); let output_data = hashmap!("trained_model".to_string() => model_output_data); diff --git a/services/management/enclave/src/service.rs b/services/management/enclave/src/service.rs index 0faa498..d48ce2f 100644 --- a/services/management/enclave/src/service.rs +++ b/services/management/enclave/src/service.rs @@ -27,7 +27,7 @@ use teaclave_service_enclave_utils::teaclave_service; use teaclave_types::Function; #[cfg(test_mode)] use teaclave_types::{FunctionInput, FunctionOutput}; -use teaclave_types::{InputDataValue, OutputDataValue, StagedTask, Task, TaskStatus}; +use teaclave_types::{FunctionInputFile, FunctionOutputFile, StagedTask, Task, TaskStatus}; use teaclave_types::{Storable, TeaclaveInputFile, TeaclaveOutputFile}; use teaclave_types::{TeaclaveServiceResponseError, TeaclaveServiceResponseResult}; use thiserror::Error; @@ -495,14 +495,14 @@ impl TeaclaveManagement for TeaclaveManagementService { .map_err(|_| TeaclaveManagementError::PermissionDenied)?; let function_arguments = task.function_arguments.clone(); - let mut input_map: HashMap<String, InputDataValue> = HashMap::new(); - let mut output_map: HashMap<String, OutputDataValue> = HashMap::new(); + let mut input_map: HashMap<String, FunctionInputFile> = HashMap::new(); + let mut output_map: HashMap<String, FunctionOutputFile> = HashMap::new(); for (data_name, data_id) in task.input_map.iter() { - let input_data: InputDataValue = if TeaclaveInputFile::match_prefix(data_id) { + let input_data: FunctionInputFile = if TeaclaveInputFile::match_prefix(data_id) { let input_file: TeaclaveInputFile = self .read_from_db(data_id.as_bytes()) .map_err(|_| TeaclaveManagementError::PermissionDenied)?; - InputDataValue::from_teaclave_input_file(&input_file) + FunctionInputFile::from_teaclave_input_file(&input_file) } else { return Err(TeaclaveManagementError::PermissionDenied.into()); }; @@ -510,14 +510,14 @@ impl TeaclaveManagement for TeaclaveManagementService { } for (data_name, data_id) in task.output_map.iter() { - let output_data: OutputDataValue = if TeaclaveOutputFile::match_prefix(data_id) { + let output_data: FunctionOutputFile = if TeaclaveOutputFile::match_prefix(data_id) { let output_file: TeaclaveOutputFile = self .read_from_db(data_id.as_bytes()) .map_err(|_| TeaclaveManagementError::PermissionDenied)?; if output_file.hash.is_some() { return Err(TeaclaveManagementError::PermissionDenied.into()); } - OutputDataValue::from_teaclave_output_file(&output_file) + FunctionOutputFile::from_teaclave_output_file(&output_file) } else { return Err(TeaclaveManagementError::PermissionDenied.into()); }; @@ -770,8 +770,8 @@ pub mod tests { let crypto_info = TeaclaveFileCryptoInfo::TeaclaveFileRootKey128( TeaclaveFileRootKey128::new(&[0; 16]).unwrap(), ); - let input_data = InputDataValue::new(url.clone(), hash, crypto_info); - let output_data = OutputDataValue::new(url, crypto_info); + let input_data = FunctionInputFile::new(url.clone(), hash, crypto_info); + let output_data = FunctionOutputFile::new(url, crypto_info); let mut input_map = HashMap::new(); input_map.insert("input".to_string(), input_data); let mut output_map = HashMap::new(); diff --git a/tests/integration/enclave/src/teaclave_worker.rs b/tests/integration/enclave/src/teaclave_worker.rs index 2ea8809..7d2feb4 100644 --- a/tests/integration/enclave/src/teaclave_worker.rs +++ b/tests/integration/enclave/src/teaclave_worker.rs @@ -1,9 +1,8 @@ use std::prelude::v1::*; use teaclave_types::{ - hashmap, read_all_bytes, ExecutorType, FunctionArguments, StagedFunction, - TeaclaveFileRootKey128, TeaclaveWorkerFileRegistry, TeaclaveWorkerInputFileInfo, - TeaclaveWorkerOutputFileInfo, + hashmap, read_all_bytes, ExecutorType, FunctionArguments, StagedFiles, StagedFunction, + StagedInputFile, StagedOutputFile, TeaclaveFileRootKey128, }; use teaclave_worker::Worker; @@ -24,26 +23,23 @@ fn test_start_worker() { let enc_output = "fixtures/functions/gbdt_training/model.enc.out"; let expected_output = "fixtures/functions/gbdt_training/expected_model.txt"; - let input_info = TeaclaveWorkerInputFileInfo::create_with_plaintext_file(plain_input).unwrap(); + let input_info = StagedInputFile::create_with_plaintext_file(plain_input).unwrap(); - let input_files = TeaclaveWorkerFileRegistry::new(hashmap!( + let input_files = StagedFiles::new(hashmap!( "training_data".to_string() => input_info)); - let output_info = - TeaclaveWorkerOutputFileInfo::new(enc_output, TeaclaveFileRootKey128::random()); + let output_info = StagedOutputFile::new(enc_output, TeaclaveFileRootKey128::random()); - let output_files = TeaclaveWorkerFileRegistry::new(hashmap!( + let output_files = StagedFiles::new(hashmap!( "trained_model".to_string() => output_info.clone())); - let request = StagedFunction { - name: "gbdt_training".to_string(), - payload: String::new(), - arguments, - input_files, - output_files, - runtime_name: "default".to_string(), - executor_type: ExecutorType::Native, - }; + let staged_function = StagedFunction::new() + .name("gbdt_training") + .arguments(arguments) + .input_files(input_files) + .output_files(output_files) + .runtime_name("default") + .executor_type(ExecutorType::Native); let worker = Worker::default(); @@ -51,7 +47,7 @@ fn test_start_worker() { assert!(capability.runtimes.contains("default")); assert!(capability.functions.contains("native-gbdt_training")); - let summary = worker.invoke_function(request).unwrap(); + let summary = worker.invoke_function(staged_function).unwrap(); assert_eq!(summary, "Trained 120 lines of data."); let result = output_info.get_plaintext().unwrap(); diff --git a/types/src/lib.rs b/types/src/lib.rs index 5f8b047..13ecacf 100644 --- a/types/src/lib.rs +++ b/types/src/lib.rs @@ -25,6 +25,10 @@ mod function; pub use function::*; mod staged_task; pub use staged_task::*; +mod staged_function; +pub use staged_function::*; +mod staged_file; +pub use staged_file::*; mod storage; pub use storage::Storable; mod task; diff --git a/types/src/staged_file.rs b/types/src/staged_file.rs new file mode 100644 index 0000000..b591b64 --- /dev/null +++ b/types/src/staged_file.rs @@ -0,0 +1,179 @@ +use crate::{TeaclaveFileCryptoInfo, TeaclaveFileRootKey128}; + +use std::collections::HashMap; +#[cfg(not(feature = "mesalock_sgx"))] +use std::fs::File; +use std::io::{self, Read, Write}; +use std::prelude::v1::*; +#[cfg(feature = "mesalock_sgx")] +use std::untrusted::fs::File; + +use protected_fs::ProtectedFile; + +#[derive(Clone, Debug, Default)] +pub struct StagedInputFile { + pub path: std::path::PathBuf, + pub crypto_info: TeaclaveFileRootKey128, +} + +#[derive(Clone, Debug, Default)] +pub struct StagedOutputFile { + pub path: std::path::PathBuf, + pub crypto_info: TeaclaveFileRootKey128, +} + +impl std::convert::From<StagedOutputFile> for StagedInputFile { + fn from(info: StagedOutputFile) -> Self { + StagedInputFile { + path: info.path, + crypto_info: info.crypto_info, + } + } +} + +impl StagedInputFile { + pub fn new( + path: impl std::convert::Into<std::path::PathBuf>, + crypto_info: TeaclaveFileRootKey128, + ) -> Self { + StagedInputFile { + path: path.into(), + crypto_info, + } + } + + pub fn get_readable_io(&self) -> anyhow::Result<Box<dyn io::Read>> { + log::debug!("path: {:?}", self.path); + log::debug!("key: {:?}", self.crypto_info.key); + let f = ProtectedFile::open_ex(&self.path, &self.crypto_info.key)?; + Ok(Box::new(f)) + } + + #[cfg(test_mode)] + pub fn create_with_plaintext_file( + path: impl AsRef<std::path::Path>, + ) -> anyhow::Result<StagedInputFile> { + let bytes = read_all_bytes(path.as_ref())?; + let dst = path.as_ref().with_extension("enc"); + Self::create_with_bytes(dst, &bytes) + } + + pub fn create_with_bytes( + path: impl AsRef<std::path::Path>, + bytes: &[u8], + ) -> anyhow::Result<StagedInputFile> { + let crypto = TeaclaveFileRootKey128::random(); + let mut f = ProtectedFile::create_ex(&path, &crypto.key)?; + f.write_all(bytes)?; + Ok(Self::new(path.as_ref(), crypto)) + } +} + +impl StagedOutputFile { + pub fn new( + path: impl std::convert::Into<std::path::PathBuf>, + crypto_info: TeaclaveFileRootKey128, + ) -> Self { + StagedOutputFile { + path: path.into(), + crypto_info, + } + } + + pub fn get_writable_io(&self) -> anyhow::Result<Box<dyn io::Write>> { + log::debug!("path: {:?}", self.path); + log::debug!("key: {:?}", self.crypto_info.key); + let f = ProtectedFile::create_ex(&self.path, &self.crypto_info.key)?; + Ok(Box::new(f)) + } + + #[cfg(test_mode)] + pub fn get_plaintext(&self) -> anyhow::Result<Vec<u8>> { + let mut content = Vec::new(); + let mut f = ProtectedFile::open_ex(&self.path, &self.crypto_info.key)?; + f.read_to_end(&mut content)?; + Ok(content) + } +} + +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) +} + +pub fn convert_encrypted_input_file( + path: impl AsRef<std::path::Path>, + crypto_info: TeaclaveFileCryptoInfo, + dst: impl AsRef<std::path::Path>, +) -> anyhow::Result<StagedInputFile> { + log::debug!("from: {:?}, to: {:?}", path.as_ref(), dst.as_ref()); + #[cfg(not(feature = "mesalock_sgx"))] + use std::fs; + #[cfg(feature = "mesalock_sgx")] + use std::untrusted::fs; + let plain_text = match crypto_info { + TeaclaveFileCryptoInfo::AesGcm128(crypto) => { + let mut bytes = read_all_bytes(path)?; + crypto.decrypt(&mut bytes)?; + bytes + } + TeaclaveFileCryptoInfo::AesGcm256(crypto) => { + let mut bytes = read_all_bytes(path)?; + crypto.decrypt(&mut bytes)?; + bytes + } + TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto) => { + fs::copy(path, dst.as_ref())?; + let dst = dst.as_ref().to_owned(); + return Ok(StagedInputFile::new(dst, crypto)); + } + TeaclaveFileCryptoInfo::Raw => read_all_bytes(path)?, + }; + StagedInputFile::create_with_bytes(dst.as_ref(), &plain_text) +} + +#[derive(Debug, Default)] +pub struct StagedFiles<T> { + pub entries: HashMap<String, T>, +} + +impl<T> StagedFiles<T> { + pub fn new(entries: HashMap<String, T>) -> Self { + StagedFiles { entries } + } +} + +impl<U, V> std::convert::TryFrom<HashMap<String, U>> for StagedFiles<V> +where + U: std::convert::TryInto<V, Error = anyhow::Error>, +{ + type Error = anyhow::Error; + fn try_from(entries: HashMap<String, U>) -> anyhow::Result<Self> { + let mut out_info: HashMap<String, V> = HashMap::new(); + entries + .into_iter() + .try_for_each(|(fid, finfo): (String, U)| -> anyhow::Result<()> { + out_info.insert(fid, finfo.try_into()?); + Ok(()) + })?; + Ok(StagedFiles { entries: out_info }) + } +} + +impl<U, V, S> std::convert::From<StagedFiles<U>> for HashMap<String, V, S> +where + V: std::convert::From<U>, + S: std::hash::BuildHasher + Default, +{ + fn from(reg: StagedFiles<U>) -> Self { + let mut out_info: HashMap<String, V, S> = HashMap::default(); + reg.entries + .into_iter() + .for_each(|(fid, finfo): (String, U)| { + out_info.insert(fid, finfo.into()); + }); + out_info + } +} diff --git a/types/src/staged_function.rs b/types/src/staged_function.rs new file mode 100644 index 0000000..cff442c --- /dev/null +++ b/types/src/staged_function.rs @@ -0,0 +1,176 @@ +use crate::{ExecutorType, StagedFiles, StagedInputFile, StagedOutputFile}; + +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::prelude::v1::*; +use std::str::FromStr; + +use anyhow::{Context, Result}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ArgumentValue { + inner: String, +} + +impl ArgumentValue { + pub fn new(value: String) -> Self { + Self { inner: value } + } + + pub fn inner(&self) -> &String { + &self.inner + } + + pub fn as_str(&self) -> &str { + &self.inner + } + + pub fn as_usize(&self) -> Result<usize> { + usize::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) + } + + pub fn as_u32(&self) -> Result<u32> { + u32::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) + } + + pub fn as_f32(&self) -> Result<f32> { + f32::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) + } + + pub fn as_f64(&self) -> Result<f64> { + f64::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) + } + + pub fn as_u8(&self) -> Result<u8> { + u8::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) + } +} + +impl std::fmt::Display for ArgumentValue { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.inner) + } +} + +#[derive(Clone, Serialize, Deserialize, Debug, Default)] +pub struct FunctionArguments { + #[serde(flatten)] + pub inner: HashMap<String, ArgumentValue>, +} + +impl<S: core::default::Default + std::hash::BuildHasher> From<FunctionArguments> + for HashMap<String, String, S> +{ + fn from(arguments: FunctionArguments) -> Self { + arguments + .inner() + .iter() + .map(|(k, v)| (k.to_owned(), v.as_str().to_owned())) + .collect() + } +} + +impl From<HashMap<String, String>> for FunctionArguments { + fn from(map: HashMap<String, String>) -> Self { + FunctionArguments::from_map(&map) + } +} + +impl FunctionArguments { + pub fn from_map<K, V>(input: &HashMap<K, V>) -> Self + where + K: std::string::ToString, + V: std::string::ToString, + { + let inner = input.iter().fold(HashMap::new(), |mut acc, (k, v)| { + acc.insert(k.to_string(), ArgumentValue::new(v.to_string())); + acc + }); + + Self { inner } + } + + pub fn inner(&self) -> &HashMap<String, ArgumentValue> { + &self.inner + } + + pub fn get(&self, key: &str) -> anyhow::Result<&ArgumentValue> { + self.inner + .get(key) + .with_context(|| format!("key not found: {}", key)) + } + + pub fn into_vec(self) -> Vec<String> { + let mut vector = Vec::new(); + + self.inner.into_iter().for_each(|(k, v)| { + vector.push(k); + vector.push(v.to_string()); + }); + + vector + } +} + +#[derive(Debug, Default)] +pub struct StagedFunction { + pub name: String, + pub payload: String, + pub arguments: FunctionArguments, + pub input_files: StagedFiles<StagedInputFile>, + pub output_files: StagedFiles<StagedOutputFile>, + pub runtime_name: String, + pub executor_type: ExecutorType, +} + +impl StagedFunction { + pub fn new() -> Self { + Self::default() + } + + pub fn name(self, name: impl ToString) -> Self { + Self { + name: name.to_string(), + ..self + } + } + + pub fn payload(self, payload: impl ToString) -> Self { + Self { + payload: payload.to_string(), + ..self + } + } + + pub fn arguments(self, arguments: FunctionArguments) -> Self { + Self { arguments, ..self } + } + + pub fn input_files(self, input_files: StagedFiles<StagedInputFile>) -> Self { + Self { + input_files, + ..self + } + } + + pub fn output_files(self, output_files: StagedFiles<StagedOutputFile>) -> Self { + Self { + output_files, + ..self + } + } + + pub fn runtime_name(self, runtime_name: impl ToString) -> Self { + Self { + runtime_name: runtime_name.to_string(), + ..self + } + } + + pub fn executor_type(self, executor_type: ExecutorType) -> Self { + Self { + executor_type, + ..self + } + } +} diff --git a/types/src/staged_task.rs b/types/src/staged_task.rs index 04a29e1..f3636a8 100644 --- a/types/src/staged_task.rs +++ b/types/src/staged_task.rs @@ -13,17 +13,17 @@ use crate::{ const STAGED_TASK_PREFIX: &str = "staged-"; // staged-task-uuid pub const QUEUE_KEY: &str = "staged-task"; -pub type FunctionInputData = HashMap<String, InputDataValue>; -pub type FunctionOutputData = HashMap<String, OutputDataValue>; +pub type FunctionInputFiles = HashMap<String, FunctionInputFile>; +pub type FunctionOutputFiles = HashMap<String, FunctionOutputFile>; #[derive(Debug, Deserialize, Serialize)] -pub struct InputDataValue { +pub struct FunctionInputFile { pub url: Url, pub hash: String, pub crypto_info: TeaclaveFileCryptoInfo, } -impl InputDataValue { +impl FunctionInputFile { pub fn new(url: Url, hash: impl ToString, crypto_info: TeaclaveFileCryptoInfo) -> Self { Self { url, @@ -42,12 +42,12 @@ impl InputDataValue { } #[derive(Debug, Deserialize, Serialize)] -pub struct OutputDataValue { +pub struct FunctionOutputFile { pub url: Url, pub crypto_info: TeaclaveFileCryptoInfo, } -impl OutputDataValue { +impl FunctionOutputFile { pub fn new(url: Url, crypto_info: TeaclaveFileCryptoInfo) -> Self { Self { url, crypto_info } } @@ -67,8 +67,8 @@ pub struct StagedTask { pub function_name: String, pub function_payload: Vec<u8>, pub function_arguments: FunctionArguments, - pub input_data: FunctionInputData, - pub output_data: FunctionOutputData, + pub input_data: FunctionInputFiles, + pub output_data: FunctionOutputFiles, } impl Storable for StagedTask { @@ -118,11 +118,11 @@ impl StagedTask { } } - pub fn input_data(self, input_data: FunctionInputData) -> Self { + pub fn input_data(self, input_data: FunctionInputFiles) -> Self { Self { input_data, ..self } } - pub fn output_data(self, output_data: FunctionOutputData) -> Self { + pub fn output_data(self, output_data: FunctionOutputFiles) -> Self { Self { output_data, ..self diff --git a/types/src/worker.rs b/types/src/worker.rs index c41f6bd..ddcc952 100644 --- a/types/src/worker.rs +++ b/types/src/worker.rs @@ -1,23 +1,8 @@ use std::collections::HashMap; use std::collections::HashSet; -use std::format; -use std::io::{self, Read, Write}; use std::prelude::v1::*; -#[cfg(not(feature = "mesalock_sgx"))] -use std::fs::File; -#[cfg(feature = "mesalock_sgx")] -use std::untrusted::fs::File; - use anyhow; -use anyhow::Context; -use anyhow::Result; - -use crate::TeaclaveFileCryptoInfo; -use crate::TeaclaveFileRootKey128; -use protected_fs::ProtectedFile; -use serde::{Deserialize, Serialize}; -use std::str::FromStr; #[derive(Debug, Copy, Clone)] pub enum ExecutorType { @@ -25,6 +10,12 @@ pub enum ExecutorType { Python, } +impl std::default::Default for ExecutorType { + fn default() -> Self { + ExecutorType::Native + } +} + impl std::convert::TryFrom<&str> for ExecutorType { type Error = anyhow::Error; @@ -47,296 +38,12 @@ impl std::fmt::Display for ExecutorType { } } -#[derive(Clone, Debug, Default)] -pub struct TeaclaveWorkerInputFileInfo { - pub path: std::path::PathBuf, - pub crypto_info: TeaclaveFileRootKey128, -} - -#[derive(Clone, Debug, Default)] -pub struct TeaclaveWorkerOutputFileInfo { - pub path: std::path::PathBuf, - pub crypto_info: TeaclaveFileRootKey128, -} - -impl std::convert::From<TeaclaveWorkerOutputFileInfo> for TeaclaveWorkerInputFileInfo { - fn from(info: TeaclaveWorkerOutputFileInfo) -> Self { - TeaclaveWorkerInputFileInfo { - path: info.path, - crypto_info: info.crypto_info, - } - } -} - -impl TeaclaveWorkerInputFileInfo { - pub fn new( - path: impl std::convert::Into<std::path::PathBuf>, - crypto_info: TeaclaveFileRootKey128, - ) -> Self { - TeaclaveWorkerInputFileInfo { - path: path.into(), - crypto_info, - } - } - - pub fn get_readable_io(&self) -> anyhow::Result<Box<dyn io::Read>> { - log::debug!("path: {:?}", self.path); - log::debug!("key: {:?}", self.crypto_info.key); - let f = ProtectedFile::open_ex(&self.path, &self.crypto_info.key)?; - Ok(Box::new(f)) - } - - #[cfg(test_mode)] - pub fn create_with_plaintext_file( - path: impl AsRef<std::path::Path>, - ) -> anyhow::Result<TeaclaveWorkerInputFileInfo> { - let bytes = read_all_bytes(path.as_ref())?; - let dst = path.as_ref().with_extension("enc"); - Self::create_with_bytes(dst, &bytes) - } - - pub fn create_with_bytes( - path: impl AsRef<std::path::Path>, - bytes: &[u8], - ) -> anyhow::Result<TeaclaveWorkerInputFileInfo> { - let crypto = TeaclaveFileRootKey128::random(); - let mut f = ProtectedFile::create_ex(&path, &crypto.key)?; - f.write_all(bytes)?; - Ok(Self::new(path.as_ref(), crypto)) - } -} - -impl TeaclaveWorkerOutputFileInfo { - pub fn new( - path: impl std::convert::Into<std::path::PathBuf>, - crypto_info: TeaclaveFileRootKey128, - ) -> Self { - TeaclaveWorkerOutputFileInfo { - path: path.into(), - crypto_info, - } - } - - pub fn get_writable_io(&self) -> anyhow::Result<Box<dyn io::Write>> { - log::debug!("path: {:?}", self.path); - log::debug!("key: {:?}", self.crypto_info.key); - let f = ProtectedFile::create_ex(&self.path, &self.crypto_info.key)?; - Ok(Box::new(f)) - } - - #[cfg(test_mode)] - pub fn get_plaintext(&self) -> anyhow::Result<Vec<u8>> { - let mut content = Vec::new(); - let mut f = ProtectedFile::open_ex(&self.path, &self.crypto_info.key)?; - f.read_to_end(&mut content)?; - Ok(content) - } -} - -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) -} - -pub fn convert_encrypted_input_file( - path: impl AsRef<std::path::Path>, - crypto_info: TeaclaveFileCryptoInfo, - dst: impl AsRef<std::path::Path>, -) -> anyhow::Result<TeaclaveWorkerInputFileInfo> { - log::debug!("from: {:?}, to: {:?}", path.as_ref(), dst.as_ref()); - #[cfg(not(feature = "mesalock_sgx"))] - use std::fs; - #[cfg(feature = "mesalock_sgx")] - use std::untrusted::fs; - let plain_text = match crypto_info { - TeaclaveFileCryptoInfo::AesGcm128(crypto) => { - let mut bytes = read_all_bytes(path)?; - crypto.decrypt(&mut bytes)?; - bytes - } - TeaclaveFileCryptoInfo::AesGcm256(crypto) => { - let mut bytes = read_all_bytes(path)?; - crypto.decrypt(&mut bytes)?; - bytes - } - TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto) => { - fs::copy(path, dst.as_ref())?; - let dst = dst.as_ref().to_owned(); - return Ok(TeaclaveWorkerInputFileInfo::new(dst, crypto)); - } - TeaclaveFileCryptoInfo::Raw => read_all_bytes(path)?, - }; - TeaclaveWorkerInputFileInfo::create_with_bytes(dst.as_ref(), &plain_text) -} - -#[derive(Debug, Default)] -pub struct TeaclaveWorkerFileRegistry<T> { - pub entries: HashMap<String, T>, -} - -impl<T> TeaclaveWorkerFileRegistry<T> { - pub fn new(entries: HashMap<String, T>) -> Self { - TeaclaveWorkerFileRegistry { entries } - } -} - -impl<U, V> std::convert::TryFrom<HashMap<String, U>> for TeaclaveWorkerFileRegistry<V> -where - U: std::convert::TryInto<V, Error = anyhow::Error>, -{ - type Error = anyhow::Error; - fn try_from(entries: HashMap<String, U>) -> anyhow::Result<Self> { - let mut out_info: HashMap<String, V> = HashMap::new(); - entries - .into_iter() - .try_for_each(|(fid, finfo): (String, U)| -> anyhow::Result<()> { - out_info.insert(fid, finfo.try_into()?); - Ok(()) - })?; - Ok(TeaclaveWorkerFileRegistry { entries: out_info }) - } -} - -impl<U, V, S> std::convert::From<TeaclaveWorkerFileRegistry<U>> for HashMap<String, V, S> -where - V: std::convert::From<U>, - S: std::hash::BuildHasher + Default, -{ - fn from(reg: TeaclaveWorkerFileRegistry<U>) -> Self { - let mut out_info: HashMap<String, V, S> = HashMap::default(); - reg.entries - .into_iter() - .for_each(|(fid, finfo): (String, U)| { - out_info.insert(fid, finfo.into()); - }); - out_info - } -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ArgumentValue { - inner: String, -} - -impl ArgumentValue { - pub fn new(value: String) -> Self { - Self { inner: value } - } - - pub fn inner(&self) -> &String { - &self.inner - } - - pub fn as_str(&self) -> &str { - &self.inner - } - - pub fn as_usize(&self) -> Result<usize> { - usize::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) - } - - pub fn as_u32(&self) -> Result<u32> { - u32::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) - } - - pub fn as_f32(&self) -> Result<f32> { - f32::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) - } - - pub fn as_f64(&self) -> Result<f64> { - f64::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) - } - - pub fn as_u8(&self) -> Result<u8> { - u8::from_str(&self.inner).with_context(|| format!("cannot parse {}", self.inner)) - } -} - -impl std::fmt::Display for ArgumentValue { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.inner) - } -} - -#[derive(Clone, Serialize, Deserialize, Debug, Default)] -pub struct FunctionArguments { - #[serde(flatten)] - pub inner: HashMap<String, ArgumentValue>, -} - -impl<S: core::default::Default + std::hash::BuildHasher> From<FunctionArguments> - for HashMap<String, String, S> -{ - fn from(arguments: FunctionArguments) -> Self { - arguments - .inner() - .iter() - .map(|(k, v)| (k.to_owned(), v.as_str().to_owned())) - .collect() - } -} - -impl From<HashMap<String, String>> for FunctionArguments { - fn from(map: HashMap<String, String>) -> Self { - FunctionArguments::from_map(&map) - } -} - -impl FunctionArguments { - pub fn from_map<K, V>(input: &HashMap<K, V>) -> Self - where - K: std::string::ToString, - V: std::string::ToString, - { - let inner = input.iter().fold(HashMap::new(), |mut acc, (k, v)| { - acc.insert(k.to_string(), ArgumentValue::new(v.to_string())); - acc - }); - - Self { inner } - } - - pub fn inner(&self) -> &HashMap<String, ArgumentValue> { - &self.inner - } - - pub fn get(&self, key: &str) -> anyhow::Result<&ArgumentValue> { - self.inner - .get(key) - .with_context(|| format!("key not found: {}", key)) - } - - pub fn into_vec(self) -> Vec<String> { - let mut vector = Vec::new(); - - self.inner.into_iter().for_each(|(k, v)| { - vector.push(k); - vector.push(v.to_string()); - }); - - vector - } -} - #[derive(Debug)] pub struct WorkerCapability { pub runtimes: HashSet<String>, pub functions: HashSet<String>, } -#[derive(Debug)] -pub struct StagedFunction { - pub name: String, - pub payload: String, - pub arguments: FunctionArguments, - pub input_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerInputFileInfo>, - pub output_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>, - pub runtime_name: String, - pub executor_type: ExecutorType, -} - #[derive(Default)] pub struct ExecutionResult { pub return_value: Vec<u8>, diff --git a/worker/src/function/context.rs b/worker/src/function/context.rs index d13d82a..7510858 100644 --- a/worker/src/function/context.rs +++ b/worker/src/function/context.rs @@ -224,10 +224,10 @@ pub mod tests { use std::str::FromStr; use teaclave_test_utils::*; use teaclave_types::hashmap; + use teaclave_types::StagedFiles; + use teaclave_types::StagedInputFile; + use teaclave_types::StagedOutputFile; use teaclave_types::TeaclaveFileRootKey128; - use teaclave_types::TeaclaveWorkerFileRegistry; - use teaclave_types::TeaclaveWorkerInputFileInfo; - use teaclave_types::TeaclaveWorkerOutputFileInfo; pub fn run_tests() -> bool { run_tests!(test_file_handle_encoding, test_rtc_api,) @@ -246,18 +246,17 @@ pub mod tests { let input = PathBuf::from_str("fixtures/functions/mesapy/input.txt").unwrap(); let output = PathBuf::from_str("fixtures/functions/mesapy/output.txt.out").unwrap(); - let input_info = TeaclaveWorkerInputFileInfo::new(input, TeaclaveFileRootKey128::random()); + let input_info = StagedInputFile::new(input, TeaclaveFileRootKey128::random()); - let output_info = - TeaclaveWorkerOutputFileInfo::new(output, TeaclaveFileRootKey128::random()); + let output_info = StagedOutputFile::new(output, TeaclaveFileRootKey128::random()); let in_fid = "in_f1"; let out_fid = "out_f1"; - let input_files = TeaclaveWorkerFileRegistry { + let input_files = StagedFiles { entries: hashmap!(in_fid.to_string() => input_info), }; - let output_files = TeaclaveWorkerFileRegistry { + let output_files = StagedFiles { entries: hashmap!(out_fid.to_string() => output_info), }; diff --git a/worker/src/function/echo.rs b/worker/src/function/echo.rs index cbd2e60..a98ee16 100644 --- a/worker/src/function/echo.rs +++ b/worker/src/function/echo.rs @@ -44,7 +44,7 @@ pub mod tests { use teaclave_types::hashmap; use teaclave_types::FunctionArguments; - use teaclave_types::TeaclaveWorkerFileRegistry; + use teaclave_types::StagedFiles; use crate::function::TeaclaveFunction; use crate::runtime::RawIoRuntime; @@ -58,8 +58,8 @@ pub mod tests { "message" => "Hello Teaclave!" )); - let input_files = TeaclaveWorkerFileRegistry::default(); - let output_files = TeaclaveWorkerFileRegistry::default(); + let input_files = StagedFiles::default(); + let output_files = StagedFiles::default(); let runtime = Box::new(RawIoRuntime::new(input_files, output_files)); let function = Echo; diff --git a/worker/src/function/gbdt_prediction.rs b/worker/src/function/gbdt_prediction.rs index ff51227..10ebdfc 100644 --- a/worker/src/function/gbdt_prediction.rs +++ b/worker/src/function/gbdt_prediction.rs @@ -102,10 +102,10 @@ pub mod tests { use teaclave_types::hashmap; use teaclave_types::FunctionArguments; + use teaclave_types::StagedFiles; + use teaclave_types::StagedInputFile; + use teaclave_types::StagedOutputFile; use teaclave_types::TeaclaveFileRootKey128; - use teaclave_types::TeaclaveWorkerFileRegistry; - use teaclave_types::TeaclaveWorkerInputFileInfo; - use teaclave_types::TeaclaveWorkerOutputFileInfo; use crate::function::TeaclaveFunction; use crate::runtime::RawIoRuntime; @@ -122,16 +122,16 @@ pub mod tests { let plain_output = "fixtures/functions/gbdt_prediction/result.txt.out"; let expected_output = "fixtures/functions/gbdt_prediction/expected_result.txt"; - let input_files = TeaclaveWorkerFileRegistry::new(hashmap!( + let input_files = StagedFiles::new(hashmap!( IN_MODEL.to_string() => - TeaclaveWorkerInputFileInfo::new(plain_if_model, TeaclaveFileRootKey128::random()), + StagedInputFile::new(plain_if_model, TeaclaveFileRootKey128::random()), IN_DATA.to_string() => - TeaclaveWorkerInputFileInfo::new(plain_if_data, TeaclaveFileRootKey128::random()) + StagedInputFile::new(plain_if_data, TeaclaveFileRootKey128::random()) )); - let output_files = TeaclaveWorkerFileRegistry::new(hashmap!( + let output_files = StagedFiles::new(hashmap!( OUT_RESULT.to_string() => - TeaclaveWorkerOutputFileInfo::new(plain_output, TeaclaveFileRootKey128::random()) + StagedOutputFile::new(plain_output, TeaclaveFileRootKey128::random()) )); let runtime = Box::new(RawIoRuntime::new(input_files, output_files)); diff --git a/worker/src/function/gbdt_training.rs b/worker/src/function/gbdt_training.rs index 8a25d93..df45c6f 100644 --- a/worker/src/function/gbdt_training.rs +++ b/worker/src/function/gbdt_training.rs @@ -137,10 +137,10 @@ pub mod tests { use teaclave_types::hashmap; use teaclave_types::FunctionArguments; + use teaclave_types::StagedFiles; + use teaclave_types::StagedInputFile; + use teaclave_types::StagedOutputFile; use teaclave_types::TeaclaveFileRootKey128; - use teaclave_types::TeaclaveWorkerFileRegistry; - use teaclave_types::TeaclaveWorkerInputFileInfo; - use teaclave_types::TeaclaveWorkerOutputFileInfo; use crate::function::TeaclaveFunction; use crate::runtime::RawIoRuntime; @@ -166,14 +166,14 @@ pub mod tests { let plain_output = "fixtures/functions/gbdt_training/training_model.txt.out"; let expected_output = "fixtures/functions/gbdt_training/expected_model.txt"; - let input_files = TeaclaveWorkerFileRegistry::new(hashmap!( + let input_files = StagedFiles::new(hashmap!( IN_DATA.to_string() => - TeaclaveWorkerInputFileInfo::new(plain_input, TeaclaveFileRootKey128::random()) + StagedInputFile::new(plain_input, TeaclaveFileRootKey128::random()) )); - let output_files = TeaclaveWorkerFileRegistry::new(hashmap!( + let output_files = StagedFiles::new(hashmap!( OUT_MODEL.to_string() => - TeaclaveWorkerOutputFileInfo::new(plain_output, TeaclaveFileRootKey128::random()) + StagedOutputFile::new(plain_output, TeaclaveFileRootKey128::random()) )); let runtime = Box::new(RawIoRuntime::new(input_files, output_files)); diff --git a/worker/src/function/mesapy.rs b/worker/src/function/mesapy.rs index b26c82c..b1d4828 100644 --- a/worker/src/function/mesapy.rs +++ b/worker/src/function/mesapy.rs @@ -109,10 +109,10 @@ pub mod tests { use crate::runtime::RawIoRuntime; use teaclave_types::hashmap; use teaclave_types::FunctionArguments; + use teaclave_types::StagedFiles; + use teaclave_types::StagedInputFile; + use teaclave_types::StagedOutputFile; use teaclave_types::TeaclaveFileRootKey128; - use teaclave_types::TeaclaveWorkerFileRegistry; - use teaclave_types::TeaclaveWorkerInputFileInfo; - use teaclave_types::TeaclaveWorkerOutputFileInfo; pub fn run_tests() -> bool { run_tests!(test_mesapy,) @@ -168,16 +168,15 @@ def entrypoint(argv): let input = "fixtures/functions/mesapy/input.txt"; let output = "fixtures/functions/mesapy/output.txt"; - let input_info = TeaclaveWorkerInputFileInfo::new(input, TeaclaveFileRootKey128::random()); + let input_info = StagedInputFile::new(input, TeaclaveFileRootKey128::random()); - let output_info = - TeaclaveWorkerOutputFileInfo::new(output, TeaclaveFileRootKey128::random()); + let output_info = StagedOutputFile::new(output, TeaclaveFileRootKey128::random()); - let input_files = TeaclaveWorkerFileRegistry { + let input_files = StagedFiles { entries: hashmap!("in_f1".to_string() => input_info), }; - let output_files = TeaclaveWorkerFileRegistry { + let output_files = StagedFiles { entries: hashmap!("out_f1".to_string() => output_info), }; let runtime = Box::new(RawIoRuntime::new(input_files, output_files)); diff --git a/worker/src/runtime/default.rs b/worker/src/runtime/default.rs index 00e1ce3..cf36793 100644 --- a/worker/src/runtime/default.rs +++ b/worker/src/runtime/default.rs @@ -22,19 +22,19 @@ use anyhow; use std::io; use super::TeaclaveRuntime; -use teaclave_types::TeaclaveWorkerFileRegistry; -use teaclave_types::TeaclaveWorkerInputFileInfo; -use teaclave_types::TeaclaveWorkerOutputFileInfo; +use teaclave_types::StagedFiles; +use teaclave_types::StagedInputFile; +use teaclave_types::StagedOutputFile; pub struct DefaultRuntime { - input_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerInputFileInfo>, - output_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>, + input_files: StagedFiles<StagedInputFile>, + output_files: StagedFiles<StagedOutputFile>, } impl DefaultRuntime { pub fn new( - input_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerInputFileInfo>, - output_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>, + input_files: StagedFiles<StagedInputFile>, + output_files: StagedFiles<StagedOutputFile>, ) -> DefaultRuntime { DefaultRuntime { input_files, diff --git a/worker/src/runtime/raw_io.rs b/worker/src/runtime/raw_io.rs index 50d5a29..8ceb457 100644 --- a/worker/src/runtime/raw_io.rs +++ b/worker/src/runtime/raw_io.rs @@ -24,19 +24,19 @@ use std::untrusted::fs::File; use anyhow; use super::TeaclaveRuntime; -use teaclave_types::TeaclaveWorkerFileRegistry; -use teaclave_types::TeaclaveWorkerInputFileInfo; -use teaclave_types::TeaclaveWorkerOutputFileInfo; +use teaclave_types::StagedFiles; +use teaclave_types::StagedInputFile; +use teaclave_types::StagedOutputFile; pub struct RawIoRuntime { - input_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerInputFileInfo>, - output_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>, + input_files: StagedFiles<StagedInputFile>, + output_files: StagedFiles<StagedOutputFile>, } impl RawIoRuntime { pub fn new( - input_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerInputFileInfo>, - output_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>, + input_files: StagedFiles<StagedInputFile>, + output_files: StagedFiles<StagedOutputFile>, ) -> RawIoRuntime { RawIoRuntime { input_files, diff --git a/worker/src/worker.rs b/worker/src/worker.rs index 84ae28c..bbcdf68 100644 --- a/worker/src/worker.rs +++ b/worker/src/worker.rs @@ -25,8 +25,8 @@ use anyhow; use serde_json; use teaclave_types::{ - ExecutorType, FunctionArguments, StagedFunction, TeaclaveWorkerFileRegistry, - TeaclaveWorkerInputFileInfo, TeaclaveWorkerOutputFileInfo, WorkerCapability, + ExecutorType, FunctionArguments, StagedFiles, StagedFunction, StagedInputFile, + StagedOutputFile, WorkerCapability, }; use crate::function::{self, TeaclaveFunction}; @@ -88,8 +88,8 @@ impl Worker { fn get_runtime( &self, name: &str, - input_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerInputFileInfo>, - output_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>, + input_files: StagedFiles<StagedInputFile>, + output_files: StagedFiles<StagedOutputFile>, ) -> anyhow::Result<Box<dyn TeaclaveRuntime + Send + Sync>> { let build_runtime = self .runtimes @@ -176,8 +176,8 @@ fn prepare_arguments( type FunctionBuilder = Box<dyn Fn() -> Box<dyn TeaclaveFunction + Send + Sync> + Send + Sync>; type RuntimeBuilder = Box< dyn Fn( - TeaclaveWorkerFileRegistry<TeaclaveWorkerInputFileInfo>, - TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>, + StagedFiles<StagedInputFile>, + StagedFiles<StagedOutputFile>, ) -> Box<dyn TeaclaveRuntime + Send + Sync> + Send + Sync, --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
