1 files changed, 9 insertions, 21 deletions
diff --git a/common/device/com/android/net/module/util/BpfMap.java b/common/device/com/android/net/module/util/BpfMap.java
index 9df2b035..d45caceb 100644
--- a/common/device/com/android/net/module/util/BpfMap.java
+++ b/common/device/com/android/net/module/util/BpfMap.java
@@ -194,9 +194,11 @@ public class BpfMap<K extends Struct, V extends Struct> implements IBpfMap<K, V>
     }
 
     private K getNextKeyInternal(@Nullable K key) throws ErrnoException {
-        final byte[] rawKey = getNextRawKey(
-                key == null ? null : key.writeToBytes());
-        if (rawKey == null) return null;
+        byte[] rawKey = new byte[mKeySize];
+
+        if (!nativeGetNextMapKey(mMapFd.getFd(),
+                                 key == null ? null : key.writeToBytes(),
+                                 rawKey)) return null;
 
         final ByteBuffer buffer = ByteBuffer.wrap(rawKey);
         buffer.order(ByteOrder.nativeOrder());
@@ -215,13 +217,6 @@ public class BpfMap<K extends Struct, V extends Struct> implements IBpfMap<K, V>
         return getNextKeyInternal(key);
     }
 
-    private byte[] getNextRawKey(@Nullable final byte[] key) throws ErrnoException {
-        byte[] nextKey = new byte[mKeySize];
-        if (nativeGetNextMapKey(mMapFd.getFd(), key, nextKey)) return nextKey;
-
-        return null;
-    }
-
     /** Get the first key of eBpf map. */
     @Override
     public K getFirstKey() throws ErrnoException {
@@ -233,30 +228,23 @@ public class BpfMap<K extends Struct, V extends Struct> implements IBpfMap<K, V>
     public boolean containsKey(@NonNull K key) throws ErrnoException {
         Objects.requireNonNull(key);
 
-        final byte[] rawValue = getRawValue(key.writeToBytes());
-        return rawValue != null;
+        byte[] rawValue = new byte[mValueSize];
+        return nativeFindMapEntry(mMapFd.getFd(), key.writeToBytes(), rawValue);
     }
 
     /** Retrieve a value from the map. Return null if there is no such key. */
     @Override
     public V getValue(@NonNull K key) throws ErrnoException {
         Objects.requireNonNull(key);
-        final byte[] rawValue = getRawValue(key.writeToBytes());
 
-        if (rawValue == null) return null;
+        byte[] rawValue = new byte[mValueSize];
+        if (!nativeFindMapEntry(mMapFd.getFd(), key.writeToBytes(), rawValue)) return null;
 
         final ByteBuffer buffer = ByteBuffer.wrap(rawValue);
         buffer.order(ByteOrder.nativeOrder());
         return Struct.parse(mValueClass, buffer);
     }
 
-    private byte[] getRawValue(final byte[] key) throws ErrnoException {
-        byte[] value = new byte[mValueSize];
-        if (nativeFindMapEntry(mMapFd.getFd(), key, value)) return value;
-
-        return null;
-    }
-
     /**
      * Iterate through the map and handle each key -> value retrieved base on the given BiConsumer.
      * The given BiConsumer may to delete the passed-in entry, but is not allowed to perform any