Merge "Add CellBroadcastService and related permissions"
diff --git a/core/java/android/os/HwParcel.java b/core/java/android/os/HwParcel.java
index cfb582e..5e8929c 100644
--- a/core/java/android/os/HwParcel.java
+++ b/core/java/android/os/HwParcel.java
@@ -23,6 +23,8 @@
 import android.annotation.TestApi;
 import android.annotation.UnsupportedAppUsage;
 
+import dalvik.annotation.optimization.FastNative;
+
 import libcore.util.NativeAllocationRegistry;
 
 import java.lang.annotation.Retention;
@@ -72,46 +74,54 @@
 
     /**
      * Writes an interface token into the parcel used to verify that
-     * a transaction has made it to the write type of interface.
+     * a transaction has made it to the right type of interface.
      *
      * @param interfaceName fully qualified name of interface message
      *     is being sent to.
      */
+    @FastNative
     public native final void writeInterfaceToken(String interfaceName);
     /**
      * Writes a boolean value to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     public native final void writeBool(boolean val);
     /**
      * Writes a byte value to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     public native final void writeInt8(byte val);
     /**
      * Writes a short value to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     public native final void writeInt16(short val);
     /**
      * Writes a int value to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     public native final void writeInt32(int val);
     /**
      * Writes a long value to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     public native final void writeInt64(long val);
     /**
      * Writes a float value to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     public native final void writeFloat(float val);
     /**
      * Writes a double value to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     public native final void writeDouble(double val);
     /**
      * Writes a String value to the end of the parcel.
@@ -120,6 +130,7 @@
      *
      * @param val to write
      */
+    @FastNative
     public native final void writeString(String val);
     /**
      * Writes a native handle (without duplicating the underlying
@@ -127,42 +138,50 @@
      *
      * @param val to write
      */
+    @FastNative
     public native final void writeNativeHandle(@Nullable NativeHandle val);
 
     /**
      * Writes an array of boolean values to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     private native final void writeBoolVector(boolean[] val);
     /**
      * Writes an array of byte values to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     private native final void writeInt8Vector(byte[] val);
     /**
      * Writes an array of short values to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     private native final void writeInt16Vector(short[] val);
     /**
      * Writes an array of int values to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     private native final void writeInt32Vector(int[] val);
     /**
      * Writes an array of long values to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     private native final void writeInt64Vector(long[] val);
     /**
      * Writes an array of float values to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     private native final void writeFloatVector(float[] val);
     /**
      * Writes an array of double values to the end of the parcel.
      * @param val to write
      */
+    @FastNative
     private native final void writeDoubleVector(double[] val);
     /**
      * Writes an array of String values to the end of the parcel.
@@ -171,6 +190,7 @@
      *
      * @param val to write
      */
+    @FastNative
     private native final void writeStringVector(String[] val);
     /**
      * Writes an array of native handles to the end of the parcel.
@@ -179,6 +199,7 @@
      *
      * @param val array of {@link NativeHandle} objects to write
      */
+    @FastNative
     private native final void writeNativeHandleVector(NativeHandle[] val);
 
     /**
@@ -299,6 +320,7 @@
      * Write a hwbinder object to the end of the parcel.
      * @param binder value to write
      */
+    @FastNative
     public native final void writeStrongBinder(IHwBinder binder);
 
     /**
@@ -314,48 +336,56 @@
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final boolean readBool();
     /**
      * Reads a byte value from the current location in the parcel.
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final byte readInt8();
     /**
      * Reads a short value from the current location in the parcel.
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final short readInt16();
     /**
      * Reads a int value from the current location in the parcel.
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final int readInt32();
     /**
      * Reads a long value from the current location in the parcel.
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final long readInt64();
     /**
      * Reads a float value from the current location in the parcel.
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final float readFloat();
     /**
      * Reads a double value from the current location in the parcel.
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final double readDouble();
     /**
      * Reads a String value from the current location in the parcel.
      * @return value parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final String readString();
     /**
      * Reads a native handle (without duplicating the underlying file
@@ -366,6 +396,7 @@
      * @return a {@link NativeHandle} instance parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final @Nullable NativeHandle readNativeHandle();
     /**
      * Reads an embedded native handle (without duplicating the underlying
@@ -379,6 +410,7 @@
      * @return a {@link NativeHandle} instance parsed from the parcel
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final @Nullable NativeHandle readEmbeddedNativeHandle(
             long parentHandle, long offset);
 
@@ -387,54 +419,63 @@
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final boolean[] readBoolVectorAsArray();
     /**
      * Reads an array of byte values from the parcel.
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final byte[] readInt8VectorAsArray();
     /**
      * Reads an array of short values from the parcel.
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final short[] readInt16VectorAsArray();
     /**
      * Reads an array of int values from the parcel.
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final int[] readInt32VectorAsArray();
     /**
      * Reads an array of long values from the parcel.
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final long[] readInt64VectorAsArray();
     /**
      * Reads an array of float values from the parcel.
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final float[] readFloatVectorAsArray();
     /**
      * Reads an array of double values from the parcel.
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final double[] readDoubleVectorAsArray();
     /**
      * Reads an array of String values from the parcel.
      * @return array of parsed values
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final String[] readStringVectorAsArray();
     /**
      * Reads an array of native handles from the parcel.
      * @return array of {@link NativeHandle} objects
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     private native final NativeHandle[] readNativeHandleAsArray();
 
     /**
@@ -537,6 +578,7 @@
      * @return binder object read from parcel or null if no binder can be read
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final IHwBinder readStrongBinder();
 
     /**
@@ -544,6 +586,7 @@
      * @return blob of size expectedSize
      * @throws IllegalArgumentException if the parcel has no more data
      */
+    @FastNative
     public native final HwBlob readBuffer(long expectedSize);
 
     /**
@@ -559,6 +602,7 @@
      * @throws NullPointerException if the transaction specified the blob to be null
      *    but nullable is false
      */
+    @FastNative
     public native final HwBlob readEmbeddedBuffer(
             long expectedSize, long parentHandle, long offset,
             boolean nullable);
@@ -567,26 +611,31 @@
      * Write a buffer into the transaction.
      * @param blob blob to write into the parcel.
      */
+    @FastNative
     public native final void writeBuffer(HwBlob blob);
     /**
      * Write a status value into the blob.
      * @param status value to write
      */
+    @FastNative
     public native final void writeStatus(int status);
     /**
      * @throws IllegalArgumentException if a success vaue cannot be read
      * @throws RemoteException if success value indicates a transaction error
      */
+    @FastNative
     public native final void verifySuccess();
     /**
      * Should be called to reduce memory pressure when this object no longer needs
      * to be written to.
      */
+    @FastNative
     public native final void releaseTemporaryStorage();
     /**
      * Should be called when object is no longer needed to reduce possible memory
      * pressure if the Java GC does not get to this object in time.
      */
+    @FastNative
     public native final void release();
 
     /**
@@ -597,6 +646,7 @@
     // Returns address of the "freeFunction".
     private static native final long native_init();
 
+    @FastNative
     private native final void native_setup(boolean allocate);
 
     static {
diff --git a/mime/Android.bp b/mime/Android.bp
index 8b2b059..23a8fbf 100644
--- a/mime/Android.bp
+++ b/mime/Android.bp
@@ -60,7 +60,7 @@
     tools: [
         "soong_zip",
     ],
-    srcs: [":mime.types"],
+    srcs: [":mime.types.minimized"],
     out: ["mimemap-res.jar"],
     cmd: "mkdir $(genDir)/res/ && cp $(in) $(genDir)/res/ && $(location soong_zip) -C $(genDir) -o $(out) -D $(genDir)/res/",
 }
@@ -73,42 +73,49 @@
     tools: [
         "soong_zip",
     ],
-    srcs: [":mime.types"],
+    srcs: [":mime.types.minimized"],
     out: ["mimemap-testing-res.jar"],
     cmd: "mkdir $(genDir)/testres/ && cp $(in) $(genDir)/testres/ && $(location soong_zip) -C $(genDir) -o $(out) -D $(genDir)/testres/",
 }
 
-// Combination of all *mime.types resources.
+// Combination of all *mime.types.minimized resources.
 filegroup {
-    name: "mime.types",
+    name: "mime.types.minimized",
     visibility: [
         "//visibility:private",
     ],
     srcs: [
-        ":debian.mime.types",
-        ":android.mime.types",
-        ":vendor.mime.types",
+        ":debian.mime.types.minimized",
+        ":android.mime.types.minimized",
+        ":vendor.mime.types.minimized",
     ],
 }
 
-filegroup {
-    name: "android.mime.types",
+java_genrule {
+    name: "android.mime.types.minimized",
     visibility: [
         "//visibility:private",
     ],
-    path: "java-res/",
+    out: ["android.mime.types"],
     srcs: [
         "java-res/android.mime.types",
     ],
+    //    strip comments            normalize whitepace       drop empty lines
+    cmd: "awk '{gsub(/#.*$$/,\"\"); $$1=$$1; print;}' $(in) | grep ' ' > $(out)",
 }
 
-filegroup {
-    name: "vendor.mime.types",
+// Unlike the other *mime.types files, vendor.mime.types gets '?' prepended to
+// every field so that its mappings will never overwrite earlier mappings by
+// the other resource files. http://b/141842825
+java_genrule {
+    name: "vendor.mime.types.minimized",
     visibility: [
         "//visibility:private",
     ],
-    path: "java-res/",
+    out: ["vendor.mime.types"],
     srcs: [
         "java-res/vendor.mime.types",
     ],
+    //    strip comments            normalize whitepace       drop empty lines   prepend ? to fields that are missing it
+    cmd: "awk '{gsub(/#.*$$/,\"\"); $$1=$$1; print;}' $(in) | grep ' '         | awk '{for(i=1;i<=NF;i++) { sub(/^\\??/, \"?\", $$i); }; print}' > $(out)",
 }
diff --git a/mime/java-res/vendor.mime.types b/mime/java-res/vendor.mime.types
index afb8f9e..06939c5 100644
--- a/mime/java-res/vendor.mime.types
+++ b/mime/java-res/vendor.mime.types
@@ -39,3 +39,6 @@
 #
 # Add your custom mappings below this line (with no "#" at the start of the line):
 
+test/mimeA extA extB extX
+?test/mimeC ?extX ?extY ?extZ
+test/mimeB extC ?extD extF
diff --git a/mime/java/android/content/type/DefaultMimeMapFactory.java b/mime/java/android/content/type/DefaultMimeMapFactory.java
index 03b685d..11d20d4 100644
--- a/mime/java/android/content/type/DefaultMimeMapFactory.java
+++ b/mime/java/android/content/type/DefaultMimeMapFactory.java
@@ -23,11 +23,9 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Function;
-import java.util.regex.Pattern;
 
 /**
  * Creates the framework default {@link MimeMap}, a bidirectional mapping
@@ -53,8 +51,6 @@
         return create(resourceName -> c.getResourceAsStream("/res/" + resourceName));
     }
 
-    private static final Pattern SPLIT_PATTERN = Pattern.compile("\\s+");
-
     /**
      * Creates a {@link MimeMap} instance whose resources are loaded from the
      * InputStreams looked up in {@code resourceSupplier}.
@@ -63,33 +59,43 @@
      */
     public static MimeMap create(Function<String, InputStream> resourceSupplier) {
         MimeMap.Builder builder = MimeMap.builder();
-        parseTypes(builder, true, resourceSupplier, "mime.types");
-        parseTypes(builder, true, resourceSupplier, "android.mime.types");
-        parseTypes(builder, false, resourceSupplier, "vendor.mime.types");
+        // The files loaded here must be in minimized format with lines of the
+        // form "mime/type ext1 ext2 ext3", i.e. no comments, no empty lines, no
+        // leading/trailing whitespace and with a single space between entries on
+        // each line.  See http://b/142267887
+        //
+        // Note: the order here matters - later entries can overwrite earlier ones
+        // (except that vendor.mime.types entries are prefixed with '?' which makes
+        // them never overwrite).
+        parseTypes(builder, resourceSupplier, "debian.mime.types");
+        parseTypes(builder, resourceSupplier, "android.mime.types");
+        parseTypes(builder, resourceSupplier, "vendor.mime.types");
         return builder.build();
     }
 
-    private static void parseTypes(MimeMap.Builder builder, boolean allowOverwrite,
+    private static void parseTypes(MimeMap.Builder builder,
             Function<String, InputStream> resourceSupplier, String resourceName) {
         try (InputStream inputStream = Objects.requireNonNull(resourceSupplier.apply(resourceName));
              BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
             String line;
+            List<String> specs = new ArrayList<>(10); // re-use for each line
             while ((line = reader.readLine()) != null) {
-                int commentPos = line.indexOf('#');
-                if (commentPos >= 0) {
-                    line = line.substring(0, commentPos);
-                }
-                line = line.trim();
-                if (line.isEmpty()) {
-                    continue;
-                }
-                List<String> specs = Arrays.asList(SPLIT_PATTERN.split(line));
-                if (!allowOverwrite) {
-                    // Pretend that the mimeType and each file extension listed in the line
-                    // carries a "?" prefix, which means that it can add new mappings but
-                    // not modify existing mappings (putIfAbsent() semantics).
-                    specs = ensurePrefix("?", specs);
-                }
+                specs.clear();
+                // Lines are of the form "mimeSpec extSpec extSpec[...]" with a single space
+                // separating them and no leading/trailing spaces and no empty lines.
+                int startIdx = 0;
+                do {
+                    int endIdx = line.indexOf(' ', startIdx);
+                    if (endIdx < 0) {
+                        endIdx = line.length();
+                    }
+                    String spec = line.substring(startIdx, endIdx);
+                    if (spec.isEmpty()) {
+                        throw new IllegalArgumentException("Malformed line: " + line);
+                    }
+                    specs.add(spec);
+                    startIdx = endIdx + 1; // skip over the space
+                } while (startIdx < line.length());
                 builder.put(specs.get(0), specs.subList(1, specs.size()));
             }
         } catch (IOException | RuntimeException e) {
@@ -97,15 +103,4 @@
         }
     }
 
-    private static List<String> ensurePrefix(String prefix, List<String> strings) {
-        List<String> result = new ArrayList<>(strings.size());
-        for (String s : strings) {
-            if (!s.startsWith(prefix)) {
-                s = prefix + s;
-            }
-            result.add(s);
-        }
-        return result;
-    }
-
 }
diff --git a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
index 44338ac..5ce215b 100644
--- a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
+++ b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
@@ -93,9 +93,11 @@
 import android.text.format.DateUtils;
 import android.util.ArraySet;
 import android.util.PrintWriterPrinter;
+import android.util.SparseArray;
 
 import com.android.internal.content.PackageHelper;
 import com.android.internal.util.ArrayUtils;
+import com.android.internal.util.IndentingPrintWriter;
 import com.android.server.LocalServices;
 import com.android.server.SystemConfig;
 
@@ -129,6 +131,7 @@
     private static final String STDIN_PATH = "-";
     /** Path where ART profiles snapshots are dumped for the shell user */
     private final static String ART_PROFILE_SNAPSHOT_DEBUG_LOCATION = "/data/misc/profman/";
+    private static final int DEFAULT_WAIT_MS = 60 * 1000;
 
     final IPackageManager mInterface;
     final private WeakHashMap<String, Resources> mResourceCache =
@@ -269,7 +272,7 @@
                 case "get-harmful-app-warning":
                     return runGetHarmfulAppWarning();
                 case "get-stagedsessions":
-                    return getStagedSessions();
+                    return runListStagedSessions();
                 case "uninstall-system-updates":
                     return uninstallSystemUpdates();
                 case "rollback-app":
@@ -341,28 +344,6 @@
         return 1;
     }
 
-    private int getStagedSessions() {
-        final PrintWriter pw = getOutPrintWriter();
-        try {
-            List<SessionInfo> stagedSessionsList =
-                    mInterface.getPackageInstaller().getStagedSessions().getList();
-            for (SessionInfo session: stagedSessionsList) {
-                pw.println("appPackageName = " + session.getAppPackageName()
-                        + "; sessionId = " + session.getSessionId()
-                        + "; isStaged = " + session.isStaged()
-                        + "; isStagedSessionReady = " + session.isStagedSessionReady()
-                        + "; isStagedSessionApplied = " + session.isStagedSessionApplied()
-                        + "; isStagedSessionFailed = " + session.isStagedSessionFailed() + ";");
-            }
-        } catch (RemoteException e) {
-            pw.println("Failure ["
-                    + e.getClass().getName() + " - "
-                    + e.getMessage() + "]");
-            return 0;
-        }
-        return 1;
-    }
-
     private int uninstallSystemUpdates() {
         final PrintWriter pw = getOutPrintWriter();
         List<String> failedUninstalls = new LinkedList<>();
@@ -535,6 +516,8 @@
                 return runListPermissionGroups();
             case "permissions":
                 return runListPermissions();
+            case "staged-sessions":
+                return runListStagedSessions();
             case "users":
                 ServiceManager.getService("user").shellCommand(
                         getInFileDescriptor(), getOutFileDescriptor(), getErrFileDescriptor(),
@@ -871,6 +854,103 @@
         return 0;
     }
 
+    private static class SessionDump {
+        boolean onlyParent; // Show parent sessions only
+        boolean onlyReady; // Show only staged sessions that are in ready state
+        boolean onlySessionId; // Show sessionId only
+    }
+
+    // Returns true if the provided flag is a session flag and given SessionDump was updated
+    private boolean setSessionFlag(String flag, SessionDump sessionDump) {
+        switch (flag) {
+            case "--only-parent":
+                sessionDump.onlyParent = true;
+                break;
+            case "--only-ready":
+                sessionDump.onlyReady = true;
+                break;
+            case "--only-sessionid":
+                sessionDump.onlySessionId = true;
+                break;
+            default:
+                return false;
+        }
+        return true;
+    }
+
+    private int runListStagedSessions() {
+        final IndentingPrintWriter pw = new IndentingPrintWriter(
+                getOutPrintWriter(), /* singleIndent */ "  ", /* wrapLength */ 120);
+
+        SessionDump sessionDump = new SessionDump();
+        String opt;
+        while ((opt = getNextOption()) != null) {
+            if (!setSessionFlag(opt, sessionDump)) {
+                pw.println("Error: Unknown option: " + opt);
+                return -1;
+            }
+        }
+
+        try {
+            List<SessionInfo> stagedSessions =
+                    mInterface.getPackageInstaller().getStagedSessions().getList();
+            printSessionList(pw, stagedSessions, sessionDump);
+        } catch (RemoteException e) {
+            pw.println("Failure ["
+                    + e.getClass().getName() + " - "
+                    + e.getMessage() + "]");
+            return -1;
+        }
+        return 1;
+    }
+
+    private void printSessionList(IndentingPrintWriter pw, List<SessionInfo> stagedSessions,
+            SessionDump sessionDump) {
+        final SparseArray<SessionInfo> sessionById = new SparseArray<>(stagedSessions.size());
+        for (SessionInfo session : stagedSessions) {
+            sessionById.put(session.getSessionId(), session);
+        }
+        for (SessionInfo session: stagedSessions) {
+            if (sessionDump.onlyReady && !session.isStagedSessionReady()) {
+                continue;
+            }
+            if (session.getParentSessionId() != SessionInfo.INVALID_ID) {
+                continue;
+            }
+            printSession(pw, session, sessionDump);
+            if (session.isMultiPackage() && !sessionDump.onlyParent) {
+                pw.increaseIndent();
+                final int[] childIds = session.getChildSessionIds();
+                for (int i = 0; i < childIds.length; i++) {
+                    final SessionInfo childSession = sessionById.get(childIds[i]);
+                    if (childSession == null) {
+                        if (sessionDump.onlySessionId) {
+                            pw.println(childIds[i]);
+                        } else {
+                            pw.println("sessionId = " + childIds[i] + "; not found");
+                        }
+                    } else {
+                        printSession(pw, childSession, sessionDump);
+                    }
+                }
+                pw.decreaseIndent();
+            }
+        }
+    }
+
+    private static void printSession(PrintWriter pw, SessionInfo session, SessionDump sessionDump) {
+        if (sessionDump.onlySessionId) {
+            pw.println(session.getSessionId());
+            return;
+        }
+        pw.println("sessionId = " + session.getSessionId()
+                + "; appPackageName = " + session.getAppPackageName()
+                + "; isStaged = " + session.isStaged()
+                + "; isReady = " + session.isStagedSessionReady()
+                + "; isApplied = " + session.isStagedSessionApplied()
+                + "; isFailed = " + session.isStagedSessionFailed() + ";");
+    }
+
     private Intent parseIntentAndUser() throws URISyntaxException {
         mTargetUser = UserHandle.USER_CURRENT;
         mBrief = false;
@@ -1078,6 +1158,45 @@
                 return 1;
             }
             abandonSession = false;
+
+            if (!params.sessionParams.isStaged || !params.waitForStagedSessionReady) {
+                pw.println("Success");
+                return 0;
+            }
+
+            long timeoutMs = params.timeoutMs <= 0
+                    ? DEFAULT_WAIT_MS
+                    : params.timeoutMs;
+            PackageInstaller.SessionInfo si = mInterface.getPackageInstaller()
+                    .getSessionInfo(sessionId);
+            long currentTime = System.currentTimeMillis();
+            long endTime = currentTime + timeoutMs;
+            // Using a loop instead of BroadcastReceiver since we can receive session update
+            // broadcast only if packageInstallerName is "android". We can't always force
+            // "android" as packageIntallerName, e.g, rollback auto implies
+            // "-i com.android.shell".
+            while (currentTime < endTime) {
+                if (si != null
+                        && (si.isStagedSessionReady() || si.isStagedSessionFailed())) {
+                    break;
+                }
+                SystemClock.sleep(Math.min(endTime - currentTime, 100));
+                currentTime = System.currentTimeMillis();
+                si = mInterface.getPackageInstaller().getSessionInfo(sessionId);
+            }
+            if (si == null) {
+                pw.println("Failure [failed to retrieve SessionInfo]");
+                return 1;
+            }
+            if (!si.isStagedSessionReady() && !si.isStagedSessionFailed()) {
+                pw.println("Failure [timed out after " + timeoutMs + " ms]");
+                return 1;
+            }
+            if (!si.isStagedSessionReady()) {
+                pw.println("Error [" + si.getStagedSessionErrorCode() + "] ["
+                        + si.getStagedSessionErrorMessage() + "]");
+                return 1;
+            }
             pw.println("Success");
             return 0;
         } finally {
@@ -2368,6 +2487,8 @@
         SessionParams sessionParams;
         String installerPackageName;
         int userId = UserHandle.USER_ALL;
+        boolean waitForStagedSessionReady = false;
+        long timeoutMs = DEFAULT_WAIT_MS;
     }
 
     private InstallParams makeInstallParams() {
@@ -2493,6 +2614,14 @@
                     }
                     sessionParams.installFlags |= PackageManager.INSTALL_ENABLE_ROLLBACK;
                     break;
+                case "--wait":
+                    params.waitForStagedSessionReady = true;
+                    try {
+                        params.timeoutMs = Long.parseLong(peekNextArg());
+                        getNextArg();
+                    } catch (NumberFormatException ignore) {
+                    }
+                    break;
                 default:
                     throw new IllegalArgumentException("Unknown option " + opt);
             }
@@ -3023,6 +3152,12 @@
         pw.println("      -d: only list dangerous permissions");
         pw.println("      -u: list only the permissions users will see");
         pw.println("");
+        pw.println("  list staged-sessions [--only-ready] [--only-sessionid] [--only-parent]");
+        pw.println("    Displays list of all staged sessions on device.");
+        pw.println("      --only-ready: show only staged sessions that are ready");
+        pw.println("      --only-sessionid: show only sessionId of each session");
+        pw.println("      --only-parent: hide all children sessions");
+        pw.println("");
         pw.println("  resolve-activity [--brief] [--components] [--query-flags FLAGS]");
         pw.println("       [--user USER_ID] INTENT");
         pw.println("    Prints the activity that resolves to the given INTENT.");
@@ -3045,7 +3180,8 @@
         pw.println("       [--referrer URI] [--abi ABI_NAME] [--force-sdk]");
         pw.println("       [--preload] [--instantapp] [--full] [--dont-kill]");
         pw.println("       [--enable-rollback]");
-        pw.println("       [--force-uuid internal|UUID] [--pkg PACKAGE] [-S BYTES] [--apex]");
+        pw.println("       [--force-uuid internal|UUID] [--pkg PACKAGE] [-S BYTES]");
+        pw.println("       [--apex] [--wait TIMEOUT]");
         pw.println("       [PATH|-]");
         pw.println("    Install an application.  Must provide the apk data to install, either as a");
         pw.println("    file path or '-' to read from stdin.  Options are:");
@@ -3075,6 +3211,9 @@
         pw.println("          3=device setup, 4=user request");
         pw.println("      --force-uuid: force install on to disk volume with given UUID");
         pw.println("      --apex: install an .apex file, not an .apk");
+        pw.println("      --wait: when performing staged install, wait TIMEOUT milliseconds");
+        pw.println("          for pre-reboot verification to complete. If TIMEOUT is not");
+        pw.println("          specified it will wait for " + DEFAULT_WAIT_MS + " milliseconds.");
         pw.println("");
         pw.println("  install-create [-lrtsfdg] [-i PACKAGE] [--user USER_ID|all|current]");
         pw.println("       [-p INHERIT_PACKAGE] [--install-location 0/1/2]");
@@ -3257,7 +3396,7 @@
         pw.println("  uninstall-system-updates");
         pw.println("    Remove updates to all system applications and fall back to their /system " +
                 "version.");
-        pw.println();
+        pw.println("");
         pw.println("  get-moduleinfo [--all | --installed] [module-name]");
         pw.println("    Displays module info. If module-name is specified only that info is shown");
         pw.println("    By default, without any argument only installed modules are shown.");