Merge changes from topic "idiag_sock_id"

* changes:
  Switch SkDestroyListener based on the flag.
  Add startSkDestroyListener flag to native_init
  Add SkDestroyListenerTest

Tethering/tests/mts/Android.bp

@@ -41,6 +41,8 @@ android_test {
         "ctstestrunner-axt",
         "junit",
         "junit-params",
+        "connectivity-net-module-utils-bpf",
+        "net-utils-device-common-bpf",
     ],
 
     jni_libs: [

Tethering/tests/mts/src/android/tethering/mts/SkDestroyListenerTest.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright (C) 2022 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.tethering.mts;
+
+import static android.system.OsConstants.AF_INET;
+import static android.system.OsConstants.AF_INET6;
+import static android.system.OsConstants.SOCK_DGRAM;
+import static android.system.OsConstants.SOCK_STREAM;
+
+import static com.android.compatibility.common.util.SystemUtil.runShellCommandOrThrow;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import android.net.TrafficStats;
+import android.os.Build;
+import android.os.Process;
+import android.system.Os;
+import android.util.Pair;
+
+import com.android.net.module.util.BpfDump;
+import com.android.net.module.util.bpf.CookieTagMapKey;
+import com.android.net.module.util.bpf.CookieTagMapValue;
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
+import com.android.testutils.DevSdkIgnoreRunner;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.FileDescriptor;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@RunWith(DevSdkIgnoreRunner.class)
+@IgnoreUpTo(Build.VERSION_CODES.S_V2)
+public class SkDestroyListenerTest {
+    private static final int COOKIE_TAG = 0x1234abcd;
+    private static final int SOCKET_COUNT = 100;
+    private static final int SOCKET_CLOSE_WAIT_MS = 200;
+    private static final String LINE_DELIMITER = "\\n";
+    private static final String DUMP_COMMAND = "dumpsys netstats --bpfRawMap --cookieTagMap";
+
+    private Map<CookieTagMapKey, CookieTagMapValue> parseBpfRawMap(final String dump) {
+        final Map<CookieTagMapKey, CookieTagMapValue> map = new HashMap<>();
+        for (final String line: dump.split(LINE_DELIMITER)) {
+            final Pair<CookieTagMapKey, CookieTagMapValue> keyValue =
+                    BpfDump.fromBase64EncodedString(CookieTagMapKey.class,
+                            CookieTagMapValue.class, line.trim());
+            map.put(keyValue.first, keyValue.second);
+        }
+        return map;
+    }
+
+    private int countTaggedSocket() {
+        final String dump = runShellCommandOrThrow(DUMP_COMMAND);
+        final Map<CookieTagMapKey, CookieTagMapValue> cookieTagMap = parseBpfRawMap(dump);
+        int count = 0;
+        for (final CookieTagMapValue value: cookieTagMap.values()) {
+            if (value.tag == COOKIE_TAG && value.uid == Process.myUid()) {
+                count++;
+            }
+        }
+        return count;
+    }
+
+    private boolean noTaggedSocket() {
+        return countTaggedSocket() == 0;
+    }
+
+    private void doTestSkDestroyListener(final int family, final int type) throws Exception {
+        assertTrue("There are tagged sockets before test", noTaggedSocket());
+
+        TrafficStats.setThreadStatsTag(COOKIE_TAG);
+        final List<FileDescriptor> fds = new ArrayList<>();
+        for (int i = 0; i < SOCKET_COUNT; i++) {
+            fds.add(Os.socket(family, type, 0 /* protocol */));
+        }
+        TrafficStats.clearThreadStatsTag();
+        assertEquals("Number of tagged socket does not match after creating sockets",
+                SOCKET_COUNT, countTaggedSocket());
+
+        for (final FileDescriptor fd: fds) {
+            Os.close(fd);
+        }
+        // Wait a bit for skDestroyListener to handle all the netlink messages.
+        Thread.sleep(SOCKET_CLOSE_WAIT_MS);
+        assertTrue("There are tagged sockets after closing sockets", noTaggedSocket());
+    }
+
+    @Test
+    public void testSkDestroyListener() throws Exception {
+        doTestSkDestroyListener(AF_INET, SOCK_STREAM);
+        doTestSkDestroyListener(AF_INET, SOCK_DGRAM);
+        doTestSkDestroyListener(AF_INET6, SOCK_STREAM);
+        doTestSkDestroyListener(AF_INET6, SOCK_DGRAM);
+    }
+}

service-t/src/com/android/server/net/NetworkStatsService.java

@@ -162,11 +162,13 @@ import com.android.net.module.util.IBpfMap;
 import com.android.net.module.util.LocationPermissionChecker;
 import com.android.net.module.util.NetworkStatsUtils;
 import com.android.net.module.util.PermissionUtils;
+import com.android.net.module.util.SharedLog;
 import com.android.net.module.util.Struct;
 import com.android.net.module.util.Struct.U32;
 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.server.BpfNetMaps;
 
 import java.io.File;
 import java.io.FileDescriptor;
@@ -454,6 +456,9 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
     @NonNull
     private final BpfInterfaceMapUpdater mInterfaceMapUpdater;
 
+    @Nullable
+    private final SkDestroyListener mSkDestroyListener;
+
     private static @NonNull Clock getDefaultClock() {
         return new BestClock(ZoneOffset.UTC, SystemClock.currentNetworkTimeClock(),
                 Clock.systemUTC());
@@ -587,6 +592,18 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         mStatsMapA = mDeps.getStatsMapA();
         mStatsMapB = mDeps.getStatsMapB();
         mAppUidStatsMap = mDeps.getAppUidStatsMap();
+
+        // TODO: Remove bpfNetMaps creation and always start SkDestroyListener
+        // Following code is for the experiment to verify the SkDestroyListener refactoring. Based
+        // on the experiment flag, BpfNetMaps starts C SkDestroyListener (existing code) or
+        // NetworkStatsService starts Java SkDestroyListener (new code).
+        final BpfNetMaps bpfNetMaps = mDeps.makeBpfNetMaps(mContext);
+        if (bpfNetMaps.isSkDestroyListenerRunning()) {
+            mSkDestroyListener = null;
+        } else {
+            mSkDestroyListener = mDeps.makeSkDestroyListener(mCookieTagMap, mHandler);
+            mHandler.post(mSkDestroyListener::start);
+        }
     }
 
     /**
@@ -782,6 +799,17 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         public boolean isDebuggable() {
             return Build.isDebuggable();
         }
+
+        /** Create a new BpfNetMaps. */
+        public BpfNetMaps makeBpfNetMaps(Context ctx) {
+            return new BpfNetMaps(ctx);
+        }
+
+        /** Create a new SkDestroyListener. */
+        public SkDestroyListener makeSkDestroyListener(
+                IBpfMap<CookieTagMapKey, CookieTagMapValue> cookieTagMap, Handler handler) {
+            return new SkDestroyListener(cookieTagMap, handler, new SharedLog(TAG));
+        }
     }
 
     /**

service-t/src/com/android/server/net/SkDestroyListener.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2022 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 com.android.server.net;
+
+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;
+
+/**
+ * Monitor socket destroy and delete entry from cookie tag bpf map.
+ */
+public class SkDestroyListener extends NetlinkMonitor {
+    private static final int SKNLGRP_INET_TCP_DESTROY = 1;
+    private static final int SKNLGRP_INET_UDP_DESTROY = 2;
+    private static final int SKNLGRP_INET6_TCP_DESTROY = 3;
+    private static final int SKNLGRP_INET6_UDP_DESTROY = 4;
+
+    // TODO: if too many sockets are closed too quickly, this can overflow the socket buffer, and
+    // some entries in mCookieTagMap will not be freed. In order to fix this it would be needed to
+    // periodically dump all sockets and remove the tag entries for sockets that have been closed.
+    // For now, set a large-enough buffer that hundreds of sockets can be closed without getting
+    // ENOBUFS and leaking mCookieTagMap entries.
+    private static final int SOCK_RCV_BUF_SIZE = 512 * 1024;
+
+    private final IBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap;
+
+    SkDestroyListener(final IBpfMap<CookieTagMapKey, CookieTagMapValue> cookieTagMap,
+            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;
+    }
+
+    @Override
+    public void processNetlinkMessage(final NetlinkMessage nlMsg, final long whenMs) {
+        if (!(nlMsg instanceof InetDiagMessage)) {
+            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);
+        }
+    }
+}

service/jni/com_android_server_BpfNetMaps.cpp

@@ -47,8 +47,8 @@ namespace android {
       ALOGE("%s failed, error code = %d", __func__, status.code()); \
   } while (0)
 
-static void native_init(JNIEnv* env, jclass clazz) {
+static void native_init(JNIEnv* env, jclass clazz, jboolean startSkDestroyListener) {
-  Status status = mTc.start();
+  Status status = mTc.start(startSkDestroyListener);
   CHECK_LOG(status);
   if (!isOk(status)) {
     uid_t uid = getuid();
@@ -201,7 +201,7 @@ static jint native_synchronizeKernelRCU(JNIEnv* env, jobject self) {
 // clang-format off
 static const JNINativeMethod gMethods[] = {
     /* name, signature, funcPtr */
-    {"native_init", "()V",
+    {"native_init", "(Z)V",
     (void*)native_init},
     {"native_addNaughtyApp", "(I)I",
     (void*)native_addNaughtyApp},

service/native/TrafficController.cpp

@@ -181,9 +181,13 @@ Status TrafficController::initMaps() {
     return netdutils::status::ok;
 }
 
-Status TrafficController::start() {
+Status TrafficController::start(bool startSkDestroyListener) {
     RETURN_IF_NOT_OK(initMaps());
 
+    if (!startSkDestroyListener) {
+        return netdutils::status::ok;
+    }
+
     auto result = makeSkDestroyListener();
     if (!isOk(result)) {
         ALOGE("Unable to create SkDestroyListener: %s", toString(result).c_str());

service/native/include/TrafficController.h

@@ -38,7 +38,7 @@ class TrafficController {
     /*
      * Initialize the whole controller
      */
-    netdutils::Status start();
+    netdutils::Status start(bool startSkDestroyListener);
 
     /*
      * Swap the stats map config from current active stats map to the idle one.

service/src/com/android/server/BpfNetMaps.java

@@ -259,10 +259,14 @@ public class BpfNetMaps {
         Log.d(TAG, "BpfNetMaps is initialized with sEnableJavaBpfMap=" + sEnableJavaBpfMap);
 
         initBpfMaps();
-        native_init();
+        native_init(!sEnableJavaBpfMap /* startSkDestroyListener */);
         sInitialized = true;
     }
 
+    public boolean isSkDestroyListenerRunning() {
+        return !sEnableJavaBpfMap;
+    }
+
     /**
      * Dependencies of BpfNetMaps, for injection in tests.
      */
@@ -919,7 +923,7 @@ public class BpfNetMaps {
         native_dump(fd, verbose);
     }
 
-    private static native void native_init();
+    private static native void native_init(boolean startSkDestroyListener);
     private native int native_addNaughtyApp(int uid);
     private native int native_removeNaughtyApp(int uid);
     private native int native_addNiceApp(int uid);

tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java

@@ -143,6 +143,7 @@ import com.android.net.module.util.Struct.U32;
 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.server.BpfNetMaps;
 import com.android.server.net.NetworkStatsService.AlertObserver;
 import com.android.server.net.NetworkStatsService.NetworkStatsSettings;
 import com.android.server.net.NetworkStatsService.NetworkStatsSettings.Config;
@@ -249,6 +250,10 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     @Mock
     private LocationPermissionChecker mLocationPermissionChecker;
     private TestBpfMap<U32, U8> mUidCounterSetMap = spy(new TestBpfMap<>(U32.class, U8.class));
+    @Mock
+    private BpfNetMaps mBpfNetMaps;
+    @Mock
+    private SkDestroyListener mSkDestroyListener;
 
     private TestBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap = new TestBpfMap<>(
             CookieTagMapKey.class, CookieTagMapValue.class);
@@ -501,6 +506,17 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
             public boolean isDebuggable() {
                 return mIsDebuggable == Boolean.TRUE;
             }
+
+            @Override
+            public BpfNetMaps makeBpfNetMaps(Context ctx) {
+                return mBpfNetMaps;
+            }
+
+            @Override
+            public SkDestroyListener makeSkDestroyListener(
+                    IBpfMap<CookieTagMapKey, CookieTagMapValue> cookieTagMap, Handler handler) {
+                return mSkDestroyListener;
+            }
         };
     }