Fix ImGui keyboard event not being handled correctly

Dalamud/Game/ClientState/Keys/InputDevicePoll.cs

@@ -0,0 +1,65 @@
+using System;
+using System.Runtime.InteropServices;
+
+using Dalamud.Hooking;
+using Dalamud.Utility;
+
+namespace Dalamud.Game.ClientState.Keys;
+
+/// <summary>
+/// This class provides events related to game inputs.
+/// </summary>
+[ServiceManager.BlockingEarlyLoadedService]
+internal sealed class InputDevicePoll : IDisposable, IServiceType
+{
+    private readonly Hook<UpdateInputDelegate> updateInputHook;
+
+    /// <summary>
+    /// Raised when the game is about to poll inputs.
+    /// </summary>
+    public event Action? OnBeforePoll;
+
+    /// <summary>
+    /// Raised when the game polled inputs.
+    /// </summary>
+    public event Action? OnAfterPoll;
+
+    [ServiceManager.ServiceConstructor]
+    public InputDevicePoll(SigScanner sigScanner)
+    {
+        // Client::System::Framework::Framework_TaskUpdateInputDevice (names from FFXIVClientStructs) calls this function
+        // 48:83EC 48   | sub rsp,48 |
+        // 48:895C24 50 | mov qword ptr ss: [rsp+50],rbx |
+        // BA 1E000000  | mov edx,1E |
+        // 48:897424 58 | mov qword ptr ss: [rsp+58],rsi |
+        // 48:897C24 40 | mov qword ptr ss: [rsp+40],rdi |
+        // 48:8BF9      | mov rdi,rcx |
+        // 48:81C1 60040000 | add rcx,460 |
+        // 0F297424 30  | movaps xmmword ptr ss: [rsp+30],xmm6 |
+        // 0F28F1       | movaps xmm6,xmm1 |
+        // E8 0158FDFF  | call ffxiv_dx11.7FF6CDD89970 |
+        var updateInputAddr =
+            sigScanner.ScanText(
+                "48?????? 48???????? BA ????0000 48???????? 48???????? 48???? 4881C1????0000 0F???????? 0F???? E8");
+        this.updateInputHook = Hook<UpdateInputDelegate>.FromAddress(updateInputAddr, this.OnUpdateInput);
+        this.updateInputHook.Enable();
+    }
+
+    [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
+    private delegate nint UpdateInputDelegate(nint a1, nint a2);
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        this.updateInputHook.Dispose();
+    }
+
+    private nint OnUpdateInput(nint a1, nint a2)
+    {
+        this.OnBeforePoll?.InvokeSafely();
+        var ret = this.updateInputHook.Original(a1, a2);
+        this.OnAfterPoll?.InvokeSafely();
+
+        return ret;
+    }
+}

Dalamud/Game/ClientState/Keys/KeyCapture.State.cs

@@ -0,0 +1,14 @@
+using System;
+
+namespace Dalamud.Game.ClientState.Keys;
+
+partial class KeyCapture
+{
+    [Flags]
+    private enum State : byte
+    {
+        CaptureAll = 0x01,
+        CaptureAllSingleFrame = 0x02,
+        RestoreAllOnNextFrame = 0x04,
+    }
+}

Dalamud/Game/ClientState/Keys/KeyCapture.cs

@@ -0,0 +1,152 @@
+using System;
+using System.Collections.Generic;
+
+using Dalamud.IoC;
+using Dalamud.IoC.Internal;
+using MonoMod.Utils;
+
+namespace Dalamud.Game.ClientState.Keys;
+
+[PluginInterface]
+[InterfaceVersion("1.0")]
+[ServiceManager.BlockingEarlyLoadedService]
+internal sealed partial class KeyCapture : IDisposable, IServiceType
+{
+    private readonly InputDevicePoll mInputDevicePoll;
+
+    private readonly SimulatedKeyState mSimulatedKeyState;
+
+    private readonly RawKeyState mRawKeyState;
+
+    private readonly List<ushort> mNextCaptureSet = new();
+
+    private readonly List<ushort> mNextRestoreSet = new();
+
+    private State mCaptureState;
+
+    [ServiceManager.ServiceConstructor]
+    private KeyCapture(InputDevicePoll devicePoll, SimulatedKeyState simulatedKeyState, RawKeyState rawKeyState)
+    {
+        this.mInputDevicePoll = devicePoll;
+        this.mSimulatedKeyState = simulatedKeyState;
+        this.mRawKeyState = rawKeyState;
+
+        this.mInputDevicePoll.OnBeforePoll += this.OnBeforeInputPoll;
+    }
+
+    /// <inheritdoc />
+    void IDisposable.Dispose()
+    {
+        this.mInputDevicePoll.OnBeforePoll -= this.OnBeforeInputPoll;
+    }
+
+    /// <summary>
+    /// Stop delivering inputs to the game indefinitely.
+    /// </summary>
+    /// <param name="doCapture">
+    /// If this is set to true then it will block all inputs.
+    /// To stop capturing, set this to false.
+    /// </param>
+    public void CaptureAll(bool doCapture = true)
+    {
+        this.mCaptureState = doCapture switch
+        {
+            true => this.mCaptureState | State.CaptureAll,
+            false => this.mCaptureState & ~State.CaptureAll,
+        };
+
+        if (!doCapture)
+        {
+            // Also queue restore so that the game can receive held inputs again.
+            this.mCaptureState |= State.RestoreAllOnNextFrame;
+        }
+    }
+
+    /// <summary>
+    /// Captures all keys for a single frame.
+    /// </summary>
+    public void CaptureAllSingleFrame()
+    {
+        this.mCaptureState |= State.CaptureAllSingleFrame;
+    }
+
+    /// <summary>
+    /// Captures a designated key for a single frame.
+    /// </summary>
+    /// <param name="vkCode">A virtual key code to capture.</param>
+    public void CaptureSingleFrame(ushort vkCode)
+    {
+        this.mNextCaptureSet.Add(vkCode);
+    }
+
+    private void OnBeforeInputPoll()
+    {
+        // Restore all keys
+        if (this.mCaptureState.HasFlag(State.RestoreAllOnNextFrame))
+        {
+            // Remove pending restore flag
+            this.mCaptureState &= ~State.RestoreAllOnNextFrame;
+
+            // Copy all simulated key states into actual buffer.
+            // Note that simulated key state is laid out very carefully so that
+            // it could just be memcpy'd in this situation.
+            this.mSimulatedKeyState.RawState.CopyTo(this.mRawKeyState.RawState);
+
+            // If restored all keys, there's no point of restoring individual keys.
+            this.mNextRestoreSet.Clear();
+        }
+
+        // Restore individual keys
+        foreach (var vkCode in this.mNextRestoreSet)
+        {
+            if (!this.mSimulatedKeyState.TryGetState(vkCode, out var state))
+            {
+                continue;
+            }
+
+            // Set key state to its original value
+            // Log.Verbose("restore: {VkCode} = {State}" vkCode, state);
+            this.mRawKeyState.SetState(vkCode, state);
+        }
+
+        this.mNextRestoreSet.Clear();
+
+        // Capture key states.
+        // Note that we process capturing **only after** finished restoring keys.
+        // This allows capturing to take higher precedence over restoring if they're queued on the same frame.
+        if (this.mCaptureState.HasFlag(State.CaptureAll))
+        {
+            // We clear game inputs here on every frame to block delivering any new inputs to the game
+            // because we didn't at DispatchMessage time.
+            //
+            // Thankfully zeroing ~1KB contiguous memory is very fast so we don't even need to touch WndProc hook at all.
+            this.mRawKeyState.RawState.Clear();
+
+            // If we captured all keys, there's no point of restoring individual keys.
+            this.mNextCaptureSet.Clear();
+        }
+
+        if (this.mCaptureState.Has(State.CaptureAllSingleFrame))
+        {
+            // Remove CaptureAllSingleFrame and Add RestoreAllOnNextFrame
+            this.mCaptureState = (this.mCaptureState | State.RestoreAllOnNextFrame) &
+                                 ~State.CaptureAllSingleFrame;
+            this.mRawKeyState.RawState.Clear();
+        }
+
+        // Capture individual keys
+        foreach (var vkCode in this.mNextCaptureSet)
+        {
+            if (!VirtualKeyExtensions.IsValidVirtualKey(vkCode))
+            {
+                continue;
+            }
+
+            // Release a key
+            this.mRawKeyState.SetState(vkCode, KeyStateFlag.None);
+            this.mNextRestoreSet.Add(vkCode);
+        }
+
+        this.mNextCaptureSet.Clear();
+    }
+}

Dalamud/Game/ClientState/Keys/KeyStateFlag.cs

@@ -0,0 +1,12 @@
+using System;
+
+namespace Dalamud.Game.ClientState.Keys;
+
+[Flags]
+public enum KeyStateFlag
+{
+    None = 0x00,
+    Down = 0x01,
+    JustPressed = 0x02,
+    JustReleased = 0x04,
+}

Dalamud/Game/ClientState/Keys/KeyStateIndex.cs

@@ -0,0 +1,68 @@
+using System;
+using System.Runtime.InteropServices;
+
+using Dalamud.Utility;
+using Serilog;
+
+namespace Dalamud.Game.ClientState.Keys;
+
+// TODO: this is probably safe to expose to plugins without major concerns
+// [PluginInterface]
+[ServiceManager.BlockingEarlyLoadedService]
+internal sealed class KeyStateIndex : IServiceType
+{
+    private const int InvalidKeyCode = 0;
+
+    private readonly unsafe byte* mVkIndex;
+
+    [ServiceManager.ServiceConstructor]
+    private unsafe KeyStateIndex(SigScanner sigScanner, ClientState clientState)
+    {
+        var moduleBaseAddress = sigScanner.Module.BaseAddress;
+        var addressResolver = clientState.AddressResolver;
+
+        this.mVkIndex = (byte*)(moduleBaseAddress + Marshal.ReadInt32(addressResolver.KeyboardStateIndexArray));
+        this.MaxValidKeyCode = this.AsSpan().Max();
+
+        Log.Verbose("VkIndex@{Address:X8}h (max={Max:X2}h)", (long)(nint)this.mVkIndex, this.MaxValidKeyCode);
+    }
+
+    /// <summary>
+    /// Gets the upper bound for the key code.
+    /// </summary>
+    public byte MaxValidKeyCode { get; }
+
+    /// <summary>
+    /// Translates a virtual key code from Windows into a key code which the game internally uses.
+    /// </summary>
+    /// <param name="vkCode">A virtual key to translate.</param>
+    /// <param name="keyCode">A reference to keyCode to receive the translated key code.</param>
+    /// <returns>
+    /// Returns true if this function successfully reads the key code, false otherwise.
+    /// </returns>
+    /// <remarks>
+    /// If this function returns false then the state of <see cref="keyCode"/> will be unspecified.
+    /// </remarks>
+    public bool TryGetKeyCode(ushort vkCode, out byte keyCode)
+    {
+        unsafe
+        {
+            keyCode = default;
+
+            if (!VirtualKeyExtensions.IsValidVirtualKey(vkCode))
+            {
+                return false;
+            }
+
+            keyCode = this.mVkIndex[vkCode];
+            if (keyCode == InvalidKeyCode)
+            {
+                return false;
+            }
+
+            return true;
+        }
+    }
+
+    private unsafe ReadOnlySpan<byte> AsSpan() => new(this.mVkIndex, VirtualKeyExtensions.MaxValidCode + 1);
+}