WIP: Add basic validate and use functionality

Author: longtomjr

rtc_auth_enclave/Cargo.lock

@@ -26,6 +26,7 @@ serde = { git = "https://github.com/mesalock-linux/serde-sgx", feature = ["deriv
 serde_json = { git = "https://github.com/mesalock-linux/serde-json-sgx", tag = "sgx_1.1.3" }
 once_cell = { git = "https://github.com/mesalock-linux/once_cell-sgx.git", tag = "sgx_1.1.3" }
 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" }
 
 secrecy = { version = "0.7.0", default-features = false }
 

rtc_auth_enclave/Cargo.toml

@@ -1,19 +1,25 @@
 use std::string::{String, ToString};
 use std::time::{SystemTime, UNIX_EPOCH};
+use std::vec;
 
 use sgx_tstd::untrusted::time::SystemTimeEx;
 
 use crate::uuid_to_string;
-use jsonwebtoken::{encode, EncodingKey, Header};
+use jsonwebtoken::{encode, Algorithm, DecodingKey, EncodingKey, Header, TokenData, Validation};
 use serde::{Deserialize, Serialize};
+use thiserror::Error;
 use uuid::Uuid;
 
+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,15 +28,15 @@ use uuid::Uuid;
 /// }
 /// ```
 #[derive(Debug, Serialize, Deserialize)]
-struct Claims {
+pub(crate) struct Claims {
     // 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?
 }
 
 impl Claims {
@@ -42,7 +48,7 @@ impl Claims {
             .as_secs();
 
         Self {
-            iss: "registree_auth_enclave".to_string(),
+            iss: CLAIMS_ISS.to_string(),
             nbf: now,
             iat: now,
             jti: token_id,
@@ -52,6 +58,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,
@@ -67,14 +110,12 @@ impl EncodedExecutionToken {
             uuid_to_string(token_id),
         );
 
-        // 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()
         };
 
@@ -85,3 +126,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/jwt.rs

@@ -11,6 +11,7 @@ extern crate sgx_tstd as std;
 
 use core::slice;
 use std::ptr;
+use std::str;
 use std::string::{String, ToString};
 
 use rtc_tenclave::crypto::RtcCrypto;
@@ -19,7 +20,9 @@ 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;
@@ -114,6 +117,32 @@ fn valid_dataset_access_key(_dataset_uuid: [u8; 16], _access_key: [u8; 24]) -> b
     true
 }
 
+#[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);
+}
+
+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/lib.rs

@@ -1,3 +1,4 @@
+use core::str::FromStr;
 use std::collections::HashMap;
 use std::io;
 use std::path::Path;
@@ -17,6 +18,8 @@ use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
 use crate::jwt;
+use crate::jwt::Claims;
+use crate::jwt::DecodedExecutionToken;
 use crate::uuid_to_string;
 
 #[derive(Serialize, Deserialize)]
@@ -62,6 +65,35 @@ 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 records) => match records.get_mut(&Uuid::from_str(&claims.jti).unwrap()) {
+            Some(mut record) if record_is_valid(claims, record).unwrap() => {
+                record.current_uses += 1;
+                store.save(&claims.dataset_uuid, &records)?;
+                Ok(true)
+            }
+            Some(_) | None => Ok(false),
+        },
+        None => Ok(false),
+    }
+}
+
+fn record_is_valid(claims: &Claims, record: &ExecutionTokenRecord) -> Result<bool, uuid::Error> {
+    let has_uses = record.allowed_uses > record.current_uses;
+    let for_correct_dataset = record.dataset_uuid == Uuid::from_str(&claims.dataset_uuid)?;
+    let for_correct_exec_module =
+        base64::encode(record.exec_module_hash) == claims.exec_module_hash;
+
+    Ok(has_uses && for_correct_dataset && for_correct_exec_module)
+}
+
 fn save_token(
     dataset_uuid: Uuid,
     token_uuid: Uuid,