Add proto dump flag to services (2/2)

Adds new PROTO flag which requests services to dump sections in proto format. Modifies PriorityDumper helper class to parse proto arguments and set asProto flags. Registers WM and AM with proto dump supprt.

Bug: 67716082

Test: frameworks/base/services/tests/runtests.py -e class "com.android.server.utils.PriorityDumpTest"
Test: adb bugreport ~/tmp/bug.zip
Test: adb shell dumpsys window --proto
Test: adb shell dumpsys activity --proto

Change-Id: Icfc6857c8a9340110a43343734a27e48d0b5a229
diff --git a/core/java/android/os/IServiceManager.java b/core/java/android/os/IServiceManager.java
index 87c65ec..2176a78 100644
--- a/core/java/android/os/IServiceManager.java
+++ b/core/java/android/os/IServiceManager.java
@@ -45,13 +45,13 @@
      * Place a new @a service called @a name into the service
      * manager.
      */
-    void addService(String name, IBinder service, boolean allowIsolated, int dumpPriority)
-                throws RemoteException;
+    void addService(String name, IBinder service, boolean allowIsolated, int dumpFlags)
+            throws RemoteException;
 
     /**
      * Return a list of all currently running services.
      */
-    String[] listServices(int dumpPriority) throws RemoteException;
+    String[] listServices(int dumpFlags) throws RemoteException;
 
     /**
      * Assign a permission controller to the service manager.  After set, this
@@ -72,9 +72,13 @@
     /*
      * Must update values in IServiceManager.h
      */
-    int DUMP_PRIORITY_CRITICAL = 1 << 0;
-    int DUMP_PRIORITY_HIGH = 1 << 1;
-    int DUMP_PRIORITY_NORMAL = 1 << 2;
-    int DUMP_PRIORITY_ALL = DUMP_PRIORITY_CRITICAL | DUMP_PRIORITY_HIGH
-            | DUMP_PRIORITY_NORMAL;
+    /* Allows services to dump sections according to priorities. */
+    int DUMP_FLAG_PRIORITY_CRITICAL = 1 << 0;
+    int DUMP_FLAG_PRIORITY_HIGH = 1 << 1;
+    int DUMP_FLAG_PRIORITY_NORMAL = 1 << 2;
+    int DUMP_FLAG_PRIORITY_ALL = DUMP_FLAG_PRIORITY_CRITICAL | DUMP_FLAG_PRIORITY_HIGH
+            | DUMP_FLAG_PRIORITY_NORMAL;
+    /* Allows services to dump sections in protobuf format. */
+    int DUMP_FLAG_PROTO = 1 << 3;
+
 }
diff --git a/core/java/android/os/ServiceManager.java b/core/java/android/os/ServiceManager.java
index f41848f..42ec315 100644
--- a/core/java/android/os/ServiceManager.java
+++ b/core/java/android/os/ServiceManager.java
@@ -83,7 +83,7 @@
      * @param service the service object
      */
     public static void addService(String name, IBinder service) {
-        addService(name, service, false, IServiceManager.DUMP_PRIORITY_NORMAL);
+        addService(name, service, false, IServiceManager.DUMP_FLAG_PRIORITY_NORMAL);
     }
 
     /**
@@ -96,7 +96,7 @@
      * to access this service
      */
     public static void addService(String name, IBinder service, boolean allowIsolated) {
-        addService(name, service, allowIsolated, IServiceManager.DUMP_PRIORITY_NORMAL);
+        addService(name, service, allowIsolated, IServiceManager.DUMP_FLAG_PRIORITY_NORMAL);
     }
 
     /**
@@ -143,7 +143,7 @@
      */
     public static String[] listServices() {
         try {
-            return getIServiceManager().listServices(IServiceManager.DUMP_PRIORITY_ALL);
+            return getIServiceManager().listServices(IServiceManager.DUMP_FLAG_PRIORITY_ALL);
         } catch (RemoteException e) {
             Log.e(TAG, "error in listServices", e);
             return null;
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index e1e53b3..6de09b1 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -49,9 +49,10 @@
 import static android.net.NetworkPolicyManager.isProcStateAllowedWhileIdleOrPowerSaveMode;
 import static android.net.NetworkPolicyManager.isProcStateAllowedWhileOnRestrictBackground;
 import static android.os.Build.VERSION_CODES.N;
-import static android.os.IServiceManager.DUMP_PRIORITY_CRITICAL;
-import static android.os.IServiceManager.DUMP_PRIORITY_HIGH;
-import static android.os.IServiceManager.DUMP_PRIORITY_NORMAL;
+import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_CRITICAL;
+import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_HIGH;
+import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_NORMAL;
+import static android.os.IServiceManager.DUMP_FLAG_PROTO;
 import static android.os.Process.BLUETOOTH_UID;
 import static android.os.Process.FIRST_APPLICATION_UID;
 import static android.os.Process.FIRST_ISOLATED_UID;
@@ -714,30 +715,36 @@
      */
     private final PriorityDump.PriorityDumper mPriorityDumper = new PriorityDump.PriorityDumper() {
         @Override
-        public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args) {
-            doDump(fd, pw, new String[] {"activities"});
+        public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args,
+                boolean asProto) {
+            if (asProto) return;
+            doDump(fd, pw, new String[]{"activities"}, asProto);
         }
 
         @Override
-        public void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args) {
-            doDump(fd, pw, new String[] {"settings"});
-            doDump(fd, pw, new String[] {"intents"});
-            doDump(fd, pw, new String[] {"broadcasts"});
-            doDump(fd, pw, new String[] {"providers"});
-            doDump(fd, pw, new String[] {"permissions"});
-            doDump(fd, pw, new String[] {"services"});
-            doDump(fd, pw, new String[] {"recents"});
-            doDump(fd, pw, new String[] {"lastanr"});
-            doDump(fd, pw, new String[] {"starter"});
-            if (mAssociations.size() > 0) {
-                doDump(fd, pw, new String[] {"associations"});
+        public void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
+            if (asProto) {
+                doDump(fd, pw, new String[0], asProto);
+            } else {
+                doDump(fd, pw, new String[]{"settings"}, asProto);
+                doDump(fd, pw, new String[]{"intents"}, asProto);
+                doDump(fd, pw, new String[]{"broadcasts"}, asProto);
+                doDump(fd, pw, new String[]{"providers"}, asProto);
+                doDump(fd, pw, new String[]{"permissions"}, asProto);
+                doDump(fd, pw, new String[]{"services"}, asProto);
+                doDump(fd, pw, new String[]{"recents"}, asProto);
+                doDump(fd, pw, new String[]{"lastanr"}, asProto);
+                doDump(fd, pw, new String[]{"starter"}, asProto);
+                if (mAssociations.size() > 0) {
+                    doDump(fd, pw, new String[]{"associations"}, asProto);
+                }
+                doDump(fd, pw, new String[]{"processes"}, asProto);
             }
-            doDump(fd, pw, new String[] {"processes"});
         }
 
         @Override
-        public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
-            doDump(fd, pw, args);
+        public void dump(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
+            doDump(fd, pw, args, asProto);
         }
     };
 
@@ -2483,15 +2490,15 @@
     public void setSystemProcess() {
         try {
             ServiceManager.addService(Context.ACTIVITY_SERVICE, this, /* allowIsolated= */ true,
-                    DUMP_PRIORITY_CRITICAL | DUMP_PRIORITY_NORMAL);
+                    DUMP_FLAG_PRIORITY_CRITICAL | DUMP_FLAG_PRIORITY_NORMAL | DUMP_FLAG_PROTO);
             ServiceManager.addService(ProcessStats.SERVICE_NAME, mProcessStats);
             ServiceManager.addService("meminfo", new MemBinder(this), /* allowIsolated= */ false,
-                    DUMP_PRIORITY_HIGH | DUMP_PRIORITY_NORMAL);
+                    DUMP_FLAG_PRIORITY_HIGH | DUMP_FLAG_PRIORITY_NORMAL);
             ServiceManager.addService("gfxinfo", new GraphicsBinder(this));
             ServiceManager.addService("dbinfo", new DbBinder(this));
             if (MONITOR_CPU_USAGE) {
                 ServiceManager.addService("cpuinfo", new CpuBinder(this),
-                        /* allowIsolated= */ false, DUMP_PRIORITY_CRITICAL);
+                        /* allowIsolated= */ false, DUMP_FLAG_PRIORITY_CRITICAL);
             }
             ServiceManager.addService("permission", new PermissionController(this));
             ServiceManager.addService("processinfo", new ProcessInfoService(this));
@@ -2544,7 +2551,9 @@
         private final PriorityDump.PriorityDumper mPriorityDumper =
                 new PriorityDump.PriorityDumper() {
             @Override
-            public void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args) {
+            public void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args,
+                    boolean asProto) {
+                if (asProto) return;
                 mActivityManagerService.dumpApplicationMemoryUsage(fd, pw, "  ", args, false, null);
             }
         };
@@ -2594,7 +2603,9 @@
         private final PriorityDump.PriorityDumper mPriorityDumper =
                 new PriorityDump.PriorityDumper() {
             @Override
-            public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args) {
+            public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args,
+                    boolean asProto) {
+                if (asProto) return;
                 if (!DumpUtils.checkDumpAndUsageStatsPermission(mActivityManagerService.mContext,
                         "cpuinfo", pw)) return;
                 synchronized (mActivityManagerService.mProcessCpuTracker) {
@@ -14789,7 +14800,7 @@
     /**
      * Wrapper function to print out debug data filtered by specified arguments.
     */
-    private void doDump(FileDescriptor fd, PrintWriter pw, String[] args) {
+    private void doDump(FileDescriptor fd, PrintWriter pw, String[] args, boolean useProto) {
         if (!DumpUtils.checkDumpAndUsageStatsPermission(mContext, TAG, pw)) return;
 
         boolean dumpAll = false;
@@ -14798,7 +14809,6 @@
         boolean dumpCheckinFormat = false;
         boolean dumpVisibleStacksOnly = false;
         boolean dumpFocusedStackOnly = false;
-        boolean useProto = false;
         String dumpPackage = null;
 
         int opti = 0;
@@ -14832,8 +14842,6 @@
             } else if ("-h".equals(opt)) {
                 ActivityManagerShellCommand.dumpHelp(pw, true);
                 return;
-            } else if ("--proto".equals(opt)) {
-                useProto = true;
             } else {
                 pw.println("Unknown argument: " + opt + "; use -h for help");
             }
diff --git a/services/core/java/com/android/server/utils/PriorityDump.java b/services/core/java/com/android/server/utils/PriorityDump.java
index 054f156..fb92c2b8 100644
--- a/services/core/java/com/android/server/utils/PriorityDump.java
+++ b/services/core/java/com/android/server/utils/PriorityDump.java
@@ -16,12 +16,19 @@
 
 package com.android.server.utils;
 
+import android.annotation.IntDef;
+
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
 
 /**
  * Helper for {@link android.os.Binder#dump(java.io.FileDescriptor, String[])} that supports the
- * {@link #PRIORITY_ARG} argument.
+ * {@link #PRIORITY_ARG} and {@link #PROTO_ARG} arguments.
  * <p>
  * Typical usage:
  *
@@ -31,13 +38,25 @@
  private final PriorityDump.PriorityDumper mPriorityDumper = new PriorityDump.PriorityDumper() {
 
      @Override
-     public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args) {
-       pw.println("Donuts in the box: 1");
+     public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
+       if (asProto) {
+         ProtoOutputStream proto = new ProtoOutputStream(fd);
+         proto.write(SpringfieldProto.DONUTS, 1);
+         proto.flush();
+       } else {
+         pw.println("Donuts in the box: 1");
+       }
      }
 
      @Override
      public void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args) {
-       pw.println("Nuclear reactor status: DANGER - MELTDOWN IMMINENT");
+        if (asProto) {
+          ProtoOutputStream proto = new ProtoOutputStream(fd);
+          proto.write(SpringfieldProto.REACTOR_STATUS, DANGER_MELTDOWN_IMMINENT);
+          proto.flush();
+        } else {
+          pw.println("Nuclear reactor status: DANGER - MELTDOWN IMMINENT");
+        }
      }
   };
 
@@ -65,6 +84,9 @@
     $ adb shell dumpsys snpp --dump-priority NORMAL
     Nuclear reactor status: DANGER - MELTDOWN IMMINENT
 
+    $ adb shell dumpsys snpp --dump-priority CRITICAL --proto
+    //binary output
+
  * </code></pre>
  *
  *
@@ -85,95 +107,146 @@
 public final class PriorityDump {
 
     public static final String PRIORITY_ARG = "--dump-priority";
+    public static final String PROTO_ARG = "--proto";
 
     private PriorityDump() {
         throw new UnsupportedOperationException();
     }
 
+    /** Enum to switch through supported priority types */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef({PRIORITY_TYPE_INVALID, PRIORITY_TYPE_CRITICAL, PRIORITY_TYPE_HIGH,
+            PRIORITY_TYPE_NORMAL})
+    private @interface PriorityType { }
+    private static final int PRIORITY_TYPE_INVALID = 0;
+    private static final int PRIORITY_TYPE_CRITICAL = 1;
+    private static final int PRIORITY_TYPE_HIGH = 2;
+    private static final int PRIORITY_TYPE_NORMAL = 3;
+
     /**
-     * Parses {@code} and call the proper {@link PriorityDumper} method when the first argument is
-     * {@code --dump-priority}, stripping the priority and its type.
+     * Parses {@code args} matching {@code --dump-priority} and/or {@code --proto}. The matching
+     * arguments are stripped.
+     * <p>
+     * If priority args are passed as an argument, it will call the appropriate method and if proto
+     * args are passed then the {@code asProto} flag is set.
      * <p>
      * For example, if called as {@code --dump-priority HIGH arg1 arg2 arg3}, it will call
-     * <code>dumper.dumpHigh(fd, pw, {"arg1", "arg2", "arg3"}) </code>
+     * <code>dumper.dumpHigh(fd, pw, {"arg1", "arg2", "arg3"}, false) </code>
      * <p>
      * If the {@code --dump-priority} is not set, it calls
-     * {@link PriorityDumper#dump(FileDescriptor, PrintWriter, String[])} passing the whole
+     * {@link PriorityDumper#dump(FileDescriptor, PrintWriter, String[], boolean)} passing the whole
      * {@code args} instead.
      */
     public static void dump(PriorityDumper dumper, FileDescriptor fd, PrintWriter pw,
             String[] args) {
-        if (args != null && args.length >= 2 && args[0].equals(PRIORITY_ARG)) {
-            final String priority = args[1];
-            switch (priority) {
-                case "CRITICAL": {
-                    dumper.dumpCritical(fd, pw, getStrippedArgs(args));
-                    return;
+        boolean asProto = false;
+        @PriorityType int priority = PRIORITY_TYPE_INVALID;
+
+        if (args == null) {
+            dumper.dump(fd, pw, args, asProto);
+            return;
+        }
+
+        String[] strippedArgs = new String[args.length];
+        int strippedCount = 0;
+        for (int argIndex = 0; argIndex < args.length; argIndex++) {
+            if (args[argIndex].equals(PROTO_ARG)) {
+                asProto = true;
+            } else if (args[argIndex].equals(PRIORITY_ARG)) {
+                if (argIndex + 1 < args.length) {
+                    argIndex++;
+                    priority = getPriorityType(args[argIndex]);
                 }
-                case "HIGH": {
-                    dumper.dumpHigh(fd, pw, getStrippedArgs(args));
-                    return;
-                }
-                case "NORMAL": {
-                    dumper.dumpNormal(fd, pw, getStrippedArgs(args));
-                    return;
-                }
+            } else {
+                strippedArgs[strippedCount++] = args[argIndex];
             }
         }
-        dumper.dump(fd, pw, args);
+
+        if (strippedCount < args.length) {
+            strippedArgs = Arrays.copyOf(strippedArgs, strippedCount);
+        }
+
+        switch (priority) {
+            case PRIORITY_TYPE_CRITICAL: {
+                dumper.dumpCritical(fd, pw, strippedArgs, asProto);
+                return;
+            }
+            case PRIORITY_TYPE_HIGH: {
+                dumper.dumpHigh(fd, pw, strippedArgs, asProto);
+                return;
+            }
+            case PRIORITY_TYPE_NORMAL: {
+                dumper.dumpNormal(fd, pw, strippedArgs, asProto);
+                return;
+            }
+            default: {
+                dumper.dump(fd, pw, strippedArgs, asProto);
+                return;
+            }
+        }
     }
 
     /**
-     * Gets an array without the {@code --dump-priority PRIORITY} prefix.
+     * Converts priority argument type to enum.
      */
-    private static String[] getStrippedArgs(String[] args) {
-        final String[] stripped = new String[args.length - 2];
-        System.arraycopy(args, 2, stripped, 0, stripped.length);
-        return stripped;
+    private static @PriorityType int getPriorityType(String arg) {
+        switch (arg) {
+            case "CRITICAL": {
+                return PRIORITY_TYPE_CRITICAL;
+            }
+            case "HIGH": {
+                return PRIORITY_TYPE_HIGH;
+            }
+            case "NORMAL": {
+                return PRIORITY_TYPE_NORMAL;
+            }
+        }
+        return PRIORITY_TYPE_INVALID;
     }
 
     /**
      * Helper for {@link android.os.Binder#dump(java.io.FileDescriptor, String[])} that supports the
-     * {@link #PRIORITY_ARG} argument.
+     * {@link #PRIORITY_ARG} and {@link #PROTO_ARG} arguments.
      *
      * @hide
      */
-    public static interface PriorityDumper {
+    public interface PriorityDumper {
 
         /**
          * Dumps only the critical section.
          */
         @SuppressWarnings("unused")
-        default void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args) {
+        default void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args,
+                boolean asProto) {
         }
 
         /**
          * Dumps only the high-priority section.
          */
         @SuppressWarnings("unused")
-        default void dumpHigh(FileDescriptor fd, PrintWriter pw, String[] args) {
+        default void dumpHigh(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
         }
 
         /**
          * Dumps only the normal section.
          */
         @SuppressWarnings("unused")
-        default void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args) {
+        default void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
         }
 
         /**
          * Dumps all sections.
          * <p>
          * This method is called when
-         * {@link PriorityDump#dump(PriorityDumper, FileDescriptor, PrintWriter, String[])} is
-         * called without priority arguments. By default, it calls the 3 {@code dumpTYPE} methods,
-         * so sub-classes just need to implement the priority types they support.
+         * {@link PriorityDump#dump(PriorityDumper, FileDescriptor, PrintWriter, String[], boolean)}
+         * is called without priority arguments. By default, it calls the 3 {@code dumpTYPE}
+         * methods, so sub-classes just need to implement the priority types they support.
          */
         @SuppressWarnings("unused")
-        default void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
-            dumpCritical(fd, pw, args);
-            dumpHigh(fd, pw, args);
-            dumpNormal(fd, pw, args);
+        default void dump(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
+            dumpCritical(fd, pw, args, asProto);
+            dumpHigh(fd, pw, args, asProto);
+            dumpNormal(fd, pw, args, asProto);
         }
     }
 }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index f0da474..e2003c9 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -394,13 +394,14 @@
 
     private final PriorityDump.PriorityDumper mPriorityDumper = new PriorityDump.PriorityDumper() {
         @Override
-        public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args) {
-            doDump(fd, pw, new String[] {"-a"});
+        public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args,
+                boolean asProto) {
+            doDump(fd, pw, new String[] {"-a"}, asProto);
         }
 
         @Override
-        public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
-            doDump(fd, pw, args);
+        public void dump(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
+            doDump(fd, pw, args, asProto);
         }
     };
 
@@ -6825,10 +6826,9 @@
         PriorityDump.dump(mPriorityDumper, fd, pw, args);
     }
 
-    private void doDump(FileDescriptor fd, PrintWriter pw, String[] args) {
+    private void doDump(FileDescriptor fd, PrintWriter pw, String[] args, boolean useProto) {
         if (!DumpUtils.checkDumpPermission(mContext, TAG, pw)) return;
         boolean dumpAll = false;
-        boolean useProto = false;
 
         int opti = 0;
         while (opti < args.length) {
@@ -6839,8 +6839,6 @@
             opti++;
             if ("-a".equals(opt)) {
                 dumpAll = true;
-            } else if ("--proto".equals(opt)) {
-                useProto = true;
             } else if ("-h".equals(opt)) {
                 pw.println("Window manager dump options:");
                 pw.println("  [-a] [-h] [cmd] ...");
diff --git a/services/java/com/android/server/SystemServer.java b/services/java/com/android/server/SystemServer.java
index d9db22e..ab2b969 100644
--- a/services/java/com/android/server/SystemServer.java
+++ b/services/java/com/android/server/SystemServer.java
@@ -125,7 +125,10 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Future;
 
-import static android.os.IServiceManager.DUMP_PRIORITY_CRITICAL;
+import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_CRITICAL;
+import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_HIGH;
+import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_NORMAL;
+import static android.os.IServiceManager.DUMP_FLAG_PROTO;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 public final class SystemServer {
@@ -827,7 +830,7 @@
                     mFactoryTestMode != FactoryTest.FACTORY_TEST_LOW_LEVEL,
                     !mFirstBoot, mOnlyCore, new PhoneWindowManager());
             ServiceManager.addService(Context.WINDOW_SERVICE, wm, /* allowIsolated= */ false,
-                    DUMP_PRIORITY_CRITICAL);
+                    DUMP_FLAG_PRIORITY_CRITICAL | DUMP_FLAG_PROTO);
             ServiceManager.addService(Context.INPUT_SERVICE, inputManager);
             traceEnd();
 
diff --git a/services/tests/servicestests/src/com/android/server/utils/PriorityDumpTest.java b/services/tests/servicestests/src/com/android/server/utils/PriorityDumpTest.java
index 8a312f6..da45d81 100644
--- a/services/tests/servicestests/src/com/android/server/utils/PriorityDumpTest.java
+++ b/services/tests/servicestests/src/com/android/server/utils/PriorityDumpTest.java
@@ -19,6 +19,7 @@
 import static com.android.server.utils.PriorityDump.dump;
 
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertArrayEquals;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Matchers.same;
 import static org.mockito.Mockito.verify;
@@ -59,13 +60,13 @@
     @Test
     public void testNullArgs() {
         dump(mDumper, mFd, mPw, null);
-        verify(mDumper).dump(same(mFd), same(mPw), eq(null));
+        verify(mDumper).dump(same(mFd), same(mPw), eq(null), /* asProto= */ eq(false));
     }
 
     @Test
     public void testNoArgs() {
         dump(mDumper, mFd, mPw, EMPTY_ARGS);
-        verify(mDumper).dump(same(mFd), same(mPw), same(EMPTY_ARGS));
+        verify(mDumper).dump(same(mFd), same(mPw), eq(EMPTY_ARGS), /* asProto= */ eq(false));
     }
 
     @Test
@@ -74,7 +75,7 @@
                 "--dumb_priority"
         };
         dump(mDumper, mFd, mPw, args);
-        verify(mDumper).dump(same(mFd), same(mPw), same(args));
+        verify(mDumper).dump(same(mFd), same(mPw), eq(args), /* asProto= */ eq(false));
     }
 
     @Test
@@ -83,7 +84,7 @@
                 "--dump-priority"
         };
         dump(mDumper, mFd, mPw, args);
-        verify(mDumper).dump(same(mFd), same(mPw), same(args));
+        verify(mDumper).dump(same(mFd), same(mPw), eq(EMPTY_ARGS), /* asProto= */ eq(false));
     }
 
     @Test
@@ -92,7 +93,7 @@
                 "--dump-priority", "SUPER_HIGH"
         };
         dump(mDumper, mFd, mPw, args);
-        verify(mDumper).dump(same(mFd), same(mPw), same(args));
+        verify(mDumper).dump(same(mFd), same(mPw), eq(EMPTY_ARGS), /* asProto= */ eq(false));
     }
 
     @Test
@@ -101,7 +102,9 @@
                 "--dump-priority", "SUPER_HIGH", "--high", "--five"
         };
         dump(mDumper, mFd, mPw, args);
-        verify(mDumper).dump(same(mFd), same(mPw), same(args));
+        verify(mDumper).dump(same(mFd), same(mPw), eq(new String[] {
+            "--high", "--five"
+        }), /* asProto= */ eq(false));
     }
 
     @Test
@@ -117,13 +120,13 @@
 
         assertSame(mFd, fakeDumper.criticalFd);
         assertSame(mPw, fakeDumper.criticalPw);
-        assertSame(args, fakeDumper.criticalArgs);
+        assertArrayEquals(args, fakeDumper.criticalArgs);
         assertSame(mFd, fakeDumper.highFd);
         assertSame(mPw, fakeDumper.highPw);
-        assertSame(args, fakeDumper.highArgs);
+        assertArrayEquals(args, fakeDumper.highArgs);
         assertSame(mFd, fakeDumper.normalFd);
         assertSame(mPw, fakeDumper.normalPw);
-        assertSame(args, fakeDumper.normalArgs);
+        assertArrayEquals(args, fakeDumper.normalArgs);
     }
 
     @Test
@@ -131,7 +134,8 @@
         dump(mDumper, mFd, mPw, new String[] {
                 "--dump-priority", "CRITICAL"
         });
-        verify(mDumper).dumpCritical(same(mFd), same(mPw), eq(EMPTY_ARGS));
+        verify(mDumper).dumpCritical(same(mFd), same(mPw), eq(EMPTY_ARGS),
+                /* asProto= */ eq(false));
     }
 
     @Test
@@ -141,7 +145,27 @@
         });
         verify(mDumper).dumpCritical(same(mFd), same(mPw), eq(new String[] {
                 "--high", "--five"
-        }));
+        }), /* asProto= */ eq(false));
+    }
+
+    @Test
+    public void testCriticalExtraArgsInMiddle() {
+        dump(mDumper, mFd, mPw, new String[] {
+                "--high", "--dump-priority", "CRITICAL", "--five"
+        });
+        verify(mDumper).dumpCritical(same(mFd), same(mPw), eq(new String[] {
+                "--high", "--five"
+        }), /* asProto= */ eq(false));
+    }
+
+    @Test
+    public void testCriticalExtraArgsAtEnd() {
+        dump(mDumper, mFd, mPw, new String[] {
+                "--high", "--five", "--dump-priority", "CRITICAL"
+        });
+        verify(mDumper).dumpCritical(same(mFd), same(mPw), eq(new String[] {
+                "--high", "--five"
+        }), /* asProto= */ eq(false));
     }
 
     @Test
@@ -149,7 +173,7 @@
         dump(mDumper, mFd, mPw, new String[] {
                 "--dump-priority", "HIGH"
         });
-        verify(mDumper).dumpHigh(same(mFd), same(mPw), eq(EMPTY_ARGS));
+        verify(mDumper).dumpHigh(same(mFd), same(mPw), eq(EMPTY_ARGS), /* asProto= */ eq(false));
     }
 
     @Test
@@ -159,7 +183,7 @@
         });
         verify(mDumper).dumpHigh(same(mFd), same(mPw), eq(new String[] {
                 "--high", "--five"
-        }));
+        }), /* asProto= */ eq(false));
     }
 
     @Test
@@ -167,17 +191,58 @@
         dump(mDumper, mFd, mPw, new String[] {
                 "--dump-priority", "NORMAL"
         });
-        verify(mDumper).dumpNormal(same(mFd), same(mPw), eq(EMPTY_ARGS));
+        verify(mDumper).dumpNormal(same(mFd), same(mPw), eq(EMPTY_ARGS), /* asProto= */ eq(false));
     }
 
     @Test
     public void testNormalExtraArgs() {
-        dump(mDumper, mFd, mPw, new String[] {
+        dump(mDumper, mFd, mPw, new String[]{
                 "--dump-priority", "NORMAL", "--high", "--five"
         });
-        verify(mDumper).dumpNormal(same(mFd), same(mPw), eq(new String[] {
+        verify(mDumper).dumpNormal(same(mFd), same(mPw), eq(new String[]{
                 "--high", "--five"
-        }));
+        }), /* asProto= */ eq(false));
+    }
+
+    @Test
+    public void testProtoArgs() {
+        dump(mDumper, mFd, mPw, new String[]{"--proto"});
+        verify(mDumper).dump(same(mFd), same(mPw), eq(EMPTY_ARGS), /* asProto= */ eq(true));
+    }
+
+    @Test
+    public void testProtoArgsWithPriorityArgs() {
+        dump(mDumper, mFd, mPw, new String[]{"--proto", "--dump-priority", "NORMAL", "--five"});
+        verify(mDumper).dumpNormal(same(mFd), same(mPw),
+                eq(new String[]{"--five"}), /* asProto= */ eq(true));
+    }
+
+    @Test
+    public void testProtoArgsWithPriorityArgsReverseOrder() {
+        dump(mDumper, mFd, mPw, new String[]{"--dump-priority", "NORMAL", "--proto", "--five"});
+        verify(mDumper).dumpNormal(same(mFd), same(mPw),
+                eq(new String[]{"--five"}), /* asProto= */ eq(true));
+    }
+
+    @Test
+    public void testProtoArgsInMiddle() {
+        dump(mDumper, mFd, mPw, new String[]{"--unknown", "--proto", "--five"});
+        verify(mDumper).dump(same(mFd), same(mPw),
+                eq(new String[]{"--unknown", "--five"}), /* asProto= */ eq(true));
+    }
+
+    @Test
+    public void testProtoArgsAtEnd() {
+        dump(mDumper, mFd, mPw, new String[]{"args", "-left", "--behind", "--proto"});
+        verify(mDumper).dump(same(mFd), same(mPw),
+                eq(new String[]{"args", "-left", "--behind"}), /* asProto= */ eq(true));
+    }
+
+    @Test
+    public void testProtoArgsWithInvalidPriorityType() {
+        dump(mDumper, mFd, mPw, new String[]{"--dump-priority", "HIGH?", "--proto"});
+        verify(mDumper).dump(same(mFd), same(mPw),
+                eq(EMPTY_ARGS), /* asProto= */ eq(true));
     }
 
     private final class FakeDumper implements PriorityDumper {
@@ -187,21 +252,22 @@
         PrintWriter criticalPw, highPw, normalPw;
 
         @Override
-        public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args) {
+        public void dumpCritical(FileDescriptor fd, PrintWriter pw, String[] args,
+                boolean asProto) {
             criticalFd = fd;
             criticalPw = pw;
             criticalArgs = args;
         }
 
         @Override
-        public void dumpHigh(FileDescriptor fd, PrintWriter pw, String[] args) {
+        public void dumpHigh(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
             highFd = fd;
             highPw = pw;
             highArgs = args;
         }
 
         @Override
-        public void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args) {
+        public void dumpNormal(FileDescriptor fd, PrintWriter pw, String[] args, boolean asProto) {
             normalFd = fd;
             normalPw = pw;
             normalArgs = args;