From d5ba44085df48507c217ef8f8dc504529f820618 Mon Sep 17 00:00:00 2001
From: Lukasz Hawrylko <lukasz.hawrylko@intel.com>
Date: Fri, 21 Feb 2020 11:07:00 +0100
Subject: [PATCH] Add support for EFI memory map parse/modification

In E820 table boot services memory is marked as available. As TBOOT is launched
after ExitBootServices() is called that behaviour looks OK. However UEFI can put
in boot services memory region some tables (like EFI_MEMORY_ATTRIBUTES_TABLE)
that Linux may want to access later. That's why TBOOT has to internally mark
this region as reserved to prevent data from overwriting. To do that, EFI memory
table has to be parsed.

TBOOT reserves some regions in memory by modifying E820 table, however in
original implementation, EFI memory map remains unchanged. Both tables are
passed to Linux kernel via boot params protocol. In this case, Linux receives
two tables describing memory regions that are contradictory. To fix that, this
patch includes also EFI memory map modification functionality to allow reserving
memory regions in the same way as in E820.

Signed-off-by: Lukasz Hawrylko <lukasz.hawrylko@intel.com>
---
 include/config.h           |   9 +-
 tboot/Makefile             |   2 +-
 tboot/common/e820.c        |  23 ++-
 tboot/common/efi_memmap.c  | 355 +++++++++++++++++++++++++++++++++++++
 tboot/common/linux.c       |  20 ++-
 tboot/common/loader.c      |  22 ++-
 tboot/common/tboot.c       |   8 +
 tboot/include/e820.h       |  38 +---
 tboot/include/efi_memmap.h | 101 +++++++++++
 tboot/txt/txt.c            |  10 ++
 tboot/txt/verify.c         |   7 +
 11 files changed, 535 insertions(+), 60 deletions(-)
 create mode 100644 tboot/common/efi_memmap.c
 create mode 100644 tboot/include/efi_memmap.h

diff --git a/include/config.h b/include/config.h
index 11c0b68..dc6ee7a 100644
--- a/include/config.h
+++ b/include/config.h
@@ -62,9 +62,14 @@
 				      TBOOT_SERIAL_LOG_SIZE)
 #define TBOOT_E820_COPY_SIZE         0x01800
 
+/* address/size for modified EFI memory map table */
+#define TBOOT_EFI_MEMMAP_COPY_ADDR   (TBOOT_E820_COPY_ADDR + \
+                                      TBOOT_E820_COPY_SIZE)
+#define TBOOT_EFI_MEMMAP_COPY_SIZE   0x08000
+
 /* address/size for modified VMM/kernel command line */
-#define TBOOT_KERNEL_CMDLINE_ADDR    (TBOOT_E820_COPY_ADDR + \
-				      TBOOT_E820_COPY_SIZE)
+#define TBOOT_KERNEL_CMDLINE_ADDR    (TBOOT_EFI_MEMMAP_COPY_ADDR + \
+                                      TBOOT_EFI_MEMMAP_COPY_SIZE)
 #define TBOOT_KERNEL_CMDLINE_SIZE    0x0400
 
 
diff --git a/tboot/Makefile b/tboot/Makefile
index a3652cc..f77201e 100644
--- a/tboot/Makefile
+++ b/tboot/Makefile
@@ -24,7 +24,7 @@ obj-y += common/vga.o common/vmac.o common/vsprintf.o common/lz.o
 obj-y += txt/acmod.o txt/errors.o txt/heap.o txt/mtrrs.o txt/txt.o
 obj-y += txt/verify.o txt/vmcs.o
 obj-y += common/tpm_12.o common/tpm_20.o 
-obj-y += common/sha256.o
+obj-y += common/sha256.o common/efi_memmap.o
 
 OBJS := $(obj-y)
 
diff --git a/tboot/common/e820.c b/tboot/common/e820.c
index a42b357..e2bff04 100644
--- a/tboot/common/e820.c
+++ b/tboot/common/e820.c
@@ -44,6 +44,7 @@
 #include <stdarg.h>
 #include <misc.h>
 #include <pci_cfgreg.h>
+#include <efi_memmap.h>
 #include <e820.h>
 #include <txt/config_regs.h>
 
@@ -593,8 +594,8 @@ bool get_ram_ranges(uint64_t *min_lo_ram, uint64_t *max_lo_ram,
      * to find it first so that we can tell when we have passed it
      */
     if ( g_min_ram > 0 ) {
-        get_highest_sized_ram(g_min_ram, 0x100000000ULL, &last_min_ram_base,
-                              &last_min_ram_size);
+        e820_get_highest_sized_ram(g_min_ram, 0x100000000ULL,
+                                   &last_min_ram_base, &last_min_ram_size);
         printk(TBOOT_DETA"highest min_ram (0x%x) region found: base=0x%Lx, size=0x%Lx\n",
                g_min_ram, last_min_ram_base, last_min_ram_size);
     }
@@ -628,6 +629,9 @@ bool get_ram_ranges(uint64_t *min_lo_ram, uint64_t *max_lo_ram,
                     printk(TBOOT_DETA"discarding RAM above reserved regions: 0x%Lx - 0x%Lx\n", base, limit);
                     if ( !e820_reserve_ram(base, limit - base) )
                         return false;
+                    if (!efi_memmap_reserve(base, limit - base)) {
+                        return false;
+                    }
                 }
             }
 
@@ -658,13 +662,13 @@ bool get_ram_ranges(uint64_t *min_lo_ram, uint64_t *max_lo_ram,
 }
 
 /* find highest (< <limit>) RAM region of at least <size> bytes */
-void get_highest_sized_ram(uint64_t size, uint64_t limit,
-                           uint64_t *ram_base, uint64_t *ram_size)
+bool e820_get_highest_sized_ram(uint64_t size, uint64_t limit,
+                                uint64_t *ram_base, uint64_t *ram_size)
 {
     uint64_t last_fit_base = 0, last_fit_size = 0;
 
     if ( ram_base == NULL || ram_size == NULL )
-        return;
+        return false;
 
     for ( unsigned int i = 0; i < g_nr_map; i++ ) {
         memory_map_t *entry = &g_copy_e820_map[i];
@@ -683,8 +687,13 @@ void get_highest_sized_ram(uint64_t size, uint64_t limit,
         }
     }
 
-    *ram_base = last_fit_base;
-    *ram_size = last_fit_size;
+    if (last_fit_size == 0) {
+        return false;
+    } else {
+        *ram_base = last_fit_base;
+        *ram_size = last_fit_size;
+        return true;
+    }
 }
 
 
diff --git a/tboot/common/efi_memmap.c b/tboot/common/efi_memmap.c
new file mode 100644
index 0000000..38c2293
--- /dev/null
+++ b/tboot/common/efi_memmap.c
@@ -0,0 +1,355 @@
+/*
+ * efi_memmap.c: EFI memory map parsing functions
+ *
+ * Copyright (c) 2006-2020, Intel Corporation
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *
+ *   * Redistributions of source code must retain the above copyright
+ *     notice, this list of conditions and the following disclaimer.
+ *   * Redistributions in binary form must reproduce the above
+ *     copyright notice, this list of conditions and the following
+ *     disclaimer in the documentation and/or other materials provided
+ *     with the distribution.
+ *   * Neither the name of the Intel Corporation nor the names of its
+ *     contributors may be used to endorse or promote products derived
+ *     from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+ * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+ * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+ * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
+ * OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ */
+
+#include <string.h>
+#include <stdbool.h>
+#include <printk.h>
+#include <uuid.h>
+#include <loader.h>
+#include <misc.h>
+#include <efi_memmap.h>
+
+static bool efi_mmap_available = false;
+static efi_memmap_t* efi_mmap = (efi_memmap_t*)TBOOT_EFI_MEMMAP_COPY_ADDR;
+
+static bool insert_after_region(uint32_t pos, uint64_t addr, uint64_t size,
+                                uint32_t type, uint64_t attr);
+static bool region_is_free(uint32_t region_type);
+
+/**
+ * @brief Copy memory map from mbi to defined memory space to allow insertion
+ *        of new entries
+ * 
+ * @param lctx loader context with mbi
+ */
+bool efi_memmap_copy(loader_ctx *lctx)
+{
+    uint32_t descr_addr, descr_ver, descr_size, mmap_size;
+    descr_addr = find_efi_memmap(lctx, &descr_size,
+                                 &descr_ver, &mmap_size);
+
+    if (descr_addr == 0 || descr_ver != EFI_MEMORY_DESCRIPTOR_VERSION) {
+        printk(TBOOT_WARN"Failed to get EFI memory map\n");
+        return false;
+    }
+
+    efi_mmap->size = mmap_size;
+    efi_mmap->descr_size = descr_size;
+    memcpy(efi_mmap->descr, (void*)descr_addr, mmap_size);
+    efi_mmap_available = true;
+
+    return true;
+}
+
+/**
+ * @brief Get address of memory map descriptors
+ * 
+ * @param descr_size return size of each descriptor
+ * @param descr_vers return descriptor version
+ * @param mmap_size  return sum of all descriptors size
+ */
+uint32_t efi_memmap_get_addr(uint32_t *descr_size, uint32_t *descr_vers,
+                             uint32_t *mmap_size)
+{
+    if (!efi_mmap_available) {
+        return 0;
+    }
+    if (descr_size != NULL) {
+        *descr_size = efi_mmap->descr_size;
+    }
+    if (descr_vers != NULL) {
+        *descr_vers = EFI_MEMORY_DESCRIPTOR_VERSION;
+    }
+    if (mmap_size != NULL) {
+        *mmap_size = efi_mmap->size;
+    }
+    return (uint32_t)efi_mmap->descr;
+}
+
+/**
+ * @brief Walk through memory map descriptors
+ * 
+ * @param prev pointer to previous descriptor, when NULL start interating from 
+ *             first one
+ */
+efi_mem_descr_t* efi_memmap_walk(efi_mem_descr_t* prev)
+{
+    if (!efi_mmap_available) {
+        printk(TBOOT_WARN"EFI memory map not available\n");
+        return NULL;
+    }
+
+    if (prev == NULL) {
+        return (efi_mem_descr_t*)efi_mmap->descr;
+    } else if ((uint32_t)prev < (uint32_t)efi_mmap->descr) {
+        /* 
+         * Should never happens, just to prevent overflow in below
+         * substraction
+          */
+        return NULL;
+    } else {
+        uint32_t next = (uint32_t)prev + efi_mmap->descr_size;
+        if (next - (uint32_t)efi_mmap->descr < efi_mmap->size) {
+            return (efi_mem_descr_t*)next;
+        } else {
+            return NULL;
+        }
+    }
+}
+
+/**
+ * @brief Mark given memory region as reserved
+ * 
+ * Region will be changed to EFI_RESERVED_TYPE, if given region already has type
+ * that indicates that it is not free, type will not be changed. Non-free means
+ * other than loader, boot, runtime and conventional memory types.
+ * 
+ * Region has to be aligned to page size, function will round non-aligned
+ * values. Base address is rounded down, length - up.
+ * 
+ * @param base   starting address
+ * @param length length of region to reserve
+ */
+bool efi_memmap_reserve(uint64_t base, uint64_t length)
+{
+    if (length == 0 || !efi_mmap_available) {
+        return true;
+    }
+
+    /* Round to page size */
+    uint64_t mask = ~((1ULL << EFI_PAGE_SHIFT) - 1ULL);
+    base &= mask;
+    if (length & ~mask) {
+        length &= mask;
+        length += (1ULL << EFI_PAGE_SHIFT);
+    }
+
+    uint64_t end = base + length;
+    efi_mem_descr_t* desc = NULL;
+    uint32_t i = 0;
+
+    while ((desc = efi_memmap_walk(desc)) != NULL) {
+        uint64_t desc_base = desc->physical_start;
+        uint64_t desc_length = desc->num_pages << EFI_PAGE_SHIFT;
+        uint64_t desc_end = desc_base + desc_length;
+
+        /* if already unusable, no need to deal with */
+        if (desc->type < EFI_LOADER_CODE ||
+                desc->type > EFI_CONVENTIONAL_MEMORY) {
+            goto cont;
+        }
+
+        /* if the range is before the current ram range, skip the ram range */
+        if (end <= desc_base) {
+            goto cont;
+        }
+        /* if the range is after the current ram range, skip the ram range */
+        if (base >= desc_end) {
+            goto cont;
+        }
+
+        /* case 1: the current ram range is within the range:
+           base, desc_base, desc_end, end */
+        if ((base <= desc_base) && (desc_end <= end)) {
+            desc->type = EFI_RESERVED_TYPE;
+        }
+        /* case 2: overlapping:
+           base, e820_base, end, e820_end */
+        else if ((desc_base >= base) && (end > desc_base) && (desc_end > end)) {
+            /* split the current ram map */
+            if (!insert_after_region(i-1, desc_base, (end - desc_base),
+                                     EFI_RESERVED_TYPE, desc->attribute)) {
+                return false;
+            }
+            /* fixup the current ram map */
+            desc = efi_memmap_walk(desc); 
+            ++i;
+            desc->physical_start = end;
+            desc->num_pages = (desc_end - end) >> EFI_PAGE_SHIFT;
+            /* no need to check more */
+            break;
+        }
+        /* case 3: overlapping:
+           desc_base, base, desc_end, end */
+        else if ((base > desc_base) && (desc_end > base) && (end >= desc_end)) {
+            /* fixup the current ram map */
+            desc->num_pages = (base - desc_base) >> EFI_PAGE_SHIFT;
+            /* split the current ram map */
+            if (!insert_after_region(i, base, (desc_end - base),
+                                     EFI_RESERVED_TYPE, desc->attribute)) {
+                return false;
+            }
+            desc = efi_memmap_walk(desc); 
+            ++i;
+        }
+        /* case 4: the range is within the current ram range:
+           desc_base, base, end, desc_end */
+        else if ((base > desc_base) && (desc_end > end)) {
+            /* fixup the current ram map */
+            desc->num_pages = (base - desc_base) >> EFI_PAGE_SHIFT;
+            /* split the current ram map */
+            if (!insert_after_region(i, base, length, EFI_RESERVED_TYPE,
+                                     desc->attribute)) {
+                return false;
+            }
+            ++i;
+            /* fixup the rest of the current ram map */
+            if (!insert_after_region(i, end, (desc_end - end), desc->type,
+                                     desc->attribute)) {
+                return false;
+            }
+            desc = efi_memmap_walk(desc);
+            desc = efi_memmap_walk(desc); 
+            ++i;
+            /* no need to check more */
+            break;
+        }
+        else {
+            printk(TBOOT_ERR"we should never get here\n");
+            return false;
+        }
+
+        cont:
+        ++i;
+    }
+
+    return true;
+}
+
+/**
+ * @brief Print whole memory map
+ */
+void efi_memmap_dump(void)
+{
+    efi_mem_descr_t* desc = NULL;
+    while ((desc = efi_memmap_walk(desc)) != NULL) {
+        printk(TBOOT_INFO" %016llx - %016llx (%-2d | 0x%llx)\n",
+               desc->physical_start,
+               desc->physical_start + (desc->num_pages << EFI_PAGE_SHIFT),
+               desc->type, desc->attribute);
+    }
+}
+
+/**
+ * @brief Find in memory map highest avaliable free space that meets given
+ *        requirements
+ * 
+ * Free space is a region in memory map of following types:
+ *   - EFI_LOADER_CODE
+ *   - EFI_LOADER_DATA
+ *   - EFI_CONVENTIONAL_MEMORY
+ * Boot services memory is excluded because it can be occupied by tables
+ * that Linux may want to access later, ex. EFI_MEMORY_ATTRIBUTES_TABLE
+ * 
+ * @param size     minimal size
+ * @param limit    highest possible address
+ * @param ram_base return address of found region
+ * @param ram_size return size of found region, bigger or equal @p size
+ */
+bool efi_memmap_get_highest_sized_ram(uint64_t size, uint64_t limit,
+                                      uint64_t *ram_base, uint64_t *ram_size)
+{
+    uint64_t last_fit_base = 0, last_fit_size = 0;
+
+    if (ram_base == NULL || ram_size == NULL || !efi_mmap_available) {
+        return false;
+    }
+
+    efi_mem_descr_t* desc = NULL;
+    while ((desc = efi_memmap_walk(desc)) != NULL) {
+        if (region_is_free(desc->type)) {
+            uint64_t base = desc->physical_start;
+            uint64_t length = desc->num_pages * (1 << EFI_PAGE_SHIFT);
+
+            /* over 4GB so use the last region that fit */
+            if ( base + length > limit )
+                break;
+            if ( size <= length ) {
+                last_fit_base = base;
+                last_fit_size = length;
+            }
+        }
+    }
+
+    printk("get_highest_sized_ram: size %llx -> base %llx, size %llx\n",
+           size, *ram_base, *ram_size);
+
+    if (last_fit_size == 0) {
+        return false;
+    } else {
+        *ram_base = last_fit_base;
+        *ram_size = last_fit_size;
+        return true;
+    }
+}
+
+static bool insert_after_region(uint32_t pos, uint64_t addr, uint64_t size,
+                                uint32_t type, uint64_t attr)
+{
+    /* no more room */
+    if (efi_mmap->size / efi_mmap->descr_size + 1 > EFI_MEMMAP_MAX_ENTRIES)
+        return false;
+
+    pos *= efi_mmap->descr_size;
+
+    /* shift (copy) everything up one entry */
+    for (uint32_t i = efi_mmap->size; i > pos; i -= efi_mmap->descr_size) {
+        uint32_t bytes = efi_mmap->descr_size;
+        void* to = efi_mmap->descr + i;
+        void* from = efi_mmap->descr + i - bytes;
+        memcpy(to, from, bytes);
+    }
+
+    efi_mem_descr_t* desc = (efi_mem_descr_t*)(efi_mmap->descr + pos +
+                            efi_mmap->descr_size);
+    memset(desc, 0, efi_mmap->descr_size);
+    desc->type = type;
+    desc->physical_start = addr;
+    desc->num_pages = size >> EFI_PAGE_SHIFT;
+    desc->attribute = attr;
+    efi_mmap->size += efi_mmap->descr_size;
+
+    return true;
+}
+
+static bool region_is_free(uint32_t region_type)
+{
+    if (region_type == EFI_LOADER_CODE || region_type == EFI_LOADER_DATA ||
+            region_type == EFI_CONVENTIONAL_MEMORY) {
+        return true;
+    } else {
+        return false;
+    }
+}
\ No newline at end of file
diff --git a/tboot/common/linux.c b/tboot/common/linux.c
index a73bd4a..05fd408 100644
--- a/tboot/common/linux.c
+++ b/tboot/common/linux.c
@@ -50,6 +50,7 @@
 #include <hash.h>
 #include <integrity.h>
 #include <processor.h>
+#include <efi_memmap.h>
 
 extern loader_ctx *g_ldr_ctx;
 
@@ -174,11 +175,13 @@ bool expand_linux_image(const void *linux_image, size_t linux_size,
             mem_limit = 0x100000000ULL;
 
         uint64_t max_ram_base, max_ram_size;
-        get_highest_sized_ram(initrd_size, mem_limit,
-                              &max_ram_base, &max_ram_size);
-        if ( max_ram_size == 0 ) {
-            printk(TBOOT_ERR"not enough RAM for initrd\n");
-            return false;
+        if (!efi_memmap_get_highest_sized_ram(initrd_size, mem_limit,
+                                              &max_ram_base, &max_ram_size)) {
+            if (!e820_get_highest_sized_ram(initrd_size, mem_limit,
+                                            &max_ram_base, &max_ram_size)) {
+                printk(TBOOT_ERR"not enough RAM for initrd\n");
+                return false;
+            }
         }
         if ( initrd_size > max_ram_size ) {
             printk(TBOOT_ERR"initrd_size is too large\n");
@@ -376,8 +379,8 @@ bool expand_linux_image(const void *linux_image, size_t linux_size,
             }
         }
 
-        efi_mmap_addr = find_efi_memmap(g_ldr_ctx, &descr_size,
-                                        &descr_vers, &mmap_size);
+        efi_mmap_addr = efi_memmap_get_addr(&descr_size, &descr_vers,
+                                            &mmap_size);
         if (!efi_mmap_addr) {
             printk(TBOOT_INFO"failed to get EFI memory map\n");
             efi->efi_memdescr_size = 0x1; // Avoid div by 0 in kernel.
@@ -394,6 +397,9 @@ bool expand_linux_image(const void *linux_image, size_t linux_size,
              */
             efi->efi_memmap_hi = 0;
 
+            printk(TBOOT_INFO"EFI memory map after modifications:\n");
+            efi_memmap_dump();
+
             printk(TBOOT_INFO "EFI memmap: memmap base: 0x%x, memmap size: 0x%x\n",
                   efi->efi_memmap, efi->efi_memmap_size);
             printk(TBOOT_INFO "EFI memmap: descr size: 0x%x, descr version: 0x%x\n",
diff --git a/tboot/common/loader.c b/tboot/common/loader.c
index cbb7def..00e0d9f 100644
--- a/tboot/common/loader.c
+++ b/tboot/common/loader.c
@@ -57,6 +57,7 @@
 #include <txt/acmod.h>
 #include <cmdline.h>
 #include <tpm.h>
+#include <efi_memmap.h>
 
 /* copy of kernel/VMM command line so that can append 'tboot=0x1234' */
 static char *new_cmdline = (char *)TBOOT_KERNEL_CMDLINE_ADDR;
@@ -800,16 +801,20 @@ static bool move_modules_to_high_memory(loader_ctx  *lctx)
         memRequired += PAGE_UP(m->mod_end - m->mod_start);
     }
 
-    /*  NOTE: the e820 map has been modified already to reserve critical
+    /*  NOTE: the memory maps have been modified already to reserve critical
         memory regions (tboot memory, etc ...). get_highest_sized_ram
         will return a range that excludes critical memory regions. */
-    get_highest_sized_ram(  memRequired, 0x100000000ULL,
-                            &max_ram_base, &max_ram_size);
-    if(!max_ram_base || !max_ram_size){
-        printk(TBOOT_INFO"ERROR No memory area found for image relocation!\n");
-        printk(TBOOT_INFO"required 0x%X\n", memRequired);
-        return false;
+    if (!efi_memmap_get_highest_sized_ram(memRequired, 0x100000000ULL,
+                                          &max_ram_base, &max_ram_size)) {
+        if (!e820_get_highest_sized_ram(memRequired, 0x100000000ULL,
+                                        &max_ram_base, &max_ram_size)) {
+            printk(TBOOT_INFO"ERROR No memory area found for image"
+                             "relocation!\n");
+            printk(TBOOT_INFO"required 0x%X\n", memRequired);
+            return false;
+        }
     }
+
     printk(TBOOT_INFO"highest suitable area @ 0x%llX (size 0x%llX)\n",
             max_ram_base, max_ram_size);
     ld_ceiling = PAGE_DOWN(max_ram_base + max_ram_size);
@@ -1383,6 +1388,9 @@ bool launch_kernel(bool is_measured_launch)
         printk(TBOOT_INFO"reserving tboot memory log (%Lx - %Lx) in e820 table\n", base, (base + size - 1));
         if ( !e820_protect_region(base, size, E820_RESERVED) )
             apply_policy(TB_ERR_FATAL);
+        if (!efi_memmap_reserve(base, size)) {
+            apply_policy(TB_ERR_FATAL);
+        }
     }
 
     /* replace map in loader context with copy */
diff --git a/tboot/common/tboot.c b/tboot/common/tboot.c
index 4518c3e..8faf4f5 100644
--- a/tboot/common/tboot.c
+++ b/tboot/common/tboot.c
@@ -71,6 +71,7 @@
 #include <cmdline.h>
 #include <tpm_20.h>
 #include <vtd.h>
+#include <efi_memmap.h>
 
 extern void _prot_to_real(uint32_t dist_addr);
 extern bool set_policy(void);
@@ -208,6 +209,9 @@ static void post_launch(void)
     printk(TBOOT_INFO"protecting tboot (%Lx - %Lx) in e820 table\n", base,      (base + size - 1));
     if ( !e820_protect_region(base, size, mem_type) )      
         apply_policy(TB_ERR_FATAL);
+    if (!efi_memmap_reserve(base, size)) {
+        apply_policy(TB_ERR_FATAL);
+    }
 
     /*
      * verify modules against policy
@@ -378,6 +382,10 @@ void begin_launch(void *addr, uint32_t magic)
     /* make copy of e820 map that we will use and adjust */
     if ( !s3_flag ) {
         if ( !copy_e820_map(g_ldr_ctx) )  apply_policy(TB_ERR_FATAL);
+        if (efi_memmap_copy(g_ldr_ctx)) {
+            printk(TBOOT_INFO"Original EFI memory map:\n");
+            efi_memmap_dump();
+        }
     }
 
     /* we need to make sure this is a (TXT-) capable platform before using */
diff --git a/tboot/include/e820.h b/tboot/include/e820.h
index efe9a79..760e1c3 100644
--- a/tboot/include/e820.h
+++ b/tboot/include/e820.h
@@ -86,42 +86,8 @@ extern void print_e820_map(void);
 extern uint32_t e820_check_region(uint64_t base, uint64_t length);
 extern bool get_ram_ranges(uint64_t *min_lo_ram, uint64_t *max_lo_ram,
                            uint64_t *min_hi_ram, uint64_t *max_hi_ram);
-extern void get_highest_sized_ram(uint64_t size, uint64_t limit,
-                                  uint64_t *ram_base, uint64_t *ram_size);
-
-/*
- * Memory map descriptor:
- */
-
-/* Memory types: */
-#define EFI_RESERVED_TYPE		 0
-#define EFI_LOADER_CODE			 1
-#define EFI_LOADER_DATA			 2
-#define EFI_BOOT_SERVICES_CODE		 3
-#define EFI_BOOT_SERVICES_DATA		 4
-#define EFI_RUNTIME_SERVICES_CODE	 5
-#define EFI_RUNTIME_SERVICES_DATA	 6
-#define EFI_CONVENTIONAL_MEMORY		 7
-#define EFI_UNUSABLE_MEMORY		 8
-#define EFI_ACPI_RECLAIM_MEMORY		 9
-#define EFI_ACPI_MEMORY_NVS		10
-#define EFI_MEMORY_MAPPED_IO		11
-#define EFI_MEMORY_MAPPED_IO_PORT_SPACE	12
-#define EFI_PAL_CODE			13
-#define EFI_MAX_MEMORY_TYPE		14
-
-/* Attribute values: */
-#define EFI_MEMORY_UC		((u64)0x0000000000000001ULL)	/* uncached */
-#define EFI_MEMORY_WC		((u64)0x0000000000000002ULL)	/* write-coalescing */
-#define EFI_MEMORY_WT		((u64)0x0000000000000004ULL)	/* write-through */
-#define EFI_MEMORY_WB		((u64)0x0000000000000008ULL)	/* write-back */
-#define EFI_MEMORY_WP		((u64)0x0000000000001000ULL)	/* write-protect */
-#define EFI_MEMORY_RP		((u64)0x0000000000002000ULL)	/* read-protect */
-#define EFI_MEMORY_XP		((u64)0x0000000000004000ULL)	/* execute-protect */
-#define EFI_MEMORY_RUNTIME	((u64)0x8000000000000000ULL)	/* range requires runtime mapping */
-#define EFI_MEMORY_DESCRIPTOR_VERSION	1
-
-#define EFI_PAGE_SHIFT		12
+extern bool e820_get_highest_sized_ram(uint64_t size, uint64_t limit,
+                                       uint64_t *ram_base, uint64_t *ram_size);
 
 #endif    /* __E820_H__ */
 
diff --git a/tboot/include/efi_memmap.h b/tboot/include/efi_memmap.h
new file mode 100644
index 0000000..1f5eed2
--- /dev/null
+++ b/tboot/include/efi_memmap.h
@@ -0,0 +1,101 @@
+/*
+ * efi_memmap.h: EFI memory map parsing functions
+ *
+ * Copyright (c) 2006-2020, Intel Corporation
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *
+ *   * Redistributions of source code must retain the above copyright
+ *     notice, this list of conditions and the following disclaimer.
+ *   * Redistributions in binary form must reproduce the above
+ *     copyright notice, this list of conditions and the following
+ *     disclaimer in the documentation and/or other materials provided
+ *     with the distribution.
+ *   * Neither the name of the Intel Corporation nor the names of its
+ *     contributors may be used to endorse or promote products derived
+ *     from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+ * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+ * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+ * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
+ * OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ */
+
+#ifndef _EFI_MEMMAP_H_
+#define _EFI_MEMMAP_H_
+
+#include <types.h>
+#include <config.h>
+
+/* Memory types: */
+#define EFI_RESERVED_TYPE               0
+#define EFI_LOADER_CODE                 1
+#define EFI_LOADER_DATA                 2
+#define EFI_BOOT_SERVICES_CODE          3
+#define EFI_BOOT_SERVICES_DATA          4
+#define EFI_RUNTIME_SERVICES_CODE       5
+#define EFI_RUNTIME_SERVICES_DATA       6
+#define EFI_CONVENTIONAL_MEMORY         7
+#define EFI_UNUSABLE_MEMORY             8
+#define EFI_ACPI_RECLAIM_MEMORY         9
+#define EFI_ACPI_MEMORY_NVS             10
+#define EFI_MEMORY_MAPPED_IO            11
+#define EFI_MEMORY_MAPPED_IO_PORT_SPACE 12
+#define EFI_PAL_CODE                    13
+#define EFI_MAX_MEMORY_TYPE             14
+
+/* Attribute values: */
+#define EFI_MEMORY_UC      ((u64)0x0000000000000001ULL)	/* uncached */
+#define EFI_MEMORY_WC      ((u64)0x0000000000000002ULL)	/* write-coalescing */
+#define EFI_MEMORY_WT      ((u64)0x0000000000000004ULL)	/* write-through */
+#define EFI_MEMORY_WB      ((u64)0x0000000000000008ULL)	/* write-back */
+#define EFI_MEMORY_WP      ((u64)0x0000000000001000ULL)	/* write-protect */
+#define EFI_MEMORY_RP      ((u64)0x0000000000002000ULL)	/* read-protect */
+#define EFI_MEMORY_XP      ((u64)0x0000000000004000ULL)	/* execute-protect */
+#define EFI_MEMORY_RUNTIME ((u64)0x8000000000000000ULL)	/* requires runtime mapping */
+#define EFI_MEMORY_DESCRIPTOR_VERSION 1
+
+#define EFI_PAGE_SHIFT 12
+#define EFI_MEMMAP_MAX_ENTRIES 682 /* limited by TBOOT_EFI_MEMMAP_COPY_SIZE */
+
+typedef struct __packed
+{
+    uint32_t size;
+    uint32_t descr_size;
+    uint8_t  descr[0]; /**< array of efi_mem_descr_t,
+                            each element has descr_size bytes */
+} efi_memmap_t;
+
+typedef struct __packed
+{
+    uint32_t type;
+    uint32_t padding;
+    uint64_t physical_start;
+    uint64_t virtual_start;
+    uint64_t num_pages;
+    uint64_t attribute;
+} efi_mem_descr_t;
+
+bool efi_memmap_copy(loader_ctx *lctx);
+uint32_t efi_memmap_get_addr(uint32_t *descr_size, uint32_t *descr_vers,
+                             uint32_t *mmap_size);
+efi_mem_descr_t* efi_memmap_walk(efi_mem_descr_t* prev);
+bool efi_memmap_is_free(uint32_t region_type);
+bool efi_memmap_reserve(uint64_t base, uint64_t length);
+bool efi_memmap_get_highest_sized_ram(uint64_t size, uint64_t limit,
+                                      uint64_t *ram_base, uint64_t *ram_size);
+void efi_memmap_dump(void);
+
+#endif
\ No newline at end of file
diff --git a/tboot/txt/txt.c b/tboot/txt/txt.c
index aaa9634..dbb4f10 100644
--- a/tboot/txt/txt.c
+++ b/tboot/txt/txt.c
@@ -57,6 +57,7 @@
 #include <cmdline.h>
 #include <acpi.h>
 #include <vtd.h>
+#include <efi_memmap.h>
 #include <txt/txt.h>
 #include <txt/config_regs.h>
 #include <txt/mtrrs.h>
@@ -1216,6 +1217,9 @@ tb_error_t txt_protect_mem_regions(void)
            (base + size - 1));
     if ( !e820_protect_region(base, size, E820_RESERVED) )
         return TB_ERR_FATAL;
+    if (!efi_memmap_reserve(base, size)) {
+        return TB_ERR_FATAL;
+    }
 
     /* SINIT */
     base = read_pub_config_reg(TXTCR_SINIT_BASE);
@@ -1224,6 +1228,9 @@ tb_error_t txt_protect_mem_regions(void)
            (base + size - 1));
     if ( !e820_protect_region(base, size, E820_RESERVED) )
         return TB_ERR_FATAL;
+    if (!efi_memmap_reserve(base, size)) {
+        return TB_ERR_FATAL;
+    }
 
     /* TXT private space */
     base = TXT_PRIV_CONFIG_REGS_BASE;
@@ -1233,6 +1240,9 @@ tb_error_t txt_protect_mem_regions(void)
            base, (base + size - 1));
     if ( !e820_protect_region(base, size, E820_RESERVED) )
         return TB_ERR_FATAL;
+    if (!efi_memmap_reserve(base, size)) {
+        return TB_ERR_FATAL;
+    }
 
     /* ensure that memory not marked as good RAM by the MDRs is RESERVED in
        the e820 table */
diff --git a/tboot/txt/verify.c b/tboot/txt/verify.c
index cb90e7e..72df3c8 100644
--- a/tboot/txt/verify.c
+++ b/tboot/txt/verify.c
@@ -54,6 +54,7 @@
 #include <integrity.h>
 #include <cmdline.h>
 #include <vtd.h>
+#include <efi_memmap.h>
 #include <txt/txt.h>
 #include <txt/smx.h>
 #include <txt/mtrrs.h>
@@ -255,6 +256,9 @@ static bool reserve_vtd_delta_mem(uint64_t min_lo_ram, uint64_t max_lo_ram,
                base, base + length);
         if ( !e820_reserve_ram(base, length) )
             return false;
+        if (!efi_memmap_reserve(base, length)) {
+           return false;
+        }
     }
     if ( max_hi_ram != (os_sinit_data->vtd_pmr_hi_base +
                         os_sinit_data->vtd_pmr_hi_size) ) {
@@ -264,6 +268,9 @@ static bool reserve_vtd_delta_mem(uint64_t min_lo_ram, uint64_t max_lo_ram,
                base, base + length);
         if ( !e820_reserve_ram(base, length) )
             return false;
+        if (!efi_memmap_reserve(base, length)) {
+           return false;
+        }
     }
 
     return true;
-- 
2.26.2

