1 files changed, 52 insertions, 15 deletions

diff --git a/aes_keywrap.c b/aes_keywrap.c
index b849304..d43a96f 100644
--- a/aes_keywrap.c
+++ b/aes_keywrap.c
@@ -51,15 +51,40 @@
 #include "hal_internal.h"
 
 /*
- * Enable use of the experimental keywrap core, if present.
+ * Enable use of the keywrap core, if present.
  */
 
-static int use_keywrap_core = 0;
+static enum {
+  unknown = -1,
+  aes_core = 0,
+  keywrap_core = 1
+} which_core = unknown;
 
-int hal_aes_use_keywrap_core(int onoff)
+static char *core_name[] = {
+  AES_CORE_NAME,
+  KEYWRAP_NAME
+};
+
+static inline hal_error_t init_aes_keywrap(void)
+{
+  if (which_core != unknown)
+    return HAL_OK;
+  else if (hal_core_find(KEYWRAP_NAME, NULL) != NULL)
+    return (which_core = keywrap_core), HAL_OK;
+  else if (hal_core_find(AES_CORE_NAME, NULL) != NULL)
+    return (which_core = aes_core), HAL_OK;
+  else
+    return HAL_ERROR_CORE_NOT_FOUND;
+}
+
+hal_error_t hal_aes_use_keywrap_core(int onoff)
 {
-  use_keywrap_core = (onoff && hal_core_find(KEYWRAP_NAME, NULL) != NULL);
-  return use_keywrap_core;
+  if (onoff && hal_core_find(KEYWRAP_NAME, NULL) != NULL)
+    return (which_core = keywrap_core), HAL_OK;
+  else if (!onoff && hal_core_find(AES_CORE_NAME, NULL) != NULL)
+    return (which_core = aes_core), HAL_OK;
+  else
+    return HAL_ERROR_CORE_NOT_FOUND;
 }
 
 
@@ -246,14 +271,16 @@ hal_error_t hal_aes_keywrap(hal_core_t *core,
   if (core) {
     const hal_core_info_t *info = hal_core_info(core);
     if (memcmp(info->name, KEYWRAP_NAME, 8) == 0)
-      use_keywrap_core = 1;
-    else if (memcmp(info->name, AES_CORE_NAME, 8) != 0)
+      which_core = keywrap_core;
+    else if (memcmp(info->name, AES_CORE_NAME, 8) == 0)
+      which_core = aes_core;
+    else
       /* I have no idea what this is */
       return HAL_ERROR_BAD_ARGUMENTS;
   }
   else {
-    const char *core_name = (use_keywrap_core ? KEYWRAP_NAME : AES_CORE_NAME);
-    if ((err = hal_core_alloc(core_name, &core, NULL)) != HAL_OK)
+    if ((err = init_aes_keywrap()) != HAL_OK ||
+        (err = hal_core_alloc(core_name[which_core], &core, NULL)) != HAL_OK)
       return err;
   }
 
@@ -277,7 +304,11 @@ hal_error_t hal_aes_keywrap(hal_core_t *core,
 
   n = calculated_C_len/8 - 1;
 
-  if (use_keywrap_core) {
+  /* Make sure the key expansion has completed. */
+  if ((err = hal_io_wait_ready(core)) != HAL_OK)
+    goto out;
+
+  if (which_core == keywrap_core) {
     err = do_keywrap_core(core, C, n);
   }
   else {
@@ -335,14 +366,16 @@ hal_error_t hal_aes_keyunwrap(hal_core_t *core,
   if (core) {
     const hal_core_info_t *info = hal_core_info(core);
     if (memcmp(info->name, KEYWRAP_NAME, 8) == 0)
-      use_keywrap_core = 1;
-    else if (memcmp(info->name, AES_CORE_NAME, 8) != 0)
+      which_core = keywrap_core;
+    else if (memcmp(info->name, AES_CORE_NAME, 8) == 0)
+      which_core = aes_core;
+    else
       /* I have no idea what this is */
       return HAL_ERROR_BAD_ARGUMENTS;
   }
   else {
-    const char *core_name = (use_keywrap_core ? KEYWRAP_NAME : AES_CORE_NAME);
-    if ((err = hal_core_alloc(core_name, &core, NULL)) != HAL_OK)
+    if ((err = init_aes_keywrap()) != HAL_OK ||
+        (err = hal_core_alloc(core_name[which_core], &core, NULL)) != HAL_OK)
       return err;
   }
 
@@ -354,7 +387,11 @@ hal_error_t hal_aes_keyunwrap(hal_core_t *core,
   if (Q != C)
     memmove(Q, C, C_len);
 
-  if (use_keywrap_core) {
+  /* Make sure the key expansion has completed. */
+  if ((err = hal_io_wait_ready(core)) != HAL_OK)
+    goto out;
+
+  if (which_core == keywrap_core) {
     err = do_keywrap_core(core, Q, n);
   }
   else {