[KA04] Expose TCP socket keepalive API

The new set of API allows applications to request keepalives
offload for established TCP sockets over wifi.

However, the application must not write to or read from the
socket after calling this method, until specific callbacks are
called.

Bug: 114151147
Test: atest FrameworksNetTests FrameworksWifiTests NetworkStackTests

Change-Id: I3880505dbc35fefa34ef6c79555458ecf5d296a4
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index fad5318..03530c4 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -68,6 +68,7 @@
 import java.lang.annotation.RetentionPolicy;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
+import java.net.Socket;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -1889,7 +1890,8 @@
      * @param callback A {@link SocketKeepalive.Callback}. Used for notifications about keepalive
      *        changes. Must be extended by applications that use this API.
      *
-     * @return A {@link SocketKeepalive} object, which can be used to control this keepalive object.
+     * @return A {@link SocketKeepalive} object that can be used to control the keepalive on the
+     *         given socket.
      **/
     public SocketKeepalive createSocketKeepalive(@NonNull Network network,
             @NonNull UdpEncapsulationSocket socket,
@@ -1918,6 +1920,8 @@
      * @param callback A {@link SocketKeepalive.Callback}. Used for notifications about keepalive
      *        changes. Must be extended by applications that use this API.
      *
+     * @return A {@link SocketKeepalive} object that can be used to control the keepalive on the
+     *         given socket.
      * @hide
      */
     @SystemApi
@@ -1933,6 +1937,34 @@
     }
 
     /**
+     * Request that keepalives be started on a TCP socket.
+     * The socket must be established.
+     *
+     * @param network The {@link Network} the socket is on.
+     * @param socket The socket that needs to be kept alive.
+     * @param executor The executor on which callback will be invoked. This implementation assumes
+     *                 the provided {@link Executor} runs the callbacks in sequence with no
+     *                 concurrency. Failing this, no guarantee of correctness can be made. It is
+     *                 the responsibility of the caller to ensure the executor provides this
+     *                 guarantee. A simple way of creating such an executor is with the standard
+     *                 tool {@code Executors.newSingleThreadExecutor}.
+     * @param callback A {@link SocketKeepalive.Callback}. Used for notifications about keepalive
+     *        changes. Must be extended by applications that use this API.
+     *
+     * @return A {@link SocketKeepalive} object that can be used to control the keepalive on the
+     *         given socket.
+     * @hide
+     */
+    @SystemApi
+    @RequiresPermission(android.Manifest.permission.PACKET_KEEPALIVE_OFFLOAD)
+    public SocketKeepalive createSocketKeepalive(@NonNull Network network,
+            @NonNull Socket socket,
+            @NonNull Executor executor,
+            @NonNull Callback callback) {
+        return new TcpSocketKeepalive(mService, network, socket, executor, callback);
+    }
+
+    /**
      * Ensure that a network route exists to deliver traffic to the specified
      * host via the specified network interface. An attempt to add a route that
      * already exists is ignored, but treated as successful.
diff --git a/core/java/android/net/IConnectivityManager.aidl b/core/java/android/net/IConnectivityManager.aidl
index 78fafeb..c926351 100644
--- a/core/java/android/net/IConnectivityManager.aidl
+++ b/core/java/android/net/IConnectivityManager.aidl
@@ -188,6 +188,9 @@
             int intervalSeconds, in Messenger messenger, in IBinder binder, String srcAddr,
             String dstAddr);
 
+    void startTcpKeepalive(in Network network, in FileDescriptor fd, int intervalSeconds,
+            in Messenger messenger, in IBinder binder);
+
     void stopKeepalive(in Network network, int slot);
 
     String getCaptivePortalServerUrl();
diff --git a/core/java/android/net/NetworkAgent.java b/core/java/android/net/NetworkAgent.java
index dfb6d6f..7bef690 100644
--- a/core/java/android/net/NetworkAgent.java
+++ b/core/java/android/net/NetworkAgent.java
@@ -177,6 +177,26 @@
      */
     public static final int EVENT_SOCKET_KEEPALIVE = BASE + 13;
 
+    // TODO: move the above 2 constants down so they are in order once merge conflicts are resolved
+    /**
+     * Sent by the KeepaliveTracker to NetworkAgent to add a packet filter.
+     *
+     * For TCP keepalive offloads, keepalive packets are sent by the firmware. However, because the
+     * remote site will send ACK packets in response to the keepalive packets, the firmware also
+     * needs to be configured to properly filter the ACKs to prevent the system from waking up.
+     * This does not happen with UDP, so this message is TCP-specific.
+     * arg1 = slot number of the keepalive to filter for.
+     * obj = the keepalive packet to send repeatedly.
+     */
+    public static final int CMD_ADD_KEEPALIVE_PACKET_FILTER = BASE + 16;
+
+    /**
+     * Sent by the KeepaliveTracker to NetworkAgent to remove a packet filter. See
+     * {@link #CMD_ADD_KEEPALIVE_PACKET_FILTER}.
+     * arg1 = slot number of the keepalive packet filter to remove.
+     */
+    public static final int CMD_REMOVE_KEEPALIVE_PACKET_FILTER = BASE + 17;
+
     /**
      * Sent by ConnectivityService to inform this network transport of signal strength thresholds
      * that when crossed should trigger a system wakeup and a NetworkCapabilities update.
@@ -312,6 +332,14 @@
                 preventAutomaticReconnect();
                 break;
             }
+            case CMD_ADD_KEEPALIVE_PACKET_FILTER: {
+                addKeepalivePacketFilter(msg);
+                break;
+            }
+            case CMD_REMOVE_KEEPALIVE_PACKET_FILTER: {
+                removeKeepalivePacketFilter(msg);
+                break;
+            }
         }
     }
 
@@ -461,6 +489,24 @@
     }
 
     /**
+     * Called by ConnectivityService to add specific packet filter to network hardware to block
+     * ACKs matching the sent keepalive packets. Implementations that support this feature must
+     * override this method.
+     */
+    protected void addKeepalivePacketFilter(Message msg) {
+        onSocketKeepaliveEvent(msg.arg1, SocketKeepalive.ERROR_HARDWARE_UNSUPPORTED);
+    }
+
+    /**
+     * Called by ConnectivityService to remove a packet filter installed with
+     * {@link #addKeepalivePacketFilter(Message)}. Implementations that support this feature
+     * must override this method.
+     */
+    protected void removeKeepalivePacketFilter(Message msg) {
+        onSocketKeepaliveEvent(msg.arg1, SocketKeepalive.ERROR_HARDWARE_UNSUPPORTED);
+    }
+
+    /**
      * Called by ConnectivityService to inform this network transport of signal strength thresholds
      * that when crossed should trigger a system wakeup and a NetworkCapabilities update.
      */
diff --git a/core/java/android/net/SocketKeepalive.java b/core/java/android/net/SocketKeepalive.java
index 7ea1bef..07728be 100644
--- a/core/java/android/net/SocketKeepalive.java
+++ b/core/java/android/net/SocketKeepalive.java
@@ -19,6 +19,7 @@
 import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.Looper;
@@ -155,7 +156,7 @@
     @NonNull private final SocketKeepalive.Callback mCallback;
     @NonNull private final Looper mLooper;
     @NonNull final Messenger mMessenger;
-    @NonNull Integer mSlot;
+    @Nullable Integer mSlot;
 
     SocketKeepalive(@NonNull IConnectivityManager service, @NonNull Network network,
             @NonNull Executor executor, @NonNull Callback callback) {
diff --git a/core/java/android/net/TcpSocketKeepalive.java b/core/java/android/net/TcpSocketKeepalive.java
new file mode 100644
index 0000000..8f6ee7b
--- /dev/null
+++ b/core/java/android/net/TcpSocketKeepalive.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright (C) 2019 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.net;
+
+import android.annotation.NonNull;
+import android.os.Binder;
+import android.os.RemoteException;
+import android.util.Log;
+
+import java.io.FileDescriptor;
+import java.net.Socket;
+import java.util.concurrent.Executor;
+
+/** @hide */
+final class TcpSocketKeepalive extends SocketKeepalive {
+
+    private final Socket mSocket;
+
+    TcpSocketKeepalive(@NonNull IConnectivityManager service,
+            @NonNull Network network,
+            @NonNull Socket socket,
+            @NonNull Executor executor,
+            @NonNull Callback callback) {
+        super(service, network, executor, callback);
+        mSocket = socket;
+    }
+
+    /**
+     * Starts keepalives. {@code mSocket} must be a connected TCP socket.
+     *
+     * - The application must not write to or read from the socket after calling this method, until
+     *   onDataReceived, onStopped, or onError are called. If it does, the keepalive will fail
+     *   with {@link #ERROR_SOCKET_NOT_IDLE}, or {@code #ERROR_INVALID_SOCKET} if the socket
+     *   experienced an error (as in poll(2) returned POLLERR); if this happens, the data received
+     *   from the socket may be invalid, and the socket can't be recovered.
+     * - If the socket has data in the send or receive buffer, then this call will fail with
+     *   {@link #ERROR_SOCKET_NOT_IDLE} and can be retried after the data has been processed.
+     *   An app could ensure this by using an application-layer protocol where it can receive
+     *   acknowledgement that it will go into keepalive mode. It could then go into keepalive
+     *   mode after having read the acknowledgement, draining the socket.
+     */
+    @Override
+    void startImpl(int intervalSec) {
+        try {
+            final FileDescriptor fd = mSocket.getFileDescriptor$();
+            mService.startTcpKeepalive(mNetwork, fd, intervalSec, mMessenger, new Binder());
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error starting packet keepalive: ", e);
+            stopLooper();
+        }
+    }
+
+    @Override
+    void stopImpl() {
+        try {
+            if (mSlot != null) {
+                mService.stopKeepalive(mNetwork, mSlot);
+            }
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error stopping packet keepalive: ", e);
+            stopLooper();
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index b98d7a1..253bdb3 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -6329,6 +6329,14 @@
     }
 
     @Override
+    public void startTcpKeepalive(Network network, FileDescriptor fd, int intervalSeconds,
+            Messenger messenger, IBinder binder) {
+        enforceKeepalivePermission();
+        mKeepaliveTracker.startTcpKeepalive(
+                getNetworkAgentInfoForNetwork(network), fd, intervalSeconds, messenger, binder);
+    }
+
+    @Override
     public void stopKeepalive(Network network, int slot) {
         mHandler.sendMessage(mHandler.obtainMessage(
                 NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE, slot, SocketKeepalive.SUCCESS, network));
diff --git a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
index d872e4d..6cff57d 100644
--- a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
+++ b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
@@ -17,6 +17,8 @@
 package com.android.server.connectivity;
 
 import static android.net.NattSocketKeepalive.NATT_PORT;
+import static android.net.NetworkAgent.CMD_ADD_KEEPALIVE_PACKET_FILTER;
+import static android.net.NetworkAgent.CMD_REMOVE_KEEPALIVE_PACKET_FILTER;
 import static android.net.NetworkAgent.CMD_START_SOCKET_KEEPALIVE;
 import static android.net.NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE;
 import static android.net.NetworkAgent.EVENT_SOCKET_KEEPALIVE;
@@ -37,6 +39,9 @@
 import android.net.NetworkAgent;
 import android.net.NetworkUtils;
 import android.net.SocketKeepalive.InvalidPacketException;
+import android.net.SocketKeepalive.InvalidSocketException;
+import android.net.TcpKeepalivePacketData;
+import android.net.TcpKeepalivePacketData.TcpSocketInfo;
 import android.net.util.IpUtils;
 import android.os.Binder;
 import android.os.Handler;
@@ -65,7 +70,7 @@
  *
  * Provides methods to stop and start keepalive requests, and keeps track of keepalives across all
  * networks. This class is tightly coupled to ConnectivityService. It is not thread-safe and its
- * methods must be called only from the ConnectivityService handler thread.
+ * handle* methods must be called only from the ConnectivityService handler thread.
  */
 public class KeepaliveTracker {
 
@@ -78,9 +83,12 @@
     private final HashMap <NetworkAgentInfo, HashMap<Integer, KeepaliveInfo>> mKeepalives =
             new HashMap<> ();
     private final Handler mConnectivityServiceHandler;
+    @NonNull
+    private final TcpKeepaliveController mTcpController;
 
     public KeepaliveTracker(Handler handler) {
         mConnectivityServiceHandler = handler;
+        mTcpController = new TcpKeepaliveController(handler);
     }
 
     /**
@@ -96,20 +104,33 @@
         private final int mUid;
         private final int mPid;
         private final NetworkAgentInfo mNai;
+        private final int mType;
+        private final FileDescriptor mFd;
 
-        /** Keepalive slot. A small integer that identifies this keepalive among the ones handled
-          * by this network. */
+        public static final int TYPE_NATT = 1;
+        public static final int TYPE_TCP = 2;
+
+        // Keepalive slot. A small integer that identifies this keepalive among the ones handled
+        // by this network.
         private int mSlot = NO_KEEPALIVE;
 
         // Packet data.
         private final KeepalivePacketData mPacket;
         private final int mInterval;
 
-        // Whether the keepalive is started or not.
-        public boolean isStarted;
+        // Whether the keepalive is started or not. The initial state is NOT_STARTED.
+        private static final int NOT_STARTED = 1;
+        private static final int STARTING = 2;
+        private static final int STARTED = 3;
+        private int mStartedState = NOT_STARTED;
 
-        public KeepaliveInfo(Messenger messenger, IBinder binder, NetworkAgentInfo nai,
-                KeepalivePacketData packet, int interval) {
+        KeepaliveInfo(@NonNull Messenger messenger,
+                @NonNull IBinder binder,
+                @NonNull NetworkAgentInfo nai,
+                @NonNull KeepalivePacketData packet,
+                int interval,
+                int type,
+                @NonNull FileDescriptor fd) {
             mMessenger = messenger;
             mBinder = binder;
             mPid = Binder.getCallingPid();
@@ -118,6 +139,8 @@
             mNai = nai;
             mPacket = packet;
             mInterval = interval;
+            mType = type;
+            mFd = fd;
 
             try {
                 mBinder.linkToDeath(this, 0);
@@ -130,32 +153,40 @@
             return mNai;
         }
 
+        private String startedStateString(final int state) {
+            switch (state) {
+                case NOT_STARTED : return "NOT_STARTED";
+                case STARTING : return "STARTING";
+                case STARTED : return "STARTED";
+            }
+            throw new IllegalArgumentException("Unknown state");
+        }
+
         public String toString() {
-            return new StringBuffer("KeepaliveInfo [")
-                    .append(" network=").append(mNai.network)
-                    .append(" isStarted=").append(isStarted)
-                    .append(" ")
-                    .append(IpUtils.addressAndPortToString(mPacket.srcAddress, mPacket.srcPort))
-                    .append("->")
-                    .append(IpUtils.addressAndPortToString(mPacket.dstAddress, mPacket.dstPort))
-                    .append(" interval=" + mInterval)
-                    .append(" packetData=" + HexDump.toHexString(mPacket.getPacket()))
-                    .append(" uid=").append(mUid).append(" pid=").append(mPid)
-                    .append(" ]")
-                    .toString();
+            return "KeepaliveInfo ["
+                    + " network=" + mNai.network
+                    + " startedState=" + startedStateString(mStartedState)
+                    + " "
+                    + IpUtils.addressAndPortToString(mPacket.srcAddress, mPacket.srcPort)
+                    + "->"
+                    + IpUtils.addressAndPortToString(mPacket.dstAddress, mPacket.dstPort)
+                    + " interval=" + mInterval
+                    + " uid=" + mUid + " pid=" + mPid
+                    + " packetData=" + HexDump.toHexString(mPacket.getPacket())
+                    + " ]";
         }
 
         /** Sends a message back to the application via its SocketKeepalive.Callback. */
         void notifyMessenger(int slot, int err) {
+            if (DBG) {
+                Log.d(TAG, "notify keepalive " + mSlot + " on " + mNai.network + " for " + err);
+            }
             KeepaliveTracker.this.notifyMessenger(mMessenger, slot, err);
         }
 
         /** Called when the application process is killed. */
         public void binderDied() {
-            // Not called from ConnectivityService handler thread, so send it a message.
-            mConnectivityServiceHandler.obtainMessage(
-                    NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE,
-                    mSlot, BINDER_DIED, mNai.network).sendToTarget();
+            stop(BINDER_DIED);
         }
 
         void unlinkDeathRecipient() {
@@ -202,7 +233,26 @@
             int error = isValid();
             if (error == SUCCESS) {
                 Log.d(TAG, "Starting keepalive " + mSlot + " on " + mNai.name());
-                mNai.asyncChannel.sendMessage(CMD_START_SOCKET_KEEPALIVE, slot, mInterval, mPacket);
+                switch (mType) {
+                    case TYPE_NATT:
+                        mNai.asyncChannel
+                                .sendMessage(CMD_START_SOCKET_KEEPALIVE, slot, mInterval, mPacket);
+                        break;
+                    case TYPE_TCP:
+                        mTcpController.startSocketMonitor(mFd, this, mSlot);
+                        mNai.asyncChannel
+                                .sendMessage(CMD_ADD_KEEPALIVE_PACKET_FILTER, slot, 0 /* Unused */,
+                                        mPacket);
+                        // TODO: check result from apf and notify of failure as needed.
+                        mNai.asyncChannel
+                                .sendMessage(CMD_START_SOCKET_KEEPALIVE, slot, mInterval, mPacket);
+                        break;
+                    default:
+                        Log.wtf(TAG, "Starting keepalive with unknown type: " + mType);
+                        handleStopKeepalive(mNai, mSlot, error);
+                        return;
+                }
+                mStartedState = STARTING;
             } else {
                 handleStopKeepalive(mNai, mSlot, error);
                 return;
@@ -216,15 +266,27 @@
                     Log.e(TAG, "Cannot stop unowned keepalive " + mSlot + " on " + mNai.network);
                 }
             }
-            if (isStarted) {
+            if (NOT_STARTED != mStartedState) {
                 Log.d(TAG, "Stopping keepalive " + mSlot + " on " + mNai.name());
-                mNai.asyncChannel.sendMessage(CMD_STOP_SOCKET_KEEPALIVE, mSlot);
+                if (mType == TYPE_NATT) {
+                    mNai.asyncChannel.sendMessage(CMD_STOP_SOCKET_KEEPALIVE, mSlot);
+                } else if (mType == TYPE_TCP) {
+                    mNai.asyncChannel.sendMessage(CMD_STOP_SOCKET_KEEPALIVE, mSlot);
+                    mNai.asyncChannel.sendMessage(CMD_REMOVE_KEEPALIVE_PACKET_FILTER, mSlot);
+                    mTcpController.stopSocketMonitor(mSlot);
+                } else {
+                    Log.wtf(TAG, "Stopping keepalive with unknown type: " + mType);
+                }
             }
             // TODO: at the moment we unconditionally return failure here. In cases where the
             // NetworkAgent is alive, should we ask it to reply, so it can return failure?
             notifyMessenger(mSlot, reason);
             unlinkDeathRecipient();
         }
+
+        void onFileDescriptorInitiatedStop(final int socketKeepaliveReason) {
+            handleStopKeepalive(mNai, mSlot, socketKeepaliveReason);
+        }
     }
 
     void notifyMessenger(Messenger messenger, int slot, int err) {
@@ -328,20 +390,38 @@
             return;
         }
 
-        if (reason == SUCCESS && !ki.isStarted) {
+        // This can be called in a number of situations :
+        // - startedState is STARTING.
+        //   - reason is SUCCESS => go to STARTED.
+        //   - reason isn't SUCCESS => it's an error starting. Go to NOT_STARTED and stop keepalive.
+        // - startedState is STARTED.
+        //   - reason is SUCCESS => it's a success stopping. Go to NOT_STARTED and stop keepalive.
+        //   - reason isn't SUCCESS => it's an error in exec. Go to NOT_STARTED and stop keepalive.
+        // The control is not supposed to ever come here if the state is NOT_STARTED. This is
+        // because in NOT_STARTED state, the code will switch to STARTING before sending messages
+        // to start, and the only way to NOT_STARTED is this function, through the edges outlined
+        // above : in all cases, keepalive gets stopped and can't restart without going into
+        // STARTING as messages are ordered. This also depends on the hardware processing the
+        // messages in order.
+        // TODO : clarify this code and get rid of mStartedState. Using a StateMachine is an
+        // option.
+        if (reason == SUCCESS && KeepaliveInfo.STARTING == ki.mStartedState) {
             // Keepalive successfully started.
             if (DBG) Log.d(TAG, "Started keepalive " + slot + " on " + nai.name());
-            ki.isStarted = true;
+            ki.mStartedState = KeepaliveInfo.STARTED;
             ki.notifyMessenger(slot, reason);
         } else {
             // Keepalive successfully stopped, or error.
-            ki.isStarted = false;
+            ki.mStartedState = KeepaliveInfo.NOT_STARTED;
             if (reason == SUCCESS) {
+                // The message indicated success stopping : don't call handleStopKeepalive.
                 if (DBG) Log.d(TAG, "Successfully stopped keepalive " + slot + " on " + nai.name());
             } else {
+                // The message indicated some error trying to start or during the course of
+                // keepalive : do call handleStopKeepalive.
+                handleStopKeepalive(nai, slot, reason);
                 if (DBG) Log.d(TAG, "Keepalive " + slot + " on " + nai.name() + " error " + reason);
             }
-            handleStopKeepalive(nai, slot, reason);
         }
     }
 
@@ -379,7 +459,47 @@
             notifyMessenger(messenger, NO_KEEPALIVE, e.error);
             return;
         }
-        KeepaliveInfo ki = new KeepaliveInfo(messenger, binder, nai, packet, intervalSeconds);
+        KeepaliveInfo ki = new KeepaliveInfo(messenger, binder, nai, packet, intervalSeconds,
+                KeepaliveInfo.TYPE_NATT, null);
+        mConnectivityServiceHandler.obtainMessage(
+                NetworkAgent.CMD_START_SOCKET_KEEPALIVE, ki).sendToTarget();
+    }
+
+    /**
+     * Called by ConnectivityService to start TCP keepalive on a file descriptor.
+     *
+     * In order to offload keepalive for application correctly, sequence number, ack number and
+     * other fields are needed to form the keepalive packet. Thus, this function synchronously
+     * puts the socket into repair mode to get the necessary information. After the socket has been
+     * put into repair mode, the application cannot access the socket until reverted to normal.
+     *
+     * See {@link android.net.SocketKeepalive}.
+     **/
+    public void startTcpKeepalive(@Nullable NetworkAgentInfo nai,
+            @NonNull FileDescriptor fd,
+            int intervalSeconds,
+            @NonNull Messenger messenger,
+            @NonNull IBinder binder) {
+        if (nai == null) {
+            notifyMessenger(messenger, NO_KEEPALIVE, ERROR_INVALID_NETWORK);
+            return;
+        }
+
+        TcpKeepalivePacketData packet = null;
+        try {
+            TcpSocketInfo tsi = TcpKeepaliveController.switchToRepairMode(fd);
+            packet = TcpKeepalivePacketData.tcpKeepalivePacket(tsi);
+        } catch (InvalidPacketException | InvalidSocketException e) {
+            try {
+                TcpKeepaliveController.switchOutOfRepairMode(fd);
+            } catch (ErrnoException e1) {
+                Log.e(TAG, "Couldn't move fd out of repair mode after failure to start keepalive");
+            }
+            notifyMessenger(messenger, NO_KEEPALIVE, e.error);
+            return;
+        }
+        KeepaliveInfo ki = new KeepaliveInfo(messenger, binder, nai, packet, intervalSeconds,
+                KeepaliveInfo.TYPE_TCP, fd);
         Log.d(TAG, "Created keepalive: " + ki.toString());
         mConnectivityServiceHandler.obtainMessage(CMD_START_SOCKET_KEEPALIVE, ki).sendToTarget();
     }
diff --git a/services/core/java/com/android/server/connectivity/TcpKeepaliveController.java b/services/core/java/com/android/server/connectivity/TcpKeepaliveController.java
index 640504f..8a9ac23 100644
--- a/services/core/java/com/android/server/connectivity/TcpKeepaliveController.java
+++ b/services/core/java/com/android/server/connectivity/TcpKeepaliveController.java
@@ -15,7 +15,6 @@
  */
 package com.android.server.connectivity;
 
-import static android.net.NetworkAgent.EVENT_SOCKET_KEEPALIVE;
 import static android.net.SocketKeepalive.DATA_RECEIVED;
 import static android.net.SocketKeepalive.ERROR_INVALID_SOCKET;
 import static android.net.SocketKeepalive.ERROR_SOCKET_NOT_IDLE;
@@ -31,10 +30,8 @@
 import android.net.TcpKeepalivePacketData.TcpSocketInfo;
 import android.net.TcpRepairWindow;
 import android.os.Handler;
-import android.os.Message;
 import android.os.MessageQueue;
 import android.os.Messenger;
-import android.os.RemoteException;
 import android.system.ErrnoException;
 import android.system.Int32Ref;
 import android.system.Os;
@@ -42,6 +39,7 @@
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.server.connectivity.KeepaliveTracker.KeepaliveInfo;
 
 import java.io.FileDescriptor;
 import java.net.InetAddress;
@@ -111,7 +109,7 @@
      * tcp/ip information.
      */
     // TODO : make this private. It's far too confusing that this gets called from outside
-    // at a time that nobody can understand, but the switch out is in this class only.
+    // at a time that nobody can understand.
     public static TcpSocketInfo switchToRepairMode(FileDescriptor fd)
             throws InvalidSocketException {
         if (DBG) Log.i(TAG, "switchToRepairMode to start tcp keepalive : " + fd);
@@ -199,7 +197,13 @@
                 trw.rcvWndScale);
     }
 
-    private static void switchOutOfRepairMode(@NonNull final FileDescriptor fd)
+    /**
+     * Switch the tcp socket out of repair mode.
+     *
+     * @param fd the fd of socket to switch back to normal.
+     */
+    // TODO : make this private.
+    public static void switchOutOfRepairMode(@NonNull final FileDescriptor fd)
             throws ErrnoException {
         Os.setsockoptInt(fd, IPPROTO_TCP, TCP_REPAIR, TCP_REPAIR_OFF);
     }
@@ -212,7 +216,7 @@
      * @param slot keepalive slot.
      */
     public void startSocketMonitor(@NonNull final FileDescriptor fd,
-            @NonNull final Messenger messenger, final int slot) {
+            @NonNull final KeepaliveInfo ki, final int slot) {
         synchronized (mListeners) {
             if (null != mListeners.get(slot)) {
                 throw new IllegalArgumentException("This slot is already taken");
@@ -226,31 +230,13 @@
                 // This can't be called twice because the queue guarantees that once the listener
                 // is unregistered it can't be called again, even for a message that arrived
                 // before it was unregistered.
-                int result;
-                try {
-                    // First move the socket out of repair mode.
-                    if (DBG) Log.d(TAG, "Moving socket out of repair mode for event : " + readyFd);
-                    switchOutOfRepairMode(readyFd);
-                    result = (0 != (events & EVENT_ERROR)) ? ERROR_INVALID_SOCKET : DATA_RECEIVED;
-                } catch (ErrnoException e) {
-                    // Could not move the socket out of repair mode. Still continue with notifying
-                    // the client
-                    Log.e(TAG, "Cannot switch socket out of repair mode", e);
-                    result = ERROR_INVALID_SOCKET;
+                final int reason;
+                if (0 != (events & EVENT_ERROR)) {
+                    reason = ERROR_INVALID_SOCKET;
+                } else {
+                    reason = DATA_RECEIVED;
                 }
-                // Prepare and send the message to the receiver.
-                final Message message = Message.obtain();
-                message.what = EVENT_SOCKET_KEEPALIVE;
-                message.arg1 = slot;
-                message.arg2 = result;
-                try {
-                    messenger.send(message);
-                } catch (RemoteException e) {
-                    // Remote process died
-                }
-                synchronized (mListeners) {
-                    mListeners.remove(slot);
-                }
+                ki.onFileDescriptorInitiatedStop(reason);
                 // The listener returns the new set of events to listen to. Because 0 means no
                 // event, the listener gets unregistered.
                 return 0;