[DK3] Send onPause/onResume keepalive callbacks

Test: CTS in the patch immediately on top of this, [DK4]
Change-Id: I208ceceb37c7977452479361f70f046fabafb37a
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 17389b4..497c105 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -2213,9 +2213,13 @@
         /** The requested keepalive was successfully started. */
         @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
         public void onStarted() {}
+        /** The keepalive was resumed after being paused by the system. */
+        public void onResumed() {}
         /** The keepalive was successfully stopped. */
         @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
         public void onStopped() {}
+        /** The keepalive was paused automatically by the system. */
+        public void onPaused() {}
         /** An error occurred. */
         @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
         public void onError(int error) {}
@@ -2314,6 +2318,18 @@
                 }
 
                 @Override
+                public void onResumed() {
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            callback.onResumed();
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
+                }
+
+                @Override
                 public void onStopped() {
                     final long token = Binder.clearCallingIdentity();
                     try {
@@ -2327,6 +2343,19 @@
                 }
 
                 @Override
+                public void onPaused() {
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            callback.onPaused();
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
+                    mExecutor.shutdown();
+                }
+
+                @Override
                 public void onError(int error) {
                     final long token = Binder.clearCallingIdentity();
                     try {
diff --git a/framework/src/android/net/ISocketKeepaliveCallback.aidl b/framework/src/android/net/ISocketKeepaliveCallback.aidl
index 1240e37..0a1de6c 100644
--- a/framework/src/android/net/ISocketKeepaliveCallback.aidl
+++ b/framework/src/android/net/ISocketKeepaliveCallback.aidl
@@ -31,4 +31,8 @@
     void onError(int error);
     /** The keepalive on a TCP socket was stopped because the socket received data. */
     void onDataReceived();
+    /** The keepalive was paused by the system because it's not necessary right now. */
+    void onPaused();
+    /** The keepalive was resumed by the system after being suspended. */
+    void onResumed();
 }
diff --git a/framework/src/android/net/SocketKeepalive.java b/framework/src/android/net/SocketKeepalive.java
index 2911ce7..00104f6 100644
--- a/framework/src/android/net/SocketKeepalive.java
+++ b/framework/src/android/net/SocketKeepalive.java
@@ -64,6 +64,12 @@
     public static final int SUCCESS = 0;
 
     /**
+     * Success when trying to suspend.
+     * @hide
+     */
+    public static final int SUCCESS_PAUSED = 1;
+
+    /**
      * No keepalive. This should only be internally as it indicates There is no keepalive.
      * It should not propagate to applications.
      * @hide
@@ -271,6 +277,18 @@
             }
 
             @Override
+            public void onResumed() {
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    mExecutor.execute(() -> {
+                        callback.onResumed();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
+            }
+
+            @Override
             public void onStopped() {
                 final long token = Binder.clearCallingIdentity();
                 try {
@@ -283,6 +301,18 @@
             }
 
             @Override
+            public void onPaused() {
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> {
+                        callback.onPaused();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
+            }
+
+            @Override
             public void onError(int error) {
                 final long token = Binder.clearCallingIdentity();
                 try {
@@ -387,8 +417,18 @@
     public static class Callback {
         /** The requested keepalive was successfully started. */
         public void onStarted() {}
+        /**
+         * The keepalive was resumed by the system after being suspended.
+         * @hide
+         **/
+        public void onResumed() {}
         /** The keepalive was successfully stopped. */
         public void onStopped() {}
+        /**
+         * The keepalive was paused by the system because it's not necessary right now.
+         * @hide
+         **/
+        public void onPaused() {}
         /** An error occurred. */
         public void onError(@ErrorCode int error) {}
         /** The keepalive on a TCP socket was stopped because the socket received data. This is
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 9dcbac8..937c93b 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -18,7 +18,7 @@
 
 import static android.net.NetworkAgent.CMD_START_SOCKET_KEEPALIVE;
 import static android.net.SocketKeepalive.ERROR_INVALID_SOCKET;
-import static android.net.SocketKeepalive.SUCCESS;
+import static android.net.SocketKeepalive.SUCCESS_PAUSED;
 import static android.provider.DeviceConfig.NAMESPACE_CONNECTIVITY;
 import static android.system.OsConstants.AF_INET;
 import static android.system.OsConstants.AF_INET6;
@@ -180,7 +180,7 @@
      *                     resumed if a TCP socket is open on the VPN.
      * See the documentation for the states for more detail.
      */
-    public class AutomaticOnOffKeepalive {
+    public class AutomaticOnOffKeepalive implements IBinder.DeathRecipient {
         @NonNull
         private final KeepaliveTracker.KeepaliveInfo mKi;
         @NonNull
@@ -237,6 +237,11 @@
             return BinderUtils.withCleanCallingIdentity(() -> PendingIntent.getBroadcast(
                     context, 0 /* requestCode */, intent, PendingIntent.FLAG_IMMUTABLE));
         }
+
+        @Override
+        public void binderDied() {
+            mConnectivityServiceHandler.post(() -> cleanupAutoOnOffKeepalive(this));
+        }
     }
 
     public AutomaticOnOffKeepaliveTracker(@NonNull Context context, @NonNull Handler handler) {
@@ -295,7 +300,7 @@
             // SUSPENDED.
             if (ki.mAutomaticOnOffState == STATE_ENABLED) {
                 ki.mAutomaticOnOffState = STATE_SUSPENDED;
-                handleSuspendKeepalive(ki.mKi);
+                handlePauseKeepalive(ki.mKi);
             }
         } else {
             handleMaybeResumeKeepalive(ki);
@@ -372,6 +377,13 @@
         mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
 
         // Add automatic on/off request into list to track its life cycle.
+        try {
+            autoKi.mKi.mCallback.asBinder().linkToDeath(autoKi, 0);
+        } catch (RemoteException e) {
+            // The underlying keepalive performs its own cleanup
+            autoKi.binderDied();
+            return;
+        }
         mAutomaticOnOffKeepalives.add(autoKi);
         if (STATE_ALWAYS_ON != autoKi.mAutomaticOnOffState) {
             startTcpPollingAlarm(autoKi.mTcpPollingAlarm);
@@ -382,10 +394,9 @@
         mKeepaliveTracker.handleStartKeepalive(ki);
     }
 
-    private void handleSuspendKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
+    private void handlePauseKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
         // TODO : mKT.handleStopKeepalive should take a KeepaliveInfo instead
-        // TODO : create a separate success code for suspend
-        mKeepaliveTracker.handleStopKeepalive(ki.getNai(), ki.getSlot(), SUCCESS);
+        mKeepaliveTracker.handleStopKeepalive(ki.getNai(), ki.getSlot(), SUCCESS_PAUSED);
     }
 
     /**
@@ -397,6 +408,8 @@
         if (autoKi.mAutomaticOnOffState != STATE_SUSPENDED) {
             final KeepaliveTracker.KeepaliveInfo ki = autoKi.mKi;
             mKeepaliveTracker.handleStopKeepalive(ki.getNai(), ki.getSlot(), reason);
+        } else {
+            mKeepaliveTracker.finalizePausedKeepalive(autoKi.mKi);
         }
 
         cleanupAutoOnOffKeepalive(autoKi);
@@ -404,10 +417,17 @@
 
     private void cleanupAutoOnOffKeepalive(@NonNull final AutomaticOnOffKeepalive autoKi) {
         ensureRunningOnHandlerThread();
-        if (null != autoKi.mTcpPollingAlarm) mAlarmManager.cancel(autoKi.mTcpPollingAlarm);
-        // Close the duplicated fd that maintains the lifecycle of socket.
+        // Close the duplicated fd that maintains the lifecycle of socket. If this fd was
+        // not duplicated this is a no-op.
         FileUtils.closeQuietly(autoKi.mFd);
-        mAutomaticOnOffKeepalives.remove(autoKi);
+        if (null != autoKi.mTcpPollingAlarm) mAlarmManager.cancel(autoKi.mTcpPollingAlarm);
+
+        // If the KI is not in the array, it's because it was already removed, or it was never
+        // added ; the only ways this can happen is if the keepalive is stopped by the app and the
+        // app dies immediately, or if the app died before the link to death could be registered.
+        if (!mAutomaticOnOffKeepalives.remove(autoKi)) return;
+
+        autoKi.mKi.mCallback.asBinder().unlinkToDeath(autoKi, 0);
     }
 
     /**
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index a512b7c..5572361 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -32,6 +32,7 @@
 import static android.net.SocketKeepalive.MIN_INTERVAL_SEC;
 import static android.net.SocketKeepalive.NO_KEEPALIVE;
 import static android.net.SocketKeepalive.SUCCESS;
+import static android.net.SocketKeepalive.SUCCESS_PAUSED;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -134,6 +135,9 @@
         public final NetworkAgentInfo mNai;
         private final int mType;
         public final FileDescriptor mFd;
+        // True if this was resumed from a previously turned off keepalive, otherwise false.
+        // This is necessary to send the correct callbacks.
+        public final boolean mResumed;
 
         public static final int TYPE_NATT = 1;
         public static final int TYPE_TCP = 2;
@@ -160,6 +164,16 @@
                 int interval,
                 int type,
                 @Nullable FileDescriptor fd) throws InvalidSocketException {
+            this(callback, nai, packet, interval, type, fd, false /* resumed */);
+        }
+
+        KeepaliveInfo(@NonNull ISocketKeepaliveCallback callback,
+                @NonNull NetworkAgentInfo nai,
+                @NonNull KeepalivePacketData packet,
+                int interval,
+                int type,
+                @Nullable FileDescriptor fd,
+                boolean resumed) throws InvalidSocketException {
             mCallback = callback;
             mPid = Binder.getCallingPid();
             mUid = Binder.getCallingUid();
@@ -169,6 +183,7 @@
             mPacket = packet;
             mInterval = interval;
             mType = type;
+            mResumed = resumed;
 
             // For SocketKeepalive, a dup of fd is kept in mFd so the source port from which the
             // keepalives are sent cannot be reused by another app even if the fd gets closed by
@@ -232,6 +247,8 @@
 
         /** Called when the application process is killed. */
         public void binderDied() {
+            // TODO b/267106526 : this is not called on the handler thread but stop() happily
+            // assumes it is, which means this is a pretty dangerous race condition.
             stop(BINDER_DIED);
         }
 
@@ -327,6 +344,10 @@
         }
 
         void start(int slot) {
+            // BINDER_DIED can happen if the binder died before the KeepaliveInfo was created and
+            // the constructor set the state to BINDER_DIED. If that's the case, the KI is already
+            // cleaned up.
+            if (BINDER_DIED == mStartedState) return;
             mSlot = slot;
             int error = isValid();
             if (error == SUCCESS) {
@@ -371,7 +392,10 @@
             // To prevent races from re-entrance of stop(), return if the state is already stopping.
             // This might happen if multiple event sources stop keepalive in a short time. Such as
             // network disconnect after user calls stop(), or tear down socket after binder died.
-            if (mStartedState == STOPPING) return;
+            // Note that it's always possible this method is called by the auto keepalive timer
+            // or any other way after the binder died, hence the check for BINDER_DIED. If the
+            // binder has died, then the KI has already been cleaned up.
+            if (mStartedState == STOPPING || mStartedState == BINDER_DIED) return;
 
             // Store the reason of stopping, and report it after the keepalive is fully stopped.
             if (mStopReason != ERROR_STOP_REASON_UNINITIALIZED) {
@@ -382,9 +406,10 @@
                     + ": " + reason);
             switch (mStartedState) {
                 case NOT_STARTED:
-                    // Remove the reference of the keepalive that meet error before starting,
+                    // Remove the reference to this keepalive that had an error before starting,
                     // e.g. invalid parameter.
                     cleanupStoppedKeepalive(mNai, mSlot);
+                    if (BINDER_DIED == reason) mStartedState = BINDER_DIED;
                     break;
                 default:
                     mStartedState = STOPPING;
@@ -422,7 +447,8 @@
          * Construct a new KeepaliveInfo from existing KeepaliveInfo with a new fd.
          */
         public KeepaliveInfo withFd(@NonNull FileDescriptor fd) throws InvalidSocketException {
-            return new KeepaliveInfo(mCallback, mNai, mPacket, mInterval, mType, fd);
+            return new KeepaliveInfo(mCallback, mNai, mPacket, mInterval, mType, fd,
+                    true /* resumed */);
         }
     }
 
@@ -523,6 +549,12 @@
             } catch (RemoteException e) {
                 Log.w(TAG, "Discarded onStop callback: " + reason);
             }
+        } else if (reason == SUCCESS_PAUSED) {
+            try {
+                ki.mCallback.onPaused();
+            } catch (RemoteException e) {
+                Log.w(TAG, "Discarded onPaused callback: " + reason);
+            }
         } else if (reason == DATA_RECEIVED) {
             try {
                 ki.mCallback.onDataReceived();
@@ -540,6 +572,25 @@
         ki.unlinkDeathRecipient();
     }
 
+    /**
+     * Finalize a paused keepalive.
+     *
+     * This will simply send the onStopped() callback after checking that this keepalive is
+     * indeed paused.
+     *
+     * @param ki the keepalive to finalize
+     */
+    public void finalizePausedKeepalive(@NonNull final KeepaliveInfo ki) {
+        if (SUCCESS_PAUSED != ki.mStopReason) {
+            throw new IllegalStateException("Keepalive is not paused");
+        }
+        try {
+            ki.mCallback.onStopped();
+        } catch (RemoteException e) {
+            Log.w(TAG, "Discarded onStopped callback while finalizing paused keepalive");
+        }
+    }
+
     public void handleCheckKeepalivesStillValid(NetworkAgentInfo nai) {
         HashMap <Integer, KeepaliveInfo> networkKeepalives = mKeepalives.get(nai);
         if (networkKeepalives != null) {
@@ -589,9 +640,14 @@
                 Log.d(TAG, "Started keepalive " + slot + " on " + nai.toShortString());
                 ki.mStartedState = KeepaliveInfo.STARTED;
                 try {
-                    ki.mCallback.onStarted();
+                    if (ki.mResumed) {
+                        ki.mCallback.onResumed();
+                    } else {
+                        ki.mCallback.onStarted();
+                    }
                 } catch (RemoteException e) {
-                    Log.w(TAG, "Discarded onStarted callback");
+                    Log.w(TAG, "Discarded " + (ki.mResumed ? "onResumed" : "onStarted")
+                            + " callback for slot " + slot);
                 }
             } else {
                 Log.d(TAG, "Failed to start keepalive " + slot + " on " + nai.toShortString()