feat(rtc_auth_enclave::token_store): access key handling

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,

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<Uuid, ExecutionTokenState>,
+}
+
+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<HashMap<Uuid, ExecutionTokenRecord>, Error = io::Error>> {
+fn kv_store<'a>() -> MutexGuard<'a, impl KvStore<ExecutionTokenSet, Error = io::Error>> {
     static TOKEN_FS_STORE: OnceCell<Mutex<FsStore<SgxFiler>>> = 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,45 +64,89 @@ 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,
 ) -> Result<String, io::Error> {
     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: <https://github.com/rust-lang/rust/issues/82766>
+        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())
+    }
 }

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<EnclaveConfig>,
+) -> ([u8; 32], [u8; 32], Vec<u8>, 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)")
+}

rtc_tenclave/src/kv_store/mod.rs

@@ -42,6 +42,27 @@ pub trait KvStore<V> {
         };
         Ok(altered)
     }
+
+    /// Mutate the value of `key`.
+    ///
+    /// This is like [`Self::mutate`], but only operates on existing values.
+    fn mutate<F>(&mut self, key: &str, mutate_fn: F) -> Result<Option<V>, 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<Option<V>, Self::Error> {
+        let loaded = self.load(key)?;
+        if loaded.is_none() {
+            self.save(key, value)?;
+        }
+        Ok(loaded)
+    }
 }
 
 #[cfg(test)]

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() {