From bce963be7ab073b77d5fac4809e6ba5e6359214e Mon Sep 17 00:00:00 2001
From: Benedict Wong <benedictwong@google.com>
Date: Thu, 30 Apr 2020 17:13:22 -0700
Subject: Add a busy-wakelock to IkeSessionStateMachine

This change adds a wakelock to ensure that the device does not go to
sleep in the middle of an exchange. Otherwise, the IKE session will be
left in a non-deterministic state

Bug: 152236790
Test: FrameworkIkeTests passing
Change-Id: Ia57bb27ac75891cb47237596580792d58bfaf693
Merged-In: Ia57bb27ac75891cb47237596580792d58bfaf693
---
 .../net/ipsec/ike/IkeLocalRequestScheduler.java        |  7 +++++--
 .../internal/net/ipsec/ike/IkeSessionStateMachine.java | 18 +++++++++++++++++-
 tests/iketests/AndroidManifest.xml                     |  2 ++
 .../net/ipsec/ike/IkeSessionStateMachineTest.java      | 12 ++++++++++++
 .../internal/net/ipsec/ike/IkeSessionTestBase.java     | 12 ++++++++++++
 5 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/java/com/android/internal/net/ipsec/ike/IkeLocalRequestScheduler.java b/src/java/com/android/internal/net/ipsec/ike/IkeLocalRequestScheduler.java
index e86df232..f8008fbe 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeLocalRequestScheduler.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeLocalRequestScheduler.java
@@ -64,15 +64,18 @@ public final class IkeLocalRequestScheduler {
      * Notifies the scheduler that the caller is ready for a new procedure
      *
      * <p>Synchronously triggers the call to onNewProcedureReady.
+     *
+     * @return whether or not a new procedure was scheduled.
      */
-    public void readyForNextProcedure() {
+    public boolean readyForNextProcedure() {
         while (!mRequestQueue.isEmpty()) {
             LocalRequest request = mRequestQueue.poll();
             if (!request.isCancelled()) {
                 mConsumer.onNewProcedureReady(request);
-                return;
+                return true;
             }
         }
+        return false;
     }
 
     /**
diff --git a/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java b/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
index 5fbcda89..46517159 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
@@ -22,6 +22,7 @@ import static android.net.ipsec.ike.exceptions.IkeProtocolException.ERROR_TYPE_I
 import static android.net.ipsec.ike.exceptions.IkeProtocolException.ERROR_TYPE_NO_ADDITIONAL_SAS;
 import static android.net.ipsec.ike.exceptions.IkeProtocolException.ERROR_TYPE_TEMPORARY_FAILURE;
 import static android.net.ipsec.ike.exceptions.IkeProtocolException.ErrorType;
+import static android.os.PowerManager.PARTIAL_WAKE_LOCK;
 
 import static com.android.internal.net.ipsec.ike.message.IkeConfigPayload.CONFIG_TYPE_REPLY;
 import static com.android.internal.net.ipsec.ike.message.IkeHeader.EXCHANGE_TYPE_INFORMATIONAL;
@@ -74,6 +75,7 @@ import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
+import android.os.PowerManager;
 import android.os.SystemClock;
 import android.system.ErrnoException;
 import android.system.Os;
@@ -368,6 +370,9 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
      */
     private final IpSecSpiGenerator mIpSecSpiGenerator;
 
+    /** Ensures that the system does not go to sleep in the middle of an exchange. */
+    private final PowerManager.WakeLock mBusyWakeLock;
+
     @VisibleForTesting
     @GuardedBy("mChildCbToSessions")
     final HashMap<ChildSessionCallback, ChildSessionStateMachine> mChildCbToSessions =
@@ -495,6 +500,10 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
             // callback instance in the future
         }
 
+        PowerManager pm = context.getSystemService(PowerManager.class);
+        mBusyWakeLock = pm.newWakeLock(PARTIAL_WAKE_LOCK, TAG + "mBusyWakeLock");
+        mBusyWakeLock.setReferenceCounted(false);
+
         mIkeSessionId = sIkeSessionIdGenerator.getAndIncrement();
         sIkeAlarmReceiver.registerIkeSession(mIkeSessionId, getHandler());
 
@@ -551,6 +560,7 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
                             sendMessageAtFrontOfQueue(CMD_EXECUTE_LOCAL_REQ, localReq);
                         });
 
+        mBusyWakeLock.acquire();
         start();
     }
 
@@ -963,6 +973,8 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
             }
             // TODO: Remove the stored ikeSessionCallback
         }
+
+        mBusyWakeLock.release();
     }
 
     private void closeAllSaRecords(boolean expectSaClosed) {
@@ -1065,7 +1077,9 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
 
         @Override
         public void enterState() {
-            mScheduler.readyForNextProcedure();
+            if (!mScheduler.readyForNextProcedure()) {
+                mBusyWakeLock.release();
+            }
 
             if (mDpdIntent == null) {
                 long remoteIkeSpi = mCurrentIkeSaRecord.getRemoteSpi();
@@ -1097,6 +1111,8 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
             // #exitState is guaranteed to be invoked when quit() or quitNow() is called
             mAlarmManager.cancel(mDpdIntent);
             logd("DPD Alarm canceled");
+
+            mBusyWakeLock.acquire();
         }
 
         @Override
diff --git a/tests/iketests/AndroidManifest.xml b/tests/iketests/AndroidManifest.xml
index 0abddbaa..bdcfdb71 100644
--- a/tests/iketests/AndroidManifest.xml
+++ b/tests/iketests/AndroidManifest.xml
@@ -22,6 +22,8 @@
     <uses-permission android:name="android.permission.INTERNET"/>
     <!--Allow tests to call ConnectivityManager#getActiveNetwork()-->
     <uses-permission android:name="android.permission.ACCESS_NETWORK_STATE"/>
+    <!--Allow tests to use wake locks-->
+    <uses-permission android:name="android.permission.WAKE_LOCK"/>
 
     <!--
        'debuggable=true' is required to properly load mockito jvmti dependencies,
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java
index f5f6765d..523040c2 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java
@@ -1180,6 +1180,7 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
         mLooper.dispatchAll();
 
         verify(mSpyCurrentIkeSocket).releaseReference(eq(mIkeSessionStateMachine));
+        verify(mMockBusyWakelock).release();
     }
 
     @Test
@@ -1435,6 +1436,8 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
 
     /** Initializes the mIkeSessionStateMachine in the IDLE state. */
     private void setupIdleStateMachine() throws Exception {
+        verify(mMockBusyWakelock).acquire();
+
         setIkeInitResults();
 
         mIkeSessionStateMachine.sendMessage(
@@ -1447,6 +1450,11 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
 
         assertTrue(
                 mIkeSessionStateMachine.getCurrentState() instanceof IkeSessionStateMachine.Idle);
+
+        verify(mMockBusyWakelock).release();
+
+        // For convenience to verify wakelocks in all other places.
+        reset(mMockBusyWakelock);
     }
 
     private void mockIkeInitAndTransitionToIkeAuth(State authState) throws Exception {
@@ -1608,6 +1616,7 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
                 mIkeSessionStateMachine.getCurrentState()
                         instanceof IkeSessionStateMachine.ChildProcedureOngoing);
         verify(mMockChildSessionStateMachine).deleteChildSession();
+        verify(mMockBusyWakelock).acquire();
     }
 
     @Test
@@ -1642,6 +1651,7 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
                 mIkeSessionStateMachine.getCurrentState()
                         instanceof IkeSessionStateMachine.ChildProcedureOngoing);
         verify(mMockChildSessionStateMachine).rekeyChildSession();
+        verify(mMockBusyWakelock).acquire();
     }
 
     @Test
@@ -1781,6 +1791,7 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
         List<IkePayload> payloadList = verifyOutInfoMsgHeaderAndGetPayloads(true /*isResp*/);
         assertEquals(1, payloadList.size());
         assertEquals(outDelPayload, ((IkeDeletePayload) payloadList.get(0)));
+        verify(mMockBusyWakelock).acquire();
     }
 
     @Test
@@ -4059,6 +4070,7 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
 
         // Verify state machine quit properly
         assertNull(mIkeSessionStateMachine.getCurrentState());
+        verify(mMockBusyWakelock).release();
     }
 
     @Test
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java
index 3728f368..566968ec 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java
@@ -18,8 +18,11 @@ package com.android.internal.net.ipsec.ike;
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 
@@ -33,6 +36,7 @@ import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.SocketKeepalive;
 import android.os.Handler;
+import android.os.PowerManager;
 
 import com.android.internal.net.ipsec.ike.testutils.MockIpSecTestUtils;
 import com.android.internal.net.ipsec.ike.utils.IkeAlarmReceiver;
@@ -50,9 +54,12 @@ public abstract class IkeSessionTestBase {
             (Inet4Address) (InetAddresses.parseNumericAddress("127.0.0.1"));
     protected static final String REMOTE_HOSTNAME = "ike.test.android.com";
 
+    protected PowerManager.WakeLock mMockBusyWakelock;
+
     protected MockIpSecTestUtils mMockIpSecTestUtils;
     protected Context mSpyContext;
     protected IpSecManager mIpSecManager;
+    protected PowerManager mPowerManager;
 
     protected ConnectivityManager mMockConnectManager;
     protected Network mMockDefaultNetwork;
@@ -74,6 +81,11 @@ public abstract class IkeSessionTestBase {
                         any(Handler.class));
         doNothing().when(mSpyContext).unregisterReceiver(any(IkeAlarmReceiver.class));
 
+        mPowerManager = mock(PowerManager.class);
+        mMockBusyWakelock = mock(PowerManager.WakeLock.class);
+        doReturn(mPowerManager).when(mSpyContext).getSystemService(eq(PowerManager.class));
+        doReturn(mMockBusyWakelock).when(mPowerManager).newWakeLock(anyInt(), anyString());
+
         mMockConnectManager = mock(ConnectivityManager.class);
         mMockDefaultNetwork = mock(Network.class);
         doReturn(mMockDefaultNetwork).when(mMockConnectManager).getActiveNetwork();
-- 
cgit v1.2.3