1 files changed, 93 insertions, 11 deletions

diff --git a/aes_keywrap.c b/aes_keywrap.c
index a3e223f..63e0cf7 100644
--- a/aes_keywrap.c
+++ b/aes_keywrap.c
@@ -53,7 +53,7 @@
  * Enable use of the experimental keywrap core, if present.
  */
 
-static int use_keywrap_core = 0;
+static int use_keywrap_core = 1;
 
 int hal_aes_use_keywrap_core(int onoff)
 {
@@ -82,17 +82,102 @@ size_t hal_aes_keywrap_ciphertext_length(const size_t plaintext_length)
 
 typedef enum { KEK_encrypting, KEK_decrypting } kek_action_t;
 
+static unsigned _kek_load = 0, _kek_skip = 0;
+
+hal_error_t hal_aes_keywrap_get_stats(unsigned *load, unsigned *skip)
+{
+  if (load == NULL || skip == NULL)
+    return HAL_ERROR_BAD_ARGUMENTS;
+
+  *load = _kek_load;
+  *skip = _kek_skip;
+
+  return HAL_OK;
+}
+
+void hal_aes_keywrap_reset_stats(void)
+{
+  _kek_load = _kek_skip = 0;
+}
+
+hal_error_t hal_aes_keywrap_zero(hal_core_t *core)
+{
+  const int free_core = (core == NULL);
+  hal_error_t err;
+
+  if (core == NULL &&
+      (err = hal_core_alloc(KEYWRAP_NAME, &core, NULL)) != HAL_OK)
+    return err;
+
+  uint8_t buf[4] = { 0, 0, 0, KEYWRAP_CTRL_ZERO };
+  err = hal_io_write(core, ADDR_CTRL, buf, sizeof(buf));
+
+  if (free_core)
+    hal_core_free(core);
+  return err;
+}
+
+hal_error_t hal_aes_keywrap_set_timeout(hal_core_t *core, uint32_t cycles)
+{
+  const int free_core = (core == NULL);
+  hal_error_t err;
+
+  if (core == NULL &&
+      (err = hal_core_alloc(KEYWRAP_NAME, &core, NULL)) != HAL_OK)
+    return err;
+
+  union {
+    uint32_t word;
+    uint8_t bytes[4];
+  } buf;
+  buf.word = htonl(cycles);
+  err = hal_io_write(core, KEYWRAP_ADDR_TIMEOUT, buf.bytes, sizeof(buf));
+
+  if (free_core)
+    hal_core_free(core);
+  return err;
+}
+
+static inline int is_kek_loaded(const hal_core_t *core)
+{
+  uint8_t buf[4];
+  
+  if (hal_io_read(core, ADDR_STATUS, buf, sizeof(buf)) == HAL_OK)
+    return (buf[3] & KEYWRAP_STATUS_LOADED);
+
+  return 0;
+}
+
 static hal_error_t load_kek(const hal_core_t *core, const uint8_t *K, const size_t K_len, const kek_action_t action)
 {
+  static size_t kek_len = 0;
   uint8_t config[4];
   hal_error_t err;
 
-  if (K == NULL)
-    return HAL_ERROR_BAD_ARGUMENTS;
+  if (K != NULL) {
+    kek_len = K_len;
+    if ((err = hal_io_write(core, AES_ADDR_KEY0, K, K_len)) != HAL_OK)
+      return err;
+  }
+  else if (!use_keywrap_core || !is_kek_loaded(core)) {
+    ++_kek_load;
+    uint8_t kek[KEK_LENGTH];
+
+    if ((err = hal_mkm_get_kek(kek, &kek_len, sizeof(kek))) == HAL_OK)
+      err = hal_io_write(core, AES_ADDR_KEY0, kek, kek_len);
+
+    memset(kek, 0, sizeof(kek));
+
+    if (err != HAL_OK)
+      return err;
+  }
+  else {
+    ++_kek_skip;
+  }
 
   memset(config, 0, sizeof(config));
 
-  switch (K_len) {
+  switch (kek_len) {
   case bitsToBytes(128):
     config[3] &= ~AES_CONFIG_KEYLEN;
     break;
@@ -120,8 +205,7 @@ static hal_error_t load_kek(const hal_core_t *core, const uint8_t *K, const size
    * Load the KEK and tell the core to expand it.
    */
 
-  if ((err = hal_io_write(core, AES_ADDR_KEY0, K, K_len))                 != HAL_OK ||
-      (err = hal_io_write(core, AES_ADDR_CONFIG, config, sizeof(config))) != HAL_OK ||
+  if ((err = hal_io_write(core, AES_ADDR_CONFIG, config, sizeof(config))) != HAL_OK ||
       (err = hal_io_init(core))                                           != HAL_OK)
     return err;
 
@@ -136,10 +220,6 @@ static hal_error_t load_kek(const hal_core_t *core, const uint8_t *K, const size
 
 static hal_error_t do_keywrap_core(const hal_core_t *core, uint8_t * const C, const size_t n)
 {
-#ifndef min
-#define min(a,b) ((a) < (b) ? (a) : (b))
-#endif
-
   hal_error_t err;
 
   hal_assert(core != NULL && C != NULL && n > 0);
@@ -246,7 +326,9 @@ hal_error_t hal_aes_keywrap(hal_core_t *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)
+    else if (memcmp(info->name, AES_CORE_NAME, 8) == 0)
+      use_keywrap_core = 0;
+    else
       /* I have no idea what this is */
       return HAL_ERROR_BAD_ARGUMENTS;
   }