From fdf64896cb30b8e0ed5cdf519c59e5b53df3bc23 Mon Sep 17 00:00:00 2001
From: AndroidX Test Team <copybara-androidxtest@google.com>
Date: Tue, 1 Oct 2024 14:40:16 -0700
Subject: [PATCH] Adding the LocalSocketShellMain for the ShellExecutor to talk
 to the ShellMain.

PiperOrigin-RevId: 681170541
---
 .../test/services/shellexecutor/BUILD         |  1 +
 .../shellexecutor/LocalSocketShellMain.kt     | 89 +++++++++++++++++++
 2 files changed, 90 insertions(+)
 create mode 100644 services/shellexecutor/java/androidx/test/services/shellexecutor/LocalSocketShellMain.kt

diff --git a/services/shellexecutor/java/androidx/test/services/shellexecutor/BUILD b/services/shellexecutor/java/androidx/test/services/shellexecutor/BUILD
index 91bf1da4d..ab0d75c07 100644
--- a/services/shellexecutor/java/androidx/test/services/shellexecutor/BUILD
+++ b/services/shellexecutor/java/androidx/test/services/shellexecutor/BUILD
@@ -60,6 +60,7 @@ kt_android_library(
     srcs = [
         "BlockingPublish.java",
         "FileObserverShellMain.kt",
+        "LocalSocketShellMain.kt",
         "ShellCommand.java",
         "ShellCommandExecutor.java",
         "ShellCommandExecutorServer.java",
diff --git a/services/shellexecutor/java/androidx/test/services/shellexecutor/LocalSocketShellMain.kt b/services/shellexecutor/java/androidx/test/services/shellexecutor/LocalSocketShellMain.kt
new file mode 100644
index 000000000..e68d4c5d0
--- /dev/null
+++ b/services/shellexecutor/java/androidx/test/services/shellexecutor/LocalSocketShellMain.kt
@@ -0,0 +1,89 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.test.services.shellexecutor
+
+import android.util.Log
+import java.io.IOException
+import java.io.InputStream
+import java.io.OutputStream
+import java.util.concurrent.Executors
+import kotlin.time.Duration.Companion.milliseconds
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.runInterruptible
+
+/** Variant of ShellMain that uses a LocalSocket to communicate with the client. */
+class LocalSocketShellMain {
+
+  suspend fun run(args: Array<String>): Int {
+    val scope = CoroutineScope(Executors.newCachedThreadPool().asCoroutineDispatcher())
+    val server = ShellCommandLocalSocketExecutorServer(scope = scope)
+    server.start()
+
+    val processArgs = args.toMutableList()
+    processArgs.addAll(
+      processArgs.size - 1,
+      listOf("-e", ShellExecSharedConstants.BINDER_KEY, server.binderKey()),
+    )
+    val pb = ProcessBuilder(processArgs.toList())
+
+    val exitCode: Int
+
+    try {
+      val process = pb.start()
+
+      val stdinCopier = scope.launch { copyStream("stdin", System.`in`, process.outputStream) }
+      val stdoutCopier = scope.launch { copyStream("stdout", process.inputStream, System.out) }
+      val stderrCopier = scope.launch { copyStream("stderr", process.errorStream, System.err) }
+
+      runInterruptible { process.waitFor() }
+      exitCode = process.exitValue()
+
+      stdinCopier.cancel() // System.`in`.close() does not force input.read() to return
+      stdoutCopier.join()
+      stderrCopier.join()
+    } finally {
+      server.stop(100.milliseconds)
+    }
+    return exitCode
+  }
+
+  suspend fun copyStream(name: String, input: InputStream, output: OutputStream) {
+    val buf = ByteArray(1024)
+    try {
+      while (true) {
+        val size = input.read(buf)
+        if (size == -1) break
+        output.write(buf, 0, size)
+      }
+      output.flush()
+    } catch (x: IOException) {
+      Log.e(TAG, "IOException on $name. Terminating.", x)
+    }
+  }
+
+  companion object {
+    private const val TAG = "LocalSocketShellMain"
+
+    @JvmStatic
+    public fun main(args: Array<String>) {
+      System.exit(runBlocking { LocalSocketShellMain().run(args) })
+    }
+  }
+}