1 files changed, 110 insertions, 25 deletions
diff --git a/DcsBoot/DcsBoot.c b/DcsBoot/DcsBoot.c
index 7eaf41b..7029a10 100644
--- a/DcsBoot/DcsBoot.c
+++ b/DcsBoot/DcsBoot.c
@@ -25,9 +25,11 @@ https://opensource.org/licenses/LGPL-3.0
 
 EFI_GUID          ImagePartGuid;
 EFI_GUID          *gEfiExecPartGuid = &ImagePartGuid;
-CHAR16            *gEfiExecCmdDefault = L"\\EFI\\Microsoft\\Boot\\Bootmgfw.efi";
+CHAR16            *gEfiExecCmdDefault = L"\\EFI\\Microsoft\\Boot\\Bootmgfw_ms.vc";
+CHAR16            *gEfiExecCmdMS = L"\\EFI\\Microsoft\\Boot\\Bootmgfw.efi";
 CHAR16            *gEfiExecCmd = NULL;
 CHAR8             gDoExecCmdMsg[256];
+CONST CHAR8* 	  g_szMsBootString = "bootmgfw.pdb";
 
 EFI_STATUS
 DoExecCmd() 
@@ -44,7 +46,10 @@ DoExecCmd()
             lockFlags = ConfigReadInt("DcsBmlLockFlags", BML_LOCK_SETVARIABLE | BML_SET_BOOTNEXT | BML_UPDATE_BOOTORDER);
             BmlLock(lockFlags);
 			res = EfiExec(NULL, gEfiExecCmd);
-			AsciiSPrint(gDoExecCmdMsg, sizeof(gDoExecCmdMsg), "\nCan't exec %s start partition %g\n", gEfiExecCmd, gEfiExecPartGuid);
+			if (EFI_ERROR(res))
+				AsciiSPrint(gDoExecCmdMsg, sizeof(gDoExecCmdMsg), "\nCan't exec %s start partition %g\n", gEfiExecCmd, gEfiExecPartGuid);
+			else
+				AsciiSPrint(gDoExecCmdMsg, sizeof(gDoExecCmdMsg), "\nDone exec %s start partition %g\n", gEfiExecCmd, gEfiExecPartGuid);
 		}	else {
 			AsciiSPrint(gDoExecCmdMsg, sizeof(gDoExecCmdMsg), "\nCan't open start partition %g\n", gEfiExecPartGuid);
 		}
@@ -54,6 +59,39 @@ DoExecCmd()
 	return res;
 }
 
+EFI_STATUS
+ExecMSWindowsLoader() {
+	
+	if (!EFI_ERROR(FileExist(NULL, gEfiExecCmdDefault)))
+		return EfiExec(NULL, gEfiExecCmdDefault);
+	else
+	{
+		if (!EFI_ERROR(FileExist(NULL, gEfiExecCmdMS)))
+		{
+			/* check if it is Microsoft one */
+			UINT8*      fileData = NULL;
+			UINTN       fileSize = 0;
+			BOOLEAN		bFound = FALSE;
+			if (!EFI_ERROR(FileLoad(NULL, gEfiExecCmdMS, &fileData, &fileSize)))
+			{
+				if ((fileSize > 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szMsBootString, AsciiStrLen(g_szMsBootString))))
+				{
+					bFound = TRUE;
+				}
+			}
+			
+			MEM_FREE(fileData);
+			
+			if (bFound)
+				return EfiExec(NULL, gEfiExecCmdMS);
+		}
+
+		ERR_PRINT(L"Could not find the original Windows loader\r\n");
+		
+		return EFI_NOT_READY;
+	}
+}
+
 //////////////////////////////////////////////////////////////////////////
 // BML
 //////////////////////////////////////////////////////////////////////////
@@ -121,6 +159,8 @@ DcsBootMain(
 	UINTN               len;
 	UINT32              attr;
 	BOOLEAN             searchOnESP = FALSE;
+	BOOLEAN             searchMsOnESP = FALSE;
+	EFI_GUID          *pEfiExecPartBackup = NULL;
 //	EFI_INPUT_KEY       key;
 
 	InitBio();
@@ -159,9 +199,33 @@ DcsBootMain(
 	EfiSetVar(L"DcsExecPartGuid", NULL, &ImagePartGuid, sizeof(EFI_GUID), EFI_VARIABLE_BOOTSERVICE_ACCESS);
 	EfiSetVar(L"DcsExecCmd", NULL, gEfiExecCmdDefault, (StrLen(gEfiExecCmdDefault) + 1) * 2, EFI_VARIABLE_BOOTSERVICE_ACCESS);
 	// 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)
+	  {
+		res = EFI_SUCCESS;
+		gST->RuntimeServices->ResetSystem(EfiResetShutdown, EFI_SUCCESS, 0, NULL);
+	  }
+	  else if (res == EFI_DCS_REBOOT_REQUESTED)
+	  {
+		res = EFI_SUCCESS;
+		gST->RuntimeServices->ResetSystem(EfiResetCold, EFI_SUCCESS, 0, NULL);
+	  }
+	  else if (res == EFI_DCS_HALT_REQUESTED)
+	  {
+		  EfiCpuHalt();
+	  }
+	  else if (res == EFI_DCS_USER_CANCELED)
+	  {
+		  /* If user cancels password prompt, call original Windows loader */
+		  res = ExecMSWindowsLoader ();
+	  }
       return res;
    }
 
@@ -169,6 +233,8 @@ DcsBootMain(
 	if (EFI_ERROR(res)) {
 		gEfiExecPartGuid = &ImagePartGuid;
 	}
+	
+	pEfiExecPartBackup = gEfiExecPartGuid;
 
 	res = EfiGetVar(L"DcsExecCmd", NULL, &gEfiExecCmd, &len, &attr);
 	if (EFI_ERROR(res)) {
@@ -177,6 +243,9 @@ DcsBootMain(
 
 	searchOnESP = CompareGuid(gEfiExecPartGuid, &ImagePartGuid) &&
 		EFI_ERROR(FileExist(NULL, gEfiExecCmd));
+		
+	searchMsOnESP = CompareGuid(gEfiExecPartGuid, &ImagePartGuid) &&
+		EFI_ERROR(FileExist(NULL, gEfiExecCmdMS));
 
     // Clear DcsExecPartGuid before execute OS to avoid problem in VirtualBox with reboot.
     EfiSetVar(L"DcsExecPartGuid", NULL, NULL, 0, EFI_VARIABLE_BOOTSERVICE_ACCESS);
@@ -187,34 +256,50 @@ DcsBootMain(
 	InitBio();
 	res = InitFS();
 
-	// Default load of bootmgfw?
-	if (searchOnESP) {
-		// gEfiExecCmd is not found on start partition. Try from ESP
-		EFI_BLOCK_IO_PROTOCOL *bio = NULL;
-		EFI_PARTITION_TABLE_HEADER *gptHdr = NULL;
-		EFI_PARTITION_ENTRY        *gptEntry = NULL;
-		HARDDRIVE_DEVICE_PATH hdp;
-		EFI_HANDLE disk;
-		if (!EFI_ERROR(res = EfiGetPartDetails(gFileRootHandle, &hdp, &disk))) {
-			if ((bio = EfiGetBlockIO(disk)) != NULL) {
-				if (!EFI_ERROR(res = GptReadHeader(bio, 1, &gptHdr)) &&
-					!EFI_ERROR(res = GptReadEntryArray(bio, gptHdr, &gptEntry))) {
-					UINT32 i;
-					for (i = 0; i < gptHdr->NumberOfPartitionEntries; ++i) {
-						if (CompareGuid(&gptEntry[i].PartitionTypeGUID, &gEfiPartTypeSystemPartGuid)) {
-							// select ESP GUID
-							CopyGuid(gEfiExecPartGuid, &gptEntry[i].UniquePartitionGUID);
-							res = DoExecCmd();
-							if(EFI_ERROR(res)) continue;
+	while (1)
+	{
+		// Default load of bootmgfw?
+		if (searchOnESP) {
+			// gEfiExecCmd is not found on start partition. Try from ESP
+			EFI_BLOCK_IO_PROTOCOL *bio = NULL;
+			EFI_PARTITION_TABLE_HEADER *gptHdr = NULL;
+			EFI_PARTITION_ENTRY        *gptEntry = NULL;
+			HARDDRIVE_DEVICE_PATH hdp;
+			EFI_HANDLE disk;
+			if (!EFI_ERROR(res = EfiGetPartDetails(gFileRootHandle, &hdp, &disk))) {
+				if ((bio = EfiGetBlockIO(disk)) != NULL) {
+					if (!EFI_ERROR(res = GptReadHeader(bio, 1, &gptHdr)) &&
+						!EFI_ERROR(res = GptReadEntryArray(bio, gptHdr, &gptEntry))) {
+						UINT32 i;
+						for (i = 0; i < gptHdr->NumberOfPartitionEntries; ++i) {
+							if (CompareGuid(&gptEntry[i].PartitionTypeGUID, &gEfiPartTypeSystemPartGuid)) {
+								// select ESP GUID
+								CopyGuid(gEfiExecPartGuid, &gptEntry[i].UniquePartitionGUID);
+								res = DoExecCmd();
+								if(EFI_ERROR(res)) continue;
+							}
 						}
 					}
 				}
 			}
+		}	else {
+			res = DoExecCmd();
 		}
-	}	else {
-		res = DoExecCmd();
+		
+		if(EFI_ERROR(res))
+		{
+			if (0 == StrCmp (gEfiExecCmd, gEfiExecCmdDefault))
+			{
+				gEfiExecCmd = gEfiExecCmdMS;
+				searchOnESP = searchMsOnESP;
+				gEfiExecPartGuid = pEfiExecPartBackup;
+			}
+			else
+				break;
+		}
+		else
+			break;
 	}
-
 	ERR_PRINT(L"%a\nStatus -  %r", gDoExecCmdMsg, res);
 	EfiCpuHalt();
    return EFI_INVALID_PARAMETER;