Merge "Implement a command handler for the ProtoLog service" into main
diff --git a/core/java/com/android/internal/protolog/ProtoLogCommandHandler.java b/core/java/com/android/internal/protolog/ProtoLogCommandHandler.java
new file mode 100644
index 0000000..3dab2e3
--- /dev/null
+++ b/core/java/com/android/internal/protolog/ProtoLogCommandHandler.java
@@ -0,0 +1,172 @@
+/*
+ * 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 com.android.internal.protolog;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.ShellCommand;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.io.PrintWriter;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+
+public class ProtoLogCommandHandler extends ShellCommand {
+    @NonNull
+    private final ProtoLogService mProtoLogService;
+    @Nullable
+    private final PrintWriter mPrintWriter;
+
+    public ProtoLogCommandHandler(@NonNull ProtoLogService protoLogService) {
+        this(protoLogService, null);
+    }
+
+    @VisibleForTesting
+    public ProtoLogCommandHandler(
+            @NonNull ProtoLogService protoLogService, @Nullable PrintWriter printWriter) {
+        this.mProtoLogService = protoLogService;
+        this.mPrintWriter = printWriter;
+    }
+
+    @Override
+    public int onCommand(String cmd) {
+        if (cmd == null) {
+            onHelp();
+            return 0;
+        }
+
+        return switch (cmd) {
+            case "groups" -> handleGroupsCommands(getNextArg());
+            case "logcat" -> handleLogcatCommands(getNextArg());
+            default -> handleDefaultCommands(cmd);
+        };
+    }
+
+    @Override
+    public void onHelp() {
+        PrintWriter pw = getOutPrintWriter();
+        pw.println("ProtoLog commands:");
+        pw.println("  help");
+        pw.println("    Print this help text.");
+        pw.println();
+        pw.println("  groups (list | status)");
+        pw.println("    list - lists all ProtoLog groups registered with ProtoLog service");
+        pw.println("    status <group> - print the status of a ProtoLog group");
+        pw.println();
+        pw.println("  logcat (enable | disable) <group>");
+        pw.println("    enable or disable ProtoLog to logcat");
+        pw.println();
+    }
+
+    @NonNull
+    @Override
+    public PrintWriter getOutPrintWriter() {
+        if (mPrintWriter != null) {
+            return mPrintWriter;
+        }
+
+        return super.getOutPrintWriter();
+    }
+
+    private int handleGroupsCommands(@Nullable String cmd) {
+        PrintWriter pw = getOutPrintWriter();
+
+        if (cmd == null) {
+            pw.println("Incomplete command. Use 'cmd protolog help' for guidance.");
+            return 0;
+        }
+
+        switch (cmd) {
+            case "list": {
+                final String[] availableGroups = mProtoLogService.getGroups();
+                if (availableGroups.length == 0) {
+                    pw.println("No ProtoLog groups registered with ProtoLog service.");
+                    return 0;
+                }
+
+                pw.println("ProtoLog groups registered with service:");
+                for (String group : availableGroups) {
+                    pw.println("- " + group);
+                }
+
+                return 0;
+            }
+            case "status": {
+                final String group = getNextArg();
+
+                if (group == null) {
+                    pw.println("Incomplete command. Use 'cmd protolog help' for guidance.");
+                    return 0;
+                }
+
+                pw.println("ProtoLog group " + group + "'s status:");
+
+                if (!Set.of(mProtoLogService.getGroups()).contains(group)) {
+                    pw.println("UNREGISTERED");
+                    return 0;
+                }
+
+                pw.println("LOG_TO_LOGCAT = " + mProtoLogService.isLoggingToLogcat(group));
+                return 0;
+            }
+            default: {
+                pw.println("Unknown command: " + cmd);
+                return -1;
+            }
+        }
+    }
+
+    private int handleLogcatCommands(@Nullable String cmd) {
+        PrintWriter pw = getOutPrintWriter();
+
+        if (cmd == null || peekNextArg() == null) {
+            pw.println("Incomplete command. Use 'cmd protolog help' for guidance.");
+            return 0;
+        }
+
+        switch (cmd) {
+            case "enable" -> {
+                mProtoLogService.enableProtoLogToLogcat(processGroups());
+                return 0;
+            }
+            case "disable" -> {
+                mProtoLogService.disableProtoLogToLogcat(processGroups());
+                return 0;
+            }
+            default -> {
+                pw.println("Unknown command: " + cmd);
+                return -1;
+            }
+        }
+    }
+
+    @NonNull
+    private String[] processGroups() {
+        if (getRemainingArgsCount() == 0) {
+            return mProtoLogService.getGroups();
+        }
+
+        final List<String> groups = new ArrayList<>();
+        while (getRemainingArgsCount() > 0) {
+            groups.add(getNextArg());
+        }
+
+        return groups.toArray(new String[0]);
+    }
+}
diff --git a/core/java/com/android/internal/protolog/ProtoLogService.java b/core/java/com/android/internal/protolog/ProtoLogService.java
index 0b13a1a..2333a06 100644
--- a/core/java/com/android/internal/protolog/ProtoLogService.java
+++ b/core/java/com/android/internal/protolog/ProtoLogService.java
@@ -33,6 +33,8 @@
 import android.annotation.SystemService;
 import android.content.Context;
 import android.os.RemoteException;
+import android.os.ResultReceiver;
+import android.os.ShellCallback;
 import android.os.SystemClock;
 import android.tracing.perfetto.DataSourceParams;
 import android.tracing.perfetto.InitArguments;
@@ -43,6 +45,7 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 
+import java.io.FileDescriptor;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -224,6 +227,14 @@
         registerGroups(client, args.getGroups(), args.getGroupsDefaultLogcatStatus());
     }
 
+    @Override
+    public void onShellCommand(@Nullable FileDescriptor in, @Nullable FileDescriptor out,
+            @Nullable FileDescriptor err, @NonNull String[] args, @Nullable ShellCallback callback,
+            @NonNull ResultReceiver resultReceiver) throws RemoteException {
+        new ProtoLogCommandHandler(this)
+                .exec(this, in, out, err, args, callback, resultReceiver);
+    }
+
     /**
      * Get the list of groups clients have registered to the protolog service.
      * @return The list of ProtoLog groups registered with this service.
diff --git a/tests/Internal/src/com/android/internal/protolog/ProtoLogCommandHandlerTest.java b/tests/Internal/src/com/android/internal/protolog/ProtoLogCommandHandlerTest.java
new file mode 100644
index 0000000..e3ec62d
--- /dev/null
+++ b/tests/Internal/src/com/android/internal/protolog/ProtoLogCommandHandlerTest.java
@@ -0,0 +1,207 @@
+/*
+ * 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 com.android.internal.protolog;
+
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.contains;
+import static org.mockito.ArgumentMatchers.endsWith;
+import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.times;
+
+import android.platform.test.annotations.Presubmit;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
+
+/**
+ * Test class for {@link ProtoLogImpl}.
+ */
+@Presubmit
+@RunWith(MockitoJUnitRunner.class)
+public class ProtoLogCommandHandlerTest {
+
+    @Mock
+    ProtoLogService mProtoLogService;
+    @Mock
+    PrintWriter mPrintWriter;
+
+    @Test
+    public void printsHelpForAllAvailableCommands() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.onHelp();
+        validateOnHelpPrinted();
+    }
+
+    @Test
+    public void printsHelpIfCommandIsNull() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.onCommand(null);
+        validateOnHelpPrinted();
+    }
+
+    @Test
+    public void handlesGroupListCommand() {
+        Mockito.when(mProtoLogService.getGroups())
+                .thenReturn(new String[] {"MY_TEST_GROUP", "MY_OTHER_GROUP"});
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "groups", "list" });
+
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("MY_TEST_GROUP"));
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("MY_OTHER_GROUP"));
+    }
+
+    @Test
+    public void handlesIncompleteGroupsCommand() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "groups" });
+
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("Incomplete command"));
+    }
+
+    @Test
+    public void handlesGroupStatusCommand() {
+        Mockito.when(mProtoLogService.getGroups()).thenReturn(new String[] {"MY_GROUP"});
+        Mockito.when(mProtoLogService.isLoggingToLogcat("MY_GROUP")).thenReturn(true);
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "groups", "status", "MY_GROUP" });
+
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("MY_GROUP"));
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("LOG_TO_LOGCAT = true"));
+    }
+
+    @Test
+    public void handlesGroupStatusCommandOfUnregisteredGroups() {
+        Mockito.when(mProtoLogService.getGroups()).thenReturn(new String[] {});
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "groups", "status", "MY_GROUP" });
+
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("MY_GROUP"));
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("UNREGISTERED"));
+    }
+
+    @Test
+    public void handlesGroupStatusCommandWithNoGroups() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "groups", "status" });
+
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("Incomplete command"));
+    }
+
+    @Test
+    public void handlesIncompleteLogcatCommand() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "logcat" });
+
+        Mockito.verify(mPrintWriter, times(1))
+                .println(contains("Incomplete command"));
+    }
+
+    @Test
+    public void handlesLogcatEnableCommand() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "logcat", "enable", "MY_GROUP" });
+        Mockito.verify(mProtoLogService).enableProtoLogToLogcat("MY_GROUP");
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "logcat", "enable", "MY_GROUP", "MY_OTHER_GROUP" });
+        Mockito.verify(mProtoLogService)
+                .enableProtoLogToLogcat("MY_GROUP", "MY_OTHER_GROUP");
+    }
+
+    @Test
+    public void handlesLogcatDisableCommand() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "logcat", "disable", "MY_GROUP" });
+        Mockito.verify(mProtoLogService).disableProtoLogToLogcat("MY_GROUP");
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "logcat", "disable", "MY_GROUP", "MY_OTHER_GROUP" });
+        Mockito.verify(mProtoLogService)
+                .disableProtoLogToLogcat("MY_GROUP", "MY_OTHER_GROUP");
+    }
+
+    @Test
+    public void handlesLogcatEnableCommandWithNoGroups() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "logcat", "enable" });
+        Mockito.verify(mPrintWriter).println(contains("Incomplete command"));
+    }
+
+    @Test
+    public void handlesLogcatDisableCommandWithNoGroups() {
+        final ProtoLogCommandHandler cmdHandler =
+                new ProtoLogCommandHandler(mProtoLogService, mPrintWriter);
+
+        cmdHandler.exec(mProtoLogService, FileDescriptor.in, FileDescriptor.out, FileDescriptor.err,
+                new String[] { "logcat", "disable" });
+        Mockito.verify(mPrintWriter).println(contains("Incomplete command"));
+    }
+
+    private void validateOnHelpPrinted() {
+        Mockito.verify(mPrintWriter, times(1)).println(endsWith("help"));
+        Mockito.verify(mPrintWriter, times(1))
+                .println(endsWith("groups (list | status)"));
+        Mockito.verify(mPrintWriter, times(1))
+                .println(endsWith("logcat (enable | disable) <group>"));
+        Mockito.verify(mPrintWriter, atLeast(0)).println(anyString());
+    }
+}