Update SuRequest handler

app/src/main/java/com/topjohnwu/magisk/core/su/SuRequestHandler.kt

@@ -13,12 +13,11 @@ import com.topjohnwu.magisk.core.magiskdb.PolicyDao
 import com.topjohnwu.magisk.core.model.MagiskPolicy
 import com.topjohnwu.magisk.core.model.toPolicy
 import com.topjohnwu.magisk.extensions.now
-import com.topjohnwu.superuser.Shell
-import com.topjohnwu.superuser.internal.UiThreadHandler
+import kotlinx.coroutines.*
 import timber.log.Timber
 import java.io.*
-import java.util.concurrent.Callable
 import java.util.concurrent.TimeUnit
+import java.util.concurrent.TimeUnit.SECONDS
 
 abstract class SuRequestHandler(
     private val packageManager: PackageManager,
@@ -33,35 +32,11 @@ abstract class SuRequestHandler(
 
     abstract fun onStart()
 
-    fun start(intent: Intent): Boolean {
+    suspend fun start(intent: Intent): Boolean {
         val name = intent.getStringExtra("socket") ?: return false
 
-        try {
-            if (Const.Version.atLeastCanary()) {
-                val server = LocalServerSocket(name)
-                val futureSocket = Shell.EXECUTOR.submit(Callable { server.accept() })
-                try {
-                    socket = futureSocket.get(1, TimeUnit.SECONDS)
-                } catch (e: Exception) {
-                    // Timeout or any IO errors
-                    throw e
-                } finally {
-                    server.close()
-                }
-            } else {
-                socket = LocalSocket()
-                socket.connect(LocalSocketAddress(name, LocalSocketAddress.Namespace.ABSTRACT))
-            }
-            output = DataOutputStream(BufferedOutputStream(socket.outputStream))
-            input = DataInputStream(BufferedInputStream(socket.inputStream))
-            val map = Shell.EXECUTOR.submit(Callable { readRequest() })
-                .runCatching { get(1, TimeUnit.SECONDS) }.getOrNull() ?: return false
-            val uid = map["uid"]?.toIntOrNull() ?: return false
-            policy = uid.toPolicy(packageManager)
-        } catch (e: Exception) {
-            Timber.e(e)
+        if (!init(name))
             return false
-        }
 
         // Never allow com.topjohnwu.magisk (could be malware)
         if (policy.packageName == BuildConfig.APPLICATION_ID)
@@ -77,10 +52,48 @@ abstract class SuRequestHandler(
                 return true
             }
         }
-        UiThreadHandler.run { onStart() }
+
+        onStart()
         return true
     }
 
+    private suspend fun <T> Deferred<T>.timedAwait() : T? {
+        return withTimeoutOrNull(SECONDS.toMillis(1)) {
+            await()
+        }
+    }
+
+    private class SocketError : IOException()
+
+    private suspend fun init(name: String) = withContext(Dispatchers.IO) {
+        try {
+            if (Const.Version.atLeastCanary()) {
+                LocalServerSocket(name).use {
+                    socket = async { it.accept() }.timedAwait() ?: throw SocketError()
+                }
+            } else {
+                socket = LocalSocket()
+                socket.connect(LocalSocketAddress(name, LocalSocketAddress.Namespace.ABSTRACT))
+            }
+            output = DataOutputStream(BufferedOutputStream(socket.outputStream))
+            input = DataInputStream(BufferedInputStream(socket.inputStream))
+            val map = async { readRequest() }.timedAwait() ?: throw SocketError()
+            val uid = map["uid"]?.toIntOrNull() ?: throw SocketError()
+            policy = uid.toPolicy(packageManager)
+            true
+        } catch (e: Exception) {
+            when (e) {
+                is IOException, is PackageManager.NameNotFoundException -> {
+                    Timber.e(e)
+                    if (::socket.isInitialized)
+                        socket.close()
+                    false
+                }
+                else -> throw e  // Unexpected error
+            }
+        }
+    }
+
     fun respond(action: Int, time: Int) {
         val until = if (time > 0)
             TimeUnit.MILLISECONDS.toSeconds(now) + TimeUnit.MINUTES.toSeconds(time.toLong())
@@ -91,7 +104,7 @@ abstract class SuRequestHandler(
         policy.until = until
         policy.uid = policy.uid % 100000 + Const.USER_ID * 100000
 
-        Shell.EXECUTOR.submit {
+        GlobalScope.launch(Dispatchers.IO) {
             try {
                 output.writeInt(policy.policy)
                 output.flush()

app/src/main/java/com/topjohnwu/magisk/ui/surequest/SuRequestActivity.kt

@@ -12,12 +12,10 @@ import com.topjohnwu.magisk.R
 import com.topjohnwu.magisk.core.su.SuCallbackHandler
 import com.topjohnwu.magisk.core.su.SuCallbackHandler.REQUEST
 import com.topjohnwu.magisk.databinding.ActivityRequestBinding
-import com.topjohnwu.magisk.extensions.subscribeK
 import com.topjohnwu.magisk.model.events.DieEvent
 import com.topjohnwu.magisk.model.events.ViewActionEvent
 import com.topjohnwu.magisk.model.events.ViewEvent
 import com.topjohnwu.magisk.ui.base.BaseUIActivity
-import io.reactivex.Single
 import org.koin.androidx.viewmodel.ext.android.viewModel
 
 open class SuRequestActivity : BaseUIActivity<SuRequestViewModel, ActivityRequestBinding>() {
@@ -39,11 +37,7 @@ open class SuRequestActivity : BaseUIActivity<SuRequestViewModel, ActivityReques
         super.onCreate(savedInstanceState)
 
         fun showRequest() {
-            Single.fromCallable {
-                viewModel.handleRequest(intent)
-            }.subscribeK {
-                if (!it) finish()
-            }
+            viewModel.handleRequest(intent)
         }
 
         fun runHandler(action: String?) {

app/src/main/java/com/topjohnwu/magisk/ui/surequest/SuRequestViewModel.kt

@@ -6,6 +6,7 @@ import android.content.pm.PackageManager
 import android.content.res.Resources
 import android.graphics.drawable.Drawable
 import android.os.CountDownTimer
+import androidx.lifecycle.viewModelScope
 import com.topjohnwu.magisk.R
 import com.topjohnwu.magisk.core.Config
 import com.topjohnwu.magisk.core.magiskdb.PolicyDao
@@ -18,6 +19,7 @@ import com.topjohnwu.magisk.model.events.DieEvent
 import com.topjohnwu.magisk.ui.base.BaseViewModel
 import com.topjohnwu.magisk.utils.KObservableField
 import com.topjohnwu.superuser.internal.UiThreadHandler
+import kotlinx.coroutines.launch
 import me.tatarka.bindingcollectionadapter2.BindingListViewAdapter
 import me.tatarka.bindingcollectionadapter2.ItemBinding
 import java.util.concurrent.TimeUnit.SECONDS
@@ -72,8 +74,11 @@ class SuRequestViewModel(
         return false
     }
 
-    fun handleRequest(intent: Intent): Boolean {
-        return handler.start(intent)
+    fun handleRequest(intent: Intent) {
+        viewModelScope.launch {
+            if (!handler.start(intent))
+                DieEvent().publish()
+        }
     }
 
     private inner class Handler : SuRequestHandler(pm, policyDB) {