Merge changes I370ecc85,If4bbf9c8

* changes:
  Use AlarmManager to trigger rekey
  Reschedule rekey in rekey state instead of TempFailureHandler

7 files changed, 162 insertions, 168 deletions

diff --git a/src/java/com/android/internal/net/ipsec/ike/AbstractSessionStateMachine.java b/src/java/com/android/internal/net/ipsec/ike/AbstractSessionStateMachine.java
index 458f2610..1ed89b97 100644
--- a/src/java/com/android/internal/net/ipsec/ike/AbstractSessionStateMachine.java
+++ b/src/java/com/android/internal/net/ipsec/ike/AbstractSessionStateMachine.java
@@ -78,6 +78,9 @@ abstract class AbstractSessionStateMachine extends StateMachine {
     // Use a value greater than the retransmit-failure timeout.
     static final long REKEY_DELETE_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(180L);
 
+    // Default delay time for retrying a request
+    static final long RETRY_INTERVAL_MS = TimeUnit.SECONDS.toMillis(15L);
+
     protected final Executor mUserCbExecutor;
     private final String mLogTag;
 
diff --git a/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachine.java b/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachine.java
index 8e244946..c92f4654 100644
--- a/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachine.java
+++ b/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachine.java
@@ -79,7 +79,6 @@ import com.android.internal.net.ipsec.ike.crypto.IkeMacPrf;
 import com.android.internal.net.ipsec.ike.exceptions.InvalidKeException;
 import com.android.internal.net.ipsec.ike.exceptions.InvalidSyntaxException;
 import com.android.internal.net.ipsec.ike.exceptions.NoValidProposalChosenException;
-import com.android.internal.net.ipsec.ike.exceptions.TemporaryFailureException;
 import com.android.internal.net.ipsec.ike.exceptions.TsUnacceptableException;
 import com.android.internal.net.ipsec.ike.message.IkeConfigPayload;
 import com.android.internal.net.ipsec.ike.message.IkeConfigPayload.ConfigAttribute;
@@ -749,12 +748,9 @@ public class ChildSessionStateMachine extends AbstractSessionStateMachine {
                                         mSkD,
                                         mChildSessionParams.isTransportMode(),
                                         true /*isLocalInit*/,
-                                        rekeyLocalRequest,
                                         buildSaLifetimeAlarmSched(
                                                 createChildResult.respSpi.getSpi()));
 
-                        mChildSmCallback.scheduleLocalRequest(rekeyLocalRequest, getRekeyTimeout());
-
                         ChildSessionConfiguration sessionConfig =
                                 buildChildSessionConfigFromResp(createChildResult, respPayloads);
                         executeUserCallback(
@@ -1292,8 +1288,7 @@ public class ChildSessionStateMachine extends AbstractSessionStateMachine {
                         ChildSessionStateMachine.this);
             } catch (SpiUnavailableException | ResourceUnavailableException e) {
                 loge("Fail to assign Child SPI. Schedule a retry for rekey Child");
-                mChildSmCallback.scheduleRetryLocalRequest(
-                        (ChildLocalRequest) mCurrentChildSaRecord.getFutureRekeyEvent());
+                mCurrentChildSaRecord.rescheduleRekey(RETRY_INTERVAL_MS);
                 transitionTo(mIdle);
             }
         }
@@ -1338,13 +1333,9 @@ public class ChildSessionStateMachine extends AbstractSessionStateMachine {
                                                 mSkD,
                                                 mChildSessionParams.isTransportMode(),
                                                 true /*isLocalInit*/,
-                                                rekeyLocalRequest,
                                                 buildSaLifetimeAlarmSched(
                                                         createChildResult.respSpi.getSpi()));
 
-                                mChildSmCallback.scheduleLocalRequest(
-                                        rekeyLocalRequest, getRekeyTimeout());
-
                                 executeUserCallback(
                                         () -> {
                                             mUserCallback.onIpSecTransformCreated(
@@ -1373,18 +1364,10 @@ public class ChildSessionStateMachine extends AbstractSessionStateMachine {
                                     resp.registeredSpi, createChildResult.exception);
                             break;
                         case CREATE_STATUS_CHILD_ERROR_RCV_NOTIFY:
-                            if (createChildResult.exception instanceof TemporaryFailureException) {
-                                loge(
-                                        "Received TEMPORARY_FAILURE for rekey Child. Retry has"
-                                                + "already been scheduled by IKE Session.");
-                            } else {
-                                loge(
-                                        "Received error notification for rekey Child. Schedule a"
-                                                + " retry");
-                                mChildSmCallback.scheduleRetryLocalRequest(
-                                        (ChildLocalRequest)
-                                                mCurrentChildSaRecord.getFutureRekeyEvent());
-                            }
+                            loge(
+                                    "Received error notification for rekey Child. Schedule a"
+                                            + " retry");
+                            mCurrentChildSaRecord.rescheduleRekey(RETRY_INTERVAL_MS);
 
                             transitionTo(mIdle);
                             break;
@@ -1551,12 +1534,9 @@ public class ChildSessionStateMachine extends AbstractSessionStateMachine {
                                         mSkD,
                                         mChildSessionParams.isTransportMode(),
                                         false /*isLocalInit*/,
-                                        rekeyLocalRequest,
                                         buildSaLifetimeAlarmSched(
                                                 createChildResult.initSpi.getSpi()));
 
-                        mChildSmCallback.scheduleLocalRequest(rekeyLocalRequest, getRekeyTimeout());
-
                         mChildSmCallback.onChildSaCreated(
                                 mRemoteInitNewChildSaRecord.getRemoteSpi(),
                                 ChildSessionStateMachine.this);
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..ec6d26a7 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
@@ -214,9 +214,6 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
     // Default fragment size in bytes.
     @VisibleForTesting static final int DEFAULT_FRAGMENT_SIZE = 1280;
 
-    // Default delay time for retrying a request
-    @VisibleForTesting static final long RETRY_INTERVAL_MS = TimeUnit.SECONDS.toMillis(15L);
-
     // Close IKE Session when all responses during this time were TEMPORARY_FAILURE(s). This
     // indicates that something has gone wrong, and we are out of sync.
     @VisibleForTesting
@@ -731,15 +728,13 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
             }
         }
 
-        /** Schedule retry when a request got rejected by TEMPORARY_FAILURE. */
-        public void handleTempFailure(LocalRequest localRequest) {
-            logd(
-                    "TempFailureHandler: Receive TEMPORARY FAILURE. Reschedule request: "
-                            + localRequest.procedureType);
-
-            // TODO: Support customized delay time when this is a rekey request and SA is going to
-            // expire soon.
-            scheduleRetry(localRequest);
+        /**
+         * Schedule temporary failure timeout.
+         *
+         * <p>Caller of this method is responsible for scheduling retry of the rejected request.
+         */
+        public void handleTempFailure() {
+            logd("TempFailureHandler: Receive TEMPORARY FAILURE");
 
             if (!mTempFailureReceived) {
                 sendEmptyMessageDelayed(TEMP_FAILURE_RETRY_TIMEOUT, TEMP_FAILURE_RETRY_TIMEOUT_MS);
@@ -773,8 +768,6 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
         // IkeSaRecord is created. Calling this method at the end of exchange will double-register
         // the SPI but it is safe because the key and value are not changed.
         mIkeSocket.registerIke(record.getLocalSpi(), this);
-
-        scheduleRekeySession(record.getFutureRekeyEvent());
     }
 
     @VisibleForTesting
@@ -1200,7 +1193,8 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
         return obtainMessage(CMD_ALARM_FIRED, mIkeSessionId, localRequestType, spiBundle);
     }
 
-    private SaLifetimeAlarmScheduler buildSaLifetimeAlarmScheduler(long remoteSpi) {
+    @VisibleForTesting
+    SaLifetimeAlarmScheduler buildSaLifetimeAlarmScheduler(long remoteSpi) {
         PendingIntent deleteSaIntent =
                 buildIkeAlarmIntent(
                         mContext,
@@ -1470,14 +1464,12 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
                     // Software keepalive alarm is fired
                     mIkeNattKeepalive.onAlarmFired();
                     return;
-                case CMD_LOCAL_REQUEST_DELETE_CHILD:
-                    // Child SA (identified by remoteChildSpi) has hit its hard lifetime
+                case CMD_LOCAL_REQUEST_DELETE_CHILD: // Hits hard lifetime; fall through
+                case CMD_LOCAL_REQUEST_REKEY_CHILD: // Hits soft lifetime
                     enqueueChildLocalRequest(message);
                     return;
-                case CMD_LOCAL_REQUEST_DELETE_IKE:
-                    // IKE SA hits its hard lifetime
-                    enqueueIkeLocalRequest(message);
-                    return;
+                case CMD_LOCAL_REQUEST_DELETE_IKE: // Hits hard lifetime; fall through
+                case CMD_LOCAL_REQUEST_REKEY_IKE: // Hits soft lifetime; fall through
                 case CMD_LOCAL_REQUEST_DPD:
                     // IKE Session has not received any protectd IKE packet for the whole DPD delay
                     enqueueIkeLocalRequest(message);
@@ -2215,7 +2207,9 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
 
         @Override
         protected void handleTempFailure() {
-            mTempFailHandler.handleTempFailure(mLocalRequestOngoing);
+            // The ChildSessionStateMachine will be responsible for rescheduling the rejected
+            // request.
+            mTempFailHandler.handleTempFailure();
         }
 
         private void transitionToIdleIfAllProceduresDone() {
@@ -2694,7 +2688,6 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
                                 mIkePrf,
                                 mIkeIntegrity == null ? 0 : mIkeIntegrity.getKeyLength(),
                                 mIkeCipher.getKeyLength(),
-                                new IkeLocalRequest(CMD_LOCAL_REQUEST_REKEY_IKE),
                                 buildSaLifetimeAlarmScheduler(mRemoteIkeSpiResource.getSpi()));
 
                 addIkeSaRecord(mCurrentIkeSaRecord);
@@ -3879,7 +3872,7 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
                                 + firstErrorNotifyPayload.notifyType
                                 + " for rekey IKE. Schedule a retry");
                 if (!isSimulRekey) {
-                    scheduleRetry(mCurrentIkeSaRecord.getFutureRekeyEvent());
+                    mCurrentIkeSaRecord.rescheduleRekey(RETRY_INTERVAL_MS);
                 }
             }
 
@@ -3941,7 +3934,6 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
                                 newIntegrity == null ? 0 : newIntegrity.getKeyLength(),
                                 newCipher.getKeyLength(),
                                 isLocalInit,
-                                new IkeLocalRequest(CMD_LOCAL_REQUEST_REKEY_IKE),
                                 buildSaLifetimeAlarmScheduler(remoteSpi));
                 addIkeSaRecord(newSaRecord);
 
@@ -3970,7 +3962,7 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
                 mRetransmitter = new EncryptedRetransmitter(buildIkeRekeyReq());
             } catch (IOException e) {
                 loge("Fail to assign IKE SPI for rekey. Schedule a retry.", e);
-                scheduleRetry(mCurrentIkeSaRecord.getFutureRekeyEvent());
+                mCurrentIkeSaRecord.rescheduleRekey(RETRY_INTERVAL_MS);
                 transitionTo(mIdle);
             }
         }
@@ -3982,7 +3974,8 @@ public class IkeSessionStateMachine extends AbstractSessionStateMachine {
 
         @Override
         protected void handleTempFailure() {
-            mTempFailHandler.handleTempFailure(mCurrentIkeSaRecord.getFutureRekeyEvent());
+            mTempFailHandler.handleTempFailure();
+            mCurrentIkeSaRecord.rescheduleRekey(RETRY_INTERVAL_MS);
         }
 
         /**
diff --git a/src/java/com/android/internal/net/ipsec/ike/SaRecord.java b/src/java/com/android/internal/net/ipsec/ike/SaRecord.java
index a0b585fc..c2d86dc9 100644
--- a/src/java/com/android/internal/net/ipsec/ike/SaRecord.java
+++ b/src/java/com/android/internal/net/ipsec/ike/SaRecord.java
@@ -31,8 +31,6 @@ import android.os.SystemClock;
 import android.util.CloseGuard;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.net.ipsec.ike.IkeLocalRequestScheduler.ChildLocalRequest;
-import com.android.internal.net.ipsec.ike.IkeLocalRequestScheduler.LocalRequest;
 import com.android.internal.net.ipsec.ike.crypto.IkeCipher;
 import com.android.internal.net.ipsec.ike.crypto.IkeMacIntegrity;
 import com.android.internal.net.ipsec.ike.crypto.IkeMacPrf;
@@ -82,10 +80,7 @@ public abstract class SaRecord implements AutoCloseable {
     private final byte[] mSkEi;
     private final byte[] mSkEr;
 
-    private final SaLifetimeAlarmScheduler mSaLifetimeAlarmScheduler;
-
-    // TODO(b/149058810): Use AlarmManager and PendingIntent to schedule rekey
-    private final LocalRequest mFutureRekeyEvent;
+    @VisibleForTesting final SaLifetimeAlarmScheduler mSaLifetimeAlarmScheduler;
 
     private final CloseGuard mCloseGuard = new CloseGuard();
 
@@ -98,7 +93,6 @@ public abstract class SaRecord implements AutoCloseable {
             byte[] skAr,
             byte[] skEi,
             byte[] skEr,
-            LocalRequest futureRekeyEvent,
             SaLifetimeAlarmScheduler saLifetimeAlarmScheduler) {
         isLocalInit = localInit;
         nonceInitiator = nonceInit;
@@ -117,9 +111,6 @@ public abstract class SaRecord implements AutoCloseable {
         mSaLifetimeAlarmScheduler = saLifetimeAlarmScheduler;
         mSaLifetimeAlarmScheduler.scheduleLifetimeExpiryAlarm(getTag());
 
-        // TODO(b/149058810): Use alarmManager to schedule rekey event and remove mFutureRekeyEvent
-        mFutureRekeyEvent = futureRekeyEvent;
-
         mCloseGuard.open("close");
     }
 
@@ -167,6 +158,11 @@ public abstract class SaRecord implements AutoCloseable {
         return isLocalInit ? mSkEr : mSkEi;
     }
 
+    /** Reschedule rekey */
+    public void rescheduleRekey(long retryDelayMs) {
+        mSaLifetimeAlarmScheduler.rescheduleRekey(retryDelayMs);
+    }
+
     /** Check that the SaRecord was closed properly. */
     @Override
     protected void finalize() throws Throwable {
@@ -178,17 +174,10 @@ public abstract class SaRecord implements AutoCloseable {
 
     @Override
     public void close() {
-        mFutureRekeyEvent.cancel();
-
         mSaLifetimeAlarmScheduler.cancelLifetimeExpiryAlarm(getTag());
     }
 
     /** Package private */
-    LocalRequest getFutureRekeyEvent() {
-        return mFutureRekeyEvent;
-    }
-
-    /** Package private */
     @VisibleForTesting
     static void setSaRecordHelper(ISaRecordHelper helper) {
         sSaRecordHelper = helper;
@@ -336,7 +325,6 @@ public abstract class SaRecord implements AutoCloseable {
                     skEr,
                     skPi,
                     skPr,
-                    ikeSaRecordConfig.futureRekeyEvent,
                     ikeSaRecordConfig.saLifetimeAlarmScheduler);
         }
 
@@ -471,8 +459,7 @@ public abstract class SaRecord implements AutoCloseable {
                         skEr,
                         inTransform,
                         outTransform,
-                        childSaRecordConfig.futureRekeyEvent,
-                        childSaRecordConfig.saLifetimeAlarmSched);
+                        childSaRecordConfig.saLifetimeAlarmScheduler);
 
             } catch (Exception e) {
                 if (initTransform != null) initTransform.close();
@@ -560,9 +547,26 @@ public abstract class SaRecord implements AutoCloseable {
                     AlarmManager.ELAPSED_REALTIME_WAKEUP,
                     SystemClock.elapsedRealtime() + mDeleteDelayMs,
                     mDeleteSaIntent);
-            getIkeLog().d(tag, "Hard lifetime expiry alarm set for " + mDeleteDelayMs + "ms");
+            mAlarmManager.setExactAndAllowWhileIdle(
+                    AlarmManager.ELAPSED_REALTIME_WAKEUP,
+                    SystemClock.elapsedRealtime() + mRekeyDelayMs,
+                    mRekeySaIntent);
+
+            getIkeLog()
+                    .d(
+                            tag,
+                            "Lifetime alarm set: Hard lifetime ("
+                                    + mDeleteDelayMs
+                                    + "ms) Soft lifetime ("
+                                    + mRekeyDelayMs
+                                    + "ms)");
+        }
 
-            // TODO: Schedule alarm for rekey
+        public void rescheduleRekey(long retryDelayMs) {
+            mAlarmManager.setExactAndAllowWhileIdle(
+                    AlarmManager.ELAPSED_REALTIME_WAKEUP,
+                    SystemClock.elapsedRealtime() + retryDelayMs,
+                    mRekeySaIntent);
         }
 
         public void cancelLifetimeExpiryAlarm(String tag) {
@@ -591,8 +595,7 @@ public abstract class SaRecord implements AutoCloseable {
         public final boolean isTransport;
         public final boolean isLocalInit;
         public final boolean hasIntegrityAlgo;
-        public final ChildLocalRequest futureRekeyEvent;
-        public final SaLifetimeAlarmScheduler saLifetimeAlarmSched;
+        public final SaLifetimeAlarmScheduler saLifetimeAlarmScheduler;
 
         ChildSaRecordConfig(
                 Context context,
@@ -607,8 +610,7 @@ public abstract class SaRecord implements AutoCloseable {
                 byte[] skD,
                 boolean isTransport,
                 boolean isLocalInit,
-                ChildLocalRequest futureRekeyEvent,
-                SaLifetimeAlarmScheduler saLifetimeAlarmSched) {
+                SaLifetimeAlarmScheduler saLifetimeAlarmScheduler) {
             this.context = context;
             this.initSpi = initSpi;
             this.respSpi = respSpi;
@@ -622,8 +624,7 @@ public abstract class SaRecord implements AutoCloseable {
             this.isTransport = isTransport;
             this.isLocalInit = isLocalInit;
             hasIntegrityAlgo = (integrityAlgo != null);
-            this.futureRekeyEvent = futureRekeyEvent;
-            this.saLifetimeAlarmSched = saLifetimeAlarmSched;
+            this.saLifetimeAlarmScheduler = saLifetimeAlarmScheduler;
         }
     }
 
@@ -663,7 +664,6 @@ public abstract class SaRecord implements AutoCloseable {
                 byte[] skEr,
                 byte[] skPi,
                 byte[] skPr,
-                LocalRequest futureRekeyEvent,
                 SaLifetimeAlarmScheduler saLifetimeAlarmScheduler) {
             super(
                     localInit,
@@ -673,7 +673,6 @@ public abstract class SaRecord implements AutoCloseable {
                     skAr,
                     skEi,
                     skEr,
-                    futureRekeyEvent,
                     saLifetimeAlarmScheduler);
 
             mInitiatorSpiResource = initSpi;
@@ -706,7 +705,6 @@ public abstract class SaRecord implements AutoCloseable {
                 IkeMacPrf prf,
                 int integrityKeyLength,
                 int encryptionKeyLength,
-                LocalRequest futureRekeyEvent,
                 SaLifetimeAlarmScheduler saLifetimeAlarmScheduler)
                 throws GeneralSecurityException {
             return sSaRecordHelper.makeFirstIkeSaRecord(
@@ -719,7 +717,6 @@ public abstract class SaRecord implements AutoCloseable {
                             integrityKeyLength,
                             encryptionKeyLength,
                             true /*isLocalInit*/,
-                            futureRekeyEvent,
                             saLifetimeAlarmScheduler));
         }
 
@@ -735,7 +732,6 @@ public abstract class SaRecord implements AutoCloseable {
                 int integrityKeyLength,
                 int encryptionKeyLength,
                 boolean isLocalInit,
-                LocalRequest futureRekeyEvent,
                 SaLifetimeAlarmScheduler saLifetimeAlarmScheduler)
                 throws GeneralSecurityException {
             return sSaRecordHelper.makeRekeyedIkeSaRecord(
@@ -750,7 +746,6 @@ public abstract class SaRecord implements AutoCloseable {
                             integrityKeyLength,
                             encryptionKeyLength,
                             isLocalInit,
-                            futureRekeyEvent,
                             saLifetimeAlarmScheduler));
         }
 
@@ -929,7 +924,6 @@ public abstract class SaRecord implements AutoCloseable {
         public final int integrityKeyLength;
         public final int encryptionKeyLength;
         public final boolean isLocalInit;
-        public final LocalRequest futureRekeyEvent;
         public final SaLifetimeAlarmScheduler saLifetimeAlarmScheduler;
 
         IkeSaRecordConfig(
@@ -939,7 +933,6 @@ public abstract class SaRecord implements AutoCloseable {
                 int integrityKeyLength,
                 int encryptionKeyLength,
                 boolean isLocalInit,
-                LocalRequest futureRekeyEvent,
                 SaLifetimeAlarmScheduler saLifetimeAlarmScheduler) {
             this.initSpi = initSpi;
             this.respSpi = respSpi;
@@ -947,7 +940,6 @@ public abstract class SaRecord implements AutoCloseable {
             this.integrityKeyLength = integrityKeyLength;
             this.encryptionKeyLength = encryptionKeyLength;
             this.isLocalInit = isLocalInit;
-            this.futureRekeyEvent = futureRekeyEvent;
             this.saLifetimeAlarmScheduler = saLifetimeAlarmScheduler;
         }
     }
@@ -979,7 +971,6 @@ public abstract class SaRecord implements AutoCloseable {
                 byte[] skEr,
                 IpSecTransform inTransform,
                 IpSecTransform outTransform,
-                ChildLocalRequest futureRekeyEvent,
                 SaLifetimeAlarmScheduler saLifetimeAlarmScheduler) {
             super(
                     localInit,
@@ -989,7 +980,6 @@ public abstract class SaRecord implements AutoCloseable {
                     skAr,
                     skEi,
                     skEr,
-                    futureRekeyEvent,
                     saLifetimeAlarmScheduler);
 
             mInboundSpi = inSpi;
@@ -1017,7 +1007,6 @@ public abstract class SaRecord implements AutoCloseable {
                 byte[] skD,
                 boolean isTransport,
                 boolean isLocalInit,
-                ChildLocalRequest futureRekeyEvent,
                 SaLifetimeAlarmScheduler saLifetimeAlarmScheduler)
                 throws GeneralSecurityException, ResourceUnavailableException,
                         SpiUnavailableException, IOException {
@@ -1037,7 +1026,6 @@ public abstract class SaRecord implements AutoCloseable {
                             skD,
                             isTransport,
                             isLocalInit,
-                            futureRekeyEvent,
                             saLifetimeAlarmScheduler));
         }
 
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachineTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachineTest.java
index 99f6d2e7..e20d0b1d 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachineTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ChildSessionStateMachineTest.java
@@ -22,6 +22,7 @@ import static android.net.ipsec.ike.exceptions.IkeProtocolException.ERROR_TYPE_T
 import static android.system.OsConstants.AF_INET;
 
 import static com.android.internal.net.TestUtils.createMockRandomFactory;
+import static com.android.internal.net.ipsec.ike.AbstractSessionStateMachine.RETRY_INTERVAL_MS;
 import static com.android.internal.net.ipsec.ike.ChildSessionStateMachine.CMD_FORCE_TRANSITION;
 import static com.android.internal.net.ipsec.ike.IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_CHILD;
 import static com.android.internal.net.ipsec.ike.IkeSessionStateMachine.IKE_EXCHANGE_SUBTYPE_DELETE_CHILD;
@@ -47,11 +48,9 @@ import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
-import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyBoolean;
 import static org.mockito.Matchers.anyInt;
-import static org.mockito.Matchers.anyLong;
 import static org.mockito.Matchers.anyObject;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.eq;
@@ -225,8 +224,9 @@ public final class ChildSessionStateMachineTest {
 
     private ArgumentMatcher<ChildLocalRequest> mRekeyChildLocalReqMatcher =
             (argument) -> {
-                return CMD_LOCAL_REQUEST_REKEY_CHILD == argument.procedureType
-                        && mMockChildSessionCallback == argument.childSessionCallback;
+                return (CMD_LOCAL_REQUEST_REKEY_CHILD == argument.procedureType
+                                && mMockChildSessionCallback == argument.childSessionCallback
+                        || CURRENT_CHILD_SA_SPI_OUT == argument.remoteSpi);
             };
 
     public ChildSessionStateMachineTest() {
@@ -385,7 +385,6 @@ public final class ChildSessionStateMachineTest {
                                 null,
                                 mock(IpSecTransform.class),
                                 mock(IpSecTransform.class),
-                                mock(ChildLocalRequest.class),
                                 mock(SaLifetimeAlarmScheduler.class)));
         doNothing().when(child).close();
         return child;
@@ -427,11 +426,7 @@ public final class ChildSessionStateMachineTest {
         assertFalse(childSaRecordConfig.isTransport);
         assertEquals(isLocalInit, childSaRecordConfig.isLocalInit);
         assertTrue(childSaRecordConfig.hasIntegrityAlgo);
-        assertEquals(
-                CMD_LOCAL_REQUEST_REKEY_CHILD, childSaRecordConfig.futureRekeyEvent.procedureType);
-        assertEquals(
-                mMockChildSessionCallback,
-                childSaRecordConfig.futureRekeyEvent.childSessionCallback);
+        assertNotNull(childSaRecordConfig.saLifetimeAlarmScheduler);
     }
 
     private void verifyNotifyUsersCreateIpSecSa(
@@ -481,8 +476,6 @@ public final class ChildSessionStateMachineTest {
 
         verify(mMockChildSessionSmCallback)
                 .onChildSaCreated(anyInt(), eq(mChildSessionStateMachine));
-        verify(mMockChildSessionSmCallback)
-                .scheduleLocalRequest(argThat(mRekeyChildLocalReqMatcher), anyLong());
         verify(mMockChildSessionSmCallback).onProcedureFinished(mChildSessionStateMachine);
         assertTrue(
                 mChildSessionStateMachine.getCurrentState()
@@ -1020,8 +1013,6 @@ public final class ChildSessionStateMachineTest {
                 .onChildSaCreated(
                         eq(mSpyLocalInitNewChildSaRecord.getRemoteSpi()),
                         eq(mChildSessionStateMachine));
-        verify(mMockChildSessionSmCallback)
-                .scheduleLocalRequest(argThat(mRekeyChildLocalReqMatcher), anyLong());
 
         verify(mMockSaRecordHelper)
                 .makeChildSaRecord(
@@ -1058,9 +1049,7 @@ public final class ChildSessionStateMachineTest {
         mLooper.dispatchAll();
 
         // Verify rekey has been rescheduled and Child Session is alive
-        verify(mMockChildSessionSmCallback)
-                .scheduleRetryLocalRequest(
-                        (ChildLocalRequest) mSpyCurrentChildSaRecord.getFutureRekeyEvent());
+        verify(mSpyCurrentChildSaRecord).rescheduleRekey(eq(RETRY_INTERVAL_MS));
         assertTrue(
                 mChildSessionStateMachine.getCurrentState()
                         instanceof ChildSessionStateMachine.Idle);
@@ -1286,8 +1275,6 @@ public final class ChildSessionStateMachineTest {
                 .onChildSaCreated(
                         eq(mSpyRemoteInitNewChildSaRecord.getRemoteSpi()),
                         eq(mChildSessionStateMachine));
-        verify(mMockChildSessionSmCallback)
-                .scheduleLocalRequest(argThat(mRekeyChildLocalReqMatcher), anyLong());
 
         verify(mMockSaRecordHelper)
                 .makeChildSaRecord(
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 1978d8a1..61f44e02 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
@@ -28,7 +28,7 @@ import static android.system.OsConstants.AF_INET;
 import static android.system.OsConstants.AF_INET6;
 
 import static com.android.internal.net.TestUtils.createMockRandomFactory;
-import static com.android.internal.net.ipsec.ike.IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_IKE;
+import static com.android.internal.net.ipsec.ike.AbstractSessionStateMachine.RETRY_INTERVAL_MS;
 import static com.android.internal.net.ipsec.ike.IkeSessionStateMachine.CMD_RECEIVE_IKE_PACKET;
 import static com.android.internal.net.ipsec.ike.IkeSessionStateMachine.IKE_EXCHANGE_SUBTYPE_DELETE_CHILD;
 import static com.android.internal.net.ipsec.ike.IkeSessionStateMachine.IKE_EXCHANGE_SUBTYPE_REKEY_CHILD;
@@ -67,6 +67,7 @@ import static org.mockito.Matchers.anyObject;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.argThat;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
@@ -641,10 +642,25 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
                 new byte[KEY_LEN_IKE_ENCR],
                 TestUtils.hexStringToByteArray(PRF_KEY_INIT_HEX_STRING),
                 TestUtils.hexStringToByteArray(PRF_KEY_RESP_HEX_STRING),
-                new IkeLocalRequest(CMD_LOCAL_REQUEST_REKEY_IKE),
                 mock(SaLifetimeAlarmScheduler.class));
     }
 
+    private void mockScheduleRekey(SaLifetimeAlarmScheduler mockSaLifetimeAlarmScheduler) {
+        IkeLocalRequest rekeyReq =
+                new IkeLocalRequest(IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_IKE);
+        doAnswer(
+                (invocation) -> {
+                        mIkeSessionStateMachine.sendMessageDelayed(
+                                IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_IKE,
+                                rekeyReq,
+                                mIkeSessionStateMachine.mIkeSessionParams
+                                        .getSoftLifetimeMsInternal());
+                        return null;
+                })
+                .when(mockSaLifetimeAlarmScheduler)
+                .scheduleLifetimeExpiryAlarm(anyString());
+    }
+
     @Before
     public void setUp() throws Exception {
         super.setUp();
@@ -1213,6 +1229,9 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
                 .thenAnswer(
                         (invocation) -> {
                             captureAndReleaseIkeSpiResource(invocation, 2);
+                            mockScheduleRekey(mSpyCurrentIkeSaRecord.mSaLifetimeAlarmScheduler);
+                            mSpyCurrentIkeSaRecord.mSaLifetimeAlarmScheduler
+                                    .scheduleLifetimeExpiryAlarm(anyString());
                             return mSpyCurrentIkeSaRecord;
                         });
     }
@@ -1225,6 +1244,9 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
                 .thenAnswer(
                         (invocation) -> {
                             captureAndReleaseIkeSpiResource(invocation, 4);
+                            mockScheduleRekey(rekeySaRecord.mSaLifetimeAlarmScheduler);
+                            rekeySaRecord.mSaLifetimeAlarmScheduler.scheduleLifetimeExpiryAlarm(
+                                    anyString());
                             return rekeySaRecord;
                         });
     }
@@ -1397,7 +1419,7 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
         assertEquals(KEY_LEN_IKE_PRF, ikeSaRecordConfig.prf.getKeyLength());
         assertEquals(KEY_LEN_IKE_INTE, ikeSaRecordConfig.integrityKeyLength);
         assertEquals(KEY_LEN_IKE_ENCR, ikeSaRecordConfig.encryptionKeyLength);
-        assertEquals(CMD_LOCAL_REQUEST_REKEY_IKE, ikeSaRecordConfig.futureRekeyEvent.procedureType);
+        assertNotNull(ikeSaRecordConfig.saLifetimeAlarmScheduler);
 
         // Validate NAT detection
         assertTrue(mIkeSessionStateMachine.mIsLocalBehindNat);
@@ -3071,17 +3093,10 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
         mIkeSessionStateMachine.sendMessage(CMD_RECEIVE_IKE_PACKET, resp);
         mLooper.dispatchAll();
 
-        // Verify IKE Session goes back to Idle
+        // Verify IKE Session goes back to Idle and retry is scheduled
+        verify(mSpyCurrentIkeSaRecord).rescheduleRekey(eq(RETRY_INTERVAL_MS));
         assertTrue(
                 mIkeSessionStateMachine.getCurrentState() instanceof IkeSessionStateMachine.Idle);
-
-        // Move time forward to trigger retry
-        mLooper.moveTimeForward(IkeSessionStateMachine.RETRY_INTERVAL_MS);
-        mLooper.dispatchAll();
-
-        assertTrue(
-                mIkeSessionStateMachine.getCurrentState()
-                        instanceof IkeSessionStateMachine.RekeyIkeLocalCreate);
     }
 
     @Test
@@ -3183,6 +3198,40 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
                         instanceof IkeSessionStateMachine.RekeyIkeLocalCreate);
     }
 
+    private void mockRescheduleRekey(IkeSaRecord spySaRecord) {
+        IkeLocalRequest rekeyReq =
+                new IkeLocalRequest(IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_IKE);
+        doAnswer(
+                (invocation) -> {
+                        mIkeSessionStateMachine.sendMessageDelayed(
+                                IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_IKE,
+                                rekeyReq,
+                                RETRY_INTERVAL_MS);
+                        return null;
+                })
+                .when(spySaRecord)
+                .rescheduleRekey(eq(RETRY_INTERVAL_MS));
+    }
+
+    @Test
+    public void testRekeyIkeLocalCreateHandleRespWithTempFailure() throws Exception {
+        setupIdleStateMachine();
+
+        // Send Rekey-Create request
+        mIkeSessionStateMachine.sendMessage(
+                IkeSessionStateMachine.CMD_EXECUTE_LOCAL_REQ,
+                new IkeLocalRequest(IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_IKE));
+        mLooper.dispatchAll();
+
+        // Mock sending TEMPORARY_FAILURE response
+        mockRcvTempFail();
+        mLooper.dispatchAll();
+
+        verify(mSpyCurrentIkeSaRecord).rescheduleRekey(eq(RETRY_INTERVAL_MS));
+        assertTrue(
+                mIkeSessionStateMachine.getCurrentState() instanceof IkeSessionStateMachine.Idle);
+    }
+
     private void mockCreateAndTransitionToRekeyDeleteLocal() {
         // Seed fake rekey data and force transition to RekeyIkeLocalDelete
         mIkeSessionStateMachine.mLocalInitNewIkeSaRecord = mSpyLocalInitIkeSaRecord;
@@ -4386,24 +4435,6 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
         mLooper.dispatchAll();
     }
 
-    @Test
-    public void testTempFailureHandlerScheduleRetry() throws Exception {
-        mockSendRekeyChildReq();
-
-        // Mock sending TEMPORARY_FAILURE response
-        mockRcvTempFail();
-
-        // Move time forward to trigger retry
-        mLooper.moveTimeForward(IkeSessionStateMachine.RETRY_INTERVAL_MS);
-        mLooper.dispatchAll();
-
-        // Verify that rekey is triggered again
-        assertTrue(
-                mIkeSessionStateMachine.getCurrentState()
-                        instanceof IkeSessionStateMachine.ChildProcedureOngoing);
-        verify(mMockChildSessionStateMachine, times(2)).rekeyChildSession();
-    }
-
     @Ignore
     public void disableTestTempFailureHandlerTimeout() throws Exception {
         long currentTime = 0;
@@ -4431,28 +4462,48 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
 
     @Test
     public void testTempFailureHandlerCancelTimer() throws Exception {
-        mockSendRekeyChildReq();
+        mockRescheduleRekey(mSpyCurrentIkeSaRecord);
+        setupIdleStateMachine();
+
+        // Send Rekey-Create request
+        mIkeSessionStateMachine.sendMessage(
+                IkeSessionStateMachine.CMD_EXECUTE_LOCAL_REQ,
+                new IkeLocalRequest(IkeSessionStateMachine.CMD_LOCAL_REQUEST_REKEY_IKE));
+        mLooper.dispatchAll();
+        verifyRetransmissionStarted();
 
         // Mock sending TEMPORARY_FAILURE response
         mockRcvTempFail();
+        mLooper.dispatchAll();
+        verify(mSpyCurrentIkeSaRecord).rescheduleRekey(eq(RETRY_INTERVAL_MS));
+        assertTrue(
+                mIkeSessionStateMachine.getCurrentState() instanceof IkeSessionStateMachine.Idle);
 
         // Move time forward to trigger retry
         mLooper.moveTimeForward(IkeSessionStateMachine.RETRY_INTERVAL_MS);
         mLooper.dispatchAll();
-        verify(mMockChildSessionStateMachine, times(2)).rekeyChildSession();
+        assertTrue(
+                mIkeSessionStateMachine.getCurrentState()
+                        instanceof IkeSessionStateMachine.RekeyIkeLocalCreate);
 
-        // Mock sending a valid response
-        ReceivedIkePacket resp =
-                makeDummyEncryptedReceivedIkePacketWithPayloadList(
-                        mSpyCurrentIkeSaRecord,
-                        EXCHANGE_TYPE_CREATE_CHILD_SA,
-                        true /*isResp*/,
-                        new LinkedList<>());
-        mIkeSessionStateMachine.sendMessage(IkeSessionStateMachine.CMD_RECEIVE_IKE_PACKET, resp);
+        // Prepare "rekeyed" SA
+        setupRekeyedIkeSa(mSpyLocalInitIkeSaRecord);
+
+        // Receive valid Rekey-Create response
+        ReceivedIkePacket dummyRekeyIkeRespReceivedPacket = makeRekeyIkeResponse();
         mIkeSessionStateMachine.sendMessage(
-                IkeSessionStateMachine.CMD_FORCE_TRANSITION, mIkeSessionStateMachine.mIdle);
+                IkeSessionStateMachine.CMD_RECEIVE_IKE_PACKET, dummyRekeyIkeRespReceivedPacket);
         mLooper.dispatchAll();
 
+        // Receive Delete response
+        ReceivedIkePacket dummyDeleteIkeRespReceivedPacket =
+                makeDeleteIkeResponse(mSpyCurrentIkeSaRecord);
+        mIkeSessionStateMachine.sendMessage(
+                IkeSessionStateMachine.CMD_RECEIVE_IKE_PACKET, dummyDeleteIkeRespReceivedPacket);
+        mLooper.dispatchAll();
+        assertTrue(
+                mIkeSessionStateMachine.getCurrentState() instanceof IkeSessionStateMachine.Idle);
+
         // Move time forward
         mLooper.moveTimeForward(IkeSessionStateMachine.TEMP_FAILURE_RETRY_TIMEOUT_MS);
         mLooper.dispatchAll();
@@ -4460,8 +4511,6 @@ public final class IkeSessionStateMachineTest extends IkeSessionTestBase {
         // Validate IKE Session is not closed
         assertTrue(
                 mIkeSessionStateMachine.getCurrentState() instanceof IkeSessionStateMachine.Idle);
-        // Validate no more retry
-        verify(mMockChildSessionStateMachine, times(2)).rekeyChildSession();
     }
 
     @Ignore
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/SaRecordTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/SaRecordTest.java
index 644ba066..b193099d 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/SaRecordTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/SaRecordTest.java
@@ -24,6 +24,7 @@ import static org.junit.Assert.assertTrue;
 import static org.mockito.AdditionalMatchers.aryEq;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyObject;
+import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -38,8 +39,6 @@ import android.net.IpSecTransform;
 import android.net.ipsec.ike.SaProposal;
 
 import com.android.internal.net.TestUtils;
-import com.android.internal.net.ipsec.ike.IkeLocalRequestScheduler.ChildLocalRequest;
-import com.android.internal.net.ipsec.ike.IkeLocalRequestScheduler.LocalRequest;
 import com.android.internal.net.ipsec.ike.SaRecord.ChildSaRecord;
 import com.android.internal.net.ipsec.ike.SaRecord.ChildSaRecordConfig;
 import com.android.internal.net.ipsec.ike.SaRecord.IIpSecTransformHelper;
@@ -148,9 +147,6 @@ public final class SaRecordTest {
     private IkeMacIntegrity mHmacSha1IntegrityMac;
     private IkeCipher mAesCbcCipher;
 
-    private LocalRequest mMockFutureRekeyIkeEvent;
-    private ChildLocalRequest mMockFutureRekeyChildEvent;
-
     private SaLifetimeAlarmScheduler mMockLifetimeAlarmScheduler;
 
     private SaRecordHelper mSaRecordHelper = new SaRecordHelper();
@@ -168,8 +164,6 @@ public final class SaRecordTest {
                                 SaProposal.ENCRYPTION_ALGORITHM_AES_CBC,
                                 SaProposal.KEY_LEN_AES_128));
 
-        mMockFutureRekeyIkeEvent = mock(LocalRequest.class);
-        mMockFutureRekeyChildEvent = mock(ChildLocalRequest.class);
         mMockLifetimeAlarmScheduler = mock(SaLifetimeAlarmScheduler.class);
     }
 
@@ -192,8 +186,7 @@ public final class SaRecordTest {
                         IKE_AUTH_ALGO_KEY_LEN,
                         IKE_ENCR_ALGO_KEY_LEN,
                         true /*isLocalInit*/,
-                        mMockFutureRekeyIkeEvent,
-                        mock(SaLifetimeAlarmScheduler.class));
+                        mMockLifetimeAlarmScheduler);
 
         int keyMaterialLen =
                 IKE_SK_D_KEY_LEN
@@ -225,10 +218,10 @@ public final class SaRecordTest {
                 TestUtils.hexStringToByteArray(IKE_SK_PRF_INIT_HEX_STRING), ikeSaRecord.getSkPi());
         assertArrayEquals(
                 TestUtils.hexStringToByteArray(IKE_SK_PRF_RESP_HEX_STRING), ikeSaRecord.getSkPr());
+        verify(mMockLifetimeAlarmScheduler).scheduleLifetimeExpiryAlarm(anyString());
 
         ikeSaRecord.close();
-
-        verify(mMockFutureRekeyIkeEvent).cancel();
+        verify(mMockLifetimeAlarmScheduler).cancelLifetimeExpiryAlarm(anyString());
     }
 
     // Test generating keying material and building IpSecTransform for making Child SA.
@@ -306,7 +299,6 @@ public final class SaRecordTest {
                         TestUtils.hexStringToByteArray(IKE_SK_D_HEX_STRING),
                         false /*isTransport*/,
                         true /*isLocalInit*/,
-                        mMockFutureRekeyChildEvent,
                         mMockLifetimeAlarmScheduler);
 
         ChildSaRecord childSaRecord =
@@ -332,8 +324,10 @@ public final class SaRecordTest {
                 TestUtils.hexStringToByteArray(FIRST_CHILD_ENCR_RESP_HEX_STRING),
                 childSaRecord.getInboundDecryptionKey());
 
+        verify(mMockLifetimeAlarmScheduler).scheduleLifetimeExpiryAlarm(anyString());
+
         childSaRecord.close();
-        verify(mMockFutureRekeyChildEvent).cancel();
+        verify(mMockLifetimeAlarmScheduler).cancelLifetimeExpiryAlarm(anyString());
 
         SaRecord.setIpSecTransformHelper(new IpSecTransformHelper());
     }