From b9548860a95af20a5b6936b1a72eee51482fa8db Mon Sep 17 00:00:00 2001
From: Motomu Utsumi <motomuman@google.com>
Date: Tue, 6 Sep 2022 16:30:05 +0900
Subject: [PATCH] Open CookieTagMap in BpfNetMaps

Upcoming CL will update BpfNetMaps to pull NetworkBpfMapInfo atom that
contains the cookieTagMap size field.

Test: m, atest BpfNetMaps
Bug: 217624062
Change-Id: I84bb0aa56e3916fc6fc0ca31ba5962f3b9be4b6f
---
 .../src/com/android/server/BpfNetMaps.java    | 27 +++++++++++++++++++
 .../com/android/server/BpfNetMapsTest.java    |  5 ++++
 2 files changed, 32 insertions(+)

diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index dc5c4c75b0..4eeef5c558 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -50,6 +50,8 @@ import com.android.net.module.util.DeviceConfigUtils;
 import com.android.net.module.util.IBpfMap;
 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 java.io.FileDescriptor;
 import java.io.IOException;
@@ -94,6 +96,8 @@ public class BpfNetMaps {
             "/sys/fs/bpf/netd_shared/map_netd_uid_owner_map";
     private static final String UID_PERMISSION_MAP_PATH =
             "/sys/fs/bpf/netd_shared/map_netd_uid_permission_map";
+    private static final String COOKIE_TAG_MAP_PATH =
+            "/sys/fs/bpf/netd_shared/map_netd_cookie_tag_map";
     private static final U32 UID_RULES_CONFIGURATION_KEY = new U32(0);
     private static final U32 CURRENT_STATS_MAP_CONFIGURATION_KEY = new U32(1);
     private static final long UID_RULES_DEFAULT_CONFIGURATION = 0;
@@ -104,6 +108,7 @@ public class BpfNetMaps {
     // BpfMap for UID_OWNER_MAP_PATH. This map is not accessed by others.
     private static IBpfMap<U32, UidOwnerValue> sUidOwnerMap = null;
     private static IBpfMap<U32, U8> sUidPermissionMap = null;
+    private static IBpfMap<CookieTagMapKey, CookieTagMapValue> sCookieTagMap = null;
 
     // LINT.IfChange(match_type)
     @VisibleForTesting public static final long NO_MATCH = 0;
@@ -153,6 +158,15 @@ public class BpfNetMaps {
         sUidPermissionMap = uidPermissionMap;
     }
 
+    /**
+     * Set cookieTagMap for test.
+     */
+    @VisibleForTesting
+    public static void setCookieTagMapForTest(
+            IBpfMap<CookieTagMapKey, CookieTagMapValue> cookieTagMap) {
+        sCookieTagMap = cookieTagMap;
+    }
+
     private static IBpfMap<U32, U32> getConfigurationMap() {
         try {
             return new BpfMap<>(
@@ -180,6 +194,15 @@ public class BpfNetMaps {
         }
     }
 
+    private static IBpfMap<CookieTagMapKey, CookieTagMapValue> getCookieTagMap() {
+        try {
+            return new BpfMap<>(COOKIE_TAG_MAP_PATH, BpfMap.BPF_F_RDWR,
+                    CookieTagMapKey.class, CookieTagMapValue.class);
+        } catch (ErrnoException e) {
+            throw new IllegalStateException("Cannot open cookie tag map", e);
+        }
+    }
+
     private static void initBpfMaps() {
         if (sConfigurationMap == null) {
             sConfigurationMap = getConfigurationMap();
@@ -209,6 +232,10 @@ public class BpfNetMaps {
         if (sUidPermissionMap == null) {
             sUidPermissionMap = getUidPermissionMap();
         }
+
+        if (sCookieTagMap == null) {
+            sCookieTagMap = getCookieTagMap();
+        }
     }
 
     /**
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index 15a2e56fcf..b7d80acbbb 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -61,6 +61,8 @@ import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.IBpfMap;
 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.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreAfter;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
@@ -117,6 +119,8 @@ public final class BpfNetMapsTest {
     private final IBpfMap<U32, UidOwnerValue> mUidOwnerMap =
             new TestBpfMap<>(U32.class, UidOwnerValue.class);
     private final IBpfMap<U32, U8> mUidPermissionMap = new TestBpfMap<>(U32.class, U8.class);
+    private final IBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap =
+            new TestBpfMap<>(CookieTagMapKey.class, CookieTagMapValue.class);
 
     @Before
     public void setUp() throws Exception {
@@ -127,6 +131,7 @@ public final class BpfNetMapsTest {
         BpfNetMaps.setConfigurationMapForTest(mConfigurationMap);
         BpfNetMaps.setUidOwnerMapForTest(mUidOwnerMap);
         BpfNetMaps.setUidPermissionMapForTest(mUidPermissionMap);
+        BpfNetMaps.setCookieTagMapForTest(mCookieTagMap);
         mBpfNetMaps = new BpfNetMaps(mContext, mNetd, mDeps);
     }