Test for getFirstMapKey of bpf maps

The upstream commit (8fe45924387be6b5c1be59a7eb330790c61d5d10 bpf:
map_get_next_key to return first key on NULL) introduce the support for
bpf maps to return the first key of the map when receive a NULL key from
userspace. This change will test this functionality in supported kernels
to make sure they have the correct backport.

Test: ./bpf_test.py
Bug: 30950746
Change-Id: I6a497a3ccb518ebb037bfa447f827d2715271291
Merged-In: I6a497a3ccb518ebb037bfa447f827d2715271291
(cherry picked from aosp commit 493cfb50cb18677628f0ae62dd8a14036a8eec0c)

2 files changed, 39 insertions, 21 deletions

diff --git a/net/test/bpf.py b/net/test/bpf.py
index ff23d79..c9ad264 100755
--- a/net/test/bpf.py
+++ b/net/test/bpf.py
@@ -195,13 +195,20 @@ def LookupMap(map_fd, key):
 
 
 def GetNextKey(map_fd, key):
-  c_key = ctypes.c_uint32(key)
+  if key is not None:
+    c_key = ctypes.c_uint32(key)
+    c_next_key = ctypes.c_uint32(0)
+    key_ptr = ctypes.addressof(c_key)
+  else:
+    key_ptr = 0;
   c_next_key = ctypes.c_uint32(0)
   attr = BpfAttrOps(
-      (map_fd, ctypes.addressof(c_key), ctypes.addressof(c_next_key), 0))
+      (map_fd, key_ptr, ctypes.addressof(c_next_key), 0))
   BpfSyscall(BPF_MAP_GET_NEXT_KEY, attr)
   return c_next_key
 
+def GetFirstKey(map_fd):
+  return GetNextKey(map_fd, None)
 
 def DeleteMap(map_fd, key):
   c_key = ctypes.c_uint32(key)
diff --git a/net/test/bpf_test.py b/net/test/bpf_test.py
index 9da3907..e7a4edb 100755
--- a/net/test/bpf_test.py
+++ b/net/test/bpf_test.py
@@ -29,7 +29,6 @@ import net_test
 import sock_diag
 
 libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
-HAVE_EBPF_SUPPORT = net_test.LINUX_VERSION >= (4, 9, 0)
 HAVE_EBPF_ACCOUNTING = net_test.LINUX_VERSION >= (4, 9, 0)
 KEY_SIZE = 8
 VALUE_SIZE = 4
@@ -146,8 +145,8 @@ INS_BPF_PARAM_STORE = [
     BpfStxMem(BPF_DW, BPF_REG_10, BPF_REG_0, key_offset),
 ]
 
-@unittest.skipUnless(HAVE_EBPF_SUPPORT,
-                     "eBPF function not fully supported")
+@unittest.skipUnless(HAVE_EBPF_ACCOUNTING,
+                     "BPF helper function is not fully supported")
 class BpfTest(net_test.NetworkTest):
 
   def setUp(self):
@@ -172,31 +171,43 @@ class BpfTest(net_test.NetworkTest):
     DeleteMap(self.map_fd, key)
     self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, key)
 
-  def testIterateMap(self):
-    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
-                            TOTAL_ENTRIES)
-    value = 1024
-    for key in xrange(1, TOTAL_ENTRIES):
-      UpdateMap(self.map_fd, key, value)
-    for key in xrange(1, TOTAL_ENTRIES):
-      self.assertEquals(value, LookupMap(self.map_fd, key).value)
-    self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, 101)
-    key = 0
+  def CheckAllMapEntry(self, nonexistent_key, totalEntries, value):
     count = 0
-    while 1:
-      if count == TOTAL_ENTRIES - 1:
+    key = nonexistent_key
+    while True:
+      if count == totalEntries:
         self.assertRaisesErrno(errno.ENOENT, GetNextKey, self.map_fd, key)
         break
       else:
         result = GetNextKey(self.map_fd, key)
         key = result.value
-        self.assertGreater(key, 0)
+        self.assertGreaterEqual(key, 0)
         self.assertEquals(value, LookupMap(self.map_fd, key).value)
         count += 1
 
-  # TODO: move this check to the begining of the class insdead.
-  @unittest.skipUnless(HAVE_EBPF_ACCOUNTING,
-                       "BPF helper function is not fully supported")
+  def testIterateMap(self):
+    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
+                            TOTAL_ENTRIES)
+    value = 1024
+    for key in xrange(0, TOTAL_ENTRIES):
+      UpdateMap(self.map_fd, key, value)
+    for key in xrange(0, TOTAL_ENTRIES):
+      self.assertEquals(value, LookupMap(self.map_fd, key).value)
+    self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, 101)
+    nonexistent_key = -1
+    self.CheckAllMapEntry(nonexistent_key, TOTAL_ENTRIES, value)
+
+  def testFindFirstMapKey(self):
+    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
+                            TOTAL_ENTRIES)
+    value = 1024
+    for key in xrange(0, TOTAL_ENTRIES):
+      UpdateMap(self.map_fd, key, value)
+    firstKey = GetFirstKey(self.map_fd)
+    key = firstKey.value
+    self.CheckAllMapEntry(key, TOTAL_ENTRIES - 1, value)
+
+
   def testRdOnlyMap(self):
     self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
                             TOTAL_ENTRIES, map_flags=BPF_F_RDONLY)