Test remote attestation in UEFI app

experimental/uefi/app/Cargo.toml

@@ -10,6 +10,7 @@ uefi = { version = "*", features = ["exts"] }
 uefi-services = "*"
 log = { version = "*" }
 oak_remote_attestation = { path = "../../../remote_attestation/rust" }
+anyhow = { version = "*", default-features = false }
 
 [dev-dependencies]
 uefi-services = { version = "*", features = ["qemu"] }

experimental/uefi/app/src/main.rs

@@ -26,7 +26,13 @@
 
 #[macro_use]
 extern crate log;
+extern crate alloc;
+extern crate anyhow;
 
+#[cfg(test)]
+use alloc::{boxed::Box, sync::Arc};
+#[cfg(test)]
+use oak_remote_attestation::handshaker::{AttestationBehavior, ClientHandshaker, ServerHandshaker};
 use uefi::{
     prelude::*,
     proto::console::serial::Serial,
@@ -130,3 +136,84 @@ fn test_simple() {
     let x = 1;
     assert_eq!(x, 1);
 }
+
+#[cfg(test)]
+const TEE_MEASUREMENT: &str = "Test TEE measurement";
+#[cfg(test)]
+const DATA: [u8; 10] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+
+#[cfg(test)]
+fn create_handshakers() -> (ClientHandshaker, ServerHandshaker) {
+    let bidirectional_attestation =
+        AttestationBehavior::create_bidirectional_attestation(&[], TEE_MEASUREMENT.as_bytes())
+            .unwrap();
+    let client_handshaker = ClientHandshaker::new(
+        bidirectional_attestation,
+        Box::new(|server_identity| {
+            if !server_identity.additional_info.is_empty() {
+                Ok(())
+            } else {
+                anyhow::bail!("No additional info provided.")
+            }
+        }),
+    );
+
+    let bidirectional_attestation =
+        AttestationBehavior::create_bidirectional_attestation(&[], TEE_MEASUREMENT.as_bytes())
+            .unwrap();
+
+    let additional_info = br"Additional Info".to_vec();
+    let server_handshaker =
+        ServerHandshaker::new(bidirectional_attestation, Arc::new(additional_info));
+
+    (client_handshaker, server_handshaker)
+}
+
+#[test_case]
+fn test_handshake() {
+    let (mut client_handshaker, mut server_handshaker) = create_handshakers();
+
+    let client_hello = client_handshaker
+        .create_client_hello()
+        .expect("Couldn't create client hello message");
+
+    let server_identity = server_handshaker
+        .next_step(&client_hello)
+        .expect("Couldn't process client hello message")
+        .expect("Empty server identity message");
+
+    let client_identity = client_handshaker
+        .next_step(&server_identity)
+        .expect("Couldn't process server identity message")
+        .expect("Empty client identity message");
+    assert!(client_handshaker.is_completed());
+
+    let result = server_handshaker
+        .next_step(&client_identity)
+        .expect("Couldn't process client identity message");
+    assert_eq!(result, None);
+    assert!(server_handshaker.is_completed());
+
+    let mut client_encryptor = client_handshaker
+        .get_encryptor()
+        .expect("Couldn't get client encryptor");
+    let mut server_encryptor = server_handshaker
+        .get_encryptor()
+        .expect("Couldn't get server encryptor");
+
+    let encrypted_client_data = client_encryptor
+        .encrypt(&DATA)
+        .expect("Couldn't encrypt client data");
+    let decrypted_client_data = server_encryptor
+        .decrypt(&encrypted_client_data)
+        .expect("Couldn't decrypt client data");
+    assert_eq!(decrypted_client_data, DATA);
+
+    let encrypted_server_data = server_encryptor
+        .encrypt(&DATA)
+        .expect("Couldn't encrypt server data");
+    let decrypted_server_data = client_encryptor
+        .decrypt(&encrypted_server_data)
+        .expect("Couldn't decrypt server data");
+    assert_eq!(decrypted_server_data, DATA);
+}

third_party/ring/Cargo.toml

@@ -33,6 +33,8 @@ include = [
     "benches/*.rs",
     "build.rs",
 
+    "stubs.c",
+
     "crypto/chacha/asm/chacha-armv4.pl",
     "crypto/chacha/asm/chacha-armv8.pl",
     "crypto/chacha/asm/chacha-x86.pl",

third_party/ring/build.rs

@@ -55,6 +55,7 @@ const RING_SRCS: &[(&[&str], &str)] = &[
     (&[X86], "crypto/chacha/asm/chacha-x86.pl"),
     (&[X86], "crypto/fipsmodule/modes/asm/ghash-x86.pl"),
 
+    (&[X86_64], "stubs.c"),
     (&[X86_64], "crypto/chacha/asm/chacha-x86_64.pl"),
     (&[X86_64], "crypto/fipsmodule/aes/asm/aesni-x86_64.pl"),
     (&[X86_64], "crypto/fipsmodule/aes/asm/vpaes-x86_64.pl"),

third_party/ring/stubs.c

@@ -0,0 +1,30 @@
+#include <stddef.h>
+
+/**
+ *
+  Stub function for win64 error handler API call inserted by nasm.
+  Stubbed as it unavailable in UEFI.
+  Ref: https://github.com/openssl/openssl/issues/12712.
+  Inspired by: https://github.com/tianocore/edk2/blob/7c0ad2c33810ead45b7919f8f8d0e282dae52e71/CryptoPkg/Library/OpensslLib/X64/ApiHooks.c
+**/
+void *
+__imp_RtlVirtualUnwind (
+  void  *Args
+  )
+{
+  return NULL;
+}
+
+/**
+  Stub function for win64 routine used for exceedomgy large variable calls.
+  Inserted MinGW, stubbed as it unavailable in UEFI.
+  Ref: https://metricpanda.com/rival-fortress-update-45-dealing-with-__chkstk-__chkstk_ms-when-cross-compiling-for-windows/
+**/
+void *
+___chkstk_ms (
+  void  *Args
+  )
+{
+  return NULL;
+}
+