From d2b42e29c66fc0f6411ce64fa7e14e1b5fcc1c3f Mon Sep 17 00:00:00 2001
From: Vaxry <vaxry@vaxry.net>
Date: Fri, 22 Mar 2024 18:45:24 +0000
Subject: [PATCH] IME: fix crashes with destroyed text-inputs

ref #5189
---
 src/helpers/WLClasses.cpp               | 36 +++++++++++++++++++++++++
 src/helpers/WLClasses.hpp               |  4 +++
 src/managers/input/InputMethodRelay.cpp | 21 ++++++++-------
 src/managers/input/InputMethodRelay.hpp |  1 +
 src/protocols/TextInputV1.hpp           |  2 --
 5 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/src/helpers/WLClasses.cpp b/src/helpers/WLClasses.cpp
index d46f7e40..1c14567d 100644
--- a/src/helpers/WLClasses.cpp
+++ b/src/helpers/WLClasses.cpp
@@ -244,3 +244,39 @@ void SKeyboard::updateXKBTranslationState(xkb_keymap* const keymap) {
     xkb_keymap_unref(NEWKEYMAP);
     xkb_context_unref(PCONTEXT);
 }
+
+void STextInput::setFocusedSurface(wlr_surface* pSurface) {
+    focusedSurface = pSurface;
+
+    hyprListener_surfaceUnmapped.removeCallback();
+    hyprListener_surfaceDestroyed.removeCallback();
+
+    if (!pSurface)
+        return;
+
+    hyprListener_surfaceUnmapped.initCallback(
+        &pSurface->events.unmap,
+        [this](void* owner, void* data) {
+            if (!focusedSurface)
+                return;
+
+            focusedSurface = nullptr;
+            hyprListener_surfaceUnmapped.removeCallback();
+            hyprListener_surfaceDestroyed.removeCallback();
+            g_pInputManager->m_sIMERelay.removeSurfaceToPTI(this);
+        },
+        this, "STextInput");
+
+    hyprListener_surfaceDestroyed.initCallback(
+        &pSurface->events.destroy,
+        [this](void* owner, void* data) {
+            if (!focusedSurface)
+                return;
+
+            focusedSurface = nullptr;
+            hyprListener_surfaceUnmapped.removeCallback();
+            hyprListener_surfaceDestroyed.removeCallback();
+            g_pInputManager->m_sIMERelay.removeSurfaceToPTI(this);
+        },
+        this, "STextInput");
+}
\ No newline at end of file
diff --git a/src/helpers/WLClasses.hpp b/src/helpers/WLClasses.hpp
index 806bf778..210405f6 100644
--- a/src/helpers/WLClasses.hpp
+++ b/src/helpers/WLClasses.hpp
@@ -295,10 +295,14 @@ struct STextInput {
     STextInputV1*      pV1Input       = nullptr;
     wlr_surface*       focusedSurface = nullptr;
 
+    void               setFocusedSurface(wlr_surface* pSurface);
+
     DYNLISTENER(textInputEnable);
     DYNLISTENER(textInputDisable);
     DYNLISTENER(textInputCommit);
     DYNLISTENER(textInputDestroy);
+    DYNLISTENER(surfaceUnmapped);
+    DYNLISTENER(surfaceDestroyed);
 };
 
 struct SIMEKbGrab {
diff --git a/src/managers/input/InputMethodRelay.cpp b/src/managers/input/InputMethodRelay.cpp
index 44ba3e42..89454d5b 100644
--- a/src/managers/input/InputMethodRelay.cpp
+++ b/src/managers/input/InputMethodRelay.cpp
@@ -143,7 +143,7 @@ void CInputMethodRelay::onNewIME(wlr_input_method_v2* pIME) {
 }
 
 wlr_surface* CInputMethodRelay::focusedSurface(STextInput* pTI) {
-    return pTI->pWlrInput ? pTI->pWlrInput->focused_surface : pTI->pV1Input->focusedSurface;
+    return pTI->pWlrInput ? pTI->pWlrInput->focused_surface : pTI->focusedSurface;
 }
 
 void CInputMethodRelay::updateInputPopup(SIMEPopup* pPopup) {
@@ -346,8 +346,8 @@ void CInputMethodRelay::createNewTextInput(wlr_text_input_v3* pInput, STextInput
 
             // v1 only, map surface to PTI
             if (PINPUT->pV1Input) {
-                wlr_surface* pSurface  = wlr_surface_from_resource((wl_resource*)data);
-                PINPUT->focusedSurface = pSurface;
+                wlr_surface* pSurface = wlr_surface_from_resource((wl_resource*)data);
+                PINPUT->setFocusedSurface(pSurface);
                 setSurfaceToPTI(pSurface, PINPUT);
                 if (m_pFocusedSurface == pSurface)
                     onTextInputEnter(pSurface);
@@ -510,8 +510,8 @@ void CInputMethodRelay::onTextInputLeave(wlr_surface* pSurface) {
         wlr_text_input_v3_send_leave(ti->pWlrInput);
     else {
         zwp_text_input_v1_send_leave(ti->pV1Input->resourceImpl);
-        ti->pV1Input->focusedSurface = nullptr;
-        ti->pV1Input->active         = false;
+        ti->setFocusedSurface(nullptr);
+        ti->pV1Input->active = false;
     }
 }
 
@@ -527,23 +527,26 @@ void CInputMethodRelay::onTextInputEnter(wlr_surface* pSurface) {
         wlr_text_input_v3_send_enter(ti->pWlrInput, pSurface);
     else {
         zwp_text_input_v1_send_enter(ti->pV1Input->resourceImpl, pSurface->resource);
-        ti->pV1Input->focusedSurface = pSurface;
-        ti->pV1Input->active         = true;
+        ti->setFocusedSurface(pSurface);
+        ti->pV1Input->active = true;
     }
 }
 
 void CInputMethodRelay::setSurfaceToPTI(wlr_surface* pSurface, STextInput* pInput) {
     if (pSurface) {
         m_mSurfaceToTextInput[pSurface] = pInput;
-        pInput->focusedSurface          = pSurface;
+        pInput->setFocusedSurface(pSurface);
     }
 }
 
 void CInputMethodRelay::removeSurfaceToPTI(STextInput* pInput) {
     if (pInput->focusedSurface) {
         m_mSurfaceToTextInput.erase(pInput->focusedSurface);
-        pInput->focusedSurface = nullptr;
+        pInput->setFocusedSurface(nullptr);
+        return;
     }
+
+    std::erase_if(m_mSurfaceToTextInput, [pInput](const auto& el) { return el.second == pInput; });
 }
 
 STextInput* CInputMethodRelay::getTextInput(wlr_surface* pSurface) {
diff --git a/src/managers/input/InputMethodRelay.hpp b/src/managers/input/InputMethodRelay.hpp
index 39465cab..6f99cae2 100644
--- a/src/managers/input/InputMethodRelay.hpp
+++ b/src/managers/input/InputMethodRelay.hpp
@@ -62,4 +62,5 @@ class CInputMethodRelay {
     friend class CHyprRenderer;
     friend class CInputManager;
     friend class CTextInputV1ProtocolManager;
+    friend struct STextInput;
 };
diff --git a/src/protocols/TextInputV1.hpp b/src/protocols/TextInputV1.hpp
index b279763d..9362a52e 100644
--- a/src/protocols/TextInputV1.hpp
+++ b/src/protocols/TextInputV1.hpp
@@ -13,8 +13,6 @@ struct STextInputV1 {
 
     wl_resource* resourceImpl = nullptr;
 
-    wlr_surface* focusedSurface = nullptr;
-
     STextInput*  pTextInput = nullptr;
 
     wl_signal    sEnable;