WIP: Validate and use exec token

rtc_auth_enclave/Cargo.toml

@@ -23,6 +23,7 @@ jsonwebtoken = { git = "https://github.com/mesalock-linux/jsonwebtoken-sgx" }
 # TODO: confirm that we have to use a forked crate here
 uuid = { git = "https://github.com/mesalock-linux/uuid-sgx", features = ["v4", "serde"] }
 base64 = { git = "https://github.com/mesalock-linux/rust-base64-sgx" }
+thiserror = { git = "https://github.com/mesalock-linux/thiserror-sgx.git", tag = "sgx_1.1.3" }
 
 # See "Cargo patch limitation workaround" in HACKING.md:
 once_cell = { git = "https://github.com/mesalock-linux/once_cell-sgx.git" }

rtc_auth_enclave/src/jwt.rs

@@ -1,19 +1,25 @@
 use std::string::{String, ToString};
 use std::time::{SystemTime, UNIX_EPOCH};
+use std::vec;
 
-use jsonwebtoken::{encode, EncodingKey, Header};
+use jsonwebtoken::{encode, Algorithm, DecodingKey, EncodingKey, Header, TokenData, Validation};
 use serde::{Deserialize, Serialize};
 use sgx_tstd::untrusted::time::SystemTimeEx;
+use thiserror::Error;
 use uuid::Uuid;
 
 use crate::uuid_to_string;
 
+const HEADER_TYP: &'static str = "ntlexec+jwt";
+const CLAIMS_ISS: &'static str = "ntls_auth_enclave";
+const CLAIMS_AUD: &'static str = "ntls_exec_enclave";
+
 /// Claims body of the JWT token
 ///
 /// Example output:
 /// ```
 /// {
-///     "iss": "registree_auth_enclave",
+///     "iss": "ntls_auth_enclave",
 ///     "nbf": 1623762799,
 ///     "iat": 1623762799,
 ///     "jti": "b300fe149d144e05aa9a9600816b42ca",
@@ -22,16 +28,16 @@ use crate::uuid_to_string;
 /// }
 /// ```
 #[derive(Debug, Serialize, Deserialize)]
-struct Claims {
+pub(crate) struct Claims {
     // Registered Claim Names: https://datatracker.ietf.org/doc/html/rfc7519#section-4.1
     // TODO: serialize to hex string? This can be mrenclave or mrsigner
     iss: String,
     nbf: u64,
     iat: u64,
-    jti: String,
+    pub(crate) jti: String,
     // TODO: Better names. use `x-ntls-mod-hash` etc?
-    exec_module_hash: String,
-    dataset_uuid: String,
+    pub(crate) exec_module_hash: String,
+    pub(crate) dataset_uuid: String, // TODO: use sub?
     dataset_size: u64,
 }
 
@@ -49,7 +55,7 @@ impl Claims {
             .as_secs();
 
         Self {
-            iss: "registree_auth_enclave".to_string(),
+            iss: CLAIMS_ISS.to_string(),
             nbf: now,
             iat: now,
             jti: token_id,
@@ -60,6 +66,43 @@ impl Claims {
     }
 }
 
+#[derive(Debug, Error)]
+pub(crate) enum DecodeError {
+    #[error("Decoding failed: {}", .0)]
+    JWT(#[from] jsonwebtoken::errors::Error),
+    #[error("Invalid typ field in the jwt header")]
+    Typ,
+}
+
+pub(crate) struct DecodedExecutionToken(TokenData<Claims>);
+
+impl DecodedExecutionToken {
+    pub(crate) fn decode(token: &str) -> Result<DecodedExecutionToken, DecodeError> {
+        let validation = Validation {
+            validate_nbf: true,
+            iss: Some(CLAIMS_ISS.to_string()),
+            algorithms: vec![Algorithm::HS256],
+
+            ..Validation::default()
+        };
+
+        let decoded = jsonwebtoken::decode::<Claims>(token, &get_decoding_key(), &validation)?;
+
+        match decoded.header.typ.as_deref() {
+            Some(HEADER_TYP) => Ok(DecodedExecutionToken(decoded)),
+            Some(_) | None => Err(DecodeError::Typ),
+        }
+    }
+
+    pub(crate) fn claims<'a>(&'a self) -> &'a Claims {
+        &self.0.claims
+    }
+
+    pub(crate) fn header<'a>(&'a self) -> &'a Header {
+        &self.0.header
+    }
+}
+
 pub(crate) struct EncodedExecutionToken {
     pub token: String,
     pub token_id: Uuid,
@@ -76,14 +119,12 @@ impl EncodedExecutionToken {
             dataset_size,
         );
 
-        // TODO: Use a signing key that corresponds to the public key
-        // in the attestation enclave held data and move to crypto module in tenclave.
-        let encoding_key = EncodingKey::from_secret("secret".as_ref());
+        let encoding_key = get_encoding_key();
         // Header size 48 characters base64
         let header = Header {
             // Explicit typing for the token type
             // SEE: https://datatracker.ietf.org/doc/html/draft-ietf-secevent-token-02#section-2.2
-            typ: Some("ntlexec+jwt".to_string()),
+            typ: Some(HEADER_TYP.to_string()),
             ..Header::default()
         };
 
@@ -94,3 +135,14 @@ impl EncodedExecutionToken {
         Self { token, token_id }
     }
 }
+
+fn get_encoding_key() -> EncodingKey {
+    // TODO: Use a signing key that corresponds to the public key
+    // in the attestation enclave held data and move to crypto module in tenclave.
+    EncodingKey::from_secret("secret".as_ref())
+}
+
+fn get_decoding_key<'a>() -> DecodingKey<'a> {
+    // TODO: Use a decoding key that can is intrinsic to this enclave instance.
+    DecodingKey::from_secret("secret".as_ref())
+}

rtc_auth_enclave/src/lib.rs

@@ -10,14 +10,21 @@ mod token_store;
 extern crate sgx_tstd as std;
 
 use core::slice;
-use std::ptr;
 use std::string::{String, ToString};
+use std::{ptr, str};
 
 use rtc_tenclave::crypto::{RtcCrypto, SodaBoxCrypto as Crypto};
 pub use rtc_tenclave::dh::*;
 #[allow(unused_imports)] // for ECALL linking
 use rtc_tenclave::enclave::enclave_create_report;
-use rtc_types::{EcallResult, EncryptedMessage, ExecReqMetadata, ExecTokenError, IssueTokenResult};
+use rtc_types::{
+    CryptoError,
+    EcallResult,
+    EncryptedMessage,
+    ExecReqMetadata,
+    ExecTokenError,
+    IssueTokenResult,
+};
 use secrecy::{ExposeSecret, Secret};
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
@@ -121,6 +128,32 @@ fn validate_dataset_access_key(
     Some(20)
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn validate_and_use_token(
+    payload_ptr: *const u8,
+    payload_len: usize,
+    metadata_ptr: *const ExecReqMetadata,
+) {
+    let payload = unsafe { slice::from_raw_parts(payload_ptr, payload_len) };
+    let metadata = unsafe { &*metadata_ptr };
+    validate_and_use_token_impl(payload, metadata).unwrap();
+}
+
+fn validate_and_use_token_impl(
+    payload: &[u8],
+    metadata: &ExecReqMetadata,
+) -> Result<(), CryptoError> {
+    let crypto = Crypto::new();
+    let message = crypto.decrypt_message(payload, &metadata.uploader_pub_key, &metadata.nonce)?;
+    // TODO: Error handling
+    let token_str = str::from_utf8(message.expose_secret()).unwrap();
+
+    match token_store::validate_and_use(token_str) {
+        Ok(true) => Ok(()),
+        Ok(_) | Err(_) => todo!(),
+    }
+}
+
 pub(crate) fn uuid_to_string(uuid: Uuid) -> String {
     let mut uuid_buf = Uuid::encode_buffer();
     uuid.to_simple().encode_lower(&mut uuid_buf).to_string()

rtc_auth_enclave/src/token_store.rs

@@ -1,3 +1,4 @@
+use core::str::FromStr;
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::io;
@@ -13,6 +14,7 @@ use sgx_tstd::sync::{SgxMutex as Mutex, SgxMutexGuard as MutexGuard};
 use sgx_tstd::untrusted::{fs as untrusted_fs, path as untrusted_path};
 use uuid::Uuid;
 
+use crate::jwt::{Claims, DecodedExecutionToken};
 use crate::{jwt, uuid_to_string};
 
 /// The set of execution tokens issued for a dataset.
@@ -112,6 +114,46 @@ pub(crate) fn issue_token(
     Ok(token)
 }
 
+pub(crate) fn validate_and_use(token: &str) -> Result<bool, io::Error> {
+    let mut store = kv_store();
+    // TODO: Error handling
+    let decoded_token = DecodedExecutionToken::decode(token).unwrap();
+    let claims = decoded_token.claims();
+
+    match store.load(&claims.dataset_uuid)? {
+        // TODO: error handling
+        Some(mut token_set) => {
+            // Consistency check:
+            assert_eq!(
+                Uuid::from_str(&claims.dataset_uuid).unwrap(),
+                token_set.dataset_uuid
+            );
+
+            let token_id = Uuid::from_str(&claims.jti).unwrap();
+            match token_set.issued_tokens.get_mut(&token_id) {
+                Some(mut token_state) if token_claim_is_valid(claims, token_state).unwrap() => {
+                    token_state.current_uses += 1;
+                    store.save(&claims.dataset_uuid, &token_set)?;
+                    Ok(true)
+                }
+                Some(_) | None => Ok(false),
+            }
+        }
+        None => Ok(false),
+    }
+}
+
+fn token_claim_is_valid(
+    claims: &Claims,
+    token_state: &ExecutionTokenState,
+) -> Result<bool, uuid::Error> {
+    let has_uses = token_state.allowed_uses > token_state.current_uses;
+    let for_correct_exec_module =
+        base64::encode(token_state.exec_module_hash) == claims.exec_module_hash;
+
+    Ok(has_uses && for_correct_exec_module)
+}
+
 /// Save a newly-issued execution token's state to the store.
 ///
 /// Fail with error for invalid `dataset_uuid`.