1 files changed, 151 insertions, 109 deletions
diff --git a/DcsCfg/DcsCfgCrypt.c b/DcsCfg/DcsCfgCrypt.c
index 0ca45ec..986824d 100644
--- a/DcsCfg/DcsCfgCrypt.c
+++ b/DcsCfg/DcsCfgCrypt.c
@@ -60,7 +60,7 @@ AskEA() {
 	CHAR16 name[128];
 	for (ea = EAGetFirst(); ea != 0; ea = EAGetNext(ea))
 	{
-		EAGetName(name, ea, 1);
+		EAGetName(name, 128, ea, 1);
 		OUT_PRINT(L"(%d) %s\n", ea, name);
 	}
 	ea = (int)AskUINTN(":", EAGetFirst());
@@ -167,12 +167,19 @@ ChangePassword(
 			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) {
@@ -417,6 +424,10 @@ RangeCrypt(
 	UINT64                  remainsOnStart;
 	UINT64                  pos;
 	UINTN                   rd;
+	BOOL                    bIsSystemEncyption = FALSE;
+
+	if (info->noIterations == get_pkcs5_iteration_count (info->pkcs5, info->volumePim, FALSE, TRUE))
+		bIsSystemEncyption = TRUE;
 
 	io = EfiGetBlockIO(disk);
 	if (!io) {
@@ -478,6 +489,13 @@ RangeCrypt(
 			if (encrypt) {
 				EncryptDataUnits(buf, (UINT64_STRUCT*)&pos, (UINT32)(rd), info);
 			}	else {
+				if (bIsSystemEncyption && (pos == start) && (0xEB52904E54465320 == BE64 (*(uint64 *) buf)))
+				{
+					// first sector is not encrypted (e.g. because of Windows repair).
+					// So we encrypt it so that decryption will lead to correct result
+					EncryptDataUnits(buf, (UINT64_STRUCT*)&pos, 1, info);
+				}
+				
 				DecryptDataUnits(buf, (UINT64_STRUCT*)&pos, (UINT32)(rd), info);
 			}
 
@@ -560,90 +578,96 @@ RangeCrypt(
 		RangeCryptProgress(size, remains, pos, remainsOnStart);
 	}
 	else if (!encrypt)
-	{
-		BOOL bIsSystemEncyption = FALSE;
-		if (info->noIterations == get_pkcs5_iteration_count (info->pkcs5, info->volumePim, FALSE, TRUE))
-			bIsSystemEncyption = TRUE;
-		
+	{		
 		if (bIsSystemEncyption)
 		{
-			/*
-			 * Case of OS decryption by Rescue Disk
-			 * Check if NTFS marker exists. If not, then probably disk affected by
-			 * bug in 1.19 Rescue Disk which caused the first 50 MB of disk to be 
-			 * decrypted in a wrong way. In this case, try to reverse the faulty decryption
-			 * and then perform correct decryption
-			 */
-			remains = size % CRYPT_BUF_SECTORS;
-			if (remains > 0)
-			{
-				/* 1.19 bug appears only when size not multiple of 50 MB */				
-				res = io->ReadBlocks(io, io->Media->MediaId, start, 512, buf);
-				if (!EFI_ERROR(res)) {
-					if (0xEB52904E54465320 != BE64 (*(uint64 *) buf)) /* NTFS */
+			res = io->ReadBlocks(io, io->Media->MediaId, start, 512, buf);
+			if (!EFI_ERROR(res)) {
+				/*
+				 * Case of OS decryption by Rescue Disk
+				 * Check if NTFS marker exists. If not, then probably disk affected by
+				 * either Windows Repair overwriting first sector or the bug in 1.19 
+				 * Rescue Disk which caused the first 50 MB of disk to be 
+				 * decrypted in a wrong way. In this case, try to reverse the faulty decryption
+				 * and then perform correct decryption
+				 */
+				if (0xEB52904E54465320 != BE64 (*(uint64 *) buf)) /* NTFS */
+				{
+					/* encrypt it to see if the first sector was unencrypted before decrypt done */
+					EncryptDataUnits(buf, (UINT64_STRUCT*)&start, 1, info);
+					
+					if (0xEB52904E54465320 == BE64 (*(uint64 *) buf)) /* NTFS */
 					{
-						if (AskConfirm("\r\nSystem already decrypted but partition can't be recognized.\r\nDid you use 1.19 Rescue Disk previously to decrypt OS?", 1)) {
-							OUT_PRINT(L"\r\nTrying to recover data corrupted by 1.19 Rescue Disk bug.");
-
-							pos = start + remains - CRYPT_BUF_SECTORS;
-							// Read
-							do {
-								res = io->ReadBlocks(io, io->Media->MediaId, pos, CRYPT_BUF_SECTORS << 9, buf);
-								if (EFI_ERROR(res)) {
-									UINT8 ar;
-									ERR_PRINT(L"Read error: %r\n", res);
-									ar = AskAR();
-									if (ar != 'R' && ar != 'r')
-										break;
-								}
-							} while (EFI_ERROR(res));
-							
-							if (EFI_ERROR(res))
-							{
-								OUT_PRINT(L"\r\nNo corrective action performed.");
+						// Write corrected first sector
+						do {
+							res = io->WriteBlocks(io, io->Media->MediaId, start, 512, buf);
+							if (EFI_ERROR(res)) {
+								UINT8 ar;
+								ERR_PRINT(L"Write error: %r\n", res);
+								ar = AskAR();
+								if (ar != 'R' && ar != 'r')
+									break;
 							}
-							else
-							{
-								UINT8* realEncryptedData = buf + ((CRYPT_BUF_SECTORS - remains) << 9);
-								BOOL bPerformWrite = FALSE;
+						} while (EFI_ERROR(res));
+						
+						if (EFI_ERROR(res))
+						{
+							OUT_PRINT(L"\r\nThe corrected first sector could not be written.");
+						}
+					}
+					else
+					{
+						/* restore original value */
+						DecryptDataUnits(buf, (UINT64_STRUCT*)&start, 1, info);
 
-								// reverse faulty decryption
-								EncryptDataUnits(buf, (UINT64_STRUCT*)&pos, (UINT32)(remains), info);
+						remains = size % CRYPT_BUF_SECTORS;
+						if (remains > 0)
+						{
+							/* 1.19 bug appears only when size not multiple of 50 MB */																
+							if (AskConfirm("\r\nSystem already decrypted but partition can't be recognized.\r\nDid you use 1.19 Rescue Disk previously to decrypt OS?", 1)) {
+								OUT_PRINT(L"\r\nTrying to recover data corrupted by 1.19 Rescue Disk bug.");
+
+								pos = start + remains - CRYPT_BUF_SECTORS;
+								// Read
+								do {
+									res = io->ReadBlocks(io, io->Media->MediaId, pos, CRYPT_BUF_SECTORS << 9, buf);
+									if (EFI_ERROR(res)) {
+										UINT8 ar;
+										ERR_PRINT(L"Read error: %r\n", res);
+										ar = AskAR();
+										if (ar != 'R' && ar != 'r')
+											break;
+									}
+								} while (EFI_ERROR(res));
 								
-								// decrypt the correct data
-								DecryptDataUnits(realEncryptedData, (UINT64_STRUCT*)&start, (UINT32)(remains), info);
-						
-								if (0xEB52904E54465320 == BE64 (*(uint64 *) realEncryptedData)) /* NTFS */
-									bPerformWrite = TRUE;
-								else
+								if (EFI_ERROR(res))
 								{
-									if (AskConfirm("\r\nDecrypted data don't contain valid partition information. Proceeed anyway?", 1))
-										bPerformWrite = TRUE;
+									OUT_PRINT(L"\r\nNo corrective action performed.");
 								}
-								
-								if (bPerformWrite)
+								else
 								{
-									// Write original encrypted data
-									do {
-										res = io->WriteBlocks(io, io->Media->MediaId, pos, (UINTN)((CRYPT_BUF_SECTORS - remains) << 9), buf);
-										if (EFI_ERROR(res)) {
-											UINT8 ar;
-											ERR_PRINT(L"Write error: %r\n", res);
-											ar = AskAR();
-											if (ar != 'R' && ar != 'r')
-												break;
-										}
-									} while (EFI_ERROR(res));
+									UINT8* realEncryptedData = buf + ((CRYPT_BUF_SECTORS - remains) << 9);
+									BOOL bPerformWrite = FALSE;
+
+									// reverse faulty decryption
+									EncryptDataUnits(buf, (UINT64_STRUCT*)&pos, (UINT32)(remains), info);
 									
-									if (EFI_ERROR(res))
+									// decrypt the correct data
+									DecryptDataUnits(realEncryptedData, (UINT64_STRUCT*)&start, (UINT32)(remains), info);
+							
+									if (0xEB52904E54465320 == BE64 (*(uint64 *) realEncryptedData)) /* NTFS */
+										bPerformWrite = TRUE;
+									else
 									{
-										OUT_PRINT(L"\r\nNo corrective action performed.");
+										if (AskConfirm("\r\nDecrypted data don't contain valid partition information. Proceeed anyway?", 1))
+											bPerformWrite = TRUE;
 									}
-									else
-									{										
-										// Write correctly decrypted data
+									
+									if (bPerformWrite)
+									{
+										// Write original encrypted data
 										do {
-											res = io->WriteBlocks(io, io->Media->MediaId, start, (UINTN) (remains << 9), realEncryptedData);
+											res = io->WriteBlocks(io, io->Media->MediaId, pos, (UINTN)((CRYPT_BUF_SECTORS - remains) << 9), buf);
 											if (EFI_ERROR(res)) {
 												UINT8 ar;
 												ERR_PRINT(L"Write error: %r\n", res);
@@ -652,30 +676,48 @@ RangeCrypt(
 													break;
 											}
 										} while (EFI_ERROR(res));
-									
+										
 										if (EFI_ERROR(res))
 										{
-											OUT_PRINT(L"\r\nFailed to write decrypted data.");
+											OUT_PRINT(L"\r\nNo corrective action performed.");
 										}
 										else
-										{
-											OUT_PRINT(L"\r\nData recovered successfully!");											
+										{										
+											// Write correctly decrypted data
+											do {
+												res = io->WriteBlocks(io, io->Media->MediaId, start, (UINTN) (remains << 9), realEncryptedData);
+												if (EFI_ERROR(res)) {
+													UINT8 ar;
+													ERR_PRINT(L"Write error: %r\n", res);
+													ar = AskAR();
+													if (ar != 'R' && ar != 'r')
+														break;
+												}
+											} while (EFI_ERROR(res));
+										
+											if (EFI_ERROR(res))
+											{
+												OUT_PRINT(L"\r\nFailed to write decrypted data.");
+											}
+											else
+											{
+												OUT_PRINT(L"\r\nData recovered successfully!");											
+											}
 										}
 									}
-								}
-								else
-								{
-									OUT_PRINT(L"\r\nNo corrective action performed.");
-								}								
-							}							
-						}
-						else
-						{
-							OUT_PRINT(L"\n\rNo corrective action attempted.");
+									else
+									{
+										OUT_PRINT(L"\r\nNo corrective action performed.");
+									}								
+								}							
+							}
+							else
+							{
+								OUT_PRINT(L"\n\rNo corrective action attempted.");
+							}
+							
 						}
-						
-					}
-					
+					}					
 				}
 			}
 			 
@@ -1566,7 +1608,7 @@ UsbScApdu(
 	EFI_STATUS res;
 	CE(InitUsb());
 	CE(UsbGetIO(gUSBHandles[UsbIndex], &UsbIo));
-	StrHexToBytes(cmd + sizeof(CCID_HEADER_OUT), &cmdLen, hexString);
+	DcsStrHexToBytes(cmd + sizeof(CCID_HEADER_OUT), &cmdLen, hexString);
 	CE(UsbScTransmit(UsbIo, cmd, cmdLen + sizeof(CCID_HEADER_OUT), resp, &respLen, &statusSc));
 	PrintBytes(resp, respLen);
 	return res;
@@ -1622,7 +1664,7 @@ UpdateDcsBoot() {
 UINTN gSecRigonCount = 0;
 
 EFI_STATUS
-SecRigionMark() 
+SecRegionMark() 
 {
 	UINT32      crc;
 	EFI_STATUS  res;
@@ -1665,7 +1707,7 @@ SecRigionMark()
 }
 
 EFI_STATUS
-SecRigionWipe()
+SecRegionWipe()
 {
 	EFI_STATUS  res;
 	CHAR8*      buf;
@@ -1719,7 +1761,7 @@ error:
 }
 
 EFI_STATUS
-SecRigionDump(
+SecRegionDump(
 	IN EFI_HANDLE   hBio,
 	IN CHAR16       *prefix
 	)
@@ -1728,9 +1770,9 @@ SecRigionDump(
 	EFI_BLOCK_IO_PROTOCOL*  bio;
 	DCS_AUTH_DATA_MARK*     adm = NULL;
 	UINT32                  crc;
-	UINT8*                  SecRegionData = NULL;
-	UINTN                   SecRegionSize = 0;
-	UINTN                   SecRegionOffset = 0;
+	UINT8*                  SecRegionDumpData = NULL;
+	UINTN                   SecRegionDumpSize = 0;
+	UINTN                   SecRegionDumpOffset = 0;
 	UINTN                   saveSize = 0;
 	UINTN                   idx = 0;
 	CHAR16                  name[128];
@@ -1755,42 +1797,42 @@ SecRigionDump(
 		res = EFI_INVALID_PARAMETER;
 	}
 
-	SecRegionSize = adm->AuthDataSize * 128 * 1024;
-	SecRegionData = MEM_ALLOC(SecRegionSize);
-	if (SecRegionData == NULL) {
+	SecRegionDumpSize = adm->AuthDataSize * 128 * 1024;
+	SecRegionDumpData = MEM_ALLOC(SecRegionDumpSize);
+	if (SecRegionDumpData == NULL) {
 		res = EFI_BUFFER_TOO_SMALL;
 		goto err;
 	}
-	CE(bio->ReadBlocks(bio, bio->Media->MediaId, 62, SecRegionSize, SecRegionData));
+	CE(bio->ReadBlocks(bio, bio->Media->MediaId, 62, SecRegionDumpSize, SecRegionDumpData));
 
 	do {
 		// EFI tables?
-		if (TablesVerify(SecRegionSize - SecRegionOffset, SecRegionData + SecRegionOffset)) {
-			EFI_TABLE_HEADER *mhdr = (EFI_TABLE_HEADER *)(SecRegionData + SecRegionOffset);
+		if (TablesVerify(SecRegionDumpSize - SecRegionDumpOffset, SecRegionDumpData + SecRegionDumpOffset)) {
+			EFI_TABLE_HEADER *mhdr = (EFI_TABLE_HEADER *)(SecRegionDumpData + SecRegionDumpOffset);
 			UINTN tblZones = (mhdr->HeaderSize + 1024 * 128 - 1) / (1024 * 128);
 			saveSize = tblZones * 1024 * 128;
 		}		else {
 			saveSize = 1024 * 128;
 		}
 		UnicodeSPrint(name, sizeof(name), L"%s%d", prefix, idx);
-		CE(FileSave(NULL, name, SecRegionData + SecRegionOffset, saveSize));
+		CE(FileSave(NULL, name, SecRegionDumpData + SecRegionDumpOffset, saveSize));
 		OUT_PRINT(L"%s saved\n", name);
 		idx += saveSize / (1024 * 128);
-		SecRegionOffset += saveSize;
-	} while (SecRegionOffset < SecRegionSize);
+		SecRegionDumpOffset += saveSize;
+	} while (SecRegionDumpOffset < SecRegionDumpSize);
 
 err:
 	if (EFI_ERROR(res)) {
 		ERR_PRINT(L"%r\n", res);
 	}
 	MEM_FREE(adm);
-	MEM_FREE(SecRegionData);
+	MEM_FREE(SecRegionDumpData);
 	return res;
 }
 
 
 EFI_STATUS
-SecRigionAdd(
+SecRegionAdd(
 	IN UINTN       regIdx
 )
 {