2 files changed, 45 insertions, 6 deletions

diff --git a/src/android/net/apf/ApfFilter.java b/src/android/net/apf/ApfFilter.java
index 663e2f1..359c859 100644
--- a/src/android/net/apf/ApfFilter.java
+++ b/src/android/net/apf/ApfFilter.java
@@ -39,6 +39,7 @@ import android.content.Intent;
 import android.content.IntentFilter;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
+import android.net.NattKeepalivePacketDataParcelable;
 import android.net.TcpKeepalivePacketDataParcelable;
 import android.net.apf.ApfGenerator.IllegalInstructionException;
 import android.net.apf.ApfGenerator.Register;
@@ -1691,13 +1692,13 @@ public class ApfFilter {
     }
 
     /**
-     * Add keepalive ack packet filter.
+     * Add TCP keepalive ack packet filter.
      * This will add a filter to drop acks to the keepalive packet passed as an argument.
      *
      * @param slot The index used to access the filter.
      * @param sentKeepalivePacket The attributes of the sent keepalive packet.
      */
-    public synchronized void addKeepalivePacketFilter(final int slot,
+    public synchronized void addTcpKeepalivePacketFilter(final int slot,
             final TcpKeepalivePacketDataParcelable sentKeepalivePacket) {
         log("Adding keepalive ack(" + slot + ")");
         if (null != mKeepaliveAcks.get(slot)) {
@@ -1711,6 +1712,18 @@ public class ApfFilter {
     }
 
     /**
+     * Add NATT keepalive packet filter.
+     * This will add a filter to drop NATT keepalive packet which is passed as an argument.
+     *
+     * @param slot The index used to access the filter.
+     * @param sentKeepalivePacket The attributes of the sent keepalive packet.
+     */
+    public synchronized void addNattKeepalivePacketFilter(final int slot,
+            final NattKeepalivePacketDataParcelable sentKeepalivePacket) {
+        Log.e(TAG, "APF add NATT keepalive filter is not implemented");
+    }
+
+    /**
      * Remove keepalive packet filter.
      *
      * @param slot The index used to access the filter.
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;
                 }