2 files changed, 10 insertions, 0 deletions
diff --git a/src/lib/Makefile.inc b/src/lib/Makefile.inc
index a3235a2062..ec57bc2d2f 100644
--- a/src/lib/Makefile.inc
+++ b/src/lib/Makefile.inc
@@ -32,6 +32,7 @@ endif
 ramstage-y += memcmp.c
 ramstage-y += memmove.c
 ramstage-y += malloc.c
+smm-$(CONFIG_SMM_TSEG) += malloc.c
 ramstage-y += delay.c
 ramstage-y += fallback_boot.c
 ramstage-y += compute_ip_checksum.c
diff --git a/src/lib/malloc.c b/src/lib/malloc.c
index d2011a10d0..43e514af58 100644
--- a/src/lib/malloc.c
+++ b/src/lib/malloc.c
@@ -1,5 +1,8 @@
 #include <stdlib.h>
 #include <console/console.h>
+#ifdef __SMM__
+#include <cpu/x86/smm.h>
+#endif
 
 #if CONFIG_DEBUG_MALLOC
 #define MALLOCDBG(x...) printk(BIOS_SPEW, x)
@@ -43,5 +46,11 @@ void *memalign(size_t boundary, size_t size)
 
 void *malloc(size_t size)
 {
+#if CONFIG_SMM_TSEG && defined(__SMM__)
+	if (!free_mem_ptr) {
+		free_mem_ptr = &_heap + smi_get_tseg_base();
+		free_mem_end_ptr = &_eheap + smi_get_tseg_base();
+	}
+#endif
 	return memalign(sizeof(u64), size);
 }