From df6fff78bbe6563fedd0e5c20ee2b51bf704c462 Mon Sep 17 00:00:00 2001
From: Katayama Hirofumi MZ <katayama.hirofumi.mz@gmail.com>
Date: Thu, 16 Sep 2021 19:29:49 +0900
Subject: [PATCH] [IMM32] Rewrite ImmAssociateContextEx (#3961)

- Rewrite ImmAssociateContextEx function.
- Modify NtUserAssociateInputContext prototype.
CORE-11700
---
 dll/win32/imm32/imm.c         | 68 ++++++++++++++++++-----------------
 sdk/include/psdk/imm.h        |  1 +
 win32ss/include/ntuser.h      |  5 +--
 win32ss/user/ntuser/ntstubs.c |  5 +--
 4 files changed, 38 insertions(+), 41 deletions(-)

diff --git a/dll/win32/imm32/imm.c b/dll/win32/imm32/imm.c
index 21832a5a3bc..6092b354d74 100644
--- a/dll/win32/imm32/imm.c
+++ b/dll/win32/imm32/imm.c
@@ -254,48 +254,50 @@ HIMC WINAPI ImmAssociateContext(HWND hWnd, HIMC hIMC)
     return old;
 }
 
-/*
- * Helper function for ImmAssociateContextEx
- */
-static BOOL CALLBACK _ImmAssociateContextExEnumProc(HWND hwnd, LPARAM lParam)
-{
-    HIMC hImc = (HIMC)lParam;
-    ImmAssociateContext(hwnd,hImc);
-    return TRUE;
-}
-
 /***********************************************************************
  *              ImmAssociateContextEx (IMM32.@)
  */
 BOOL WINAPI ImmAssociateContextEx(HWND hWnd, HIMC hIMC, DWORD dwFlags)
 {
-    TRACE("(%p, %p, 0x%x):\n", hWnd, hIMC, dwFlags);
+    HWND hwndFocus;
+    PWND pFocusWnd;
+    HIMC hOldIMC = NULL;
+    DWORD dwValue;
 
-    if (!hWnd)
+    TRACE("(%p, %p, 0x%lX)\n", hWnd, hIMC, dwFlags);
+
+    if (!g_psi || !(g_psi->dwSRVIFlags & SRVINFO_IMM32))
         return FALSE;
 
-    switch (dwFlags)
-    {
-    case 0:
-        ImmAssociateContext(hWnd,hIMC);
-        return TRUE;
-    case IACE_DEFAULT:
-    {
-        HIMC defaultContext = get_default_context( hWnd );
-        if (!defaultContext) return FALSE;
-        ImmAssociateContext(hWnd,defaultContext);
-        return TRUE;
-    }
-    case IACE_IGNORENOCONTEXT:
-        if (GetPropW(hWnd,szwWineIMCProperty))
-            ImmAssociateContext(hWnd,hIMC);
-        return TRUE;
-    case IACE_CHILDREN:
-        EnumChildWindows(hWnd,_ImmAssociateContextExEnumProc,(LPARAM)hIMC);
-        return TRUE;
-    default:
-        FIXME("Unknown dwFlags 0x%x\n",dwFlags);
+    if (hIMC && !(dwFlags & IACE_DEFAULT) && Imm32IsCrossThreadAccess(hIMC))
         return FALSE;
+
+    hwndFocus = (HWND)NtUserQueryWindow(hWnd, QUERY_WINDOW_FOCUS);
+    pFocusWnd = ValidateHwndNoErr(hwndFocus);
+    if (pFocusWnd)
+        hOldIMC = pFocusWnd->hImc;
+
+    dwValue = NtUserAssociateInputContext(hWnd, hIMC, dwFlags);
+    switch (dwValue)
+    {
+        case 0:
+            return TRUE;
+
+        case 1:
+            pFocusWnd = ValidateHwndNoErr(hwndFocus);
+            if (pFocusWnd)
+            {
+                hIMC = pFocusWnd->hImc;
+                if (hIMC != hOldIMC)
+                {
+                    ImmSetActiveContext(hwndFocus, hOldIMC, FALSE);
+                    ImmSetActiveContext(hwndFocus, hIMC, TRUE);
+                }
+            }
+            return TRUE;
+
+        default:
+            return FALSE;
     }
 }
 
diff --git a/sdk/include/psdk/imm.h b/sdk/include/psdk/imm.h
index f28b8e69f3c..6232ee115ff 100644
--- a/sdk/include/psdk/imm.h
+++ b/sdk/include/psdk/imm.h
@@ -627,6 +627,7 @@ BOOL WINAPI ImmConfigureIMEW(_In_ HKL, _In_ HWND, _In_ DWORD, _In_ LPVOID);
 #define ImmConfigureIME WINELIB_NAME_AW(ImmConfigureIME)
 
 HIMC WINAPI ImmCreateContext(void);
+BOOL WINAPI ImmSetActiveContext(HWND hwnd, HIMC hIMC, BOOL fFlag);
 BOOL WINAPI ImmDestroyContext(_In_ HIMC hIMC);
 BOOL WINAPI ImmDisableIME(_In_ DWORD idThread);
 BOOL WINAPI ImmEnumInputContext(_In_ DWORD, _In_ IMCENUMPROC, _In_ LPARAM);
diff --git a/win32ss/include/ntuser.h b/win32ss/include/ntuser.h
index 636392365fb..6d463ddb69a 100644
--- a/win32ss/include/ntuser.h
+++ b/win32ss/include/ntuser.h
@@ -1317,10 +1317,7 @@ C_ASSERT(sizeof(CLIENTIMC) == 0x34);
 
 DWORD
 NTAPI
-NtUserAssociateInputContext(
-    DWORD dwUnknown1,
-    DWORD dwUnknown2,
-    DWORD dwUnknown3);
+NtUserAssociateInputContext(HWND hWnd, HIMC hIMC, DWORD dwFlags);
 
 NTSTATUS
 NTAPI
diff --git a/win32ss/user/ntuser/ntstubs.c b/win32ss/user/ntuser/ntstubs.c
index 4de9ea36a58..6a1b3c4673e 100644
--- a/win32ss/user/ntuser/ntstubs.c
+++ b/win32ss/user/ntuser/ntstubs.c
@@ -11,10 +11,7 @@ DBG_DEFAULT_CHANNEL(UserMisc);
 
 DWORD
 APIENTRY
-NtUserAssociateInputContext(
-    DWORD dwUnknown1,
-    DWORD dwUnknown2,
-    DWORD dwUnknown3)
+NtUserAssociateInputContext(HWND hWnd, HIMC hIMC, DWORD dwFlags)
 {
     STUB
     return 0;
-- 
2.17.1