1 files changed, 30 insertions, 4 deletions

diff --git a/src/android/net/ip/IpClient.java b/src/android/net/ip/IpClient.java
index 96e09fa..dc74c04 100644
--- a/src/android/net/ip/IpClient.java
+++ b/src/android/net/ip/IpClient.java
@@ -29,6 +29,7 @@ import android.net.INetd;
 import android.net.IpPrefix;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
+import android.net.NattKeepalivePacketDataParcelable;
 import android.net.NetworkStackIpMemoryStore;
 import android.net.ProvisioningConfigurationParcelable;
 import android.net.ProxyInfo;
@@ -371,6 +372,10 @@ public class IpClient extends StateMachine {
     private boolean mMulticastFiltering;
     private long mStartTimeMillis;
 
+    /* This must match the definition in KeepaliveTracker.KeepaliveInfo */
+    private static final int TYPE_NATT = 1;
+    private static final int TYPE_TCP = 2;
+
     /**
      * Reading the snapshot is an asynchronous operation initiated by invoking
      * Callback.startReadPacketFilter() and completed when the WiFi Service responds with an
@@ -553,6 +558,11 @@ public class IpClient extends StateMachine {
             IpClient.this.addKeepalivePacketFilter(slot, pkt);
         }
         @Override
+        public void addNattKeepalivePacketFilter(int slot, NattKeepalivePacketDataParcelable pkt) {
+            checkNetworkStackCallingPermission();
+            IpClient.this.addNattKeepalivePacketFilter(slot, pkt);
+        }
+        @Override
         public void removeKeepalivePacketFilter(int slot) {
             checkNetworkStackCallingPermission();
             IpClient.this.removeKeepalivePacketFilter(slot);
@@ -691,11 +701,20 @@ public class IpClient extends StateMachine {
     }
 
     /**
-     * Called by WifiStateMachine to add keepalive packet filter before setting up
+     * Called by WifiStateMachine to add TCP keepalive packet filter before setting up
      * keepalive offload.
      */
     public void addKeepalivePacketFilter(int slot, @NonNull TcpKeepalivePacketDataParcelable pkt) {
-        sendMessage(CMD_ADD_KEEPALIVE_PACKET_FILTER_TO_APF, slot, 0 /* Unused */, pkt);
+        sendMessage(CMD_ADD_KEEPALIVE_PACKET_FILTER_TO_APF, slot, TYPE_TCP, pkt);
+    }
+
+    /**
+     *  Called by WifiStateMachine to add NATT keepalive packet filter before setting up
+     *  keepalive offload.
+     */
+    public void addNattKeepalivePacketFilter(int slot,
+            @NonNull NattKeepalivePacketDataParcelable pkt) {
+        sendMessage(CMD_ADD_KEEPALIVE_PACKET_FILTER_TO_APF, slot, TYPE_NATT, pkt);
     }
 
     /**
@@ -1607,9 +1626,16 @@ public class IpClient extends StateMachine {
 
                 case CMD_ADD_KEEPALIVE_PACKET_FILTER_TO_APF: {
                     final int slot = msg.arg1;
+                    final int type = msg.arg2;
+
                     if (mApfFilter != null) {
-                        mApfFilter.addKeepalivePacketFilter(slot,
-                                (TcpKeepalivePacketDataParcelable) msg.obj);
+                        if (type == TYPE_NATT) {
+                            mApfFilter.addNattKeepalivePacketFilter(slot,
+                                    (NattKeepalivePacketDataParcelable) msg.obj);
+                        } else {
+                            mApfFilter.addTcpKeepalivePacketFilter(slot,
+                                    (TcpKeepalivePacketDataParcelable) msg.obj);
+                        }
                     }
                     break;
                 }