BpfBitmap: some minor optimizations

Test: TreeHugger, atest BpfBitmapTest
Flag: EXEMPT mainline
Signed-off-by: Maciej Żenczykowski <maze@google.com>
Change-Id: I5acf5bcb31777e503e58cef733d617a53006d2b9

staticlibs/device/com/android/net/module/util/BpfBitmap.java

@@ -22,6 +22,8 @@ import android.system.ErrnoException;
 import androidx.annotation.NonNull;
 import androidx.annotation.RequiresApi;
 
+import dalvik.annotation.optimization.CriticalNative;
+
  /**
  *
  * Generic bitmap class for use with BPF programs. Corresponds to a BpfMap
@@ -30,7 +32,12 @@ import androidx.annotation.RequiresApi;
  */
 @RequiresApi(Build.VERSION_CODES.S)
 public class BpfBitmap {
-    private BpfMap<Struct.S32, Struct.S64> mBpfMap;
+    static {
+        System.loadLibrary(JniUtil.getJniLibraryName(BpfBitmap.class.getPackage()));
+    }
+
+    private final BpfMap<Struct.S32, Struct.S64> mBpfMap;
+    private final int mMapFd;
 
     /**
      * Create a BpfBitmap map wrapper with "path" of filesystem.
@@ -39,21 +46,16 @@ public class BpfBitmap {
      */
     public BpfBitmap(@NonNull String path) throws ErrnoException {
         mBpfMap = new BpfMap<>(path, Struct.S32.class, Struct.S64.class);
+        mMapFd = mBpfMap.getFd();
     }
 
-    /**
-     * Retrieves the value from BpfMap for the given key.
-     *
-     * @param key The key in the map corresponding to the value to return.
-     */
-    private long getBpfMapValue(Struct.S32 key) throws ErrnoException  {
-        Struct.S64 curVal = mBpfMap.getValue(key);
-        if (curVal != null) {
-            return curVal.val;
-        } else {
-            return 0;
-        }
-    }
+    // Returns > 0 if bit is set, 0 if not set, < 0 on error (negative errno).
+    @CriticalNative
+    private static native int nativeGet(int fd, int index);
+
+    // Returns 0 on success, < 0 on error (negative errno).
+    @CriticalNative
+    private static native int nativeSet(int fd, int index, boolean set);
 
     /**
      * Retrieves the bit for the given index in the bitmap.
@@ -63,8 +65,26 @@ public class BpfBitmap {
     public boolean get(int index) throws ErrnoException  {
         if (index < 0) return false;
 
-        Struct.S32 key = new Struct.S32(index >> 6);
-        return ((getBpfMapValue(key) >>> (index & 63)) & 1L) != 0;
+        final int ret = nativeGet(mMapFd, index);
+        if (ret < 0) {
+            throw new ErrnoException("nativeGet", -ret);
+        }
+        return ret > 0;
+    }
+
+    /**
+     * Change the specified index in the bitmap to set value.
+     *
+     * @param index Position to (un)set in bitmap.
+     * @param set Boolean indicating to set or unset index.
+     */
+    public void set(int index, boolean set) throws ErrnoException {
+        if (index < 0) throw new IllegalArgumentException("Index out of bounds.");
+
+        final int ret = nativeSet(mMapFd, index, set);
+        if (ret < 0) {
+            throw new ErrnoException("nativeSet", -ret);
+        }
     }
 
     /**
@@ -85,22 +105,6 @@ public class BpfBitmap {
         set(index, false);
     }
 
-    /**
-     * Change the specified index in the bitmap to set value.
-     *
-     * @param index Position to unset in bitmap.
-     * @param set Boolean indicating to set or unset index.
-     */
-    public void set(int index, boolean set) throws ErrnoException {
-        if (index < 0) throw new IllegalArgumentException("Index out of bounds.");
-
-        Struct.S32 key = new Struct.S32(index >> 6);
-        long mask = (1L << (index & 63));
-        long val = getBpfMapValue(key);
-        if (set) val |= mask; else val &= ~mask;
-        mBpfMap.updateEntry(key, new Struct.S64(val));
-    }
-
     /**
      * Clears the map. The map may already be empty.
      *
@@ -118,7 +122,8 @@ public class BpfBitmap {
     public boolean isEmpty() throws ErrnoException {
         Struct.S32 key = mBpfMap.getFirstKey();
         while (key != null) {
-            if (getBpfMapValue(key) != 0) {
+            Struct.S64 val = mBpfMap.getValue(key);
+            if (val != null && val.val != 0) {
                 return false;
             }
             key = mBpfMap.getNextKey(key);

staticlibs/device/com/android/net/module/util/BpfMap.java

@@ -71,6 +71,11 @@ public class BpfMap<K extends Struct, V extends Struct> implements IBpfMap<K, V>
     private final int mKeySize;
     private final int mValueSize;
 
+    // The following is (ab)used by BpfBitmap.java
+    /* package */ int getFd() {
+        return mMapFd.getFd();
+    }
+
     private static ConcurrentHashMap<Pair<String, Integer>, ParcelFileDescriptor> sFdCache =
             new ConcurrentHashMap<>();
 

staticlibs/native/bpfmapjni/Android.bp

@@ -20,7 +20,7 @@ package {
 cc_library_static {
     name: "libnet_utils_device_common_bpfjni",
     srcs: [
-        "com_android_net_module_util_BpfMap.cpp",
+        "com_android_net_module_util_Bpf.cpp",
         "com_android_net_module_util_TcUtils.cpp",
     ],
     header_libs: [

staticlibs/native/bpfmapjni/com_android_net_module_util_Bpf.cpp

@@ -0,0 +1,221 @@
+/*
+ * Copyright (C) 2020 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.
+ */
+
+#include <errno.h>
+#include <linux/pfkeyv2.h>
+#include <stdint.h>
+#include <sys/socket.h>
+#include <jni.h>
+#include <nativehelper/JNIHelp.h>
+#include <nativehelper/ScopedLocalRef.h>
+
+#include "nativehelper/scoped_primitive_array.h"
+#include "nativehelper/scoped_utf_chars.h"
+
+#include <android-base/unique_fd.h>
+#include "BpfSyscallWrappers.h"
+#include "bpf/KernelUtils.h"
+
+namespace android {
+
+using ::android::base::unique_fd;
+
+static jint com_android_net_module_util_BpfMap_nativeBpfFdGet(JNIEnv *env, jclass clazz,
+        jstring path, jint mode, jint keySize, jint valueSize) {
+    ScopedUtfChars pathname(env, path);
+
+    unique_fd fd;
+    switch (mode) {
+      case 0:
+        fd.reset(bpf::mapRetrieveRW(pathname.c_str()));
+        break;
+      case BPF_F_RDONLY:
+        fd.reset(bpf::mapRetrieveRO(pathname.c_str()));
+        break;
+      case BPF_F_WRONLY:
+        fd.reset(bpf::mapRetrieveWO(pathname.c_str()));
+        break;
+      case BPF_F_RDONLY|BPF_F_WRONLY:
+        fd.reset(bpf::mapRetrieveExclusiveRW(pathname.c_str()));
+        break;
+      default:
+        errno = EINVAL;
+        break;
+    }
+
+    if (!fd.ok()) {
+        jniThrowErrnoException(env, "nativeBpfFdGet", errno);
+        return -1;
+    }
+
+    if (bpf::isAtLeastKernelVersion(4, 14, 0)) {
+        // These likely fail with -1 and set errno to EINVAL on <4.14
+        if (bpf::bpfGetFdKeySize(fd) != keySize) {
+            jniThrowErrnoException(env, "nativeBpfFdGet KeySize", EBADFD);
+            return -1;
+        }
+        if (bpf::bpfGetFdValueSize(fd) != valueSize) {
+            jniThrowErrnoException(env, "nativeBpfFdGet ValueSize", EBADFD);
+            return -1;
+        }
+    }
+
+    return fd.release();
+}
+
+static void com_android_net_module_util_BpfMap_nativeWriteToMapEntry(JNIEnv *env, jobject self,
+        jint fd, jbyteArray key, jbyteArray value, jint flags) {
+    ScopedByteArrayRO keyRO(env, key);
+    ScopedByteArrayRO valueRO(env, value);
+
+    int ret = bpf::writeToMapEntry(static_cast<int>(fd), keyRO.get(), valueRO.get(),
+            static_cast<int>(flags));
+
+    if (ret) jniThrowErrnoException(env, "nativeWriteToMapEntry", errno);
+}
+
+static jboolean throwIfNotEnoent(JNIEnv *env, const char* functionName, int ret, int err) {
+    if (ret == 0) return true;
+
+    if (err != ENOENT) jniThrowErrnoException(env, functionName, err);
+    return false;
+}
+
+static jboolean com_android_net_module_util_BpfMap_nativeDeleteMapEntry(JNIEnv *env, jobject self,
+        jint fd, jbyteArray key) {
+    ScopedByteArrayRO keyRO(env, key);
+
+    // On success, zero is returned.  If the element is not found, -1 is returned and errno is set
+    // to ENOENT.
+    int ret = bpf::deleteMapEntry(static_cast<int>(fd), keyRO.get());
+
+    return throwIfNotEnoent(env, "nativeDeleteMapEntry", ret, errno);
+}
+
+static jboolean com_android_net_module_util_BpfMap_nativeGetNextMapKey(JNIEnv *env, jobject self,
+        jint fd, jbyteArray key, jbyteArray nextKey) {
+    // If key is found, the operation returns zero and sets the next key pointer to the key of the
+    // next element.  If key is not found, the operation returns zero and sets the next key pointer
+    // to the key of the first element.  If key is the last element, -1 is returned and errno is
+    // set to ENOENT.  Other possible errno values are ENOMEM, EFAULT, EPERM, and EINVAL.
+    ScopedByteArrayRW nextKeyRW(env, nextKey);
+    int ret;
+    if (key == nullptr) {
+        // Called by getFirstKey. Find the first key in the map.
+        ret = bpf::getNextMapKey(static_cast<int>(fd), nullptr, nextKeyRW.get());
+    } else {
+        ScopedByteArrayRO keyRO(env, key);
+        ret = bpf::getNextMapKey(static_cast<int>(fd), keyRO.get(), nextKeyRW.get());
+    }
+
+    return throwIfNotEnoent(env, "nativeGetNextMapKey", ret, errno);
+}
+
+static jboolean com_android_net_module_util_BpfMap_nativeFindMapEntry(JNIEnv *env, jobject self,
+        jint fd, jbyteArray key, jbyteArray value) {
+    ScopedByteArrayRO keyRO(env, key);
+    ScopedByteArrayRW valueRW(env, value);
+
+    // If an element is found, the operation returns zero and stores the element's value into
+    // "value".  If no element is found, the operation returns -1 and sets errno to ENOENT.
+    int ret = bpf::findMapEntry(static_cast<int>(fd), keyRO.get(), valueRW.get());
+
+    return throwIfNotEnoent(env, "nativeFindMapEntry", ret, errno);
+}
+
+static jint com_android_net_module_util_BpfMap_nativeSynchronizeKernelRCU() {
+    const int pfSocket = socket(AF_KEY, SOCK_RAW | SOCK_CLOEXEC, PF_KEY_V2);
+    if (pfSocket < 0) return -errno;
+    // On Linux close() will always close the fd, any error it returns is a previous pending error.
+    if (close(pfSocket)) return -errno;  // in practice cannot fail
+    return 0;
+}
+
+static jint com_android_net_module_util_BpfBitmap_nativeGet(jint fd, jint index) {
+    if (index < 0) return -EINVAL;
+
+    const uint32_t key = index >> 6;
+    const uint32_t subkey = index & 63;
+    uint64_t value = 0;
+
+    // findMapEntry returns 0 on success, -1 on error.
+    // If the entry does not exist, it's not an error, the value is just 0.
+    if (bpf::findMapEntry(fd, &key, &value) && errno != ENOENT) return -errno;
+
+    return (value >> subkey) & 1;
+}
+
+static jint com_android_net_module_util_BpfBitmap_nativeSet(jint fd, jint index, jboolean set) {
+    if (index < 0) return -EINVAL;
+
+    const uint32_t key = index >> 6;
+    const uint32_t subkey = index & 63;
+    uint64_t value = 0;
+
+    // Read the existing value. It's okay if it doesn't exist, value will be 0.
+    if (bpf::findMapEntry(fd, &key, &value) && errno != ENOENT) return -errno;
+
+    const uint64_t mask = 1uLL << subkey;
+    if (set) {
+        value |= mask;
+    } else {
+        value &= ~mask;
+    }
+
+    // Write the updated value back. BPF_ANY will create or update as needed.
+    if (bpf::writeToMapEntry(fd, &key, &value, BPF_ANY)) return -errno;
+
+    return 0;
+}
+
+/*
+ * JNI registration.
+ */
+static const JNINativeMethod gBpfMapMethods[] = {
+    /* name, signature, funcPtr */
+    { "nativeBpfFdGet", "(Ljava/lang/String;III)I",
+        (void*) com_android_net_module_util_BpfMap_nativeBpfFdGet },
+    { "nativeWriteToMapEntry", "(I[B[BI)V",
+        (void*) com_android_net_module_util_BpfMap_nativeWriteToMapEntry },
+    { "nativeDeleteMapEntry", "(I[B)Z",
+        (void*) com_android_net_module_util_BpfMap_nativeDeleteMapEntry },
+    { "nativeGetNextMapKey", "(I[B[B)Z",
+        (void*) com_android_net_module_util_BpfMap_nativeGetNextMapKey },
+    { "nativeFindMapEntry", "(I[B[B)Z",
+        (void*) com_android_net_module_util_BpfMap_nativeFindMapEntry },
+    // CriticalNative
+    { "nativeSynchronizeKernelRCU", "()I",
+        (void*) com_android_net_module_util_BpfMap_nativeSynchronizeKernelRCU },
+
+};
+
+int register_com_android_net_module_util_BpfMap(JNIEnv* env, char const* class_name) {
+    return jniRegisterNativeMethods(env, class_name, gBpfMapMethods, NELEM(gBpfMapMethods));
+}
+
+static const JNINativeMethod gBpfBitmapMethods[] = {
+    /* name, signature, funcPtr */
+    // CriticalNative
+    { "nativeGet", "(II)I", (void*)com_android_net_module_util_BpfBitmap_nativeGet },
+    // CriticalNative
+    { "nativeSet", "(IIZ)I", (void*)com_android_net_module_util_BpfBitmap_nativeSet },
+};
+
+int register_com_android_net_module_util_BpfBitmap(JNIEnv* env, char const* class_name) {
+    return jniRegisterNativeMethods(env, class_name, gBpfBitmapMethods, NELEM(gBpfBitmapMethods));
+}
+
+}; // namespace android