Merge "QcomModulePkg: Fix fastboot getvar all command access invalid data"
diff --git a/AndroidBoot.mk b/AndroidBoot.mk
index aaa072a..0e879be 100644
--- a/AndroidBoot.mk
+++ b/AndroidBoot.mk
@@ -49,6 +49,14 @@
 	BOARD_BOOTLOADER_PRODUCT_NAME := QC_Reference_Phone
 endif
 
+ifeq ($(BOARD_ABL_SAFESTACK_DISABLE),true)
+	ABL_SAFESTACK := false
+else
+	ABL_SAFESTACK := true
+endif
+
+SAFESTACK_SUPPORTED_CLANG_VERSION = 6.0
+
 # For most platform, abl needed always be built
 # in aarch64 arthitecture to run.
 # Specify BOOTLOADER_ARCH if needed to built with
@@ -92,6 +100,8 @@
 		CLANG_BIN=$(CLANG_BIN) \
 		CLANG_PREFIX=$(CLANG35_PREFIX)\
 		ABL_USE_SDLLVM=$(ABL_USE_SDLLVM) \
+		ABL_SAFESTACK=$(ABL_SAFESTACK) \
+		SAFESTACK_SUPPORTED_CLANG_VERSION=$(SAFESTACK_SUPPORTED_CLANG_VERSION) \
 		CLANG_GCC_TOOLCHAIN=$(CLANG35_GCC_TOOLCHAIN)\
 		TARGET_ARCHITECTURE=$(TARGET_ARCHITECTURE) \
 		BOARD_BOOTLOADER_PRODUCT_NAME=$(BOARD_BOOTLOADER_PRODUCT_NAME)
diff --git a/QcomModulePkg/Application/LinuxLoader/LinuxLoader.c b/QcomModulePkg/Application/LinuxLoader/LinuxLoader.c
index f5a6bf8..db37872 100644
--- a/QcomModulePkg/Application/LinuxLoader/LinuxLoader.c
+++ b/QcomModulePkg/Application/LinuxLoader/LinuxLoader.c
@@ -51,6 +51,34 @@
 STATIC BOOLEAN BootIntoFastboot = FALSE;
 STATIC BOOLEAN BootIntoRecovery = FALSE;
 
+STATIC VOID* UnSafeStackPtr;
+
+STATIC EFI_STATUS __attribute__ ( (no_sanitize ("safe-stack")))
+AllocateUnSafeStackPtr (VOID)
+{
+
+  EFI_STATUS Status = EFI_SUCCESS;
+
+  UnSafeStackPtr = AllocatePool (BOOT_LOADER_MAX_UNSAFE_STACK_SIZE);
+  if (UnSafeStackPtr == NULL) {
+    DEBUG ((EFI_D_ERROR, "Failed to Allocate memory for UnSafeStack \n"));
+    Status = EFI_OUT_OF_RESOURCES;
+    return Status;
+  }
+
+  UnSafeStackPtr += BOOT_LOADER_MAX_UNSAFE_STACK_SIZE;
+
+  return Status;
+}
+
+//This function is to return the Unsafestack ptr address
+VOID** __attribute__ ( (no_sanitize ("safe-stack")))
+__safestack_pointer_address (VOID)
+{
+
+  return (VOID**) &UnSafeStackPtr;
+}
+
 // This function is used to Deactivate MDTP by entering recovery UI
 
 STATIC EFI_STATUS MdtpDisable (VOID)
@@ -111,7 +139,7 @@
 
  **/
 
-EFI_STATUS EFIAPI
+EFI_STATUS EFIAPI  __attribute__ ( (no_sanitize ("safe-stack")))
 LinuxLoaderEntry (IN EFI_HANDLE ImageHandle, IN EFI_SYSTEM_TABLE *SystemTable)
 {
   EFI_STATUS Status;
@@ -123,6 +151,13 @@
 
   DEBUG ((EFI_D_INFO, "Loader Build Info: %a %a\n", __DATE__, __TIME__));
 
+  Status = AllocateUnSafeStackPtr ();
+  if (Status != EFI_SUCCESS) {
+    DEBUG ((EFI_D_ERROR, "Unable to Allocate memory for Unsafe Stack: %r\n",
+            Status));
+    goto stack_guard_update_default;
+  }
+
   StackGuardChkSetup ();
 
   BootStatsSetTimeStamp (BS_BL_START);
diff --git a/QcomModulePkg/Application/LinuxLoader/LinuxLoader.inf b/QcomModulePkg/Application/LinuxLoader/LinuxLoader.inf
index 3de6230..7a1d3e7 100644
--- a/QcomModulePkg/Application/LinuxLoader/LinuxLoader.inf
+++ b/QcomModulePkg/Application/LinuxLoader/LinuxLoader.inf
@@ -23,6 +23,9 @@
 [Sources]
 	LinuxLoader.c
 
+[BuildOptions]
+	GCC:*_*_*_CC_FLAGS = $(LLVM_ENABLE_SAFESTACK) $(LLVM_SAFESTACK_USE_PTR) $(LLVM_SAFESTACK_COLORING)
+
 [BuildOptions.AARCH64]
 	GCC:*_*_*_CC_FLAGS = $(SDLLVM_COMPILE_ANALYZE) $(SDLLVM_ANALYZE_REPORT)
 	GCC:*_*_*_CC_FLAGS = $(UBSAN_UEFI_GCC_FLAG_UNDEFINED)
diff --git a/QcomModulePkg/Include/Library/Board.h b/QcomModulePkg/Include/Library/Board.h
index ad99a33..1ea8004 100644
--- a/QcomModulePkg/Include/Library/Board.h
+++ b/QcomModulePkg/Include/Library/Board.h
@@ -89,8 +89,8 @@
 UfsGetSetBootLun (UINT32 *UfsBootlun, BOOLEAN IsGet);
 BOOLEAN BoardPlatformFusion (VOID);
 UINT32 BoardPlatformRawChipId (VOID);
-EFI_STATUS
-GetRamPartitions (RamPartitionEntry **RamPartitions, UINT32 *NumPartitions);
-VOID
-GetPageSize (UINT32 *PageSize);
+EFI_STATUS GetRamPartitions (RamPartitionEntry **RamPartitions,
+                  UINT32 *NumPartitions);
+EFI_STATUS GetGranuleSize (UINT32 *MinPasrGranuleSize);
+VOID GetPageSize (UINT32 *PageSize);
 #endif
diff --git a/QcomModulePkg/Include/Library/LinuxLoaderLib.h b/QcomModulePkg/Include/Library/LinuxLoaderLib.h
index 577cec1..1e34256 100644
--- a/QcomModulePkg/Include/Library/LinuxLoaderLib.h
+++ b/QcomModulePkg/Include/Library/LinuxLoaderLib.h
@@ -110,6 +110,10 @@
 #define BOOT_DEV_MAX_LEN 32
 #define BOOT_DEV_NAME_SIZE_MAX 10
 
+/* Allocate unsafe stack size of 128KB to address worst case,
+ * which is same as the  normal stack size */
+#define BOOT_LOADER_MAX_UNSAFE_STACK_SIZE (1 << 17)
+
 /* Any data specific to additional attributes can be added here. */
 typedef struct {
   EFI_GUID *RootDeviceType; /* GUID Selecting the root device type */
diff --git a/QcomModulePkg/Include/Protocol/EFIRamPartition.h b/QcomModulePkg/Include/Protocol/EFIRamPartition.h
index 1824fd2..56ff17a 100644
--- a/QcomModulePkg/Include/Protocol/EFIRamPartition.h
+++ b/QcomModulePkg/Include/Protocol/EFIRamPartition.h
@@ -1,4 +1,4 @@
-/* Copyright (c) 2015-2016, The Linux Foundation. All rights reserved.
+/* Copyright (c) 2015-2016, 2018 The Linux Foundation. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -38,7 +38,7 @@
 /**
   Protocol version.
 */
-#define EFI_RAMPARTITION_PROTOCOL_REVISION 0x0000000000010000
+#define EFI_RAMPARTITION_PROTOCOL_REVISION 0x0000000000010001
 /** @} */ /* end_addtogroup efi_ramPartition_constants */
 
 /*  Protocol GUID definition */
@@ -119,6 +119,28 @@
     OUT UINT32 *HighestBankBit);
 
 /* ============================================================================
+**  Function : EFI_RamPartition_GetMinPasrSize
+** ============================================================================
+*/
+/** @ingroup EFI_RamPartition_GetMinPasrSize
+  @par Summary
+  Gets the MinPasrSize
+
+  @param[in]   This         Pointer to the EFI_RAMPARTITION_PROTOCOL instance.
+  @param[out]  MinPasrSize  Pointer to MinPasrSize
+
+  @return
+  EFI_SUCCESS        -- Function completed successfully.
+  EFI_PROTOCOL_ERROR -- Error occurred during the operation.
+*/
+typedef
+EFI_STATUS
+(EFIAPI *EFI_RAMPARTITION_GETMINPASRSIZE)(
+   IN EFI_RAMPARTITION_PROTOCOL *This,
+   OUT UINT32                *MinPasrSize
+   );
+
+/* ============================================================================
 **  Function : EFI_RamPartition_GetRamPartitions
 ** ============================================================================
 */
@@ -154,6 +176,7 @@
   EFI_RAMPARTITION_GETRAMPARTITIONVERSION GetRamPartitionVersion;
   EFI_RAMPARTITION_GETHIGHESTBANKBIT GetHighestBankBit;
   EFI_RAMPARTITION_GETRAMPARTITIONS GetRamPartitions;
+  EFI_RAMPARTITION_GETMINPASRSIZE GetMinPasrSize;
 };
 
 #endif /* __EFIRAMPARTITION_H__ */
diff --git a/QcomModulePkg/Library/BootLib/Board.c b/QcomModulePkg/Library/BootLib/Board.c
index cc83bd3..022f630 100644
--- a/QcomModulePkg/Library/BootLib/Board.c
+++ b/QcomModulePkg/Library/BootLib/Board.c
@@ -79,6 +79,33 @@
 }
 
 EFI_STATUS
+GetGranuleSize (UINT32 *MinPasrGranuleSize)
+{
+  EFI_STATUS Status = EFI_NOT_FOUND;
+  EFI_RAMPARTITION_PROTOCOL *pRamPartProtocol = NULL;
+
+  Status = gBS->LocateProtocol (&gEfiRamPartitionProtocolGuid, NULL,
+                                (VOID **)&pRamPartProtocol);
+  if (EFI_ERROR (Status) ||
+        (pRamPartProtocol == NULL)) {
+    DEBUG ((EFI_D_ERROR,
+            "Locate EFI_RAMPARTITION_Protocol failed, Status = %r \n",
+            Status));
+    return Status;
+  }
+
+  Status = pRamPartProtocol->GetMinPasrSize (pRamPartProtocol,
+                                             MinPasrGranuleSize);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((EFI_D_ERROR,
+            "Failed to get MinPasrSize, Status = %r\n",
+            Status));
+    return Status;
+  }
+  return Status;
+}
+
+EFI_STATUS
 BaseMem (UINT64 *BaseMemory)
 {
   EFI_STATUS Status = EFI_NOT_FOUND;
diff --git a/QcomModulePkg/Library/BootLib/BootLib.inf b/QcomModulePkg/Library/BootLib/BootLib.inf
index 8e6202e..981c9b7 100644
--- a/QcomModulePkg/Library/BootLib/BootLib.inf
+++ b/QcomModulePkg/Library/BootLib/BootLib.inf
@@ -36,6 +36,9 @@
 	VERSION_STRING                 = 1.0
 	LIBRARY_CLASS                  = BootLib
 
+[BuildOptions]
+	GCC:*_*_*_CC_FLAGS = $(LLVM_ENABLE_SAFESTACK) $(LLVM_SAFESTACK_USE_PTR) $(LLVM_SAFESTACK_COLORING)
+
 [BuildOptions.AARCH64]
 	GCC:*_*_*_CC_FLAGS = -O2 $(SDLLVM_COMPILE_ANALYZE) $(SDLLVM_ANALYZE_REPORT)
 	GCC:*_*_*_CC_FLAGS = $(UBSAN_UEFI_GCC_FLAG_UNDEFINED)
diff --git a/QcomModulePkg/Library/BootLib/LinuxLoaderLib.c b/QcomModulePkg/Library/BootLib/LinuxLoaderLib.c
index 51c2b63..0f49783 100644
--- a/QcomModulePkg/Library/BootLib/LinuxLoaderLib.c
+++ b/QcomModulePkg/Library/BootLib/LinuxLoaderLib.c
@@ -29,6 +29,7 @@
 #include "LinuxLoaderLib.h"
 #include "AutoGen.h"
 #include <Library/BootLinux.h>
+#include <FastbootLib/FastbootCmds.h>
 
 /* Volume Label size 11 chars, round off to 16 */
 #define VOLUME_LABEL_SIZE 16
@@ -500,6 +501,9 @@
   CHAR8 *ImageBuffer = NULL;
   UINT32 DivMsgBufSize;
   UINT32 WriteBlockSize;
+  UINT64 WriteUnitSize = MAX_WRITE_SIZE;
+  INT64 LeftSize = 0;
+  UINT32 WriteSize = 0;
 
   if ((BlockIo == NULL) ||
     (Image == NULL)) {
@@ -536,18 +540,33 @@
   }
 
   DivMsgBufSize = (Size / WriteBlockSize) * WriteBlockSize;
-
+  WriteUnitSize = ROUND_TO_PAGE (WriteUnitSize, WriteBlockSize - 1);
   if (DivMsgBufSize) {
-    Status = BlockIo->WriteBlocks (BlockIo,
-                                   BlockIo->Media->MediaId,
-                                   Offset,
-                                   DivMsgBufSize,
-                                   Image);
-    if (Status != EFI_SUCCESS) {
-      DEBUG ((EFI_D_ERROR, "Write the divisible Image failed :%r\n", Status));
-      return Status;
+    /* The big image buffer may take a long flashing time which will block
+       parallel usb image download. It will cause the fastboot  protocol host
+       side timeout. So split the image into small writing units  to let usb
+       have chance to champ in and doing work in parallel.
+      */
+    if (!IsUsbTimerStarted ()) {
+      WriteUnitSize = DivMsgBufSize;
     }
-    Offset += DivMsgBufSize / BlockIo->Media->BlockSize;
+
+    LeftSize = DivMsgBufSize;
+    while (LeftSize > 0) {
+      WriteSize = LeftSize > WriteUnitSize? WriteUnitSize : LeftSize;
+      Status = BlockIo->WriteBlocks (BlockIo,
+                                     BlockIo->Media->MediaId,
+                                     Offset,
+                                     WriteSize,
+                                     Image + DivMsgBufSize - LeftSize);
+
+      if (Status != EFI_SUCCESS) {
+        DEBUG ((EFI_D_ERROR, "Write the divisible Image failed :%r\n", Status));
+        return Status;
+      }
+      Offset += WriteSize / BlockIo->Media->BlockSize;
+      LeftSize -= WriteSize;
+    }
   }
 
   if (Size - DivMsgBufSize > 0) {
@@ -675,8 +694,10 @@
   QCOM_VERIFIEDBOOT_PROTOCOL *VbIntf;
 
   /* If verified boot is not enabled, return SUCCESS */
-  if (!VerifiedBootEnbled ())
+  if (!VerifiedBootEnbled () ||
+     (GetAVBVersion () == AVB_LE )) {
     return EFI_SUCCESS;
+  }
 
   Status = gBS->LocateProtocol (&gEfiQcomVerifiedBootProtocolGuid, NULL,
                                 (VOID **)&VbIntf);
diff --git a/QcomModulePkg/Library/BootLib/UpdateDeviceTree.c b/QcomModulePkg/Library/BootLib/UpdateDeviceTree.c
index c90315a..aa713ed 100644
--- a/QcomModulePkg/Library/BootLib/UpdateDeviceTree.c
+++ b/QcomModulePkg/Library/BootLib/UpdateDeviceTree.c
@@ -219,6 +219,35 @@
 }
 
 STATIC
+VOID
+UpdateGranuleInfo (VOID *fdt)
+{
+  EFI_STATUS Status = EFI_SUCCESS;
+  UINT32 GranuleNodeOffset;
+  UINT32 GranuleSize;
+  INT32 Ret;
+
+  Status = GetGranuleSize (&GranuleSize);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((EFI_D_ERROR,
+            "Update Granule Size failed!!! Status = %r\r\n",
+            Status));
+    return;
+  }
+
+  GranuleNodeOffset = fdt_path_offset (fdt, "/mem-offline");
+  if (GranuleNodeOffset < 0) {
+    DEBUG ((EFI_D_ERROR, "WARNING: Could not find mem-offline node.\n"));
+    return;
+  }
+
+  Ret = fdt_setprop_u32 (fdt, GranuleNodeOffset, "granule", GranuleSize);
+  if (Ret) {
+    DEBUG ((EFI_D_ERROR, "WARNING: Granule size update failed.\n"));
+  }
+}
+
+STATIC
 EFI_STATUS
 QueryMemoryCellSize (IN VOID *Fdt, OUT UINT32 *MemoryCellLen)
 {
@@ -338,11 +367,16 @@
 
   /* Get Available memory from partition table */
   Status = AddMemMap (fdt, MemNodeOffset, BootWith32Bit);
-  if (EFI_ERROR (Status))
+  if (EFI_ERROR (Status)) {
     DEBUG ((EFI_D_ERROR,
             "Invalid memory configuration, check memory partition table: %r\n",
             Status));
+    goto out;
+  }
 
+  UpdateGranuleInfo (fdt);
+
+out:
   return Status;
 }
 
diff --git a/QcomModulePkg/Library/FastbootLib/FastbootCmds.c b/QcomModulePkg/Library/FastbootLib/FastbootCmds.c
index a5cb334..36540cf 100644
--- a/QcomModulePkg/Library/FastbootLib/FastbootCmds.c
+++ b/QcomModulePkg/Library/FastbootLib/FastbootCmds.c
@@ -186,6 +186,13 @@
   VOID *Data;
 } CmdInfo;
 
+STATIC BOOLEAN UsbTimerStarted;
+
+BOOLEAN IsUsbTimerStarted (VOID)
+{
+  return UsbTimerStarted;
+}
+
 #ifdef DISABLE_PARALLEL_DOWNLOAD_FLASH
 BOOLEAN IsDisableParallelDownloadFlash (VOID)
 {
@@ -1352,6 +1359,7 @@
     gBS->CloseEvent (UsbTimerEvent);
     UsbTimerEvent = NULL;
   }
+  UsbTimerStarted = FALSE;
 }
 #else
 STATIC VOID StopUsbTimer (VOID)
@@ -1666,6 +1674,7 @@
         IsFlashComplete = TRUE;
         StopUsbTimer ();
       } else {
+        UsbTimerStarted = TRUE;
         FastbootOkay ("");
       }
     }
@@ -2396,7 +2405,6 @@
          return;
     }
     FastbootOkay ("");
-    RebootDevice (RECOVERY_MODE);
   }
 }
 #endif
diff --git a/QcomModulePkg/Library/FastbootLib/FastbootCmds.h b/QcomModulePkg/Library/FastbootLib/FastbootCmds.h
index 7ec19fe..0b2c2a1 100644
--- a/QcomModulePkg/Library/FastbootLib/FastbootCmds.h
+++ b/QcomModulePkg/Library/FastbootLib/FastbootCmds.h
@@ -50,6 +50,7 @@
 #define MAX_DOWNLOAD_SIZE_STR "536870912"
 #endif
 
+#define MAX_WRITE_SIZE (1024 * 1024)
 #define MAX_BUFFER_SIZE MAX_DOWNLOAD_SIZE
 #define MAX_RSP_SIZE 64
 #define ERASE_BUFF_SIZE 256 * 1024
@@ -174,4 +175,5 @@
 UpdateDevInfo (CHAR16 *Pname, CHAR8 *ImgVersion);
 VOID
 GetDevInfo (DeviceInfo **DevinfoPtr);
+BOOLEAN IsUsbTimerStarted (VOID);
 #endif
diff --git a/QcomModulePkg/Library/FastbootLib/FastbootLib.inf b/QcomModulePkg/Library/FastbootLib/FastbootLib.inf
index b4968b4..75d88bb 100644
--- a/QcomModulePkg/Library/FastbootLib/FastbootLib.inf
+++ b/QcomModulePkg/Library/FastbootLib/FastbootLib.inf
@@ -52,6 +52,7 @@
 [BuildOptions]
   GCC:*_*_*_CC_FLAGS = $(UBSAN_UEFI_GCC_FLAG_UNDEFINED)
   GCC:*_*_*_CC_FLAGS = $(UBSAN_UEFI_GCC_FLAG_ALIGNMENT)
+  GCC:*_*_*_CC_FLAGS = $(LLVM_ENABLE_SAFESTACK) $(LLVM_SAFESTACK_USE_PTR) $(LLVM_SAFESTACK_COLORING)
 
 [Packages]
   EmbeddedPkg/EmbeddedPkg.dec
diff --git a/QcomModulePkg/Library/StackCanary/StackCanary.inf b/QcomModulePkg/Library/StackCanary/StackCanary.inf
index 6de63d2..fd7ff19 100644
--- a/QcomModulePkg/Library/StackCanary/StackCanary.inf
+++ b/QcomModulePkg/Library/StackCanary/StackCanary.inf
@@ -1,5 +1,5 @@
 #/*
-# * Copyright (c) 2017, The Linux Foundation. All rights reserved.
+# * Copyright (c) 2017-2018, The Linux Foundation. All rights reserved.
 # *
 # * Redistribution and use in source and binary forms, with or without
 # * modification, are permitted provided that the following conditions are
@@ -39,9 +39,13 @@
 [Sources]
 	StackCanary.c
 
+[BuildOptions]
+	GCC:*_*_*_CC_FLAGS = $(LLVM_ENABLE_SAFESTACK) $(LLVM_SAFESTACK_USE_PTR) $(LLVM_SAFESTACK_COLORING)
+
 [BuildOptions.AARCH64]
 	GCC:*_*_*_CC_FLAGS = -fno-stack-protector
 	GCC:*_*_*_CC_FLAGS = $(SDLLVM_COMPILE_ANALYZE) $(SDLLVM_ANALYZE_REPORT)
+
 [BuildOptions.ARM]
 	GCC:*_*_*_CC_FLAGS = -fno-stack-protector
 
diff --git a/QcomModulePkg/Library/UbsanLib/UbsanLib.inf b/QcomModulePkg/Library/UbsanLib/UbsanLib.inf
index b6edfe2..c4660b9 100644
--- a/QcomModulePkg/Library/UbsanLib/UbsanLib.inf
+++ b/QcomModulePkg/Library/UbsanLib/UbsanLib.inf
@@ -1,5 +1,5 @@
 #/* @file
-#  Copyright (c) 2017, The Linux Foundation. All rights reserved.
+#  Copyright (c) 2017-2018, The Linux Foundation. All rights reserved.
 #  Portions Copyright (c) 2011-2013, ARM Limited. All rights reserved.
 #
 #  This program and the accompanying materials
@@ -23,6 +23,9 @@
 [Sources.common]
   UbsanLib.c
 
+[BuildOptions]
+  GCC:*_*_*_CC_FLAGS = $(LLVM_ENABLE_SAFESTACK) $(LLVM_SAFESTACK_USE_PTR) $(LLVM_SAFESTACK_COLORING)
+
 [Packages]
   QcomModulePkg/QcomModulePkg.dec
   ArmPkg/ArmPkg.dec
diff --git a/QcomModulePkg/Library/avb/AvbLib.inf b/QcomModulePkg/Library/avb/AvbLib.inf
index bfb3a75..be4d6ee 100644
--- a/QcomModulePkg/Library/avb/AvbLib.inf
+++ b/QcomModulePkg/Library/avb/AvbLib.inf
@@ -38,6 +38,7 @@
 
 [BuildOptions]
 	GCC:*_*_*_CC_FLAGS = -DAVB_COMPILATION -DAVB_ENABLE_DEBUG
+	GCC:*_*_*_CC_FLAGS = $(LLVM_ENABLE_SAFESTACK) $(LLVM_SAFESTACK_USE_PTR) $(LLVM_SAFESTACK_COLORING)
 
 [BuildOptions.AARCH64]
 	GCC:*_*_*_CC_FLAGS = $(SDLLVM_COMPILE_ANALYZE) $(SDLLVM_ANALYZE_REPORT)
diff --git a/QcomModulePkg/Library/zlib/zlib.inf b/QcomModulePkg/Library/zlib/zlib.inf
index b25af58..5bbb559 100644
--- a/QcomModulePkg/Library/zlib/zlib.inf
+++ b/QcomModulePkg/Library/zlib/zlib.inf
@@ -6,6 +6,8 @@
   VERSION_STRING                 = 1.0
   LIBRARY_CLASS                  = Zlib
 
+[BuildOptions]
+  GCC:*_*_*_CC_FLAGS = $(LLVM_ENABLE_SAFESTACK) $(LLVM_SAFESTACK_USE_PTR) $(LLVM_SAFESTACK_COLORING)
 
 [BuildOptions.AARCH64]
   GCC:*_*_*_CC_FLAGS = -O2 -DZ_SOLO
diff --git a/makefile b/makefile
index beae335..8751fba 100644
--- a/makefile
+++ b/makefile
@@ -33,6 +33,27 @@
 ABL_FV_ELF := $(BOOTLOADER_OUT)/../../abl.elf
 SHELL:=/bin/bash
 
+# This function is to check version compatibility, used to control features based on the compiler version. \
+Arguments should be return value, current version and supported version in order. \
+It sets return value to true if the current version is equal or greater than the supported version.
+define check_version_compatibility
+	$(eval CURR_VERSION := $(shell $(2)/clang --version |& grep -i "clang version" |& sed 's/[^0-9.]//g'))
+	$(eval CURR_VERSION_MAJOR := $(shell echo $(CURR_VERSION) |& cut -d. -f1))
+	$(eval CURR_VERSION_MINOR := $(shell echo $(CURR_VERSION) |& cut -d. -f2))
+	$(eval SUPPORTED_VERSION := $(3))
+	$(eval SUPPORTED_VERSION_MAJOR := $(shell echo $(SUPPORTED_VERSION) |& cut -d. -f1))
+	$(eval SUPPORTED_VERSION_MINOR := $(shell echo $(SUPPORTED_VERSION) |& cut -d. -f2))
+
+	ifeq ($(shell expr $(CURR_VERSION_MAJOR) \> $(SUPPORTED_VERSION_MAJOR)), 1)
+		$(1) := true
+	endif
+	ifeq ($(shell expr $(CURR_VERSION_MAJOR) \= $(SUPPORTED_VERSION_MAJOR)), 1)
+		ifeq ($(shell expr $(CURR_VERSION_MINOR) \>= $(SUPPORTED_VERSION_MINOR)), 1)
+			$(1) := true
+		endif
+	endif
+endef
+
 # UEFI UBSAN Configuration
 # ENABLE_UEFI_UBSAN := true
 
@@ -73,6 +94,24 @@
 export SDLLVM_COMPILE_ANALYZE := $(SDLLVM_COMPILE_ANALYZE)
 export SDLLVM_ANALYZE_REPORT := $(SDLLVM_ANALYZE_REPORT)
 
+CLANG_SUPPORTS_SAFESTACK := false
+$(eval $(call check_version_compatibility, CLANG_SUPPORTS_SAFESTACK, $(CLANG_BIN), $(SAFESTACK_SUPPORTED_CLANG_VERSION)))
+
+ifeq "$(ABL_SAFESTACK)" "true"
+	ifeq "$(CLANG_SUPPORTS_SAFESTACK)" "true"
+		LLVM_ENABLE_SAFESTACK := -fsanitize=safe-stack
+		LLVM_SAFESTACK_USE_PTR := -mllvm -safestack-use-pointer-address
+		LLVM_SAFESTACK_COLORING := -mllvm -safe-stack-coloring=true
+	endif
+else
+	LLVM_ENABLE_SAFESTACK :=
+	LLVM_SAFESTACK_USE_PTR :=
+	LLVM_SAFESTACK_COLORING :=
+endif
+export LLVM_ENABLE_SAFESTACK := $(LLVM_ENABLE_SAFESTACK)
+export LLVM_SAFESTACK_USE_PTR := $(LLVM_SAFESTACK_USE_PTR)
+export LLVM_SAFESTACK_COLORING := $(LLVM_SAFESTACK_COLORING)
+
 .PHONY: all cleanall
 
 all: ABL_FV_ELF