5 files changed, 183 insertions, 41 deletions
diff --git a/DcsRe/DcsRe.c b/DcsRe/DcsRe.c
index 7f4096f..769cb88 100644
--- a/DcsRe/DcsRe.c
+++ b/DcsRe/DcsRe.c
@@ -21,8 +21,10 @@ https://opensource.org/licenses/LGPL-3.0
 #include "common/Tcdefs.h"
 
 #ifdef _M_X64
+#define ARCHdot L"x64."
 #define ARCHdotEFI L"x64.efi"
 #else
+#define ARCHdot L"IA32."
 #define ARCHdotEFI L"IA32.efi"
 #endif
 
@@ -48,48 +50,56 @@ SelectEfiVolume()
 	EFI_FILE     *file;
 	EFI_FILE     **efiVolumes;
 	UINTN        efiVolumesCount = 0;
-	EFI_HANDLE   startHandle;
 	if (EfiBootVolume != NULL) return;
-	res = EfiGetStartDevice(&startHandle);
-	if (EFI_ERROR(res)) {
-		ERR_PRINT(L"GetStartDevice %r", res);
-		return;
-	}
+
 	efiVolumes = MEM_ALLOC(sizeof(EFI_FILE*) * gFSCount);
 	for (i = 0; i < gFSCount; ++i) {
+		if (gFSHandles[i] == gFileRootHandle)
+			continue;
 		res = FileOpenRoot(gFSHandles[i], &file);
-		if(EFI_ERROR(res)) continue;
-		if (!EFI_ERROR(FileExist(file, L"EFI\\Boot\\boot" ARCHdotEFI))) {
+		if(EFI_ERROR(res)) { ERR_PRINT(L"FileOpenRoot %r\n", res); continue;}
+		if (	!EFI_ERROR(FileExist(file, L"EFI\\Boot\\boot" ARCHdotEFI))
+			||	!EFI_ERROR(FileExist(file, L"EFI\\Microsoft\\Boot\\bootmgfw.efi"))
+			||	!EFI_ERROR(FileExist(file, L"EFI\\Microsoft\\Boot\\bootmgfw_ms.vc"))
+			) 
+		{
 			efiVolumesCount++;
 			efiVolumes[i] = file;
-			if (gFSHandles[i] != startHandle) {
-				EfiBootVolumeIndex = i;
-				EfiBootVolume = file;
-			}
+			EfiBootVolumeIndex = i;
+			EfiBootVolume = file;
 		}	else {
 			FileClose(file);
 		}
 	}
-
-	for (i = 0; i < gFSCount; ++i) {
-		OUT_PRINT(L"%H%d)%N ", i);
-		if (efiVolumes[i] != NULL) {
-			if (gFSHandles[i] == startHandle) {
-				OUT_PRINT(L"%V [Boot Rescue] %N");
-			}
-			else {
+	
+	if (efiVolumesCount > 1)
+	{
+		for (i = 0; i < gFSCount; ++i) {
+			OUT_PRINT(L"%H%d)%N ", i);
+			if (efiVolumes[i] != NULL) {
 				OUT_PRINT(L"%V [Boot] %N");
 			}
+			EfiPrintDevicePath(gFSHandles[i]);
+			OUT_PRINT(L"\n");
 		}
-		EfiPrintDevicePath(gFSHandles[i]);
-		OUT_PRINT(L"\n");
-	}
 
-	do {
-		EfiBootVolumeIndex = AskUINTN("Select EFI boot volume:", EfiBootVolumeIndex);
-		if (EfiBootVolumeIndex >= gFSCount) continue;
-		EfiBootVolume = efiVolumes[EfiBootVolumeIndex];
-	} while (EfiBootVolume == NULL);
+		do {
+			EfiBootVolumeIndex = AskUINTN("Select EFI boot volume:", EfiBootVolumeIndex);
+			if (EfiBootVolumeIndex >= gFSCount) continue;
+			EfiBootVolume = efiVolumes[EfiBootVolumeIndex];
+		} while (EfiBootVolume == NULL);
+		
+		/* free unused descriptors */
+		for (i = 0; i < gFSCount; ++i) {
+			if (efiVolumes[i] != NULL && efiVolumes[i] != EfiBootVolume) {
+				FileClose(efiVolumes[i]);
+			}
+		}
+
+		OUT_PRINT (L"\n");
+	}
+	
+	
 	MEM_FREE(efiVolumes);
 }
 
@@ -108,9 +118,7 @@ ActionShell(IN VOID* ctx) {
 
 EFI_STATUS
 ActionDcsBoot(IN VOID* ctx) {
-	SelectEfiVolume();
-	if (EfiBootVolume == NULL) return EFI_NOT_READY;
-	return EfiExec(gFSHandles[EfiBootVolumeIndex], L"EFI\\VeraCrypt\\DcsBoot.efi");
+	return EfiExec(gFileRootHandle, L"EFI\\VeraCrypt\\DcsBoot.efi");
 }
 
 CHAR16* DcsBootBins[] = {
@@ -128,13 +136,77 @@ EFI_STATUS
 ActionRestoreDcsLoader(IN VOID* ctx) {
 	EFI_STATUS res = EFI_NOT_READY;
 	UINTN i;
+	CONST CHAR8* g_szMsBootString = "bootmgfw.pdb";
+	CONST CHAR16* g_szVcBootString = L"VeraCrypt";
 	SelectEfiVolume();
 	if (EfiBootVolume == NULL) return EFI_NOT_READY;
+	
+	DirectoryCreate (EfiBootVolume, L"EFI\\VeraCrypt");
+	
 	for (i = 0; i < sizeof(DcsBootBins) / sizeof(CHAR16*); ++i) {
 		res = FileCopy(NULL, DcsBootBins[i], EfiBootVolume, DcsBootBins[i], 1024 * 1024);
 		if (EFI_ERROR(res)) return res;
 	}
-	return res;
+	/* restore standard boot file */
+	if (!EFI_ERROR(FileExist(EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI)))
+	{
+		/* check if it is Microsoft one or ours */
+		UINT8*      fileData = NULL;
+		UINTN       fileSize = 0;
+		res = EFI_SUCCESS;
+		if (!EFI_ERROR(FileLoad(EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, &fileData, &fileSize)))
+		{
+			if ((fileSize > 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szMsBootString, AsciiStrLen(g_szMsBootString))))
+			{
+				res = FileCopy(EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, EfiBootVolume, L"\\EFI\\Boot\\original_boot" ARCHdot L"vc_backup", 1024 * 1024);
+				if (!EFI_ERROR(res))
+					res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, 1024 * 1024);				
+			}
+			else if ((fileSize <= 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szVcBootString, StrLen (g_szVcBootString) * 2)))
+			{
+				res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, 1024 * 1024);
+			}
+			MEM_FREE(fileData);
+			
+			if (EFI_ERROR(res)) return res;
+		}		
+	}
+	else if (!EFI_ERROR(FileExist(EfiBootVolume, L"\\EFI\\Boot\\original_boot" ARCHdot L"vc_backup")))
+	{
+		res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, 1024 * 1024);
+		if (EFI_ERROR(res)) return res;
+	}
+	
+	if (!EFI_ERROR(FileExist(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi")))
+	{
+		/* check if it is Microsoft one */
+		UINT8*      fileData = NULL;
+		UINTN       fileSize = 0;
+		res = EFI_SUCCESS;
+		if (!EFI_ERROR(FileLoad(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi", &fileData, &fileSize)))
+		{
+			if ((fileSize > 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szMsBootString, AsciiStrLen(g_szMsBootString))))
+			{
+				res = FileCopy(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi", EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw_ms.vc", 1024 * 1024);				
+			}
+
+			MEM_FREE(fileData);
+			
+			if (EFI_ERROR(res)) return res;
+		}
+
+		res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw.efi", 1024 * 1024);
+		if (EFI_ERROR(res)) return res;
+	}
+	else if (!EFI_ERROR(FileExist(EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw_ms.vc")))		
+	{
+		res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw.efi", 1024 * 1024);
+		if (EFI_ERROR(res)) return res;
+	}
+	
+	OUT_PRINT (L"\nVeraCrypt Loader restored to disk successfully\n\n");
+	
+	return EFI_SUCCESS;
 }
 
 CHAR16* sDcsBootEfi = L"EFI\\VeraCrypt\\DcsBoot.efi";
diff --git a/Include/Library/CommonLib.h b/Include/Library/CommonLib.h
index 479c5c1..d0db41f 100644
--- a/Include/Library/CommonLib.h
+++ b/Include/Library/CommonLib.h
@@ -80,6 +80,13 @@ PrepareMemory(
    IN UINTN    len,
    OUT VOID**  mem
    );
+   
+EFI_STATUS
+MemoryHasPattern (
+	CONST VOID* buffer,
+	UINTN bufferLen,
+	CONST VOID* pattern,
+	UINTN patternLen);
 
 //////////////////////////////////////////////////////////////////////////
 // handles
@@ -651,6 +658,12 @@ EFI_STATUS
 InitFS();
 
 EFI_STATUS
+DirectoryCreate(
+   IN    EFI_FILE*   root,
+   IN    CHAR16*     name
+   );
+
+EFI_STATUS
 FileOpenRoot(
    IN    EFI_HANDLE rootHandle,
    OUT   EFI_FILE** rootFile);
@@ -685,7 +698,7 @@ EFI_STATUS
 FileWrite(
    IN       EFI_FILE*   f,
    IN       VOID*       data,
-   IN OUT   UINTN*      bytes,
+   IN OUT   UINTN       bytes,
    IN OUT   UINT64*     position);
 
 UINTN
diff --git a/Library/CommonLib/EfiFile.c b/Library/CommonLib/EfiFile.c
index fdc999c..4839eae 100644
--- a/Library/CommonLib/EfiFile.c
+++ b/Library/CommonLib/EfiFile.c
@@ -42,6 +42,22 @@ InitFS() {
 }
 
 EFI_STATUS
+DirectoryCreate(
+   IN    EFI_FILE*   root,
+   IN    CHAR16*     name
+   )
+{
+   EFI_FILE*      file;
+   EFI_STATUS     res;
+   if (!name) { return EFI_INVALID_PARAMETER; }
+
+   res = FileOpen(root, name, &file, EFI_FILE_MODE_READ | EFI_FILE_MODE_CREATE | EFI_FILE_MODE_WRITE, EFI_FILE_DIRECTORY);
+   if (EFI_ERROR(res)) return res;
+   FileClose(file);
+   return res;
+}
+
+EFI_STATUS
 FileOpenRoot(
    IN    EFI_HANDLE rootHandle,
    OUT   EFI_FILE** rootFile)
@@ -131,12 +147,14 @@ EFI_STATUS
 FileWrite(
    IN       EFI_FILE*   f, 
    IN       VOID*       data,
-   IN OUT   UINTN*      bytes, 
+   IN OUT   UINTN      bytes, 
    IN OUT   UINT64*     position) 
 {
    EFI_STATUS res;
+   UINTN remaining;
+   UINT8* pbData = (UINT8*) data;
 
-   if (!f || !data || !bytes) { 
+   if (!f || !data) { 
       return EFI_INVALID_PARAMETER; 
    }
    if (position != NULL) {
@@ -145,7 +163,20 @@ FileWrite(
          return res;
       }
    }
-   res = f->Write(f, bytes, data);
+   remaining = bytes;
+   res = f->Write(f, &bytes, pbData);
+   if (!EFI_ERROR(res)) {
+	   remaining -= bytes;
+	   pbData += bytes;
+	   bytes = remaining;
+	   while ((remaining > 0) && !EFI_ERROR(res))
+	   {
+		   res = f->Write(f, &bytes, pbData);
+		   remaining -= bytes;
+		   pbData += bytes;
+		   bytes = remaining;
+ 	   }
+   }
    if (position != NULL) {
       f->GetPosition(f, position);
    }
@@ -265,12 +296,11 @@ FileSave(
 {
    EFI_FILE*      file;
    EFI_STATUS     res;
-   UINTN          sz = size;
    if (!data || !name) { return EFI_INVALID_PARAMETER; }
    FileDelete(root, name);
    res = FileOpen(root, name, &file, EFI_FILE_MODE_READ | EFI_FILE_MODE_CREATE | EFI_FILE_MODE_WRITE, 0);
    if (EFI_ERROR(res)) return res;
-   res = FileWrite(file, data, &sz, NULL);
+   res = FileWrite(file, data, size, NULL);
    FileClose(file);
    return res;
 }
@@ -356,7 +386,7 @@ FileCopy(
 		datasz = remains > bufSz ? bufSz : remains;
 		res =FileRead(srcfile, data, &datasz, NULL);
 		if (EFI_ERROR(res)) goto copyerr;
-		res = FileWrite(dstfile, data, &datasz, NULL);
+		res = FileWrite(dstfile, data, datasz, NULL);
 		if (EFI_ERROR(res)) goto copyerr;
 		remains -= datasz;
 	} while (remains > 0);
diff --git a/Library/CommonLib/EfiMem.c b/Library/CommonLib/EfiMem.c
index d9386c0..872d3de 100644
--- a/Library/CommonLib/EfiMem.c
+++ b/Library/CommonLib/EfiMem.c
@@ -72,3 +72,30 @@ PrepareMemory(
    *mem = buf;
    return status;
 }
+
+//////////////////////////////////////////////////////////////////////////
+// Memory misc
+//////////////////////////////////////////////////////////////////////////
+EFI_STATUS MemoryHasPattern (
+	CONST VOID* buffer,
+	UINTN bufferLen,
+	CONST VOID* pattern,
+	UINTN patternLen)
+{
+	EFI_STATUS status = EFI_NOT_FOUND;
+	if (patternLen <= bufferLen)
+	{
+		UINTN i;
+		CONST UINT8* memPtr = (CONST UINT8*) buffer;
+		for (i = 0; i <= (bufferLen - patternLen); ++i)
+		{
+			if (CompareMem (&memPtr[i], pattern, patternLen) == 0)
+			{
+				status = EFI_SUCCESS;
+				break;
+			}
+		}
+	}
+
+	return status;
+}
diff --git a/Library/DcsCfgLib/GptEdit.c b/Library/DcsCfgLib/GptEdit.c
index 5f5d0b0..0b74bb3 100644
--- a/Library/DcsCfgLib/GptEdit.c
+++ b/Library/DcsCfgLib/GptEdit.c
@@ -289,13 +289,13 @@ DeListSaveToFile() {
 				UINTN pad;
 				len = (UINTN)DeList->DE[i].Length;
 				pad = (((len + 511) >> 9) << 9) - len;
-				res = FileWrite(file, DeData[i], &len, NULL);
+				res = FileWrite(file, DeData[i], len, NULL);
 				if (EFI_ERROR(res)) {
 					ERR_PRINT(L"Write: %r\n", res);
 					goto error;
 				}
 				if (pad > 0) {
-					res = FileWrite(file, pad512buf, &pad, NULL);
+					res = FileWrite(file, pad512buf, pad, NULL);
 					if (EFI_ERROR(res)) {
 						ERR_PRINT(L"Write: %r\n", res);
 						goto error;