From 2d4d3ab4ed0ee86d46fe2b44ca1e56a3d2916f68 Mon Sep 17 00:00:00 2001
From: Mounir IDRASSI <mounir.idrassi@idrix.fr>
Date: Fri, 8 Nov 2019 23:06:25 +0100
Subject: Fix "ActionFailed" not working and add "ActionCancelled" to customize
 handling of user hitting ESC on password prompt

---
 DcsBoot/DcsBoot.c                   | 6 +++++-
 DcsInt/DcsInt.c                     | 5 +++++
 Include/Library/CommonLib.h         | 1 +
 Library/VeraCryptLib/DcsVeraCrypt.c | 3 +++
 Library/VeraCryptLib/DcsVeraCrypt.h | 1 +
 5 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/DcsBoot/DcsBoot.c b/DcsBoot/DcsBoot.c
index 080b052..7029a10 100644
--- a/DcsBoot/DcsBoot.c
+++ b/DcsBoot/DcsBoot.c
@@ -201,7 +201,11 @@ DcsBootMain(
 	// Authorize
 	gBS->SetWatchdogTimer(0, 0, 0, NULL);
 	res = EfiExec(NULL, L"\\EFI\\VeraCrypt\\DcsInt.dcs");
-   if (EFI_ERROR(res)) {
+   if (EFI_ERROR(res) && (res != EFI_DCS_POSTEXEC_REQUESTED)) {
+
+      // Clear DcsExecPartGuid before execute OS to avoid problem in VirtualBox with reboot.
+      EfiSetVar(L"DcsExecPartGuid", NULL, NULL, 0, EFI_VARIABLE_BOOTSERVICE_ACCESS);
+      EfiSetVar(L"DcsExecCmd", NULL, NULL, 0, EFI_VARIABLE_BOOTSERVICE_ACCESS);
       // ERR_PRINT(L"\nDcsInt.efi %r\n",res);
 	  if (res == EFI_DCS_SHUTDOWN_REQUESTED)
 	  {
diff --git a/DcsInt/DcsInt.c b/DcsInt/DcsInt.c
index 8133726..7dd98ec 100644
--- a/DcsInt/DcsInt.c
+++ b/DcsInt/DcsInt.c
@@ -804,6 +804,7 @@ enum OnExitTypes{
 	OnExitAuthFaild = 1,
 	OnExitAuthNotFound,
 	OnExitAuthTimeout,
+	OnExitAuthCancelled,
 	OnExitSuccess
 };
 
@@ -976,6 +977,8 @@ OnExit(
 		if (fileStr != NULL) {
 			EfiSetVar(L"DcsExecCmd", NULL, fileStr, (StrLen(fileStr) + 1) * 2, EFI_VARIABLE_BOOTSERVICE_ACCESS);
 		}
+
+		retValue = EFI_DCS_POSTEXEC_REQUESTED;
 		goto exit;
 	}
 
@@ -1185,6 +1188,8 @@ UefiMain(
 	if (EFI_ERROR(res)) {
 		if (res == EFI_TIMEOUT)
 			return OnExit(gOnExitTimeout, OnExitAuthTimeout, res);
+		else if (res == EFI_DCS_USER_CANCELED)
+			return OnExit(gOnExitCancelled, OnExitAuthCancelled, res);
 		else
 			return OnExit(gOnExitFailed, OnExitAuthFaild, res);
 	}
diff --git a/Include/Library/CommonLib.h b/Include/Library/CommonLib.h
index 0bfb631..2cf10d3 100644
--- a/Include/Library/CommonLib.h
+++ b/Include/Library/CommonLib.h
@@ -32,6 +32,7 @@ https://opensource.org/licenses/LGPL-3.0
 #define EFI_DCS_REBOOT_REQUESTED	ENCODE_ERROR(0xDC50002)
 #define EFI_DCS_HALT_REQUESTED		ENCODE_ERROR(0xDC50003)
 #define EFI_DCS_USER_CANCELED		ENCODE_ERROR(0xDC50004)
+#define EFI_DCS_POSTEXEC_REQUESTED	ENCODE_ERROR(0xDC50005)
 
 //////////////////////////////////////////////////////////////////////////
 // Check error 
diff --git a/Library/VeraCryptLib/DcsVeraCrypt.c b/Library/VeraCryptLib/DcsVeraCrypt.c
index 4cf0e05..c3e8a39 100644
--- a/Library/VeraCryptLib/DcsVeraCrypt.c
+++ b/Library/VeraCryptLib/DcsVeraCrypt.c
@@ -82,6 +82,7 @@ CHAR8* gOnExitFailed = NULL;
 CHAR8* gOnExitSuccess = NULL;
 CHAR8* gOnExitNotFound = NULL;
 CHAR8* gOnExitTimeout = NULL;
+CHAR8* gOnExitCancelled = NULL;
 
 //////////////////////////////////////////////////////////////////////////
 // Authorize
@@ -184,6 +185,8 @@ VCAuthLoadConfig()
 	ConfigReadString("ActionFailed", "Exit", gOnExitFailed, MAX_MSG);
 	VCCONFIG_ALLOC(gOnExitTimeout, MAX_MSG);
 	ConfigReadString("ActionTimeout", "Shutdown", gOnExitTimeout, MAX_MSG);
+	VCCONFIG_ALLOC(gOnExitCancelled, MAX_MSG);
+	ConfigReadString("ActionCancelled", "Exit", gOnExitCancelled, MAX_MSG);
 
 	strTemp = MEM_ALLOC(MAX_MSG);
 	ConfigReadString("PartitionGuidOS", "", strTemp, MAX_MSG);
diff --git a/Library/VeraCryptLib/DcsVeraCrypt.h b/Library/VeraCryptLib/DcsVeraCrypt.h
index 152a335..1f25ae9 100644
--- a/Library/VeraCryptLib/DcsVeraCrypt.h
+++ b/Library/VeraCryptLib/DcsVeraCrypt.h
@@ -75,6 +75,7 @@ extern CHAR8* gOnExitFailed;
 extern CHAR8* gOnExitSuccess;
 extern CHAR8* gOnExitNotFound;
 extern CHAR8* gOnExitTimeout;
+extern CHAR8* gOnExitCancelled;
 
 void
 VCAuthAsk();
-- 
cgit v1.2.3