From ecb29e3b3f1a67800922a79c009e80a0eb45ef91 Mon Sep 17 00:00:00 2001
From: Mounir IDRASSI <mounir.idrassi@idrix.fr>
Date: Thu, 21 Mar 2019 20:58:50 +0100
Subject: Enhance Rescue Disk implementation of restoring VeraCrypt loader.

---
 Library/CommonLib/EfiFile.c | 42 ++++++++++++++++++++++++++++++++++++------
 Library/CommonLib/EfiMem.c  | 27 +++++++++++++++++++++++++++
 Library/DcsCfgLib/GptEdit.c |  4 ++--
 3 files changed, 65 insertions(+), 8 deletions(-)

(limited to 'Library')

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
@@ -41,6 +41,22 @@ InitFS() {
    return res;
 }
 
+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,
@@ -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;
-- 
cgit v1.2.3