[Pal/Linux-SGX] EDMM hybrid allocation

Documentation/manifest-syntax.rst

@@ -483,6 +483,21 @@ before enclave creation (because it involves more enclave exits and syscalls).
 .. note::
    Support for EDMM first appeared in Linux 6.0.
 
+EDMM heap pre-allocated size
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+::
+
+    sgx.edmm_heap_prealloc_size = "[SIZE]"
+    (default: "0")
+
+When ``sgx.edmm_enable`` is enabled, users can precisely set the amount of heap
+to pre-allocate by setting this option. For example, when size is set to "64M",
+Gramine will pre-allocate 64M of enclave memory (with read-write-execute
+permissions, similar to SGXv1) rather than allocating dynamically. Higher values
+in this option may improve run time but could worsen startup time, and
+vice versa.
+
 Enclave size
 ^^^^^^^^^^^^
 
@@ -824,8 +839,10 @@ predictable.
 Please note that using this option makes sense only when the :term:`EPC` is
 large enough to hold the whole heap area.
 
-This option is invalid (i.e. must be ``false``) if specified together with
-``sgx.edmm_enable``, as there are no heap pages to pre-fault.
+This option is invalid (i.e. must be ``false``) when only ``sgx.edmm_enable`` is
+specified, as there are no heap pages to pre-fault. But if used together with
+``sgx.edmm_heap_prealloc_size``, then the pre-allocated heap will also be
+pre-faulted.
 
 Enabling per-thread and process-wide SGX stats
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

libos/test/regression/manifest.template

@@ -24,6 +24,7 @@ fs.mounts = [
 sgx.max_threads = 16
 sgx.debug = true
 sgx.edmm_enable = {{ 'true' if env.get('EDMM', '0') == '1' else 'false' }}
+sgx.edmm_heap_prealloc_size = "{{ '64M' if env.get('EDMM', '0') == '1' else '0' }}"
 
 sgx.allowed_files = [
   "file:tmp/",

pal/src/host/linux-sgx/enclave_ecalls.c

@@ -101,6 +101,7 @@ void handle_ecall(long ecall_index, void* ecall_args, void* exit_target, void* e
                        COPY_UNTRUSTED_VALUE(&start_args->rpc_queue),
                        COPY_UNTRUSTED_VALUE(&start_args->dns_host_conf),
                        COPY_UNTRUSTED_VALUE(&start_args->edmm_enabled),
+                       COPY_UNTRUSTED_VALUE(&start_args->edmm_heap_prealloc_size),
                        COPY_UNTRUSTED_VALUE(&start_args->reserved_mem_ranges),
                        COPY_UNTRUSTED_VALUE(&start_args->reserved_mem_ranges_size));
     } else {

pal/src/host/linux-sgx/enclave_edmm.c

@@ -30,6 +30,37 @@ static void sgx_emodpe(uint64_t addr, uint64_t prot) {
     /* `EMODPE` does not return errors, it can only fault. */
 }
 
+/* Updates page count such that the request is fully below the pre-allocated heap. If `count` is
+ * updated to 0, then the entire request overlaps with pre-allocated heap.
+ *
+ * Partial overlap illustration:
+                +----------------------+ --> heap_max
+                |                      |
+addr + size <-- |  Pre-allocated heap  |
+                |                      |
+                +----------------------+ --> edmm_heap_prealloc_start
+                |                      |     (heap_max - edmm_heap_prealloc_size)
+                |  Dynamically         |
+       addr <-- |  allocated heap      |
+                |                      |
+                +----------------------+
+*/
+static void exclude_preallocated_pages(uint64_t addr, size_t* count) {
+    size_t size = *count * PAGE_SIZE;
+    uint64_t edmm_heap_prealloc_start = (uint64_t)g_pal_linuxsgx_state.heap_max -
+                                        g_pal_linuxsgx_state.edmm_heap_prealloc_size;
+
+    if (addr >= edmm_heap_prealloc_start) {
+        /* full overlap: entire request lies in the pre-allocated region */
+        *count = 0;
+    } else if (addr + size > edmm_heap_prealloc_start) {
+        /* partial overlap: update count to skip the pre-allocated region */
+        *count = (edmm_heap_prealloc_start - addr) / PAGE_SIZE;
+    } else {
+        /* no overlap: don't update count */
+    }
+}
+
 int sgx_edmm_add_pages(uint64_t addr, size_t count, uint64_t prot) {
     int ret;
 
@@ -38,6 +69,20 @@ int sgx_edmm_add_pages(uint64_t addr, size_t count, uint64_t prot) {
         prot |= SGX_SECINFO_FLAGS_R;
     }
 
+    if (g_pal_linuxsgx_state.edmm_heap_prealloc_size > 0) {
+        size_t original_count = count;
+        exclude_preallocated_pages(addr, &count);
+
+        size_t preallocated_count = original_count - count;
+        if (preallocated_count != 0) {
+            memset((void*)(addr + count * PAGE_SIZE), 0, preallocated_count * PAGE_SIZE);
+            if (count == 0) {
+                /* Entire request is in pre-allocated range */
+                return 0;
+            }
+        }
+    }
+
     for (size_t i = 0; i < count; i++) {
         /* SGX2 HW requires initial page permissions to be RW. */
         ret = sgx_eaccept(addr + i * PAGE_SIZE, (SGX_PAGE_TYPE_REG << SGX_SECINFO_FLAGS_TYPE_SHIFT)
@@ -84,6 +129,12 @@ int sgx_edmm_add_pages(uint64_t addr, size_t count, uint64_t prot) {
 }
 
 int sgx_edmm_remove_pages(uint64_t addr, size_t count) {
+    if (g_pal_linuxsgx_state.edmm_heap_prealloc_size > 0) {
+        exclude_preallocated_pages(addr, &count);
+        if (count == 0)
+            return 0;
+    }
+
     int ret = ocall_edmm_modify_pages_type(addr, count, SGX_PAGE_TYPE_TRIM);
     if (ret < 0) {
         return unix_to_pal_error(ret);
@@ -114,6 +165,12 @@ int sgx_edmm_remove_pages(uint64_t addr, size_t count) {
 }
 
 int sgx_edmm_set_page_permissions(uint64_t addr, size_t count, uint64_t prot) {
+    if (g_pal_linuxsgx_state.edmm_heap_prealloc_size > 0) {
+        exclude_preallocated_pages(addr, &count);
+        if (count == 0)
+            return 0;
+    }
+
     if (prot & SGX_SECINFO_FLAGS_W) {
         /* HW limitation. */
         prot |= SGX_SECINFO_FLAGS_R;

pal/src/host/linux-sgx/host_ecalls.c

@@ -9,8 +9,8 @@
 int ecall_enclave_start(char* libpal_uri, char* args, size_t args_size, char* env,
                         size_t env_size, int parent_stream_fd, sgx_target_info_t* qe_targetinfo,
                         struct pal_topo_info* topo_info, struct pal_dns_host_conf* dns_conf,
-                        bool edmm_enabled, void* reserved_mem_ranges,
-                        size_t reserved_mem_ranges_size) {
+                        bool edmm_enabled, size_t edmm_heap_prealloc_size,
+                        void* reserved_mem_ranges, size_t reserved_mem_ranges_size) {
     g_rpc_queue = NULL;
 
     if (g_pal_enclave.rpc_thread_num > 0) {
@@ -34,6 +34,7 @@ int ecall_enclave_start(char* libpal_uri, char* args, size_t args_size, char* en
         .topo_info                = topo_info,
         .dns_host_conf            = dns_conf,
         .edmm_enabled             = edmm_enabled,
+        .edmm_heap_prealloc_size  = edmm_heap_prealloc_size,
         .reserved_mem_ranges      = reserved_mem_ranges,
         .reserved_mem_ranges_size = reserved_mem_ranges_size,
         .rpc_queue                = g_rpc_queue,

pal/src/host/linux-sgx/host_ecalls.h

@@ -9,8 +9,8 @@
 int ecall_enclave_start(char* libpal_uri, char* args, size_t args_size, char* env, size_t env_size,
                         int parent_stream_fd, sgx_target_info_t* qe_targetinfo,
                         struct pal_topo_info* topo_info, struct pal_dns_host_conf* host_conf,
-                        bool edmm_enabled, void* reserved_mem_ranges,
-                        size_t reserved_mem_ranges_size);
+                        bool edmm_enabled, size_t edmm_heap_prealloc_size,
+                        void* reserved_mem_ranges, size_t reserved_mem_ranges_size);
 
 int ecall_thread_start(void);
 

pal/src/host/linux-sgx/host_internal.h

@@ -47,6 +47,7 @@ struct pal_enclave {
     unsigned long rpc_thread_num;
     unsigned long ssa_frame_size;
     bool edmm_enabled;
+    size_t edmm_heap_prealloc_size;
     enum sgx_attestation_type attestation_type;
     char* libpal_uri; /* Path to the PAL binary */
 

pal/src/host/linux-sgx/host_main.c

@@ -490,8 +490,20 @@ static int initialize_enclave(struct pal_enclave* enclave, const char* manifest_
 
         if (areas[i].skip_eextend && enclave->edmm_enabled) {
             assert(areas[i].data_src == ZERO);
-            /* If EDMM is enabled, no need to add non-measured zero pages. */
-            continue;
+            if (enclave->edmm_heap_prealloc_size == 0) {
+                /* If EDMM is enabled and no pre-allocated heap is requested, then skip adding
+                 non-measured zero pages. */
+                continue;
+            } else {
+                if (enclave->edmm_heap_prealloc_size > areas[i].size) {
+                    log_error("'sgx.edmm_heap_prealloc_size' must be less than total heap size "
+                              "0x%lx", areas[i].size);
+                    ret = -EINVAL;
+                    goto out;
+                }
+                areas[i].addr = areas[i].addr + areas[i].size - enclave->edmm_heap_prealloc_size;
+                areas[i].size = enclave->edmm_heap_prealloc_size;
+            }
         }
 
         void* data = NULL;
@@ -682,6 +694,12 @@ static int parse_loader_config(char* manifest, struct pal_enclave* enclave_info,
         goto out;
     }
 
+    if (!enclave_info->size || !IS_POWER_OF_2(enclave_info->size)) {
+        log_error("Enclave size not a power of two (an SGX-imposed requirement)");
+        ret = -EINVAL;
+        goto out;
+    }
+
     ret = toml_bool_in(manifest_root, "sgx.edmm_enable", /*defaultval=*/false,
                        &enclave_info->edmm_enabled);
     if (ret < 0) {
@@ -690,8 +708,22 @@ static int parse_loader_config(char* manifest, struct pal_enclave* enclave_info,
         goto out;
     }
 
-    if (!enclave_info->size || !IS_POWER_OF_2(enclave_info->size)) {
-        log_error("Enclave size not a power of two (an SGX-imposed requirement)");
+    ret = toml_sizestring_in(manifest_root, "sgx.edmm_heap_prealloc_size", /*defaultval=*/0,
+                             &enclave_info->edmm_heap_prealloc_size);
+    if (ret < 0) {
+        log_error("Cannot parse 'sgx.edmm_heap_prealloc_size'");
+        ret = -EINVAL;
+        goto out;
+    }
+
+    if (!enclave_info->edmm_enabled && enclave_info->edmm_heap_prealloc_size > 0) {
+        log_error("'sgx.edmm_heap_prealloc_size' must be used together with 'sgx.edmm_enable'!");
+        ret = -EINVAL;
+        goto out;
+    }
+
+    if (!IS_ALIGNED(enclave_info->edmm_heap_prealloc_size, g_page_size)) {
+        log_error("'sgx.edmm_heap_prealloc_size' must be 4K (page) aligned");
         ret = -EINVAL;
         goto out;
     }
@@ -1083,7 +1115,8 @@ static int load_enclave(struct pal_enclave* enclave, char* args, size_t args_siz
     /* start running trusted PAL */
     ecall_enclave_start(enclave->libpal_uri, args, args_size, env, env_size, parent_stream_fd,
                         &qe_targetinfo, &topo_info, &dns_conf, enclave->edmm_enabled,
-                        reserved_mem_ranges, reserved_mem_ranges_size);
+                        enclave->edmm_heap_prealloc_size, reserved_mem_ranges,
+                        reserved_mem_ranges_size);
 
     unmap_tcs();
     DO_SYSCALL(munmap, alt_stack, ALT_STACK_SIZE);

pal/src/host/linux-sgx/pal_ecall_types.h

@@ -29,6 +29,7 @@ struct ecall_enclave_start {
     struct pal_topo_info*     topo_info;
     struct pal_dns_host_conf* dns_host_conf;
     unsigned char             edmm_enabled;
+    size_t                    edmm_heap_prealloc_size;
     void*                     reserved_mem_ranges;
     size_t                    reserved_mem_ranges_size;
 

pal/src/host/linux-sgx/pal_linux.h

@@ -34,6 +34,7 @@ extern struct pal_linuxsgx_state {
     /* enclave information */
     bool enclave_initialized;        /* thread creation ECALL is allowed only after this is set */
     bool edmm_enabled;
+    size_t edmm_heap_prealloc_size;
     sgx_target_info_t qe_targetinfo; /* received from untrusted host, use carefully */
     sgx_report_body_t enclave_info;  /* cached self-report result, trusted */
 
@@ -72,7 +73,8 @@ noreturn void pal_linux_main(void* uptr_libpal_uri, size_t libpal_uri_len, void*
                              size_t args_size, void* uptr_env, size_t env_size,
                              int parent_stream_fd, void* uptr_qe_targetinfo, void* uptr_topo_info,
                              void* uptr_rpc_queue, void* uptr_dns_conf, bool edmm_enabled,
-                             void* urts_reserved_mem_ranges, size_t urts_reserved_mem_ranges_size);
+                             size_t edmm_heap_prealloc_size, void* urts_reserved_mem_ranges,
+                             size_t urts_reserved_mem_ranges_size);
 void pal_start_thread(void);
 
 extern char __text_start, __text_end, __data_start, __data_end;

pal/src/host/linux-sgx/pal_main.c

@@ -528,8 +528,16 @@ static int print_warnings_on_insecure_configs(PAL_HANDLE parent_process) {
 
 __attribute_no_sanitize_address
 static void do_preheat_enclave(void) {
-    for (uint8_t* i = g_pal_linuxsgx_state.heap_min; i < (uint8_t*)g_pal_linuxsgx_state.heap_max;
-             i += g_page_size) {
+    /* Heap allocation requests are serviced starting from highest heap address. So when
+     * sgx.edmm_heap_prealloc_size is turned on, preheat from the top of the heap until
+     * sgx.edmm_heap_prealloc_size. */
+    uint8_t* start = (uint8_t*)g_pal_linuxsgx_state.heap_min;
+    if (g_pal_linuxsgx_state.edmm_heap_prealloc_size > 0) {
+        start = (uint8_t*)g_pal_linuxsgx_state.heap_max -
+                g_pal_linuxsgx_state.edmm_heap_prealloc_size;
+    }
+
+    for (uint8_t* i = start; i < (uint8_t*)g_pal_linuxsgx_state.heap_max; i += g_page_size) {
         READ_ONCE(*(size_t*)i);
     }
 }
@@ -541,7 +549,8 @@ noreturn void pal_linux_main(void* uptr_libpal_uri, size_t libpal_uri_len, void*
                              size_t args_size, void* uptr_env, size_t env_size,
                              int parent_stream_fd, void* uptr_qe_targetinfo, void* uptr_topo_info,
                              void* uptr_rpc_queue, void* uptr_dns_conf, bool edmm_enabled,
-                             void* urts_reserved_mem_ranges, size_t urts_reserved_mem_ranges_size) {
+                             size_t edmm_heap_prealloc_size, void* urts_reserved_mem_ranges,
+                             size_t urts_reserved_mem_ranges_size) {
     /* All our arguments are coming directly from the host. We are responsible to check them. */
     int ret;
 
@@ -569,6 +578,24 @@ noreturn void pal_linux_main(void* uptr_libpal_uri, size_t libpal_uri_len, void*
     g_pal_linuxsgx_state.heap_max = GET_ENCLAVE_TCB(heap_max);
     g_pal_linuxsgx_state.edmm_enabled = edmm_enabled;
 
+    if (!edmm_enabled && edmm_heap_prealloc_size > 0) {
+        log_error("'sgx.edmm_heap_prealloc_size' must be used together with 'sgx.edmm_enable'!");
+        ocall_exit(1, /*is_exitgroup=*/true);
+    }
+
+    size_t total_heap_size = g_pal_linuxsgx_state.heap_max - g_pal_linuxsgx_state.heap_min;
+    if (edmm_heap_prealloc_size > total_heap_size) {
+        log_error("'sgx.edmm_heap_prealloc_size' must be less than total heap size 0x%lx",
+                   total_heap_size);
+        ocall_exit(1, /*is_exitgroup=*/true);
+    }
+
+    if (!IS_ALIGNED(edmm_heap_prealloc_size, g_page_size)) {
+        log_error("edmm_heap_prealloc_size must be 4K (page) aligned");
+        ocall_exit(1, /*is_exitgroup=*/true);
+    }
+    g_pal_linuxsgx_state.edmm_heap_prealloc_size = edmm_heap_prealloc_size;
+
     /* No need for adding any initial memory ranges - they are all outside of the available memory
      * set below. */
     g_pal_public_state.memory_address_start = g_pal_linuxsgx_state.heap_min;
@@ -696,6 +723,20 @@ noreturn void pal_linux_main(void* uptr_libpal_uri, size_t libpal_uri_len, void*
         ocall_exit(1, /*is_exitgroup=*/true);
     }
 
+    size_t edmm_heap_prealloc_size_manifest;
+    ret = toml_sizestring_in(g_pal_public_state.manifest_root, "sgx.edmm_heap_prealloc_size",
+                             /*defaultval=*/0, &edmm_heap_prealloc_size_manifest);
+    if (ret < 0) {
+        log_error("Cannot parse 'sgx.edmm_heap_prealloc_size'");
+        ocall_exit(1, /*is_exitgroup=*/true);
+    }
+
+    if (edmm_heap_prealloc_size_manifest != g_pal_linuxsgx_state.edmm_heap_prealloc_size) {
+        log_error("edmm_heap_prealloc_size_manifest(=%ld) != edmm_heap_prealloc_size(=%ld)",
+                   edmm_heap_prealloc_size_manifest, g_pal_linuxsgx_state.edmm_heap_prealloc_size);
+        ocall_exit(1, /*is_exitgroup=*/true);
+    }
+
     int64_t rpc_thread_num;
     ret = toml_int_in(g_pal_public_state.manifest_root, "sgx.insecure__rpc_thread_num",
                       /*defaultval=*/0, &rpc_thread_num);
@@ -729,8 +770,9 @@ noreturn void pal_linux_main(void* uptr_libpal_uri, size_t libpal_uri_len, void*
         ocall_exit(1, /*is_exitgroup=*/true);
     }
     if (preheat_enclave) {
-        if (g_pal_linuxsgx_state.edmm_enabled) {
-            log_error("'sgx.preheat_enclave' manifest option makes no sense with EDMM enabled!");
+        if (g_pal_linuxsgx_state.edmm_enabled && !g_pal_linuxsgx_state.edmm_heap_prealloc_size) {
+            log_error("'sgx.preheat_enclave' manifest option makes no sense with EDMM enabled and "
+                      "'sgx.edmm_heap_prealloc_size' set to zero!");
             ocall_exit(1, /*is_exitgroup=*/true);
         }
         do_preheat_enclave();

python/graminelibos/manifest.py

@@ -104,6 +104,7 @@ def __init__(self, manifest_str):
         sgx.setdefault('require_amx', False)
         sgx.setdefault('require_exinfo', False)
         sgx.setdefault('enable_stats', False)
+        sgx.setdefault('edmm_heap_prealloc_size', '0')
 
         if not isinstance(sgx['trusted_files'], list):
             raise ValueError("Unsupported trusted files syntax, more info: " +