From 0684820afa26cf19efe3f7327e1992c4c0768137 Mon Sep 17 00:00:00 2001
From: Szczepan Zalega <szczepan@nitrokey.com>
Date: Thu, 29 Jun 2023 17:43:20 +0200
Subject: [PATCH] fixup Fix tests

---
 src/virt.rs             | 27 +++++++++++++++++++++++++--
 tests/hmacsha256p256.rs |  6 +++---
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/virt.rs b/src/virt.rs
index 94f6cc4..7649e0b 100644
--- a/src/virt.rs
+++ b/src/virt.rs
@@ -16,6 +16,9 @@ use crate::streaming::ChunkedExtension;
 #[cfg(feature = "manage")]
 use trussed::types::{Location, Path};
 
+#[cfg(feature = "hmacsha256p256")]
+use crate::hmacsha256p256::HmacSha256P256Extension;
+
 #[derive(Default, Debug)]
 pub struct Dispatcher {
     backend: StagingBackend,
@@ -34,6 +37,14 @@ pub enum ExtensionIds {
     Chunked,
     #[cfg(feature = "manage")]
     Manage,
+    #[cfg(feature = "hmacsha256p256")]
+    HmacShaP256,
+}
+
+#[cfg(feature = "hmacsha256p256")]
+impl ExtensionId<HmacSha256P256Extension> for Dispatcher {
+    type Id = ExtensionIds;
+    const ID: ExtensionIds = ExtensionIds::HmacShaP256;
 }
 
 #[cfg(feature = "wrap-key-to-file")]
@@ -63,6 +74,8 @@ impl From<ExtensionIds> for u8 {
             ExtensionIds::Chunked => 1,
             #[cfg(feature = "manage")]
             ExtensionIds::Manage => 2,
+            #[cfg(feature = "hmacsha256p256")]
+            ExtensionIds::HmacShaP256 => 2,
         }
     }
 }
@@ -77,6 +90,8 @@ impl TryFrom<u8> for ExtensionIds {
             1 => Ok(Self::Chunked),
             #[cfg(feature = "manage")]
             2 => Ok(Self::Manage),
+            #[cfg(feature = "hmacsha256p256")]
+            3 => Ok(Self::HmacShaP256),
             _ => Err(Error::FunctionNotSupported),
         }
     }
@@ -122,7 +137,16 @@ impl ExtensionDispatch for Dispatcher {
                     resources,
                 )
             }
-
+            #[cfg(feature = "hmacsha256p256")]
+            ExtensionIds::HmacShaP256 => {
+                ExtensionImpl::<HmacSha256P256Extension>::extension_request_serialized(
+                    &mut self.backend,
+                    &mut ctx.core,
+                    &mut ctx.backends,
+                    request,
+                    resources,
+                )
+            }
             #[cfg(feature = "chunked")]
             ExtensionIds::Chunked => {
                 ExtensionImpl::<ChunkedExtension>::extension_request_serialized(
@@ -133,7 +157,6 @@ impl ExtensionDispatch for Dispatcher {
                     resources,
                 )
             }
-
             #[cfg(feature = "manage")]
             ExtensionIds::Manage => ExtensionImpl::<ManageExtension>::extension_request_serialized(
                 &mut self.backend,
diff --git a/tests/hmacsha256p256.rs b/tests/hmacsha256p256.rs
index 62f5c73..757d576 100644
--- a/tests/hmacsha256p256.rs
+++ b/tests/hmacsha256p256.rs
@@ -17,22 +17,22 @@ use trussed_staging::hmacsha256p256::HmacSha256P256Client;
 
 #[test]
 fn hmac_inject_any() {
+    use trussed::types::Message;
     with_ram_client("staging-tests", |mut client| {
         let client = &mut client;
 
         let key = syscall!(client.inject_any_key(
-            b"12345678123456781234567812345678",
+            Message::from_slice(b"12345678123456781234567812345678").unwrap(),
             Volatile,
             Kind::P256
         ))
         .key
         .unwrap();
 
-        let pk = syscall!(client.derive_p256_public_key(key, Location::Volatile)).key;
+        let _pk = syscall!(client.derive_p256_public_key(key, Location::Volatile)).key;
 
         let signature =
             syscall!(client.sign(Mechanism::P256, key, &[], SignatureSerialization::Raw)).signature;
         assert!(signature.len() > 0);
-        todo!();
     });
 }