Update SkDestroyListener to take consumer for inetDiagMessage processing

This is a no-op refactoring (except for logging) in preparation for an
upcoming CL(aosp/3502915) that will use SkDestroyListener to close quic
connections for destroyed quic sockets.

Test: atest FrameworksNetTests
Bug: 311792075
Change-Id: I41e2a16075fc9242cae19f84d1255358333744ab
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 27f039f..03c64f5 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -188,6 +188,8 @@
 import com.android.net.module.util.Struct.U8;
 import com.android.net.module.util.bpf.CookieTagMapKey;
 import com.android.net.module.util.bpf.CookieTagMapValue;
+import com.android.net.module.util.netlink.InetDiagMessage;
+import com.android.net.module.util.netlink.StructInetDiagSockId;
 import com.android.networkstack.apishim.BroadcastOptionsShimImpl;
 import com.android.networkstack.apishim.ConstantsShim;
 import com.android.networkstack.apishim.common.UnsupportedApiLevelException;
@@ -215,6 +217,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 /**
  * Collect and persist detailed network statistics, and provide this data to
@@ -725,7 +728,14 @@
             mTrafficStatsUidCache = null;
         }
 
-        mSkDestroyListener = mDeps.makeSkDestroyListener(mCookieTagMap, mHandler);
+        mSkDestroyListener = mDeps.makeSkDestroyListener((message) -> {
+            final StructInetDiagSockId sockId = message.inetDiagMsg.id;
+            try {
+                mCookieTagMap.deleteEntry(new CookieTagMapKey(sockId.cookie));
+            } catch (ErrnoException e) {
+                Log.e(TAG, "Failed to delete CookieTagMap entry for " + sockId.cookie  + ": " + e);
+            }
+        }, mHandler);
         mHandler.post(mSkDestroyListener::start);
     }
 
@@ -947,10 +957,10 @@
         }
 
         /** Create a new SkDestroyListener. */
-        public SkDestroyListener makeSkDestroyListener(
-                IBpfMap<CookieTagMapKey, CookieTagMapValue> cookieTagMap, Handler handler) {
-            return new SkDestroyListener(
-                    cookieTagMap, handler, new SharedLog(MAX_SOCKET_DESTROY_LISTENER_LOGS, TAG));
+        public SkDestroyListener makeSkDestroyListener(Consumer<InetDiagMessage> consumer,
+                Handler handler) {
+            return new SkDestroyListener(consumer, handler,
+                    new SharedLog(MAX_SOCKET_DESTROY_LISTENER_LOGS, TAG));
         }
 
         /**
diff --git a/service-t/src/com/android/server/net/SkDestroyListener.java b/service-t/src/com/android/server/net/SkDestroyListener.java
index a6cc2b5..789e9d6 100644
--- a/service-t/src/com/android/server/net/SkDestroyListener.java
+++ b/service-t/src/com/android/server/net/SkDestroyListener.java
@@ -19,18 +19,14 @@
 import static android.system.OsConstants.NETLINK_INET_DIAG;
 
 import android.os.Handler;
-import android.system.ErrnoException;
 
-import com.android.net.module.util.IBpfMap;
 import com.android.net.module.util.SharedLog;
-import com.android.net.module.util.bpf.CookieTagMapKey;
-import com.android.net.module.util.bpf.CookieTagMapValue;
 import com.android.net.module.util.ip.NetlinkMonitor;
 import com.android.net.module.util.netlink.InetDiagMessage;
 import com.android.net.module.util.netlink.NetlinkMessage;
-import com.android.net.module.util.netlink.StructInetDiagSockId;
 
 import java.io.PrintWriter;
+import java.util.function.Consumer;
 
 /**
  * Monitor socket destroy and delete entry from cookie tag bpf map.
@@ -48,17 +44,17 @@
     // ENOBUFS and leaking mCookieTagMap entries.
     private static final int SOCK_RCV_BUF_SIZE = 512 * 1024;
 
-    private final IBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap;
+    private final Consumer<InetDiagMessage> mSkDestroyCallback;
 
-    SkDestroyListener(final IBpfMap<CookieTagMapKey, CookieTagMapValue> cookieTagMap,
-            final Handler handler, final SharedLog log) {
+    SkDestroyListener(final Consumer<InetDiagMessage> consumer, final Handler handler,
+            final SharedLog log) {
         super(handler, log, "SkDestroyListener", NETLINK_INET_DIAG,
                 1 << (SKNLGRP_INET_TCP_DESTROY - 1)
                         | 1 << (SKNLGRP_INET_UDP_DESTROY - 1)
                         | 1 << (SKNLGRP_INET6_TCP_DESTROY - 1)
                         | 1 << (SKNLGRP_INET6_UDP_DESTROY - 1),
                 SOCK_RCV_BUF_SIZE);
-        mCookieTagMap = cookieTagMap;
+        mSkDestroyCallback = consumer;
     }
 
     @Override
@@ -67,12 +63,7 @@
             mLog.e("Received non InetDiagMessage");
             return;
         }
-        final StructInetDiagSockId sockId = ((InetDiagMessage) nlMsg).inetDiagMsg.id;
-        try {
-            mCookieTagMap.deleteEntry(new CookieTagMapKey(sockId.cookie));
-        } catch (ErrnoException e) {
-            mLog.e("Failed to delete CookieTagMap entry for " + sockId.cookie  + ": " + e);
-        }
+        mSkDestroyCallback.accept((InetDiagMessage) nlMsg);
     }
 
     /**
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 581c709..d93e709 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -168,6 +168,7 @@
 import com.android.net.module.util.Struct.U8;
 import com.android.net.module.util.bpf.CookieTagMapKey;
 import com.android.net.module.util.bpf.CookieTagMapValue;
+import com.android.net.module.util.netlink.InetDiagMessage;
 import com.android.server.connectivity.ConnectivityResources;
 import com.android.server.net.NetworkStatsService.AlertObserver;
 import com.android.server.net.NetworkStatsService.NetworkStatsSettings;
@@ -210,6 +211,7 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 
 /**
  * Tests for {@link NetworkStatsService}.
@@ -605,8 +607,8 @@
         }
 
         @Override
-        public SkDestroyListener makeSkDestroyListener(
-                IBpfMap<CookieTagMapKey, CookieTagMapValue> cookieTagMap, Handler handler) {
+        public SkDestroyListener makeSkDestroyListener(Consumer<InetDiagMessage> consumer,
+                Handler handler) {
             return mSkDestroyListener;
         }
 
diff --git a/tests/unit/java/com/android/server/net/SkDestroyListenerTest.kt b/tests/unit/java/com/android/server/net/SkDestroyListenerTest.kt
index 18785e5..dec71f8 100644
--- a/tests/unit/java/com/android/server/net/SkDestroyListenerTest.kt
+++ b/tests/unit/java/com/android/server/net/SkDestroyListenerTest.kt
@@ -54,7 +54,7 @@
         doReturn(sharedLog).`when`(sharedLog).forSubComponent(any())
 
         val handler = Handler(handlerThread.looper)
-        val skDestroylistener = SkDestroyListener(null /* cookieTagMap */, handler, sharedLog)
+        val skDestroylistener = SkDestroyListener({} /* consumer */, handler, sharedLog)
         val pw = PrintWriter(System.out)
         skDestroylistener.dump(pw)