From f48d41da2c675337f84ab5cb7168a692b9a3e92e Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <cwinter@microsoft.com>
Date: Fri, 21 Aug 2020 13:40:03 +0000
Subject: [PATCH] Clean up disk encryption key config handling

Signed-off-by: Christoph M. Wintersteiger <cwinter@microsoft.com>
---
 src/enclave/enclave_init.c          | 43 ++++++++++++++++++++++++++++-
 src/include/enclave/enclave_state.h |  2 ++
 src/lkl/setup.c                     | 13 ++-------
 3 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/src/enclave/enclave_init.c b/src/enclave/enclave_init.c
index c88c0b71d..72e37fd7a 100644
--- a/src/enclave/enclave_init.c
+++ b/src/enclave/enclave_init.c
@@ -19,6 +19,45 @@ extern struct mpmcq __scheduler_queue;
 _Noreturn void __dls3(elf64_stack_t* conf, void* tos);
 extern void init_sysconf(long nproc_conf, long nproc_onln);
 
+static void get_disk_keys()
+{
+    /* Here or earlier: get keys from remote key provider or auxv. */
+    const sgxlkl_enclave_config_t* cfg = sgxlkl_enclave_state.config;
+    sgxlkl_enclave_disk_state_t* disk_states = sgxlkl_enclave_state.disk_state;
+    if (cfg->root.key)
+    {
+        uint8_t* key = cfg->root.key;
+        size_t len = cfg->root.key_len;
+        disk_states[0].key = oe_malloc(sizeof(uint8_t) * len);
+        memcpy(disk_states[0].key, key, len);
+    }
+    for (size_t i = 0; i < cfg->num_mounts; i++)
+    {
+        if (cfg->mounts[i].key)
+        {
+            uint8_t* key = cfg->mounts[i].key;
+            size_t len = cfg->mounts[i].key_len;
+            disk_states[i + 1].key = oe_malloc(sizeof(uint8_t) * len);
+            memcpy(disk_states[i + 1].key, key, len);
+        }
+    }
+}
+
+static void wipe_disk_keys()
+{
+    sgxlkl_enclave_disk_state_t* disk_state = sgxlkl_enclave_state.disk_state;
+    for (size_t i = 0; i < sgxlkl_enclave_state.num_disk_state; i++)
+    {
+        if (disk_state[i].key)
+        {
+            memset(disk_state[i].key, 0, disk_state[i].key_len);
+            oe_free(disk_state[i].key);
+        }
+        disk_state[i].key = NULL;
+        disk_state[i].key_len = 0;
+    }
+}
+
 static void find_and_mount_disks()
 {
     const sgxlkl_enclave_config_t* cfg = sgxlkl_enclave_state.config;
@@ -30,7 +69,7 @@ static void find_and_mount_disks()
     estate->disk_state = oe_calloc(n, sizeof(sgxlkl_enclave_disk_state_t));
     estate->num_disk_state = n;
 
-    // root disk index
+    // root disk index 0
     estate->disk_state[0].host_disk_index = 0;
 
     for (int i = 0; i < cfg->num_mounts; i++)
@@ -57,7 +96,9 @@ static void find_and_mount_disks()
                 cfg_disk->destination);
     }
 
+    get_disk_keys();
     lkl_mount_disks(&cfg->root, cfg->mounts, cfg->num_mounts, cfg->cwd);
+    wipe_disk_keys();
 }
 
 static void init_wireguard_peers()
diff --git a/src/include/enclave/enclave_state.h b/src/include/enclave/enclave_state.h
index dc8d7d0dd..ca2e1761f 100644
--- a/src/include/enclave/enclave_state.h
+++ b/src/include/enclave/enclave_state.h
@@ -20,6 +20,8 @@ typedef struct sgxlkl_enclave_disk_state
     int fd;                 /* File descriptor of the disk */
     size_t capacity;        /* Capacity of the disk */
     bool mounted;           /* Tracks whether the disk has been mounted */
+    uint8_t* key;           /* Encryption key */
+    size_t key_len;         /* Length of encryption key */
 } sgxlkl_enclave_disk_state_t;
 
 typedef struct
diff --git a/src/lkl/setup.c b/src/lkl/setup.c
index 55b44845a..e51d92129 100644
--- a/src/lkl/setup.c
+++ b/src/lkl/setup.c
@@ -408,13 +408,6 @@ static void* lkl_activate_crypto_disk_thread(struct lkl_crypt_device* lkl_cd)
 
     crypt_free(cd);
 
-    // The key is only needed during activation, so don't keep it around
-    // afterwards and free up space.
-    memset(lkl_cd->disk_config.key, 0, lkl_cd->disk_config.key_len);
-
-    lkl_cd->disk_config.key = NULL;
-    lkl_cd->disk_config.key_len = 0;
-
     return 0;
 }
 #endif
@@ -621,13 +614,11 @@ static void lkl_mount_disk(
     lkl_cd.readonly = disk->readonly;
     lkl_cd.disk_config = *disk;
 
-    (void)lkl_cd;
-
-    if (disk->create && disk->fresh_key)
+    if (disk->create && disk->fresh_key && !disk->key)
     {
         disk->key_len = CREATED_DISK_KEY_LENGTH / 8;
         SGXLKL_VERBOSE("Generating random disk encryption key\n");
-        disk->key = malloc(disk->key_len);
+        disk->key = oe_malloc(disk->key_len);
         if (disk->key == NULL)
             sgxlkl_fail("Could not allocate memory for disk encryption key\n");
         for (size_t i = 0; i < disk->key_len; i++)