diff --git a/rtc_auth_enclave/src/lib.rs b/rtc_auth_enclave/src/lib.rs index a036e0e2..082446f4 100644 --- a/rtc_auth_enclave/src/lib.rs +++ b/rtc_auth_enclave/src/lib.rs @@ -93,6 +93,7 @@ fn issue_execution_token_impl( { let token = token_store::issue_token( Uuid::from_bytes(message.dataset_uuid), + message.dataset_access_key, message.exec_module_hash, message.number_of_uses, dataset_size, diff --git a/rtc_auth_enclave/src/token_store.rs b/rtc_auth_enclave/src/token_store.rs index 6c5e83ed..2923ddf0 100644 --- a/rtc_auth_enclave/src/token_store.rs +++ b/rtc_auth_enclave/src/token_store.rs @@ -1,3 +1,4 @@ +use std::collections::hash_map::Entry; use std::collections::HashMap; use std::io; use std::path::Path; @@ -14,16 +15,41 @@ use uuid::Uuid; use crate::{jwt, uuid_to_string}; +/// The set of execution tokens issued for a dataset. #[derive(Serialize, Deserialize)] -struct ExecutionTokenRecord { +struct ExecutionTokenSet { + dataset_uuid: Uuid, // XXX(Pi): This may be redundant? Remove, or keep for self-integrity checking? + + /// The dataset's access key. + access_key: [u8; 24], + + /// The dataset's unsealed size in bytes. + dataset_size: u64, + + /// Usage state of the issued execution tokens, by JWT ID (`jti`). + issued_tokens: HashMap, +} + +impl ExecutionTokenSet { + fn new(dataset_uuid: Uuid, access_key: [u8; 24], dataset_size: u64) -> ExecutionTokenSet { + ExecutionTokenSet { + dataset_uuid, + access_key, + dataset_size, + issued_tokens: HashMap::new(), + } + } +} + +/// Usage state of a single execution token. +#[derive(Serialize, Deserialize)] +struct ExecutionTokenState { exec_module_hash: [u8; 32], - dataset_uuid: Uuid, allowed_uses: u32, current_uses: u32, } -fn kv_store<'a>( -) -> MutexGuard<'a, impl KvStore, Error = io::Error>> { +fn kv_store<'a>() -> MutexGuard<'a, impl KvStore> { static TOKEN_FS_STORE: OnceCell>> = OnceCell::new(); let store = TOKEN_FS_STORE.get_or_init(|| { // TODO: Evaluate if this make sense, and what the possible attack vectors can be from relying on the @@ -38,9 +64,36 @@ fn kv_store<'a>( store.lock().expect("FS store mutex poisoned") } +/// Save a new dataset access key and associated metadata to the store. +/// +/// This must be called before [`issue_token`] can be called. +/// +/// # Panics +/// +/// If `dataset_uuid` already exists in the store. (This should not happen.) +#[allow(dead_code)] // TODO +pub(crate) fn save_access_key( + dataset_uuid: Uuid, + access_key: [u8; 24], + dataset_size: u64, +) -> Result<(), io::Error> { + let mut store = kv_store(); + let dataset_uuid_string = uuid_to_string(dataset_uuid); + let empty_token_set = ExecutionTokenSet::new(dataset_uuid, access_key, dataset_size); + + match store.try_insert(&dataset_uuid_string, &empty_token_set)? { + None => Ok(()), + Some(_existing) => panic!( + "token_store::save_access_key: access key for dateset_uuid={:?} already saved (this should not happen)", + dataset_uuid, + ) + } +} + // Returns exec token hash pub(crate) fn issue_token( dataset_uuid: Uuid, + access_key: [u8; 24], exec_module_hash: [u8; 32], number_of_allowed_uses: u32, dataset_size: u64, @@ -48,35 +101,52 @@ pub(crate) fn issue_token( let EncodedExecutionToken { token, token_id } = EncodedExecutionToken::new(exec_module_hash, dataset_uuid, dataset_size); - save_token( - dataset_uuid, - token_id, + let token_state = ExecutionTokenState { exec_module_hash, - number_of_allowed_uses, - )?; + allowed_uses: number_of_allowed_uses, + current_uses: 0u32, + }; + + save_token(dataset_uuid, access_key, token_id, token_state)?; Ok(token) } +/// Save a newly-issued execution token's state to the store. +/// +/// Fail with error for invalid `dataset_uuid`. +/// +/// # Panics +/// +/// If `token_uuid` was already issued. fn save_token( dataset_uuid: Uuid, - token_uuid: Uuid, - exec_module_hash: [u8; 32], - number_of_allowed_uses: u32, + access_key: [u8; 24], + token_id: Uuid, + token_state: ExecutionTokenState, ) -> Result<(), io::Error> { let mut store = kv_store(); let dataset_uuid_string = uuid_to_string(dataset_uuid); - let new_record = ExecutionTokenRecord { - dataset_uuid, - exec_module_hash, - allowed_uses: number_of_allowed_uses, - current_uses: 0u32, - }; - store.alter(&dataset_uuid_string, |records| { - let mut records = records.unwrap_or_else(HashMap::new); - records.insert(token_uuid, new_record); - Some(records) - })?; - Ok(()) + let mut token_set = store + .load(&dataset_uuid_string)? + // TODO(Pi): Use something better than the io NotFound here? + .ok_or_else(|| io::ErrorKind::NotFound)?; + + // Update if the access key matches. + if token_set.access_key == access_key { + // TODO: Use [`HashMap::try_insert`] once stable. + // Unstable tracking issue: + match token_set.issued_tokens.entry(token_id) { + Entry::Occupied(_entry) => panic!( + "token_store::save_token: token_uuid={:?} already issued (this should not happen)", + token_id, + ), + Entry::Vacant(entry) => entry.insert(token_state), + }; + store.save(&dataset_uuid_string, &token_set)?; + Ok(()) + } else { + Err(io::ErrorKind::NotFound.into()) + } } diff --git a/rtc_data_service/tests/ecalls/issue_execution_token.rs b/rtc_data_service/tests/ecalls/issue_execution_token.rs index 4e015ec6..8ca2887f 100644 --- a/rtc_data_service/tests/ecalls/issue_execution_token.rs +++ b/rtc_data_service/tests/ecalls/issue_execution_token.rs @@ -2,6 +2,7 @@ use std::convert::TryInto; use std::str::FromStr; use rtc_types::ExecReqMetadata; +use rtc_uenclave::{EnclaveConfig, RtcAuthEnclave}; use serde::{Deserialize, Serialize}; use sgx_types::sgx_target_info_t; @@ -15,10 +16,9 @@ pub struct ExecReqData { number_of_uses: u32, } -#[test] -fn test_issue_execution_token_success() { - let enclave = helpers::init_auth_enclave(); - +fn make_request( + enclave: &RtcAuthEnclave, +) -> ([u8; 32], [u8; 32], Vec, ExecReqMetadata) { let enclave_pubkey = enclave .create_report(&sgx_target_info_t::default()) .unwrap() @@ -53,15 +53,26 @@ fn test_issue_execution_token_success() { ) .unwrap(); - let result = enclave - .issue_execution_token( - &ciphertext[CRYPTO_BOX_BOXZEROBYTES..], - ExecReqMetadata { - uploader_pub_key: pubkey, - nonce, - }, - ) - .unwrap(); + let payload = ciphertext[CRYPTO_BOX_BOXZEROBYTES..].to_vec(); + + let metadata = ExecReqMetadata { + uploader_pub_key: pubkey, + nonce, + }; + + (enclave_pubkey, privkey, payload, metadata) +} + +// FIXME: The success case currently fails because there's no actual data / access key. +// Complete the test once that's available. +#[allow(dead_code)] +// #[test] +fn test_issue_execution_token_success() { + let enclave = helpers::init_auth_enclave(); + + let (enclave_pubkey, privkey, payload, metadata) = make_request(&enclave); + + let result = enclave.issue_execution_token(&payload, metadata).unwrap(); let mut m = vec![0_u8; result.ciphertext.len() + CRYPTO_BOX_BOXZEROBYTES]; @@ -89,3 +100,17 @@ fn test_issue_execution_token_success() { // TODO: Assert that decrypted value is a valid JWT } + +// Lookup failure: Invalid dataset UUID and access key. +#[test] +fn test_lookup_failure() { + let enclave = helpers::init_auth_enclave(); + + let (_enclave_pubkey, _privkey, payload, metadata) = make_request(&enclave); + + let err = enclave + .issue_execution_token(&payload, metadata) + .unwrap_err(); + + assert_eq!(format!("{:?}", err), "RtcEnclave(IO)") +} diff --git a/rtc_tenclave/src/kv_store/mod.rs b/rtc_tenclave/src/kv_store/mod.rs index 2edd5901..302df0c6 100644 --- a/rtc_tenclave/src/kv_store/mod.rs +++ b/rtc_tenclave/src/kv_store/mod.rs @@ -42,6 +42,27 @@ pub trait KvStore { }; Ok(altered) } + + /// Mutate the value of `key`. + /// + /// This is like [`Self::mutate`], but only operates on existing values. + fn mutate(&mut self, key: &str, mutate_fn: F) -> Result, Self::Error> + where + F: FnOnce(V) -> V, + { + self.alter(key, |opt_v| opt_v.map(mutate_fn)) + } + + /// Insert a value for `key`, if absent. If `key` already has a value, do nothing. + /// + /// Return the key's prior value (`None` if `value` was inserted) + fn try_insert(&mut self, key: &str, value: &V) -> Result, Self::Error> { + let loaded = self.load(key)?; + if loaded.is_none() { + self.save(key, value)?; + } + Ok(loaded) + } } #[cfg(test)] diff --git a/rtc_tenclave/src/kv_store/tests.rs b/rtc_tenclave/src/kv_store/tests.rs index 1198f783..130750a9 100644 --- a/rtc_tenclave/src/kv_store/tests.rs +++ b/rtc_tenclave/src/kv_store/tests.rs @@ -10,10 +10,36 @@ use tempfile::TempDir; use super::fs::std_filer::StdFiler; use super::fs::FsStore; -use super::in_memory::{InMemoryJsonStore, InMemoryStore}; +use super::in_memory::{InMemoryJsonStore, InMemoryStore, Never}; use super::inspect::InspectStore; use super::KvStore; +#[test] +fn test_mutate() -> Result<(), Never> { + let mut store = InMemoryStore::default(); + + assert_eq!(store.mutate("missing", |n| n + 1)?, None); + + store.save("existing", &2)?; + assert_eq!(store.mutate("existing", |n| n + 1)?, Some(3)); + assert_eq!(store.load("existing")?, Some(3)); + + Ok(()) +} +#[test] +fn test_try_insert() -> Result<(), Never> { + let mut store = InMemoryStore::default(); + + assert_eq!(store.try_insert("missing", &42)?, None); + assert_eq!(store.load("missing")?, Some(42)); + + store.save("existing", &5)?; + assert_eq!(store.try_insert("existing", &42)?, Some(5)); + assert_eq!(store.load("existing")?, Some(5)); + + Ok(()) +} + /// Verify that executing a sequence of store operations matches a simple model. #[test] fn prop_store_ops_match_model() {