Merge "CAS: Add AIDL Support"
diff --git a/core/api/current.txt b/core/api/current.txt
index 097dc44..a0c132f 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -17088,8 +17088,8 @@
     method public boolean flush(android.hardware.SensorEventListener);
     method public static float getAltitude(float, float);
     method public static void getAngleChange(float[], float[], float[]);
-    method public android.hardware.Sensor getDefaultSensor(int);
-    method public android.hardware.Sensor getDefaultSensor(int, boolean);
+    method @Nullable public android.hardware.Sensor getDefaultSensor(int);
+    method @Nullable public android.hardware.Sensor getDefaultSensor(int, boolean);
     method public java.util.List<android.hardware.Sensor> getDynamicSensorList(int);
     method public static float getInclination(float[]);
     method public static float[] getOrientation(float[], float[]);
diff --git a/core/java/android/hardware/SensorManager.java b/core/java/android/hardware/SensorManager.java
index dec424c..6d8c4a9 100644
--- a/core/java/android/hardware/SensorManager.java
+++ b/core/java/android/hardware/SensorManager.java
@@ -16,6 +16,7 @@
 
 package android.hardware;
 
+import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.annotation.SystemService;
 import android.compat.annotation.UnsupportedAppUsage;
@@ -496,7 +497,7 @@
      * @see #getSensorList(int)
      * @see Sensor
      */
-    public Sensor getDefaultSensor(int type) {
+    public @Nullable Sensor getDefaultSensor(int type) {
         // TODO: need to be smarter, for now, just return the 1st sensor
         List<Sensor> l = getSensorList(type);
         boolean wakeUpSensor = false;
@@ -544,7 +545,7 @@
      *         and the application has the necessary permissions, or null otherwise.
      * @see Sensor#isWakeUpSensor()
      */
-    public Sensor getDefaultSensor(int type, boolean wakeUp) {
+    public @Nullable Sensor getDefaultSensor(int type, boolean wakeUp) {
         List<Sensor> l = getSensorList(type);
         for (Sensor sensor : l) {
             if (sensor.isWakeUpSensor() == wakeUp) {
diff --git a/core/java/android/hardware/camera2/impl/CameraMetadataNative.java b/core/java/android/hardware/camera2/impl/CameraMetadataNative.java
index 468e604..5abe677 100644
--- a/core/java/android/hardware/camera2/impl/CameraMetadataNative.java
+++ b/core/java/android/hardware/camera2/impl/CameraMetadataNative.java
@@ -88,8 +88,8 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Map;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 
@@ -1027,6 +1027,19 @@
         return fixedFaceRectangles;
     }
 
+    private boolean setLensShadingMap(LensShadingMap lensShadingMap) {
+        if (lensShadingMap == null) {
+            return false;
+        }
+        float[] lsmArray = new float[lensShadingMap.getGainFactorCount()];
+        lensShadingMap.copyGainFactors(lsmArray, 0);
+        setBase(CaptureResult.STATISTICS_LENS_SHADING_MAP, lsmArray);
+
+        Size s = new Size(lensShadingMap.getRowCount(), lensShadingMap.getColumnCount());
+        setBase(CameraCharacteristics.LENS_INFO_SHADING_MAP_SIZE, s);
+        return true;
+    }
+
     private LensShadingMap getLensShadingMap() {
         float[] lsmArray = getBase(CaptureResult.STATISTICS_LENS_SHADING_MAP);
         Size s = get(CameraCharacteristics.LENS_INFO_SHADING_MAP_SIZE);
@@ -1851,6 +1864,13 @@
                 metadata.setAERegions(value);
             }
         });
+        sSetCommandMap.put(CaptureResult.STATISTICS_LENS_SHADING_CORRECTION_MAP.getNativeKey(),
+                new SetCommand() {
+                    @Override
+                    public <T> void setValue(CameraMetadataNative metadata, T value) {
+                        metadata.setLensShadingMap((LensShadingMap) value);
+                    }
+                });
     }
 
     private boolean setAvailableFormats(int[] value) {
diff --git a/core/tests/coretests/src/android/hardware/camera2/OWNERS b/core/tests/coretests/src/android/hardware/camera2/OWNERS
new file mode 100644
index 0000000..f48a95c
--- /dev/null
+++ b/core/tests/coretests/src/android/hardware/camera2/OWNERS
@@ -0,0 +1 @@
+include platform/frameworks/av:/camera/OWNERS
diff --git a/core/tests/coretests/src/android/hardware/camera2/impl/CaptureMetadataNativeTest.java b/core/tests/coretests/src/android/hardware/camera2/impl/CaptureMetadataNativeTest.java
new file mode 100644
index 0000000..a38c040
--- /dev/null
+++ b/core/tests/coretests/src/android/hardware/camera2/impl/CaptureMetadataNativeTest.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2023 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 android.hardware.camera2.impl;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.hardware.camera2.CaptureResult;
+import android.hardware.camera2.params.LensShadingMap;
+import android.util.Size;
+
+import androidx.test.filters.SmallTest;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.Arrays;
+
+@SmallTest
+@RunWith(JUnit4.class)
+/** Tests for {@link CameraMetadataNative} class. */
+public class CaptureMetadataNativeTest {
+
+    @Test
+    public void setLensShadingMap() {
+        final Size s = new Size(10, 10);
+        // 4 x rows x columns
+        final float[] elements = new float[400];
+        Arrays.fill(elements, 42);
+
+        final LensShadingMap lensShadingMap =
+                new LensShadingMap(elements, s.getHeight(), s.getWidth());
+        CameraMetadataNative captureResults = new CameraMetadataNative();
+        captureResults.set(CaptureResult.STATISTICS_LENS_SHADING_CORRECTION_MAP, lensShadingMap);
+
+        final LensShadingMap output =
+                captureResults.get(CaptureResult.STATISTICS_LENS_SHADING_CORRECTION_MAP);
+
+        assertThat(output).isEqualTo(lensShadingMap);
+    }
+}
diff --git a/libs/hwui/Android.bp b/libs/hwui/Android.bp
index c022087..2800e31 100644
--- a/libs/hwui/Android.bp
+++ b/libs/hwui/Android.bp
@@ -634,7 +634,7 @@
 cc_defaults {
     name: "hwui_test_defaults",
     defaults: ["hwui_defaults"],
-    test_suites: ["device-tests"],
+    test_suites: ["general-tests"],
     header_libs: ["libandroid_headers_private"],
     target: {
         android: {
diff --git a/libs/hwui/AndroidTest.xml b/libs/hwui/AndroidTest.xml
index 381fb9f..911315f 100644
--- a/libs/hwui/AndroidTest.xml
+++ b/libs/hwui/AndroidTest.xml
@@ -16,22 +16,22 @@
 <configuration description="Config for hwuimicro">
     <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
         <option name="cleanup" value="true" />
-        <option name="push" value="hwui_unit_tests->/data/nativetest/hwui_unit_tests" />
-        <option name="push" value="hwuimicro->/data/benchmarktest/hwuimicro" />
-        <option name="push" value="hwuimacro->/data/benchmarktest/hwuimacro" />
+        <option name="push" value="hwui_unit_tests->/data/local/tmp/nativetest/hwui_unit_tests" />
+        <option name="push" value="hwuimicro->/data/local/tmp/benchmarktest/hwuimicro" />
+        <option name="push" value="hwuimacro->/data/local/tmp/benchmarktest/hwuimacro" />
     </target_preparer>
     <option name="test-suite-tag" value="apct" />
     <test class="com.android.tradefed.testtype.GTest" >
-        <option name="native-test-device-path" value="/data/nativetest" />
+        <option name="native-test-device-path" value="/data/local/tmp/nativetest" />
         <option name="module-name" value="hwui_unit_tests" />
     </test>
     <test class="com.android.tradefed.testtype.GoogleBenchmarkTest" >
-        <option name="native-benchmark-device-path" value="/data/benchmarktest" />
+        <option name="native-benchmark-device-path" value="/data/local/tmp/benchmarktest" />
         <option name="benchmark-module-name" value="hwuimicro" />
         <option name="file-exclusion-filter-regex" value=".*\.config$" />
     </test>
     <test class="com.android.tradefed.testtype.GoogleBenchmarkTest" >
-        <option name="native-benchmark-device-path" value="/data/benchmarktest" />
+        <option name="native-benchmark-device-path" value="/data/local/tmp/benchmarktest" />
         <option name="benchmark-module-name" value="hwuimacro" />
         <option name="file-exclusion-filter-regex" value=".*\.config$" />
     </test>
diff --git a/libs/hwui/TEST_MAPPING b/libs/hwui/TEST_MAPPING
index b1719a9..03682e8 100644
--- a/libs/hwui/TEST_MAPPING
+++ b/libs/hwui/TEST_MAPPING
@@ -5,6 +5,9 @@
     },
     {
       "name": "CtsAccelerationTestCases"
+    },
+    {
+      "name": "hwui_unit_tests"
     }
   ],
   "imports": [
diff --git a/libs/hwui/tests/unit/CacheManagerTests.cpp b/libs/hwui/tests/unit/CacheManagerTests.cpp
index edd3e4e..fc84abb 100644
--- a/libs/hwui/tests/unit/CacheManagerTests.cpp
+++ b/libs/hwui/tests/unit/CacheManagerTests.cpp
@@ -32,7 +32,8 @@
     return cacheUsage;
 }
 
-RENDERTHREAD_SKIA_PIPELINE_TEST(CacheManager, trimMemory) {
+// TOOD(258700630): fix this test and re-enable
+RENDERTHREAD_SKIA_PIPELINE_TEST(CacheManager, DISABLED_trimMemory) {
     int32_t width = DeviceInfo::get()->getWidth();
     int32_t height = DeviceInfo::get()->getHeight();
     GrDirectContext* grContext = renderThread.getGrContext();
diff --git a/libs/hwui/tests/unit/GraphicsStatsServiceTests.cpp b/libs/hwui/tests/unit/GraphicsStatsServiceTests.cpp
index 098b4cc..c2d23e6 100644
--- a/libs/hwui/tests/unit/GraphicsStatsServiceTests.cpp
+++ b/libs/hwui/tests/unit/GraphicsStatsServiceTests.cpp
@@ -14,17 +14,18 @@
  * limitations under the License.
  */
 
+#include <android-base/macros.h>
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
-
-#include "protos/graphicsstats.pb.h"
-#include "service/GraphicsStatsService.h"
-
 #include <stdio.h>
 #include <stdlib.h>
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <unistd.h>
 
+#include "protos/graphicsstats.pb.h"
+#include "service/GraphicsStatsService.h"
+
 using namespace android;
 using namespace android::uirenderer;
 
@@ -49,12 +50,14 @@
 
 // No code left untested
 TEST(GraphicsStats, findRootPath) {
-#ifdef __LP64__
-    std::string expected = "/data/nativetest64/hwui_unit_tests";
-#else
-    std::string expected = "/data/nativetest/hwui_unit_tests";
-#endif
-    EXPECT_EQ(expected, findRootPath());
+    // Different tools/infrastructure seem to push this to different locations. It shouldn't really
+    // matter where the binary is, so add new locations here as needed. This test still seems good
+    // as it's nice to understand the possibility space, and ensure findRootPath continues working
+    // as expected.
+    std::string acceptableLocations[] = {"/data/nativetest/hwui_unit_tests",
+                                         "/data/nativetest64/hwui_unit_tests",
+                                         "/data/local/tmp/nativetest/hwui_unit_tests/" ABI_STRING};
+    EXPECT_THAT(acceptableLocations, ::testing::Contains(findRootPath()));
 }
 
 TEST(GraphicsStats, saveLoad) {
diff --git a/libs/incident/Android.bp b/libs/incident/Android.bp
index ff1714d..a996700 100644
--- a/libs/incident/Android.bp
+++ b/libs/incident/Android.bp
@@ -60,6 +60,7 @@
         ":libincident_aidl",
         "src/IncidentReportArgs.cpp",
     ],
+    min_sdk_version: "29",
 }
 
 cc_library_shared {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/ImageTransformState.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/ImageTransformState.java
index 76025ab7..0446165 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/ImageTransformState.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/ImageTransformState.java
@@ -53,7 +53,9 @@
             return true;
         }
         if (otherState instanceof ImageTransformState) {
-            return mIcon != null && mIcon.sameAs(((ImageTransformState) otherState).getIcon());
+            final Icon otherIcon = ((ImageTransformState) otherState).mIcon;
+            return mIcon == otherIcon || (mIcon != null && otherIcon != null && mIcon.sameAs(
+                    otherIcon));
         }
         return false;
     }
diff --git a/services/core/Android.bp b/services/core/Android.bp
index 9268fc0..bf71e6b 100644
--- a/services/core/Android.bp
+++ b/services/core/Android.bp
@@ -154,8 +154,8 @@
         "android.hardware.health-translate-java",
         "android.hardware.light-V1-java",
         "android.hardware.tv.cec-V1.1-java",
-        "android.hardware.tv.cec-V1-java",
-        "android.hardware.tv.hdmi-V1-java",
+        "android.hardware.tv.hdmi.cec-V1-java",
+        "android.hardware.tv.hdmi.connection-V1-java",
         "android.hardware.weaver-V1.0-java",
         "android.hardware.biometrics.face-V1.0-java",
         "android.hardware.biometrics.fingerprint-V2.3-java",
diff --git a/services/core/java/com/android/server/hdmi/HdmiCecController.java b/services/core/java/com/android/server/hdmi/HdmiCecController.java
index bf0052d..dbba9ee 100644
--- a/services/core/java/com/android/server/hdmi/HdmiCecController.java
+++ b/services/core/java/com/android/server/hdmi/HdmiCecController.java
@@ -19,16 +19,16 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.hardware.hdmi.HdmiPortInfo;
-import android.hardware.tv.cec.CecMessage;
-import android.hardware.tv.cec.IHdmiCec;
-import android.hardware.tv.cec.IHdmiCecCallback;
 import android.hardware.tv.cec.V1_0.HotplugEvent;
 import android.hardware.tv.cec.V1_0.IHdmiCec.getPhysicalAddressCallback;
 import android.hardware.tv.cec.V1_0.OptionKey;
 import android.hardware.tv.cec.V1_0.Result;
 import android.hardware.tv.cec.V1_0.SendMessageResult;
-import android.hardware.tv.hdmi.IHdmi;
-import android.hardware.tv.hdmi.IHdmiCallback;
+import android.hardware.tv.hdmi.cec.CecMessage;
+import android.hardware.tv.hdmi.cec.IHdmiCec;
+import android.hardware.tv.hdmi.cec.IHdmiCecCallback;
+import android.hardware.tv.hdmi.connection.IHdmiConnection;
+import android.hardware.tv.hdmi.connection.IHdmiConnectionCallback;
 import android.icu.util.IllformedLocaleException;
 import android.icu.util.ULocale;
 import android.os.Binder;
@@ -178,7 +178,7 @@
         if (controller != null) {
             return controller;
         }
-        HdmiLogger.warning("Unable to use CEC and HDMI AIDL HALs");
+        HdmiLogger.warning("Unable to use CEC and HDMI Connection AIDL HALs");
 
         controller = createWithNativeWrapper(service, new NativeWrapperImpl11(), atomWriter);
         if (controller != null) {
@@ -872,14 +872,14 @@
     private static final class NativeWrapperImplAidl
             implements NativeWrapper, IBinder.DeathRecipient {
         private IHdmiCec mHdmiCec;
-        private IHdmi mHdmi;
+        private IHdmiConnection mHdmiConnection;
         @Nullable private HdmiCecCallback mCallback;
 
         private final Object mLock = new Object();
 
         @Override
         public String nativeInit() {
-            return connectToHal() ? mHdmiCec.toString() + " " + mHdmi.toString() : null;
+            return connectToHal() ? mHdmiCec.toString() + " " + mHdmiConnection.toString() : null;
         }
 
         boolean connectToHal() {
@@ -896,15 +896,15 @@
                 HdmiLogger.error("Couldn't link to death : ", e);
             }
 
-            mHdmi =
-                    IHdmi.Stub.asInterface(
-                            ServiceManager.getService(IHdmi.DESCRIPTOR + "/default"));
-            if (mHdmi == null) {
-                HdmiLogger.error("Could not initialize HDMI AIDL HAL");
+            mHdmiConnection =
+                    IHdmiConnection.Stub.asInterface(
+                            ServiceManager.getService(IHdmiConnection.DESCRIPTOR + "/default"));
+            if (mHdmiConnection == null) {
+                HdmiLogger.error("Could not initialize HDMI Connection AIDL HAL");
                 return false;
             }
             try {
-                mHdmi.asBinder().linkToDeath(this, 0);
+                mHdmiConnection.asBinder().linkToDeath(this, 0);
             } catch (RemoteException e) {
                 HdmiLogger.error("Couldn't link to death : ", e);
             }
@@ -915,8 +915,8 @@
         public void binderDied() {
             // One of the services died, try to reconnect to both.
             mHdmiCec.asBinder().unlinkToDeath(this, 0);
-            mHdmi.asBinder().unlinkToDeath(this, 0);
-            HdmiLogger.error("HDMI or CEC service died, reconnecting");
+            mHdmiConnection.asBinder().unlinkToDeath(this, 0);
+            HdmiLogger.error("HDMI Connection or CEC service died, reconnecting");
             connectToHal();
             // Reconnect the callback
             if (mCallback != null) {
@@ -935,7 +935,7 @@
             }
             try {
                 // Create an AIDL callback that can callback onHotplugEvent
-                mHdmi.setCallback(new HdmiCallbackAidl(callback));
+                mHdmiConnection.setCallback(new HdmiConnectionCallbackAidl(callback));
             } catch (RemoteException e) {
                 HdmiLogger.error("Couldn't initialise tv.hdmi callback : ", e);
             }
@@ -1052,10 +1052,11 @@
         @Override
         public HdmiPortInfo[] nativeGetPortInfos() {
             try {
-                android.hardware.tv.hdmi.HdmiPortInfo[] hdmiPortInfos = mHdmi.getPortInfo();
+                android.hardware.tv.hdmi.connection.HdmiPortInfo[] hdmiPortInfos =
+                        mHdmiConnection.getPortInfo();
                 HdmiPortInfo[] hdmiPortInfo = new HdmiPortInfo[hdmiPortInfos.length];
                 int i = 0;
-                for (android.hardware.tv.hdmi.HdmiPortInfo portInfo : hdmiPortInfos) {
+                for (android.hardware.tv.hdmi.connection.HdmiPortInfo portInfo : hdmiPortInfos) {
                     hdmiPortInfo[i] =
                             new HdmiPortInfo(
                                     portInfo.portId,
@@ -1076,7 +1077,7 @@
         @Override
         public boolean nativeIsConnected(int port) {
             try {
-                return mHdmi.isConnected(port);
+                return mHdmiConnection.isConnected(port);
             } catch (RemoteException e) {
                 HdmiLogger.error("Failed to get connection info : ", e);
                 return false;
@@ -1580,10 +1581,10 @@
         }
     }
 
-    private static final class HdmiCallbackAidl extends IHdmiCallback.Stub {
+    private static final class HdmiConnectionCallbackAidl extends IHdmiConnectionCallback.Stub {
         private final HdmiCecCallback mHdmiCecCallback;
 
-        HdmiCallbackAidl(HdmiCecCallback hdmiCecCallback) {
+        HdmiConnectionCallbackAidl(HdmiCecCallback hdmiCecCallback) {
             mHdmiCecCallback = hdmiCecCallback;
         }
 
@@ -1594,12 +1595,12 @@
 
         @Override
         public synchronized String getInterfaceHash() throws android.os.RemoteException {
-            return IHdmiCallback.Stub.HASH;
+            return IHdmiConnectionCallback.Stub.HASH;
         }
 
         @Override
         public int getInterfaceVersion() throws android.os.RemoteException {
-            return IHdmiCallback.Stub.VERSION;
+            return IHdmiConnectionCallback.Stub.VERSION;
         }
     }
 
diff --git a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
index 4d525da..9b42cfc 100644
--- a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
+++ b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
@@ -36,6 +36,7 @@
 import android.net.ConnectivityManager;
 import android.net.Network;
 import android.net.NetworkCapabilities;
+import android.net.NetworkRequest;
 import android.os.Handler;
 import android.os.PowerManager;
 import android.os.SystemClock;
@@ -103,12 +104,12 @@
 
     /**
      * Number of boots until we consider the escrow data to be stale for the purposes of metrics.
-     * <p>
-     * If the delta between the current boot number and the boot number stored when the mechanism
+     *
+     * <p>If the delta between the current boot number and the boot number stored when the mechanism
      * was armed is under this number and the escrow mechanism fails, we report it as a failure of
      * the mechanism.
-     * <p>
-     * If the delta over this number and escrow fails, we will not report the metric as failed
+     *
+     * <p>If the delta over this number and escrow fails, we will not report the metric as failed
      * since there most likely was some other issue if the device rebooted several times before
      * getting to the escrow restore code.
      */
@@ -120,8 +121,11 @@
      */
     private static final int DEFAULT_LOAD_ESCROW_DATA_RETRY_COUNT = 3;
     private static final int DEFAULT_LOAD_ESCROW_DATA_RETRY_INTERVAL_SECONDS = 30;
+
     // 3 minutes. It's enough for the default 3 retries with 30 seconds interval
-    private static final int DEFAULT_WAKE_LOCK_TIMEOUT_MILLIS = 180_000;
+    private static final int DEFAULT_LOAD_ESCROW_BASE_TIMEOUT_MILLIS = 180_000;
+    // 5 seconds. An extension of the overall RoR timeout to account for overhead.
+    private static final int DEFAULT_LOAD_ESCROW_TIMEOUT_EXTENSION_MILLIS = 5000;
 
     @IntDef(prefix = {"ERROR_"}, value = {
             ERROR_NONE,
@@ -133,6 +137,7 @@
             ERROR_PROVIDER_MISMATCH,
             ERROR_KEYSTORE_FAILURE,
             ERROR_NO_NETWORK,
+            ERROR_TIMEOUT_EXHAUSTED,
     })
     @Retention(RetentionPolicy.SOURCE)
     @interface RebootEscrowErrorCode {
@@ -147,6 +152,7 @@
     static final int ERROR_PROVIDER_MISMATCH = 6;
     static final int ERROR_KEYSTORE_FAILURE = 7;
     static final int ERROR_NO_NETWORK = 8;
+    static final int ERROR_TIMEOUT_EXHAUSTED = 9;
 
     private @RebootEscrowErrorCode int mLoadEscrowDataErrorCode = ERROR_NONE;
 
@@ -168,6 +174,15 @@
     /** Notified when mRebootEscrowReady changes. */
     private RebootEscrowListener mRebootEscrowListener;
 
+    /** Set when unlocking reboot escrow times out. */
+    private boolean mRebootEscrowTimedOut = false;
+
+    /**
+     * Set when {@link #loadRebootEscrowDataWithRetry} is called to ensure the function is only
+     * called once.
+     */
+    private boolean mLoadEscrowDataWithRetry = false;
+
     /**
      * Hold this lock when checking or generating the reboot escrow key.
      */
@@ -192,6 +207,7 @@
 
     PowerManager.WakeLock mWakeLock;
 
+    private ConnectivityManager.NetworkCallback mNetworkCallback;
 
     interface Callbacks {
         boolean isUserSecure(int userId);
@@ -246,6 +262,11 @@
                     "server_based_ror_enabled", false);
         }
 
+        public boolean waitForInternet() {
+            return DeviceConfig.getBoolean(
+                    DeviceConfig.NAMESPACE_OTA, "wait_for_internet_ror", false);
+        }
+
         public boolean isNetworkConnected() {
             final ConnectivityManager connectivityManager =
                     mContext.getSystemService(ConnectivityManager.class);
@@ -263,6 +284,38 @@
                             NetworkCapabilities.NET_CAPABILITY_VALIDATED);
         }
 
+        /**
+         * Request network with internet connectivity with timeout.
+         *
+         * @param networkCallback callback to be executed if connectivity manager exists.
+         * @return true if success
+         */
+        public boolean requestNetworkWithInternet(
+                ConnectivityManager.NetworkCallback networkCallback) {
+            final ConnectivityManager connectivityManager =
+                    mContext.getSystemService(ConnectivityManager.class);
+            if (connectivityManager == null) {
+                return false;
+            }
+            NetworkRequest request =
+                    new NetworkRequest.Builder()
+                            .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                            .build();
+
+            connectivityManager.requestNetwork(
+                    request, networkCallback, getLoadEscrowTimeoutMillis());
+            return true;
+        }
+
+        public void stopRequestingNetwork(ConnectivityManager.NetworkCallback networkCallback) {
+            final ConnectivityManager connectivityManager =
+                    mContext.getSystemService(ConnectivityManager.class);
+            if (connectivityManager == null) {
+                return;
+            }
+            connectivityManager.unregisterNetworkCallback(networkCallback);
+        }
+
         public Context getContext() {
             return mContext;
         }
@@ -318,6 +371,16 @@
                     DEFAULT_LOAD_ESCROW_DATA_RETRY_INTERVAL_SECONDS);
         }
 
+        @VisibleForTesting
+        public int getLoadEscrowTimeoutMillis() {
+            return DEFAULT_LOAD_ESCROW_BASE_TIMEOUT_MILLIS;
+        }
+
+        @VisibleForTesting
+        public int getWakeLockTimeoutMillis() {
+            return getLoadEscrowTimeoutMillis() + DEFAULT_LOAD_ESCROW_TIMEOUT_EXTENSION_MILLIS;
+        }
+
         public void reportMetric(boolean success, int errorCode, int serviceType, int attemptCount,
                 int escrowDurationInSeconds, int vbmetaDigestStatus,
                 int durationSinceBootCompleteInSeconds) {
@@ -351,13 +414,37 @@
         mKeyStoreManager = injector.getKeyStoreManager();
     }
 
-    private void onGetRebootEscrowKeyFailed(List<UserInfo> users, int attemptCount) {
+    /** Wrapper function to set error code serialized through handler, */
+    private void setLoadEscrowDataErrorCode(@RebootEscrowErrorCode int value, Handler handler) {
+        if (mInjector.waitForInternet()) {
+            mInjector.post(
+                    handler,
+                    () -> {
+                        mLoadEscrowDataErrorCode = value;
+                    });
+        } else {
+            mLoadEscrowDataErrorCode = value;
+        }
+    }
+
+    /** Wrapper function to compare and set error code serialized through handler. */
+    private void compareAndSetLoadEscrowDataErrorCode(
+            @RebootEscrowErrorCode int expectedValue,
+            @RebootEscrowErrorCode int newValue,
+            Handler handler) {
+        if (expectedValue == mLoadEscrowDataErrorCode) {
+            setLoadEscrowDataErrorCode(newValue, handler);
+        }
+    }
+
+    private void onGetRebootEscrowKeyFailed(
+            List<UserInfo> users, int attemptCount, Handler retryHandler) {
         Slog.w(TAG, "Had reboot escrow data for users, but no key; removing escrow storage.");
         for (UserInfo user : users) {
             mStorage.removeRebootEscrow(user.id);
         }
 
-        onEscrowRestoreComplete(false, attemptCount);
+        onEscrowRestoreComplete(false, attemptCount, retryHandler);
     }
 
     void loadRebootEscrowDataIfAvailable(Handler retryHandler) {
@@ -380,39 +467,130 @@
         mWakeLock = mInjector.getWakeLock();
         if (mWakeLock != null) {
             mWakeLock.setReferenceCounted(false);
-            mWakeLock.acquire(DEFAULT_WAKE_LOCK_TIMEOUT_MILLIS);
+            mWakeLock.acquire(mInjector.getWakeLockTimeoutMillis());
+        }
+
+        if (mInjector.waitForInternet()) {
+            // Timeout to stop retrying same as the wake lock timeout.
+            mInjector.postDelayed(
+                    retryHandler,
+                    () -> {
+                        mRebootEscrowTimedOut = true;
+                    },
+                    mInjector.getLoadEscrowTimeoutMillis());
+
+            mInjector.post(
+                    retryHandler,
+                    () -> loadRebootEscrowDataOnInternet(retryHandler, users, rebootEscrowUsers));
+            return;
         }
 
         mInjector.post(retryHandler, () -> loadRebootEscrowDataWithRetry(
                 retryHandler, 0, users, rebootEscrowUsers));
     }
 
-    void scheduleLoadRebootEscrowDataOrFail(Handler retryHandler, int attemptNumber,
-            List<UserInfo> users, List<UserInfo> rebootEscrowUsers) {
+    void scheduleLoadRebootEscrowDataOrFail(
+            Handler retryHandler,
+            int attemptNumber,
+            List<UserInfo> users,
+            List<UserInfo> rebootEscrowUsers) {
         Objects.requireNonNull(retryHandler);
 
         final int retryLimit = mInjector.getLoadEscrowDataRetryLimit();
         final int retryIntervalInSeconds = mInjector.getLoadEscrowDataRetryIntervalSeconds();
 
-        if (attemptNumber < retryLimit) {
+        if (attemptNumber < retryLimit && !mRebootEscrowTimedOut) {
             Slog.i(TAG, "Scheduling loadRebootEscrowData retry number: " + attemptNumber);
             mInjector.postDelayed(retryHandler, () -> loadRebootEscrowDataWithRetry(
-                    retryHandler, attemptNumber, users, rebootEscrowUsers),
+                            retryHandler, attemptNumber, users, rebootEscrowUsers),
                     retryIntervalInSeconds * 1000);
             return;
         }
 
+        if (mInjector.waitForInternet()) {
+            if (mRebootEscrowTimedOut) {
+                Slog.w(TAG, "Failed to load reboot escrow data within timeout");
+                compareAndSetLoadEscrowDataErrorCode(
+                        ERROR_NONE, ERROR_TIMEOUT_EXHAUSTED, retryHandler);
+            } else {
+                Slog.w(
+                        TAG,
+                        "Failed to load reboot escrow data after " + attemptNumber + " attempts");
+                compareAndSetLoadEscrowDataErrorCode(
+                        ERROR_NONE, ERROR_RETRY_COUNT_EXHAUSTED, retryHandler);
+            }
+            onGetRebootEscrowKeyFailed(users, attemptNumber, retryHandler);
+            return;
+        }
+
         Slog.w(TAG, "Failed to load reboot escrow data after " + attemptNumber + " attempts");
         if (mInjector.serverBasedResumeOnReboot() && !mInjector.isNetworkConnected()) {
             mLoadEscrowDataErrorCode = ERROR_NO_NETWORK;
         } else {
             mLoadEscrowDataErrorCode = ERROR_RETRY_COUNT_EXHAUSTED;
         }
-        onGetRebootEscrowKeyFailed(users, attemptNumber);
+        onGetRebootEscrowKeyFailed(users, attemptNumber, retryHandler);
     }
 
-    void loadRebootEscrowDataWithRetry(Handler retryHandler, int attemptNumber,
-            List<UserInfo> users, List<UserInfo> rebootEscrowUsers) {
+    void loadRebootEscrowDataOnInternet(
+            Handler retryHandler, List<UserInfo> users, List<UserInfo> rebootEscrowUsers) {
+
+        // HAL-Based RoR does not require network connectivity.
+        if (!mInjector.serverBasedResumeOnReboot()) {
+            loadRebootEscrowDataWithRetry(
+                    retryHandler, /* attemptNumber = */ 0, users, rebootEscrowUsers);
+            return;
+        }
+
+        mNetworkCallback =
+                new ConnectivityManager.NetworkCallback() {
+                    @Override
+                    public void onAvailable(Network network) {
+                        compareAndSetLoadEscrowDataErrorCode(
+                                ERROR_NO_NETWORK, ERROR_NONE, retryHandler);
+
+                        if (!mLoadEscrowDataWithRetry) {
+                            mLoadEscrowDataWithRetry = true;
+                            // Only kickoff retry mechanism on first onAvailable call.
+                            loadRebootEscrowDataWithRetry(
+                                    retryHandler,
+                                    /* attemptNumber = */ 0,
+                                    users,
+                                    rebootEscrowUsers);
+                        }
+                    }
+
+                    @Override
+                    public void onUnavailable() {
+                        Slog.w(TAG, "Failed to connect to network within timeout");
+                        compareAndSetLoadEscrowDataErrorCode(
+                                ERROR_NONE, ERROR_NO_NETWORK, retryHandler);
+                        onGetRebootEscrowKeyFailed(users, /* attemptCount= */ 0, retryHandler);
+                    }
+
+                    @Override
+                    public void onLost(Network lostNetwork) {
+                        // TODO(b/231660348): If network is lost, wait for network to become
+                        // available again.
+                        Slog.w(TAG, "Network lost, still attempting to load escrow key.");
+                        compareAndSetLoadEscrowDataErrorCode(
+                                ERROR_NONE, ERROR_NO_NETWORK, retryHandler);
+                    }
+                };
+
+        // Fallback to retrying without waiting for internet on failure.
+        boolean success = mInjector.requestNetworkWithInternet(mNetworkCallback);
+        if (!success) {
+            loadRebootEscrowDataWithRetry(
+                    retryHandler, /* attemptNumber = */ 0, users, rebootEscrowUsers);
+        }
+    }
+
+    void loadRebootEscrowDataWithRetry(
+            Handler retryHandler,
+            int attemptNumber,
+            List<UserInfo> users,
+            List<UserInfo> rebootEscrowUsers) {
         // Fetch the key from keystore to decrypt the escrow data & escrow key; this key is
         // generated before reboot. Note that we will clear the escrow key even if the keystore key
         // is null.
@@ -423,7 +601,7 @@
 
         RebootEscrowKey escrowKey;
         try {
-            escrowKey = getAndClearRebootEscrowKey(kk);
+            escrowKey = getAndClearRebootEscrowKey(kk, retryHandler);
         } catch (IOException e) {
             Slog.i(TAG, "Failed to load escrow key, scheduling retry.", e);
             scheduleLoadRebootEscrowDataOrFail(retryHandler, attemptNumber + 1, users,
@@ -438,12 +616,12 @@
                         ? RebootEscrowProviderInterface.TYPE_SERVER_BASED
                         : RebootEscrowProviderInterface.TYPE_HAL;
                 if (providerType != mStorage.getInt(REBOOT_ESCROW_KEY_PROVIDER, -1, USER_SYSTEM)) {
-                    mLoadEscrowDataErrorCode = ERROR_PROVIDER_MISMATCH;
+                    setLoadEscrowDataErrorCode(ERROR_PROVIDER_MISMATCH, retryHandler);
                 } else {
-                    mLoadEscrowDataErrorCode = ERROR_LOAD_ESCROW_KEY;
+                    setLoadEscrowDataErrorCode(ERROR_LOAD_ESCROW_KEY, retryHandler);
                 }
             }
-            onGetRebootEscrowKeyFailed(users, attemptNumber + 1);
+            onGetRebootEscrowKeyFailed(users, attemptNumber + 1, retryHandler);
             return;
         }
 
@@ -454,10 +632,10 @@
             allUsersUnlocked &= restoreRebootEscrowForUser(user.id, escrowKey, kk);
         }
 
-        if (!allUsersUnlocked && mLoadEscrowDataErrorCode == ERROR_NONE) {
-            mLoadEscrowDataErrorCode = ERROR_UNLOCK_ALL_USERS;
+        if (!allUsersUnlocked) {
+            compareAndSetLoadEscrowDataErrorCode(ERROR_NONE, ERROR_UNLOCK_ALL_USERS, retryHandler);
         }
-        onEscrowRestoreComplete(allUsersUnlocked, attemptNumber + 1);
+        onEscrowRestoreComplete(allUsersUnlocked, attemptNumber + 1, retryHandler);
     }
 
     private void clearMetricsStorage() {
@@ -497,7 +675,8 @@
                 .REBOOT_ESCROW_RECOVERY_REPORTED__VBMETA_DIGEST_STATUS__MISMATCH;
     }
 
-    private void reportMetricOnRestoreComplete(boolean success, int attemptCount) {
+    private void reportMetricOnRestoreComplete(
+            boolean success, int attemptCount, Handler retryHandler) {
         int serviceType = mInjector.serverBasedResumeOnReboot()
                 ? FrameworkStatsLog.REBOOT_ESCROW_RECOVERY_REPORTED__TYPE__SERVER_BASED
                 : FrameworkStatsLog.REBOOT_ESCROW_RECOVERY_REPORTED__TYPE__HAL;
@@ -511,52 +690,69 @@
         }
 
         int vbmetaDigestStatus = getVbmetaDigestStatusOnRestoreComplete();
-        if (!success && mLoadEscrowDataErrorCode == ERROR_NONE) {
-            mLoadEscrowDataErrorCode = ERROR_UNKNOWN;
+        if (!success) {
+            compareAndSetLoadEscrowDataErrorCode(ERROR_NONE, ERROR_UNKNOWN, retryHandler);
         }
 
-        Slog.i(TAG, "Reporting RoR recovery metrics, success: " + success + ", service type: "
-                + serviceType + ", error code: " + mLoadEscrowDataErrorCode);
+        Slog.i(
+                TAG,
+                "Reporting RoR recovery metrics, success: "
+                        + success
+                        + ", service type: "
+                        + serviceType
+                        + ", error code: "
+                        + mLoadEscrowDataErrorCode);
         // TODO(179105110) report the duration since boot complete.
-        mInjector.reportMetric(success, mLoadEscrowDataErrorCode, serviceType, attemptCount,
-                escrowDurationInSeconds, vbmetaDigestStatus, -1);
+        mInjector.reportMetric(
+                success,
+                mLoadEscrowDataErrorCode,
+                serviceType,
+                attemptCount,
+                escrowDurationInSeconds,
+                vbmetaDigestStatus,
+                -1);
 
-        mLoadEscrowDataErrorCode = ERROR_NONE;
+        setLoadEscrowDataErrorCode(ERROR_NONE, retryHandler);
     }
 
-    private void onEscrowRestoreComplete(boolean success, int attemptCount) {
+    private void onEscrowRestoreComplete(boolean success, int attemptCount, Handler retryHandler) {
         int previousBootCount = mStorage.getInt(REBOOT_ESCROW_ARMED_KEY, -1, USER_SYSTEM);
 
         int bootCountDelta = mInjector.getBootCount() - previousBootCount;
         if (success || (previousBootCount != -1 && bootCountDelta <= BOOT_COUNT_TOLERANCE)) {
-            reportMetricOnRestoreComplete(success, attemptCount);
+            reportMetricOnRestoreComplete(success, attemptCount, retryHandler);
         }
-
         // Clear the old key in keystore. A new key will be generated by new RoR requests.
         mKeyStoreManager.clearKeyStoreEncryptionKey();
         // Clear the saved reboot escrow provider
         mInjector.clearRebootEscrowProvider();
         clearMetricsStorage();
 
+        if (mNetworkCallback != null) {
+            mInjector.stopRequestingNetwork(mNetworkCallback);
+        }
+
         if (mWakeLock != null) {
             mWakeLock.release();
         }
     }
 
-    private RebootEscrowKey getAndClearRebootEscrowKey(SecretKey kk) throws IOException {
+    private RebootEscrowKey getAndClearRebootEscrowKey(SecretKey kk, Handler retryHandler)
+            throws IOException {
         RebootEscrowProviderInterface rebootEscrowProvider =
                 mInjector.createRebootEscrowProviderIfNeeded();
         if (rebootEscrowProvider == null) {
-            Slog.w(TAG,
+            Slog.w(
+                    TAG,
                     "Had reboot escrow data for users, but RebootEscrowProvider is unavailable");
-            mLoadEscrowDataErrorCode = ERROR_NO_PROVIDER;
+            setLoadEscrowDataErrorCode(ERROR_NO_PROVIDER, retryHandler);
             return null;
         }
 
         // Server based RoR always need the decryption key from keystore.
         if (rebootEscrowProvider.getType() == RebootEscrowProviderInterface.TYPE_SERVER_BASED
                 && kk == null) {
-            mLoadEscrowDataErrorCode = ERROR_KEYSTORE_FAILURE;
+            setLoadEscrowDataErrorCode(ERROR_KEYSTORE_FAILURE, retryHandler);
             return null;
         }
 
@@ -870,6 +1066,9 @@
         pw.print("mRebootEscrowListener=");
         pw.println(mRebootEscrowListener);
 
+        pw.print("mLoadEscrowDataErrorCode=");
+        pw.println(mLoadEscrowDataErrorCode);
+
         boolean keySet;
         synchronized (mKeyGenerationLock) {
             keySet = mPendingRebootEscrowKey != null;
diff --git a/services/core/java/com/android/server/pm/dex/ArtStatsLogUtils.java b/services/core/java/com/android/server/pm/dex/ArtStatsLogUtils.java
index 1407530..e056151 100644
--- a/services/core/java/com/android/server/pm/dex/ArtStatsLogUtils.java
+++ b/services/core/java/com/android/server/pm/dex/ArtStatsLogUtils.java
@@ -298,7 +298,9 @@
                     dexMetadataType,
                     apkType,
                     ISA_MAP.getOrDefault(isa,
-                            ArtStatsLog.ART_DATUM_REPORTED__ISA__ART_ISA_UNKNOWN));
+                            ArtStatsLog.ART_DATUM_REPORTED__ISA__ART_ISA_UNKNOWN),
+                    ArtStatsLog.ART_DATUM_REPORTED__GC__ART_GC_COLLECTOR_TYPE_UNKNOWN,
+                    ArtStatsLog.ART_DATUM_REPORTED__UFFD_SUPPORT__ART_UFFD_SUPPORT_UNKNOWN);
         }
     }
 
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index ed4ba0d..7f8f406 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -3202,8 +3202,7 @@
             }
             final PowerGroup powerGroup = mPowerGroups.get(groupId);
             wakefulness = powerGroup.getWakefulnessLocked();
-            if ((wakefulness == WAKEFULNESS_DREAMING || wakefulness == WAKEFULNESS_DOZING) &&
-                    powerGroup.isSandmanSummonedLocked() && powerGroup.isReadyLocked()) {
+            if (powerGroup.isSandmanSummonedLocked() && powerGroup.isReadyLocked()) {
                 startDreaming = canDreamLocked(powerGroup) || canDozeLocked(powerGroup);
                 powerGroup.setSandmanSummonedLocked(/* isSandmanSummoned= */ false);
             } else {
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
index b01c1c8..ce6bd6c 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
@@ -50,6 +50,8 @@
 import android.content.ContextWrapper;
 import android.content.pm.UserInfo;
 import android.hardware.rebootescrow.IRebootEscrow;
+import android.net.ConnectivityManager;
+import android.net.Network;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.RemoteException;
@@ -72,6 +74,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.function.Consumer;
 
 import javax.crypto.SecretKey;
 import javax.crypto.spec.SecretKeySpec;
@@ -113,6 +116,7 @@
     private RebootEscrowManager mService;
     private SecretKey mAesKey;
     private MockInjector mMockInjector;
+    private Handler mHandler;
 
     public interface MockableRebootEscrowInjected {
         int getBootCount();
@@ -132,6 +136,9 @@
         private final RebootEscrowKeyStoreManager mKeyStoreManager;
         private boolean mServerBased;
         private RebootEscrowProviderInterface mRebootEscrowProviderInUse;
+        private ConnectivityManager.NetworkCallback mNetworkCallback;
+        private Consumer<ConnectivityManager.NetworkCallback> mNetworkConsumer;
+        private boolean mWaitForInternet;
 
         MockInjector(Context context, UserManager userManager,
                 IRebootEscrow rebootEscrow,
@@ -142,6 +149,7 @@
             mRebootEscrow = rebootEscrow;
             mServiceConnection = null;
             mServerBased = false;
+            mWaitForInternet = false;
             RebootEscrowProviderHalImpl.Injector halInjector =
                     new RebootEscrowProviderHalImpl.Injector() {
                         @Override
@@ -164,6 +172,7 @@
             mServiceConnection = serviceConnection;
             mRebootEscrow = null;
             mServerBased = true;
+            mWaitForInternet = false;
             RebootEscrowProviderServerBasedImpl.Injector injector =
                     new RebootEscrowProviderServerBasedImpl.Injector(serviceConnection) {
                         @Override
@@ -199,11 +208,33 @@
         }
 
         @Override
+        public boolean waitForInternet() {
+            return mWaitForInternet;
+        }
+
+        public void setWaitForNetwork(boolean waitForNetworkEnabled) {
+            mWaitForInternet = waitForNetworkEnabled;
+        }
+
+        @Override
         public boolean isNetworkConnected() {
             return false;
         }
 
         @Override
+        public boolean requestNetworkWithInternet(
+                ConnectivityManager.NetworkCallback networkCallback) {
+            mNetworkCallback = networkCallback;
+            mNetworkConsumer.accept(networkCallback);
+            return true;
+        }
+
+        @Override
+        public void stopRequestingNetwork(ConnectivityManager.NetworkCallback networkCallback) {
+            mNetworkCallback = null;
+        }
+
+        @Override
         public RebootEscrowProviderInterface createRebootEscrowProviderIfNeeded() {
             mRebootEscrowProviderInUse = mDefaultRebootEscrowProvider;
             return mRebootEscrowProviderInUse;
@@ -242,6 +273,12 @@
         }
 
         @Override
+        public int getLoadEscrowTimeoutMillis() {
+            // Timeout in 3 seconds.
+            return 3000;
+        }
+
+        @Override
         public String getVbmetaDigest(boolean other) {
             return other ? "" : "fake digest";
         }
@@ -291,6 +328,9 @@
         mMockInjector = new MockInjector(mContext, mUserManager, mRebootEscrow,
                 mKeyStoreManager, mStorage, mInjected);
         mService = new RebootEscrowManager(mMockInjector, mCallbacks, mStorage);
+        HandlerThread thread = new HandlerThread("RebootEscrowManagerTest");
+        thread.start();
+        mHandler = new Handler(thread.getLooper());
     }
 
     private void setServerBasedRebootEscrowProvider() throws Exception {
@@ -462,7 +502,7 @@
 
     @Test
     public void loadRebootEscrowDataIfAvailable_NothingAvailable_Success() throws Exception {
-        mService.loadRebootEscrowDataIfAvailable(null);
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
     }
 
     @Test
@@ -499,7 +539,7 @@
                 eq(20), eq(0) /* vbmeta status */, anyInt());
         when(mRebootEscrow.retrieveKey()).thenAnswer(invocation -> keyByteCaptor.getValue());
 
-        mService.loadRebootEscrowDataIfAvailable(null);
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
         verify(mRebootEscrow).retrieveKey();
         assertTrue(metricsSuccessCaptor.getValue());
         verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
@@ -531,9 +571,16 @@
         // pretend reboot happens here
         when(mInjected.getBootCount()).thenReturn(1);
         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
-        doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture(),
-                eq(0) /* error code */, eq(2) /* Server based */, eq(1) /* attempt count */,
-                anyInt(), eq(0) /* vbmeta status */, anyInt());
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        eq(0) /* error code */,
+                        eq(2) /* Server based */,
+                        eq(1) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
 
         when(mServiceConnection.unwrap(any(), anyLong()))
                 .thenAnswer(invocation -> invocation.getArgument(0));
@@ -569,15 +616,23 @@
         when(mInjected.getBootCount()).thenReturn(1);
         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
-        doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture(),
-                metricsErrorCodeCaptor.capture(), eq(2) /* Server based */,
-                eq(1) /* attempt count */, anyInt(), eq(0) /* vbmeta status */, anyInt());
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(1) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
 
         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(RemoteException.class);
         mService.loadRebootEscrowDataIfAvailable(null);
         verify(mServiceConnection).unwrap(any(), anyLong());
         assertFalse(metricsSuccessCaptor.getValue());
-        assertEquals(Integer.valueOf(RebootEscrowManager.ERROR_LOAD_ESCROW_KEY),
+        assertEquals(
+                Integer.valueOf(RebootEscrowManager.ERROR_LOAD_ESCROW_KEY),
                 metricsErrorCodeCaptor.getValue());
     }
 
@@ -606,18 +661,24 @@
         when(mInjected.getBootCount()).thenReturn(1);
         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
-        doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture(),
-                metricsErrorCodeCaptor.capture(), eq(2) /* Server based */,
-                eq(2) /* attempt count */, anyInt(), eq(0) /* vbmeta status */, anyInt());
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(2) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(IOException.class);
 
-        HandlerThread thread = new HandlerThread("RebootEscrowManagerTest");
-        thread.start();
-        mService.loadRebootEscrowDataIfAvailable(new Handler(thread.getLooper()));
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
         // Sleep 5s for the retry to complete
         Thread.sleep(5 * 1000);
         assertFalse(metricsSuccessCaptor.getValue());
-        assertEquals(Integer.valueOf(RebootEscrowManager.ERROR_NO_NETWORK),
+        assertEquals(
+                Integer.valueOf(RebootEscrowManager.ERROR_NO_NETWORK),
                 metricsErrorCodeCaptor.getValue());
     }
 
@@ -645,16 +706,22 @@
         // pretend reboot happens here
         when(mInjected.getBootCount()).thenReturn(1);
         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
-        doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture(),
-                anyInt(), anyInt(), eq(2) /* attempt count */, anyInt(), anyInt(), anyInt());
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        anyInt(),
+                        anyInt(),
+                        eq(2) /* attempt count */,
+                        anyInt(),
+                        anyInt(),
+                        anyInt());
 
         when(mServiceConnection.unwrap(any(), anyLong()))
                 .thenThrow(new IOException())
                 .thenAnswer(invocation -> invocation.getArgument(0));
 
-        HandlerThread thread = new HandlerThread("RebootEscrowManagerTest");
-        thread.start();
-        mService.loadRebootEscrowDataIfAvailable(new Handler(thread.getLooper()));
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
         // Sleep 5s for the retry to complete
         Thread.sleep(5 * 1000);
         verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
@@ -663,6 +730,447 @@
     }
 
     @Test
+    public void loadRebootEscrowDataIfAvailable_serverBasedWaitForInternet_success()
+            throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        eq(0) /* error code */,
+                        eq(2) /* Server based */,
+                        eq(1) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
+
+        // load escrow data
+        when(mServiceConnection.unwrap(any(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        Network mockNetwork = mock(Network.class);
+        mMockInjector.mNetworkConsumer =
+                (callback) -> {
+                    callback.onAvailable(mockNetwork);
+                };
+
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        verify(mServiceConnection).unwrap(any(), anyLong());
+        assertTrue(metricsSuccessCaptor.getValue());
+        verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
+    public void loadRebootEscrowDataIfAvailable_serverBasedWaitForInternetRemoteException_Failure()
+            throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(1) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
+
+        // load escrow data
+        when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(RemoteException.class);
+        Network mockNetwork = mock(Network.class);
+        mMockInjector.mNetworkConsumer =
+                (callback) -> {
+                    callback.onAvailable(mockNetwork);
+                };
+
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        verify(mServiceConnection).unwrap(any(), anyLong());
+        assertFalse(metricsSuccessCaptor.getValue());
+        assertEquals(
+                Integer.valueOf(RebootEscrowManager.ERROR_LOAD_ESCROW_KEY),
+                metricsErrorCodeCaptor.getValue());
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
+    public void loadRebootEscrowDataIfAvailable_waitForInternet_networkUnavailable()
+            throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(0) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
+
+        // Network is not available within timeout.
+        mMockInjector.mNetworkConsumer = ConnectivityManager.NetworkCallback::onUnavailable;
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        assertFalse(metricsSuccessCaptor.getValue());
+        assertEquals(
+                Integer.valueOf(RebootEscrowManager.ERROR_NO_NETWORK),
+                metricsErrorCodeCaptor.getValue());
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
+    public void loadRebootEscrowDataIfAvailable_waitForInternet_networkLost() throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(2) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
+
+        // Network is available, then lost.
+        when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(new IOException());
+        Network mockNetwork = mock(Network.class);
+        mMockInjector.mNetworkConsumer =
+                (callback) -> {
+                    callback.onAvailable(mockNetwork);
+                    callback.onLost(mockNetwork);
+                };
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        // Sleep 5s for the retry to complete
+        Thread.sleep(5 * 1000);
+        assertFalse(metricsSuccessCaptor.getValue());
+        assertEquals(
+                Integer.valueOf(RebootEscrowManager.ERROR_NO_NETWORK),
+                metricsErrorCodeCaptor.getValue());
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
+    public void loadRebootEscrowDataIfAvailable_waitForInternet_networkAvailableWithDelay()
+            throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(1) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
+
+        // load escrow data
+        when(mServiceConnection.unwrap(any(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        // network available after 1 sec
+        Network mockNetwork = mock(Network.class);
+        mMockInjector.mNetworkConsumer =
+                (callback) -> {
+                    try {
+                        Thread.sleep(1000);
+                    } catch (InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
+                    callback.onAvailable(mockNetwork);
+                };
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        verify(mServiceConnection).unwrap(any(), anyLong());
+        assertTrue(metricsSuccessCaptor.getValue());
+        verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
+    public void loadRebootEscrowDataIfAvailable_waitForInternet_timeoutExhausted()
+            throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(1) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
+
+        // load reboot escrow data
+        when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(IOException.class);
+        Network mockNetwork = mock(Network.class);
+        // wait past timeout
+        mMockInjector.mNetworkConsumer =
+                (callback) -> {
+                    try {
+                        Thread.sleep(3500);
+                    } catch (InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
+                    callback.onAvailable(mockNetwork);
+                };
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        verify(mServiceConnection).unwrap(any(), anyLong());
+        assertFalse(metricsSuccessCaptor.getValue());
+        assertEquals(
+                Integer.valueOf(RebootEscrowManager.ERROR_TIMEOUT_EXHAUSTED),
+                metricsErrorCodeCaptor.getValue());
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
+    public void loadRebootEscrowDataIfAvailable_serverBasedWaitForNetwork_retryCountExhausted()
+            throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        metricsErrorCodeCaptor.capture(),
+                        eq(2) /* Server based */,
+                        eq(2) /* attempt count */,
+                        anyInt(),
+                        eq(0) /* vbmeta status */,
+                        anyInt());
+
+        when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(new IOException());
+        Network mockNetwork = mock(Network.class);
+        mMockInjector.mNetworkConsumer =
+                (callback) -> {
+                    callback.onAvailable(mockNetwork);
+                };
+
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        // Sleep 5s for the retry to complete
+        Thread.sleep(5 * 1000);
+        verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
+        assertFalse(metricsSuccessCaptor.getValue());
+        assertEquals(
+                Integer.valueOf(RebootEscrowManager.ERROR_RETRY_COUNT_EXHAUSTED),
+                metricsErrorCodeCaptor.getValue());
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
+    public void loadRebootEscrowDataIfAvailable_ServerBasedWaitForInternet_RetrySuccess()
+            throws Exception {
+        setServerBasedRebootEscrowProvider();
+        mMockInjector.setWaitForNetwork(true);
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        doNothing()
+                .when(mInjected)
+                .reportMetric(
+                        metricsSuccessCaptor.capture(),
+                        anyInt(),
+                        anyInt(),
+                        eq(2) /* attempt count */,
+                        anyInt(),
+                        anyInt(),
+                        anyInt());
+
+        when(mServiceConnection.unwrap(any(), anyLong()))
+                .thenThrow(new IOException())
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        Network mockNetwork = mock(Network.class);
+        mMockInjector.mNetworkConsumer =
+                (callback) -> {
+                    callback.onAvailable(mockNetwork);
+                };
+
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
+        // Sleep 5s for the retry to complete
+        Thread.sleep(5 * 1000);
+        verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
+        assertTrue(metricsSuccessCaptor.getValue());
+        verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
+        assertNull(mMockInjector.mNetworkCallback);
+    }
+
+    @Test
     public void loadRebootEscrowDataIfAvailable_TooManyBootsInBetween_NoMetrics() throws Exception {
         when(mInjected.getBootCount()).thenReturn(0);
 
@@ -687,7 +1195,7 @@
         when(mInjected.getBootCount()).thenReturn(10);
         when(mRebootEscrow.retrieveKey()).thenReturn(new byte[32]);
 
-        mService.loadRebootEscrowDataIfAvailable(null);
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
         verify(mRebootEscrow).retrieveKey();
         verify(mInjected, never()).reportMetric(anyBoolean(), anyInt(), anyInt(), anyInt(),
                 anyInt(), anyInt(), anyInt());
@@ -715,7 +1223,7 @@
         when(mInjected.getBootCount()).thenReturn(10);
         when(mRebootEscrow.retrieveKey()).thenReturn(new byte[32]);
 
-        mService.loadRebootEscrowDataIfAvailable(null);
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
         verify(mInjected, never()).reportMetric(anyBoolean(), anyInt(), anyInt(), anyInt(),
                 anyInt(), anyInt(), anyInt());
     }
@@ -753,7 +1261,7 @@
         // Trigger a vbmeta digest mismatch
         mStorage.setString(RebootEscrowManager.REBOOT_ESCROW_KEY_VBMETA_DIGEST,
                 "non sense value", USER_SYSTEM);
-        mService.loadRebootEscrowDataIfAvailable(null);
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
         verify(mInjected).reportMetric(eq(true), eq(0) /* error code */, eq(1) /* HAL based */,
                 eq(1) /* attempt count */, anyInt(), eq(2) /* vbmeta status */, anyInt());
         assertEquals(mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_KEY_VBMETA_DIGEST,
@@ -790,7 +1298,7 @@
                 eq(1) /* attempt count */, anyInt(), anyInt(), anyInt());
 
         when(mRebootEscrow.retrieveKey()).thenAnswer(invocation -> null);
-        mService.loadRebootEscrowDataIfAvailable(null);
+        mService.loadRebootEscrowDataIfAvailable(mHandler);
         verify(mRebootEscrow).retrieveKey();
         assertFalse(metricsSuccessCaptor.getValue());
         assertEquals(Integer.valueOf(RebootEscrowManager.ERROR_LOAD_ESCROW_KEY),