From ffd42de7f1f1370425ea051f4602588065ae83fb Mon Sep 17 00:00:00 2001
From: rhunk <101876869+rhunk@users.noreply.github.com>
Date: Fri, 29 Dec 2023 01:01:28 +0100
Subject: [PATCH] fix(core): SuspendLocationUpdates

---
 .../impl/global/SuspendLocationUpdates.kt     | 39 +++++++++++++++++--
 1 file changed, 35 insertions(+), 4 deletions(-)

diff --git a/core/src/main/kotlin/me/rhunk/snapenhance/core/features/impl/global/SuspendLocationUpdates.kt b/core/src/main/kotlin/me/rhunk/snapenhance/core/features/impl/global/SuspendLocationUpdates.kt
index 767a13e3c..292105348 100644
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/features/impl/global/SuspendLocationUpdates.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/features/impl/global/SuspendLocationUpdates.kt
@@ -10,6 +10,7 @@ import me.rhunk.snapenhance.core.ui.ViewAppearanceHelper
 import me.rhunk.snapenhance.core.util.hook.HookStage
 import me.rhunk.snapenhance.core.util.hook.hook
 import me.rhunk.snapenhance.core.util.ktx.getId
+import java.util.WeakHashMap
 
 //TODO: bridge shared preferences
 class SuspendLocationUpdates : BridgeFileFeature(
@@ -17,15 +18,39 @@ class SuspendLocationUpdates : BridgeFileFeature(
     loadParams = FeatureLoadParams.INIT_SYNC or FeatureLoadParams.ACTIVITY_CREATE_SYNC,
     bridgeFileType = BridgeFileType.SUSPEND_LOCATION_STATE
 ) {
+    private val streamSendHandlerInstanceMap = WeakHashMap<Any, () -> Unit>()
     private val isEnabled get() = context.config.global.suspendLocationUpdates.get()
+
     override fun init() {
         if (!isEnabled) return
         reload()
 
-        context.classCache.unifiedGrpcService.hook("bidiStreamingCall", HookStage.BEFORE) { param ->
-            val uri = param.arg<String>(0)
-            if (uri == "/snapchat.valis.Valis/Communicate" && exists("true")) {
-                param.setResult(null)
+        findClass("com.snapchat.client.grpc.ClientStreamSendHandler\$CppProxy").hook("send", HookStage.BEFORE) { param ->
+            if (param.nullableThisObject<Any>() !in streamSendHandlerInstanceMap) return@hook
+            if (!exists("true")) return@hook
+            param.setResult(null)
+        }
+
+        context.classCache.unifiedGrpcService.apply {
+            hook("unaryCall", HookStage.BEFORE) { param ->
+                val uri = param.arg<String>(0)
+                if (exists("true") && uri == "/snapchat.valis.Valis/SendClientUpdate") {
+                    param.setResult(null)
+                }
+            }
+
+            hook("bidiStreamingCall", HookStage.AFTER) { param ->
+                val uri = param.arg<String>(0)
+                if (uri != "/snapchat.valis.Valis/Communicate") return@hook
+                param.getResult()?.let { instance ->
+                    streamSendHandlerInstanceMap[instance] = {
+                        runCatching {
+                            instance::class.java.methods.first { it.name == "closeStream" }.invoke(instance)
+                        }.onFailure {
+                            context.log.error("Failed to close stream send handler instance", it)
+                        }
+                    }
+                }
             }
         }
     }
@@ -47,6 +72,12 @@ class SuspendLocationUpdates : BridgeFileFeature(
                         ViewGroup.LayoutParams.WRAP_CONTENT
                     )
                     setOnCheckedChangeListener { _, isChecked ->
+                        if (isChecked) {
+                            streamSendHandlerInstanceMap.entries.removeIf { (_, closeStream) ->
+                                closeStream()
+                                true
+                            }
+                        }
                         setState("true", isChecked)
                     }
                 })