From c8b9a73916b727d973cf51f1d7ebac3b9670d06b Mon Sep 17 00:00:00 2001
From: Mounir IDRASSI <mounir.idrassi@idrix.fr>
Date: Sun, 4 Aug 2024 10:25:12 +0200
Subject: Windows: Fix rare cases being stuck in Secure Desktop after it is
 used for password entry

Cause seems to be IME as documented in KeePass project (https://keepass.info/help/kb/sec_desk.html#ime). We use the same approach as KeePass to disable IME in Secure Desktop.
This commit also add few changes:
  - we switch to secure desktop only if SetThreadDesktop succeeds
  - we call SwitchDesktop to switch to original desktop only if we actually succeeded in displaying secure desktop
---
 src/Common/Dlgcode.c | 56 +++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/src/Common/Dlgcode.c b/src/Common/Dlgcode.c
index ce86c9da..488b050f 100644
--- a/src/Common/Dlgcode.c
+++ b/src/Common/Dlgcode.c
@@ -13921,20 +13921,33 @@ static unsigned int __stdcall SecureDesktopThread( LPVOID lpThreadParameter )
 	StringCbCopy(SecureDesktopName, sizeof (SecureDesktopName), pParam->szDesktopName);
 	pParam->hDesk = hSecureDesk;
 
-	// wait for SwitchDesktop to succeed before using it for current thread
-	while (true)
+	bNewDesktopSet = SetThreadDesktop (hSecureDesk);
+
+	if (bNewDesktopSet)
 	{
-		if (SwitchDesktop (hSecureDesk))
+		// call ImmDisableIME　from imm32.dll to disable IME since it can create issue with secure desktop
+		// cf: https://keepass.info/help/kb/sec_desk.html#ime
+		HMODULE hImmDll = LoadLibraryEx (L"imm32.dll", NULL, LOAD_LIBRARY_SEARCH_SYSTEM32);
+		if (hImmDll)
 		{
-			break;
+			typedef BOOL (WINAPI *ImmDisableIME_t)(DWORD);
+			ImmDisableIME_t ImmDisableIME = (ImmDisableIME_t) GetProcAddress (hImmDll, "ImmDisableIME");
+			if (ImmDisableIME)
+			{
+				ImmDisableIME (0);
+			}
 		}
-		Sleep (SECUREDESKTOP_MONOTIR_PERIOD);
-	}
 
-	bNewDesktopSet = SetThreadDesktop (hSecureDesk);
+		// wait for SwitchDesktop to succeed before using it for current thread
+		while (true)
+		{
+			if (SwitchDesktop (hSecureDesk))
+			{
+				break;
+			}
+			Sleep (SECUREDESKTOP_MONOTIR_PERIOD);
+		}
 
-	if (bNewDesktopSet)
-	{
 		// create the thread that will ensure that VeraCrypt secure desktop has always user input
 		// this is done only if the stop event is created successfully
 		HANDLE hStopEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
@@ -13964,6 +13977,12 @@ static unsigned int __stdcall SecureDesktopThread( LPVOID lpThreadParameter )
 		}
 
 		pParam->bDlgDisplayed = TRUE;
+
+		// free imm32.dll handle
+		if (hImmDll)
+		{
+			FreeLibrary (hImmDll);
+		}
 	}
 	else
 	{
@@ -14084,19 +14103,20 @@ INT_PTR SecureDesktopDialogBoxParam(
 					// dialog box was indeed displayed in Secure Desktop
 					retValue = param.retValue;
 					bSuccess = TRUE;
+
+					// switch back to the original desktop
+					while (!SwitchDesktop (hOriginalDesk))
+					{
+						Sleep (SECUREDESKTOP_MONOTIR_PERIOD);
+					}
+
+					SetThreadDesktop (hOriginalDesk);
 				}
-			}
 
-			if (param.hDesk)
-			{	
-				while (!SwitchDesktop (hOriginalDesk))
+				if (param.hDesk)
 				{
-					Sleep (SECUREDESKTOP_MONOTIR_PERIOD);
+					CloseDesktop (param.hDesk);
 				}
-
-				SetThreadDesktop (hOriginalDesk);
-
-				CloseDesktop (param.hDesk);
 			}
 
 			// get the new list of ctfmon.exe processes in order to find the ID of the
-- 
cgit v1.2.3