1 files changed, 91 insertions, 35 deletions
diff --git a/DcsInt/DcsInt.c b/DcsInt/DcsInt.c
index f7c3965..a8fe493 100644
--- a/DcsInt/DcsInt.c
+++ b/DcsInt/DcsInt.c
@@ -84,7 +84,7 @@ UINTN                   SecRegionOffset = 0;
 PCRYPTO_INFO            SecRegionCryptInfo = NULL;
 
 VOID
-CleanSensitiveData()
+CleanSensitiveData(BOOLEAN bClearBootParams)
 {
 	if (SecRegionCryptInfo != NULL) {
 		MEM_BURN(SecRegionCryptInfo, sizeof(*SecRegionCryptInfo));
@@ -97,6 +97,10 @@ CleanSensitiveData()
 	if (SecRegionData != NULL) {
 		MEM_BURN(SecRegionData, SecRegionSize);
 	}
+	
+	if (bootParams != NULL && bClearBootParams) {
+		MEM_BURN(bootParams, sizeof(*bootParams));
+	}
 
 	if (gAutoPassword != NULL) {
 		MEM_BURN(gAutoPassword, MAX_PASSWORD);
@@ -105,7 +109,7 @@ CleanSensitiveData()
 
 void HaltPrint(const CHAR16* Msg)
 {
-	CleanSensitiveData();
+	CleanSensitiveData(TRUE);
 	Print(L"%s - system Halted\n", Msg);
 	EfiCpuHalt();
 }
@@ -160,29 +164,33 @@ PrepareBootParams(
 	IN PCRYPTO_INFO   cryptoInfo)
 {
 	BootArguments           *bootArgs;
-	if (bootParams == NULL) return EFI_UNSUPPORTED;
-	bootArgs = &bootParams->BootArgs;
-	TC_SET_BOOT_ARGUMENTS_SIGNATURE(bootArgs->Signature);
-	bootArgs->BootLoaderVersion = VERSION_NUM;
-	bootArgs->CryptoInfoOffset = (uint16)(FIELD_OFFSET(BOOT_PARAMS, BootCryptoInfo));
-	bootArgs->CryptoInfoLength = (uint16)(sizeof(BOOT_CRYPTO_HEADER) + 2 + sizeof(SECREGION_BOOT_PARAMS));
-	bootArgs->HeaderSaltCrc32 = gHeaderSaltCrc32;
-	CopyMem(&bootArgs->BootPassword, &gAuthPassword, sizeof(gAuthPassword));
-	bootArgs->HiddenSystemPartitionStart = 0;
-	bootArgs->DecoySystemPartitionStart = 0;
-	bootArgs->BootDriveSignature = bootDriveSignature;
-	bootArgs->Flags = (uint32)(gAuthPim << 16);
-	bootArgs->BootArgumentsCrc32 = GetCrc32((byte *)bootArgs, (int)((byte *)&bootArgs->BootArgumentsCrc32 - (byte *)bootArgs));
-	bootParams->BootCryptoInfo.ea = (uint16)cryptoInfo->ea;
-	bootParams->BootCryptoInfo.mode = (uint16)cryptoInfo->mode;
-	bootParams->BootCryptoInfo.pkcs5 = (uint16)cryptoInfo->pkcs5;
-	SetSecRegionParamsMemory();
+	EFI_STATUS              status;
+	if (bootParams == NULL) status = EFI_UNSUPPORTED;
+	else {
+		bootArgs = &bootParams->BootArgs;
+		TC_SET_BOOT_ARGUMENTS_SIGNATURE(bootArgs->Signature);
+		bootArgs->BootLoaderVersion = VERSION_NUM;
+		bootArgs->CryptoInfoOffset = (uint16)(FIELD_OFFSET(BOOT_PARAMS, BootCryptoInfo));
+		bootArgs->CryptoInfoLength = (uint16)(sizeof(BOOT_CRYPTO_HEADER) + 2 + sizeof(SECREGION_BOOT_PARAMS));
+		bootArgs->HeaderSaltCrc32 = gHeaderSaltCrc32;
+		CopyMem(&bootArgs->BootPassword, &gAuthPassword, sizeof(gAuthPassword));
+		bootArgs->HiddenSystemPartitionStart = 0;
+		bootArgs->DecoySystemPartitionStart = 0;
+		bootArgs->BootDriveSignature = bootDriveSignature;
+		bootArgs->Flags = (uint32)(gAuthPim << 16);
+		bootArgs->BootArgumentsCrc32 = GetCrc32((byte *)bootArgs, (int)((byte *)&bootArgs->BootArgumentsCrc32 - (byte *)bootArgs));
+		bootParams->BootCryptoInfo.ea = (uint16)cryptoInfo->ea;
+		bootParams->BootCryptoInfo.mode = (uint16)cryptoInfo->mode;
+		bootParams->BootCryptoInfo.pkcs5 = (uint16)cryptoInfo->pkcs5;
+		SetSecRegionParamsMemory();
+		status = EFI_SUCCESS;
+	}
 
 	// Clean auth data
 	MEM_BURN(&gAuthPassword, sizeof(gAuthPassword));
 	MEM_BURN(&gAuthPim, sizeof(gAuthPim));
 
-	return EFI_SUCCESS;
+	return status;
 }
 
 void GetIntersection(uint64 start1, uint32 length1, uint64 start2, uint64 end2, uint64 *intersectStart, uint32 *intersectLength)
@@ -563,12 +571,19 @@ SecRegionChangePwd() {
 		ZeroMem(&confirmPassword, sizeof(newPassword));
 		VCAskPwd(AskPwdNew, &newPassword);
 		if (gAuthPwdCode == AskPwdRetCancel) {
-			return EFI_NOT_READY;
+			return EFI_DCS_USER_CANCELED;
+		}
+		if (gAuthPwdCode == AskPwdRetTimeout) {
+			return EFI_TIMEOUT;
 		}
 		VCAskPwd(AskPwdConfirm, &confirmPassword);
 		if (gAuthPwdCode == AskPwdRetCancel) {
 			MEM_BURN(&newPassword, sizeof(newPassword));
-			return EFI_NOT_READY;
+			return EFI_DCS_USER_CANCELED;
+		}
+		if (gAuthPwdCode == AskPwdRetTimeout) {
+			MEM_BURN(&newPassword, sizeof(newPassword));
+			return EFI_TIMEOUT;
 		}
 		if (newPassword.Length == confirmPassword.Length) {
 			if (CompareMem(newPassword.Text, confirmPassword.Text, confirmPassword.Length) == 0) {
@@ -627,7 +642,7 @@ SecRegionChangePwd() {
 		if (key.UnicodeChar == 'r') {
 			MEM_BURN(&newPassword, sizeof(newPassword));
 			MEM_BURN(&confirmPassword, sizeof(confirmPassword));
-			CleanSensitiveData();
+			CleanSensitiveData(TRUE);
 			gST->RuntimeServices->ResetSystem(EfiResetCold, EFI_SUCCESS, 0, NULL);
 		}
 	}
@@ -675,7 +690,10 @@ SecRegionTryDecrypt()
 		SecRegionOffset = 0;
 		VCAuthAsk();
 		if (gAuthPwdCode == AskPwdRetCancel) {
-			return EFI_NOT_READY;
+			return EFI_DCS_USER_CANCELED;
+		}
+		if (gAuthPwdCode == AskPwdRetTimeout) {
+			return EFI_TIMEOUT;
 		}
 		OUT_PRINT(L"%a", gAuthStartMsg);
 		do {
@@ -698,6 +716,10 @@ SecRegionTryDecrypt()
 			break;
 		}	else {
 			ERR_PRINT(L"%a", gAuthErrorMsg);
+			// clear previous failed authentication information
+			MEM_BURN(&gAuthPassword, sizeof(gAuthPassword));
+			if (gAuthPimRqt)
+				MEM_BURN(&gAuthPim, sizeof(gAuthPim));
 		}
 		retry--;
 	} while (vcres != 0 && retry > 0);
@@ -793,6 +815,8 @@ SecRegionTryDecrypt()
 enum OnExitTypes{
 	OnExitAuthFaild = 1,
 	OnExitAuthNotFound,
+	OnExitAuthTimeout,
+	OnExitAuthCancelled,
 	OnExitSuccess
 };
 
@@ -820,7 +844,7 @@ AsciiStrNStr(
 			++posp;
 			++pos2;
 		}
-		if (*pos2 == 0) return NULL;
+		if (*pos2 == 0 && *posp) return NULL;
 		if (*posp == 0) return pos1;
 		++pos1;
 	}
@@ -866,10 +890,17 @@ OnExit(
 	CHAR8* delayStr = NULL;
 	EFI_GUID *guid = NULL;
 	CHAR16  *fileStr  = NULL;
+	
+	if (EFI_ERROR(retValue))
+	{
+		CleanSensitiveData(TRUE);
+	}
+	
 	if (action == NULL) return retValue;
+
 	if (OnExitGetParam(action, "guid", &guidStr, NULL)) {
 		EFI_GUID tmp;
-		if (AsciiStrToGuid(&tmp, guidStr)) {
+		if (DcsAsciiStrToGuid(&tmp, guidStr)) {
 			guid = MEM_ALLOC(sizeof(EFI_GUID));
 			CopyMem(guid, &tmp, sizeof(EFI_GUID));
 		}
@@ -905,29 +936,43 @@ OnExit(
 	}
 
 	if (AsciiStrNStr(action, "halt") == action) {
-		EfiCpuHalt();
+		retValue = EFI_DCS_HALT_REQUESTED;
+	}
+
+	else if (AsciiStrNStr(action, "shutdown") == action) {
+		retValue = EFI_DCS_SHUTDOWN_REQUESTED;
+	}
+	
+	else if (AsciiStrNStr(action, "reboot") == action) {
+		retValue = EFI_DCS_REBOOT_REQUESTED;
 	}
 
-	if (AsciiStrNStr(action, "exec") == action) {
+	else if (AsciiStrNStr(action, "exec") == action) {
 		if (guid != NULL) {
 			EFI_STATUS res;
 			EFI_HANDLE h;
 			res = EfiFindPartByGUID(guid, &h);
 			if (EFI_ERROR(res)) {
 				ERR_PRINT(L"\nCan't find start partition\n");
-				EfiCpuHalt();
+				CleanSensitiveData(TRUE);
+				retValue = EFI_DCS_HALT_REQUESTED;
+				goto exit;
 			}
 			// Try to exec
 			if (fileStr != NULL) {				
 				res = EfiExec(h, fileStr);
 				if (EFI_ERROR(res)) {
 					ERR_PRINT(L"\nStart %s - %r\n", fileStr, res);
-					EfiCpuHalt();
+					CleanSensitiveData(TRUE);
+					retValue = EFI_DCS_HALT_REQUESTED;
+					goto exit;
 				}
 			}
 			else {
 				ERR_PRINT(L"\nNo EFI execution path specified. Halting!\n");
-				EfiCpuHalt();
+				CleanSensitiveData(TRUE);
+				retValue = EFI_DCS_HALT_REQUESTED;
+				goto exit;
 			}
 		}		
 
@@ -937,17 +982,19 @@ OnExit(
 		goto exit;
 	}
 
-	if (AsciiStrNStr(action, "postexec") == action) {
+	else if (AsciiStrNStr(action, "postexec") == action) {
 		if (guid != NULL) {
 			EfiSetVar(L"DcsExecPartGuid", NULL, &guid, sizeof(EFI_GUID), EFI_VARIABLE_BOOTSERVICE_ACCESS);
 		}
 		if (fileStr != NULL) {
 			EfiSetVar(L"DcsExecCmd", NULL, fileStr, (StrLen(fileStr) + 1) * 2, EFI_VARIABLE_BOOTSERVICE_ACCESS);
 		}
+
+		retValue = EFI_DCS_POSTEXEC_REQUESTED;
 		goto exit;
 	}
 
-	if (AsciiStrStr(action, "exit") == action) {
+	else if (AsciiStrStr(action, "exit") == action) {
 		goto exit;
 	}
 
@@ -973,7 +1020,7 @@ VirtualNotifyEvent(
 	)
 {
 	// Clean all sensible info and keys before transfer to OS
-	CleanSensitiveData();
+	CleanSensitiveData(FALSE);
 }
 
 //////////////////////////////////////////////////////////////////////////
@@ -1151,7 +1198,16 @@ UefiMain(
 	gST->ConIn->Reset(gST->ConIn, FALSE);
 
 	if (EFI_ERROR(res)) {
-		return OnExit(gOnExitFailed, OnExitAuthFaild, res);
+		// clear buffers with potential authentication data
+		MEM_BURN(&gAuthPassword, sizeof(gAuthPassword));
+		MEM_BURN(&gAuthPim, sizeof(gAuthPim));
+
+		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);
 	}
 
 	res = PrepareBootParams(BootDriveSignature, SecRegionCryptInfo);