Fix leaks in TLSWebViewClient (#5445)

* Introduce SafeKeyChainAliasCallback

* Add missing request ignore when activity is null

Author: TimoPtr

app/src/main/kotlin/io/homeassistant/companion/android/util/TLSWebViewClient.kt

@@ -10,19 +10,18 @@ import android.webkit.ClientCertRequest
 import android.webkit.WebView
 import android.webkit.WebViewClient
 import androidx.annotation.RequiresApi
-import androidx.appcompat.app.AppCompatActivity
-import androidx.lifecycle.lifecycleScope
 import io.homeassistant.companion.android.common.data.keychain.KeyChainRepository
-import java.security.Principal
+import java.lang.ref.WeakReference
 import java.security.PrivateKey
 import java.security.cert.CertificateException
 import java.security.cert.X509Certificate
-import javax.inject.Inject
-import javax.inject.Named
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
+import timber.log.Timber
 
-open class TLSWebViewClient @Inject constructor(@Named("keyChainRepository") private var keyChainRepository: KeyChainRepository) : WebViewClient() {
-
+open class TLSWebViewClient(private var keyChainRepository: KeyChainRepository) : WebViewClient() {
     var isTLSClientAuthNeeded = false
         private set
 
@@ -39,17 +38,14 @@ open class TLSWebViewClient @Inject constructor(@Named("keyChainRepository") pri
         if (context == null) {
             return null
         } else if (context is ContextWrapper) {
-            return if (context is Activity) {
-                context
-            } else {
-                getActivity(context.baseContext)
-            }
+            return context as? Activity ?: getActivity(context.baseContext)
         }
         return null
     }
 
     @RequiresApi(Build.VERSION_CODES.M)
     override fun onReceivedClientCertRequest(view: WebView, request: ClientCertRequest) {
+        Timber.d("onReceivedClientCertRequest invoked looking for cert in local storage or ask the user for it")
         // Let the WebViewActivity know the endpoint requires TLS Client Auth
         isTLSClientAuthNeeded = true
 
@@ -75,42 +71,37 @@ open class TLSWebViewClient @Inject constructor(@Named("keyChainRepository") pri
                         // If no key is available, then the user must be prompt for a key
                         // The whole operation is wrapped in the selectPrivateKey method but caution as it must occurs outside of the main thread
                         // see: https://developer.android.com/reference/android/security/KeyChain#getPrivateKey(android.content.Context,%20java.lang.String)
-                        selectClientCert(activity, request.principals, request)
+                        selectClientCert(activity, request)
                     }
                 }
+            } else {
+                request.ignore()
             }
         }
     }
 
     @RequiresApi(Build.VERSION_CODES.M)
-    private fun selectClientCert(activity: Activity, principals: Array<Principal>?, request: ClientCertRequest) {
-        require(activity is AppCompatActivity)
-
-        val kcac = KeyChainAliasCallback { alias ->
-            if (alias != null) {
-                activity.lifecycleScope.launch {
-                    // Load the key and the chain
-                    keyChainRepository.load(activity.applicationContext, alias)
-
-                    key = keyChainRepository.getPrivateKey()
-                    chain = keyChainRepository.getCertificateChain()
-
-                    // If we got the key and the cert
-                    if (key == null || chain == null) {
-                        // Either the user didn't choose a key or no key was available
-                        hasUserDeniedAccess = true
-                    }
-
+    private fun selectClientCert(activity: Activity, request: ClientCertRequest) {
+        // prompt the user for a key
+        KeyChain.choosePrivateKeyAlias(
+            activity,
+            SafeKeyChainAliasCallback(keyChainRepository, activity.applicationContext) { key, chain ->
+                if (key == null || chain == null) {
+                    hasUserDeniedAccess = true
+                    request.ignore()
+                } else {
                     checkChainValidity()
+                    this.key = key
+                    this.chain = chain
                     request.proceed(key, chain)
                 }
-            } else {
-                request.proceed(key, chain)
-            }
-        }
-
-        // prompt the user for a key
-        KeyChain.choosePrivateKeyAlias(activity, kcac, arrayOf<String>(), principals, null, null)
+            },
+            request.keyTypes,
+            request.principals,
+            request.host,
+            request.port,
+            null,
+        )
     }
 
     private fun checkChainValidity() {
@@ -125,3 +116,39 @@ open class TLSWebViewClient @Inject constructor(@Named("keyChainRepository") pri
         }
     }
 }
+
+/**
+ * Addresses a potential memory leak with [KeyChain.choosePrivateKeyAlias].
+ *
+ * [KeyChain.choosePrivateKeyAlias] holds a strong reference to its callback even after
+ * invocation. To prevent this callback from leaking its capturing context (e.g., a WebView),
+ * this wrapper stores the actual result consumer ([onResult]) in a [WeakReference].
+ *
+ * If the consumer (e.g., WebView) is destroyed before the user selects a key,
+ * the [WeakReference] will allow it to be garbage collected, and the result will not be
+ * delivered to the (now-gone) consumer. The user's selection is intended to be
+ * handled independently by [KeyChainRepository.load] within its coroutine scope,
+ * ensuring the choice is persisted even if the initial UI component is gone.
+ */
+private class SafeKeyChainAliasCallback(
+    private var keyChainRepository: KeyChainRepository,
+    context: Context,
+    onResult: (key: PrivateKey?, chain: Array<X509Certificate>?) -> Unit,
+) : KeyChainAliasCallback {
+    private val ioScope: CoroutineScope = CoroutineScope(Dispatchers.IO + Job())
+    private val context = context.applicationContext
+    private val onResult = WeakReference(onResult)
+
+    override fun alias(alias: String?) {
+        if (alias != null) {
+            ioScope.launch {
+                keyChainRepository.load(context, alias)
+                val key = keyChainRepository.getPrivateKey()
+                val chain = keyChainRepository.getCertificateChain()
+                onResult.get()?.invoke(key, chain)
+            }
+        } else {
+            onResult.get()?.invoke(null, null)
+        }
+    }
+}