Add remote attestation test to the UEFI app

Author: Juliette Pretot

experimental/uefi/app/Cargo.lock

@@ -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/Cargo.toml

@@ -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);
+}