Integration tests for LocalSocketShellMain.

PiperOrigin-RevId: 681170539

Author: copybara-androidxtest

services/CHANGELOG.md

@@ -11,6 +11,15 @@
 
 **New Features**
 
+* LocalSocketProtocol: a replacement for SpeakEasy.
+* ShellCommandLocalSocketClient: client that speaks LocalSocketProtocol.
+* ShellCommandLocalSocketExecutorServer: server that speaks LocalSocketProtocol.
+* LocalSocketShellMain: a replacement for ShellMain. Using this in place of
+  ShellMain avoids the use of the SpeakEasy protocol, so androidx.test.services
+  can be freely killed and restarted by the operating system without breaking
+  tests. The ShellExecutorFactory will automatically figure out which protocol
+  to use based on the binder key passed down by the variant of ShellMain.
+
 **Breaking Changes**
 
 **API Changes**

services/shellexecutor/java/androidx/test/services/shellexecutor/BUILD

@@ -29,15 +29,43 @@ kt_android_library(
     ],
 )
 
+proto_library(
+    name = "local_socket_protocol_pb",
+    srcs = ["local_socket_protocol.proto"],
+)
+
+java_lite_proto_library(
+    name = "local_socket_protocol_pb_java_proto_lite",
+    visibility = [
+        "//services/shellexecutor/javatests/androidx/test/services/shellexecutor:__subpackages__",
+    ],
+    deps = [":local_socket_protocol_pb"],
+)
+
+kt_android_library(
+    name = "local_socket_protocol",
+    srcs = ["LocalSocketProtocol.kt"],
+    visibility = [
+        "//services/shellexecutor/javatests/androidx/test/services/shellexecutor:__subpackages__",
+    ],
+    deps = [
+        ":local_socket_protocol_pb_java_proto_lite",
+        "@com_google_protobuf//:protobuf_javalite",
+        "@maven//:org_jetbrains_kotlinx_kotlinx_coroutines_core",
+    ],
+)
+
 kt_android_library(
     name = "exec_server",
     srcs = [
         "BlockingPublish.java",
         "FileObserverShellMain.kt",
+        "LocalSocketShellMain.kt",
         "ShellCommand.java",
         "ShellCommandExecutor.java",
         "ShellCommandExecutorServer.java",
         "ShellCommandFileObserverExecutorServer.kt",
+        "ShellCommandLocalSocketExecutorServer.kt",
         "ShellExecSharedConstants.java",
         "ShellMain.java",
     ],
@@ -46,6 +74,8 @@ kt_android_library(
     deps = [
         ":coroutine_file_observer",
         ":file_observer_protocol",
+        ":local_socket_protocol",
+        ":local_socket_protocol_pb_java_proto_lite",
         "//services/speakeasy/java/androidx/test/services/speakeasy:protocol",
         "//services/speakeasy/java/androidx/test/services/speakeasy/client",
         "//services/speakeasy/java/androidx/test/services/speakeasy/client:tool_connection",
@@ -62,17 +92,21 @@ kt_android_library(
         "ShellCommand.java",
         "ShellCommandClient.java",
         "ShellCommandFileObserverClient.kt",
+        "ShellCommandLocalSocketClient.kt",
         "ShellExecSharedConstants.java",
         "ShellExecutor.java",
         "ShellExecutorFactory.java",
         "ShellExecutorFileObserverImpl.kt",
         "ShellExecutorImpl.java",
+        "ShellExecutorLocalSocketImpl.kt",
     ],
     idl_srcs = ["Command.aidl"],
     visibility = [":export"],
     deps = [
         ":coroutine_file_observer",
         ":file_observer_protocol",
+        ":local_socket_protocol",
+        ":local_socket_protocol_pb_java_proto_lite",
         "//services/speakeasy/java/androidx/test/services/speakeasy:protocol",
         "//services/speakeasy/java/androidx/test/services/speakeasy/client",
         "//services/speakeasy/java/androidx/test/services/speakeasy/client:tool_connection",

services/shellexecutor/java/androidx/test/services/shellexecutor/LocalSocketProtocol.kt

@@ -0,0 +1,145 @@
+/*
+ * 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.net.LocalSocket
+import android.net.LocalSocketAddress
+import android.util.Log
+import androidx.test.services.shellexecutor.LocalSocketProtocolProto.RunCommandRequest
+import androidx.test.services.shellexecutor.LocalSocketProtocolProto.RunCommandResponse
+import com.google.protobuf.ByteString
+import java.io.IOException
+import java.net.URLDecoder
+import java.net.URLEncoder
+import kotlin.time.Duration
+
+/**
+ * Protocol for ShellCommandLocalSocketClient to talk to ShellCommandLocalSocketExecutorServer.
+ *
+ * Since androidx.test.services already includes the protobuf runtime, we aren't paying much extra
+ * for adding some more protos to ship back and forth, which is vastly easier to deal with than
+ * PersistableBundles (which don't even support ByteArray types).
+ *
+ * A conversation consists of a single RunCommandRequest from the client followed by a stream of
+ * RunCommandResponses from the server; the final response has an exit code.
+ */
+object LocalSocketProtocol {
+  /** Composes a RunCommandRequest and sends it over the LocalSocket. */
+  fun LocalSocket.sendRequest(
+    secret: String,
+    argv: List<String>,
+    env: Map<String, String>? = null,
+    timeout: Duration,
+  ) {
+    val builder = RunCommandRequest.newBuilder()
+    builder.setSecret(secret)
+    builder.addAllArgv(argv)
+    env?.forEach { (k, v) -> builder.putEnvironment(k, v) }
+    if (timeout.isInfinite() || timeout.isNegative() || timeout == Duration.ZERO) {
+      builder.setTimeoutMs(0) // <= 0 means no timeout
+    } else {
+      builder.setTimeoutMs(timeout.inWholeMilliseconds)
+    }
+    builder.build().writeDelimitedTo(outputStream)
+  }
+
+  /** Reads a RunCommandRequest from the LocalSocket. */
+  fun LocalSocket.readRequest(): RunCommandRequest {
+    return RunCommandRequest.parseDelimitedFrom(inputStream)!!
+  }
+
+  /** Composes a RunCommandResponse and sends it over the LocalSocket. */
+  fun LocalSocket.sendResponse(
+    buffer: ByteArray? = null,
+    size: Int = 0,
+    exitCode: Int? = null,
+  ): Boolean {
+    val builder = RunCommandResponse.newBuilder()
+    buffer?.let {
+      val bufferSize = if (size > 0) size else it.size
+      builder.buffer = ByteString.copyFrom(it, 0, bufferSize)
+    }
+    // Since we're currently stuck on a version of protobuf where we don't have hasExitCode(), we
+    // use a magic value to indicate that exitCode is not set. When we upgrade to a newer version
+    // of protobuf, we can obsolete this.
+    if (exitCode != null) {
+      builder.exitCode = exitCode
+    } else {
+      builder.exitCode = HAS_NOT_EXITED
+    }
+
+    try {
+      builder.build().writeDelimitedTo(outputStream)
+    } catch (x: IOException) {
+      // Sadly, the only way to discover that the client cut the connection is an exception that
+      // can only be distinguished by its text.
+      if (x.message.equals("Broken pipe")) {
+        Log.i(TAG, "LocalSocket stream closed early")
+      } else {
+        Log.w(TAG, "LocalSocket write failed", x)
+      }
+      return false
+    }
+    return true
+  }
+
+  /** Reads a RunCommandResponse from the LocalSocket. */
+  fun LocalSocket.readResponse(): RunCommandResponse? {
+    return RunCommandResponse.parseDelimitedFrom(inputStream)
+  }
+
+  /**
+   * Is this the end of the stream?
+   *
+   * Once we upgrade to a newer version of protobuf, we can switch to hasExitCode().
+   */
+  fun RunCommandResponse.hasExited() = exitCode != HAS_NOT_EXITED
+
+  /**
+   * Builds a binder key, given the server address and secret. Binder keys should be opaque outside
+   * this directory.
+   *
+   * The address can contain spaces, and since it gets passed through a command line, we need to
+   * encode it so it doesn't get split by argv. java.net.URLEncoder is conveniently available on all
+   * SDK versions.
+   */
+  @JvmStatic
+  fun LocalSocketAddress.asBinderKey(secret: String) = buildString {
+    append(":")
+    append(URLEncoder.encode(name, "UTF-8")) // Will convert any : to %3A
+    append(":")
+    append(URLEncoder.encode(secret, "UTF-8"))
+    append(":")
+  }
+
+  /** Extracts the address from a binder key. */
+  @JvmStatic
+  fun addressFromBinderKey(binderKey: String) =
+    LocalSocketAddress(URLDecoder.decode(binderKey.split(":")[1], "UTF-8"))
+
+  /** Extracts the secret from a binder key. */
+  @JvmStatic
+  fun secretFromBinderKey(binderKey: String) = URLDecoder.decode(binderKey.split(":")[2], "UTF-8")
+
+  /** Is this a valid binder key? */
+  @JvmStatic
+  fun isBinderKey(maybeKey: String) =
+    maybeKey.startsWith(':') && maybeKey.endsWith(':') && maybeKey.split(":").size == 4
+
+  const val TAG = "LocalSocketProtocol"
+  private const val HAS_NOT_EXITED = 0xCA7F00D
+}

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) })
+    }
+  }
+}