1 files changed, 56 insertions, 17 deletions
diff --git a/src/Common/Dlgcode.c b/src/Common/Dlgcode.c
index 2fd18606..76748fe5 100644
--- a/src/Common/Dlgcode.c
+++ b/src/Common/Dlgcode.c
@@ -7166,8 +7166,9 @@ BOOL CALLBACK KeyfileGeneratorDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LP
 	{
 	case WM_INITDIALOG:
 		{
 			HWND hComboBox = GetDlgItem (hwndDlg, IDC_PRF_ID);
+			HWND hSizeUnit = GetDlgItem (hwndDlg, IDC_KEYFILES_SIZE_UNIT);
 			HCRYPTPROV hRngProv = NULL;
 
 			VirtualLock (randPool, sizeof(randPool));
 			VirtualLock (lastRandPool, sizeof(lastRandPool));
@@ -7195,8 +7196,18 @@ BOOL CALLBACK KeyfileGeneratorDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LP
 					AddComboPair (hComboBox, HashGetName(hid), hid);
 			}
 			SelectAlgo (hComboBox, &hash_algo);
 
+			// populate keyfiles size unit combo
+			SendMessage (hSizeUnit, CB_RESETCONTENT, 0, 0);
+			AddComboPair (hSizeUnit,  GetString ("BYTES"), 0);
+			AddComboPair (hSizeUnit, GetString ("KB"), 1);
+			AddComboPair (hSizeUnit, GetString ("MB"), 2);
+			AddComboPair (hSizeUnit, GetString ("GB"), 3);
+
+			// set default keyfiles size unit
+			SendMessage (hSizeUnit, CB_SETCURSEL, 0, 0);
+
 			SetCheckBox (hwndDlg, IDC_DISPLAY_POOL_CONTENTS, bDisplayPoolContents);
 			hEntropyBar = GetDlgItem (hwndDlg, IDC_ENTROPY_BAR);
 			SendMessage (hEntropyBar, PBM_SETRANGE32, 0, maxEntropyLevel);
 			SendMessage (hEntropyBar, PBM_SETSTEP, 1, 0);
@@ -7302,8 +7313,9 @@ BOOL CALLBACK KeyfileGeneratorDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LP
 
 		if (lw == IDC_KEYFILES_RANDOM_SIZE)
 		{
 			EnableWindow(GetDlgItem (hwndDlg, IDC_KEYFILES_SIZE), !GetCheckBox (hwndDlg, IDC_KEYFILES_RANDOM_SIZE));
+			EnableWindow(GetDlgItem (hwndDlg, IDC_KEYFILES_SIZE_UNIT), !GetCheckBox (hwndDlg, IDC_KEYFILES_RANDOM_SIZE));
 		}
 
 		if (lw == IDC_GENERATE_AND_SAVE_KEYFILE)
 		{
@@ -7312,9 +7324,12 @@ BOOL CALLBACK KeyfileGeneratorDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LP
 			wchar_t szDirName[TC_MAX_PATH];
 			wchar_t szFileName [2*TC_MAX_PATH + 16];
 			unsigned char *keyfile = NULL;
 			int fhKeyfile = -1, status;
-			long keyfilesCount = 0, keyfilesSize = 0, i;
+			long keyfilesCount = 0, i;
+			unsigned long long keyfilesSize = 0, remainingBytes = 0;
+			int selectedUnitIndex, selectedUnitFactor, loopIndex, rndBytesLength;
+			DWORD dwLastError = 0;
 			wchar_t* fileExtensionPtr = 0;
 			wchar_t szSuffix[32];
 			BOOL bRandomSize = GetCheckBox (hwndDlg, IDC_KEYFILES_RANDOM_SIZE);
 
@@ -7334,14 +7349,25 @@ BOOL CALLBACK KeyfileGeneratorDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LP
 				if (!GetWindowText(GetDlgItem (hwndDlg, IDC_KEYFILES_SIZE), szNumber, ARRAYSIZE(szNumber)))
 					szNumber[0] = 0;
 
 				keyfilesSize = wcstoul(szNumber, NULL, 0);
-				if (keyfilesSize < 64 || keyfilesSize > 1024*1024)
+				// multiply by the unit factor
+				selectedUnitIndex = ComboBox_GetCurSel (GetDlgItem (hwndDlg, IDC_KEYFILES_SIZE_UNIT));
+				if (selectedUnitIndex != CB_ERR)
+				{
+					selectedUnitFactor = (CK_SLOT_ID) ComboBox_GetItemData (GetDlgItem (hwndDlg, IDC_KEYFILES_SIZE_UNIT), selectedUnitIndex);
+					for (loopIndex = 0; loopIndex < selectedUnitFactor; loopIndex++)
+						keyfilesSize *= 1024ULL;
+				}
+
+				if (keyfilesSize < 64)
 				{
 					Warning("KEYFILE_INCORRECT_SIZE", hwndDlg);
 					SendMessage(hwndDlg, WM_NEXTDLGCTL, (WPARAM) GetDlgItem (hwndDlg, IDC_KEYFILES_SIZE), TRUE);
 					return 1;
 				}
+
+				remainingBytes = keyfilesSize;
 			}
 
 			if (!GetWindowText(GetDlgItem (hwndDlg, IDC_KEYFILES_BASE_NAME), szFileBaseName, TC_MAX_PATH))
 				szFileBaseName[0] = 0;
@@ -7371,9 +7397,9 @@ BOOL CALLBACK KeyfileGeneratorDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LP
 				StringCbCat(szDirName, sizeof(szDirName), L"\\");
 
 			WaitCursor();
 
-			keyfile = (unsigned char*) TCalloc( bRandomSize? KEYFILE_MAX_READ_LEN : keyfilesSize );
+			keyfile = (unsigned char*) TCalloc(KEYFILE_MAX_READ_LEN);
 
 			for (i= 0; i < keyfilesCount; i++)
 			{
 				StringCbCopyW(szFileName, sizeof(szFileName), szDirName);
@@ -7434,34 +7460,47 @@ BOOL CALLBACK KeyfileGeneratorDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LP
 						NormalCursor();
 						return 1;
 					}
 					
-					/* since keyfilesSize < 1024 * 1024, we mask with 0x000FFFFF */
-					keyfilesSize = (long) (((unsigned long) keyfilesSize) & 0x000FFFFF);
+					/* since random keyfilesSize < 1024 * 1024, we mask with 0x000FFFFF */
+					keyfilesSize = (unsigned long long) (((unsigned long) keyfilesSize) & 0x000FFFFF);
 
 					keyfilesSize %= ((KEYFILE_MAX_READ_LEN - 64) + 1);
 					keyfilesSize += 64;
+
+					remainingBytes = keyfilesSize;
 				}
 
-				/* Generate the keyfile */ 				
-				if (!RandgetBytesFull (hwndDlg, keyfile, keyfilesSize, TRUE, TRUE))
-				{
-					_close (fhKeyfile);
-					DeleteFile (szFileName);
-					TCfree(keyfile);
-					NormalCursor();
-					return 1;
-				}				
+				do {
+					rndBytesLength = (int) min (remainingBytes, (unsigned long long) KEYFILE_MAX_READ_LEN);
+
+					/* Generate the keyfile */ 				
+					if (!RandgetBytesFull (hwndDlg, keyfile, rndBytesLength, TRUE, TRUE))
+					{
+						_close (fhKeyfile);
+						DeleteFile (szFileName);
+						TCfree(keyfile);
+						NormalCursor();
+						return 1;
+					}				
+
+					/* Write the keyfile */
+					status = _write (fhKeyfile, keyfile, rndBytesLength);
+				} while (status != -1 && (remainingBytes -= (unsigned long long) rndBytesLength) > 0);
+
+				/* save last error code */
+				if (status == -1)
+					dwLastError = GetLastError();
 
-				/* Write the keyfile */
-				status = _write (fhKeyfile, keyfile, keyfilesSize);
-				burn (keyfile, keyfilesSize);
+				burn (keyfile, KEYFILE_MAX_READ_LEN);
 				_close (fhKeyfile);
 
 				if (status == -1)
 				{
 					TCfree(keyfile);
 					NormalCursor();
+					/* restore last error code */
+					SetLastError(dwLastError);
 					handleWin32Error (hwndDlg, SRC_POS);
 					return 1;
 				}				
 			}