New version of key wrapping using stateful CVM guest attestation + TPM

cgpuvm-attest-cli/src/main.rs

@@ -12,5 +12,6 @@ pub fn main() {
     let Ok(token) = attest(s.as_bytes(), 0xffff, maa_url) else {
         panic!("Failed to get MAA token")
     };
+
     println!("Got MAA token: {}", String::from_utf8(token).unwrap());
 }

cgpuvm-attest/src/err.rs

@@ -2,8 +2,10 @@ use thiserror::Error;
 
 #[derive(Error, Debug)]
 pub enum AttestError {
+    #[error("Failed to initialize guest attestation library")]
+    Initialization,
     #[error("Failed to convert endpoint URL to CString")]
     Convertion,
     #[error("CVM guest attestation library returned error: {0}")]
-    MAAToken(i32),
+    LibraryError(i32),
 }

cgpuvm-attest/src/lib.rs

@@ -1,7 +1,7 @@
 pub mod err;
 
 use err::AttestError;
-use libc::{c_char, c_int, size_t};
+use libc::{c_char, c_int, size_t, c_void};
 use std::ffi::CString;
 
 type Res<T> = Result<T, Box<dyn std::error::Error>>;
@@ -15,6 +15,29 @@ extern "C" {
         jwt_len: *mut size_t,
         endpoint_url: *const c_char,
     ) -> c_int;
+
+    fn ga_create(
+      st: *mut *mut c_void
+    ) -> c_int;
+
+    fn ga_free(
+      st: *mut c_void
+    );
+
+    fn ga_get_token(
+      st: *mut c_void,
+      app_data: *const u8,
+      pcr: u32,
+      jwt: *mut u8,
+      jwt_len: *mut size_t,
+      endpoint_url: *const c_char
+    ) -> c_int;
+
+    fn ga_decrypt(
+      st: *mut c_void,
+      cipher: *mut u8,
+      len: *mut size_t
+    ) -> c_int;
 }
 
 pub fn attest(data: &[u8], pcrs: u32, endpoint_url: &str) -> Res<Vec<u8>> {
@@ -23,17 +46,83 @@ pub fn attest(data: &[u8], pcrs: u32, endpoint_url: &str) -> Res<Vec<u8>> {
             let mut dstlen = 32 * 1024;
             let mut dst = Vec::with_capacity(dstlen);
             let pdst = dst.as_mut_ptr();
-
             let url_ptr = endpoint_url_cstring.as_ptr();
 
             let ret = get_attestation_token(data.as_ptr(), pcrs, pdst, &mut dstlen, url_ptr);
             if ret == 0 {
                 dst.set_len(dstlen);
                 Ok(dst)
             } else {
-                Err(Box::new(AttestError::MAAToken(ret)))
+                Err(Box::new(AttestError::LibraryError(ret)))
             }
         },
         _e => Err(Box::new(AttestError::Convertion)),
     }
 }
+
+pub struct AttestationClient {
+  st: *mut c_void
+}
+
+impl AttestationClient {
+  pub fn new() -> Res<AttestationClient> {
+    let mut c = AttestationClient { st: std::ptr::null_mut() };
+    unsafe {
+      let rc = ga_create(&mut c.st);
+
+      if rc == 0 {
+        return Ok(c);
+      }
+
+      return Err(Box::new(AttestError::Initialization));
+    }
+  }
+
+  pub fn attest(&mut self, data: &[u8], pcrs: u32, endpoint_url: &str) -> Res<Vec<u8>> {
+    match CString::new(endpoint_url) {
+      Ok(endpoint_url_cstring) =>
+        unsafe {
+            let url_ptr = endpoint_url_cstring.as_ptr();
+            let mut dstlen = 32 * 1024;
+            let mut dst = Vec::with_capacity(dstlen);
+            let pdst = dst.as_mut_ptr();
+            let rc = ga_get_token(self.st, data.as_ptr(), pcrs, pdst, &mut dstlen, url_ptr);
+
+            if rc == 0 {
+              dst.set_len(dstlen);
+              Ok(dst)
+            } else {
+              Err(Box::new(AttestError::LibraryError(rc)))
+            }
+        },
+      _ => Err(Box::new(AttestError::Convertion)),
+    }
+  }
+
+  pub fn decrypt(&mut self, data: &[u8]) -> Res<Vec<u8>> {
+    unsafe {
+      let mut buf = Vec::from(data);
+      let mut len = data.len();
+      let rc = ga_decrypt(self.st, buf.as_mut_ptr(), &mut len);
+
+      if rc == 0 {
+        buf.set_len(len);
+        Ok(buf)
+      } else {
+        Err(Box::new(AttestError::LibraryError(rc)))
+      }
+    }
+  }
+}
+
+impl Drop for AttestationClient {
+  fn drop(&mut self) {
+    unsafe {
+      ga_free(self.st);
+    }
+  }
+}
+
+unsafe impl Send for AttestationClient {
+
+}

ohttp-client/src/main.rs

@@ -240,7 +240,11 @@ async fn get_kms_config(kms_url: String, cert: &str) -> Res<String> {
 
     loop {
         // Make the GET request
-        let response = client.get(url.clone()).send().await?.error_for_status()?;
+        let response = client
+            .get(url.clone())
+            .send()
+            .await?
+            .error_for_status()?;
 
         // We may have to wait for receipt to be ready
         match response.status().as_u16() {
@@ -256,12 +260,12 @@ async fn get_kms_config(kms_url: String, cert: &str) -> Res<String> {
                 } else {
                     Err("Max retries reached, giving up. Cannot reach key management service")?;
                 }
-            }
+            },
             200 => {
                 let body = response.text().await?;
                 assert!(!body.is_empty());
-                return Ok(body);
-            }
+                return Ok(body)
+            },
             e => {
                 Err(format!("KMS returned unexpected {} status code.", e))?;
             }

ohttp-server/Cargo.toml

@@ -12,6 +12,7 @@ rust-hpke = ["ohttp/rust-hpke"]
 [dependencies]
 env_logger = {version = "0.10", default-features = false}
 hex = "0.4"
+base64 = "0.22.1"
 lazy_static = "1.4"
 moka = { version = "0.12", features = ["future"] }
 tokio = { version = "1", features = ["full"] }

ohttp-server/src/err.rs

@@ -18,4 +18,10 @@ pub enum ServerError {
     KMSUnreachable,
     #[error("Private key missing from SKR response")]
     PrivateKeyMissing,
+    #[error("CVM guest attestation library initialization failure")]
+    AttestationLibraryInit,
+    #[error("Guest attestation library failed to decrypt HPKE private key")]
+    TPMDecryptionFailure,
+    #[error("SKR for the requested KID has failed in the past 60 seconds, waiting to retry.")]
+    CachedSKRError,
 }