upload android base code part6

android/kernel/tests/net/test/bpf_test.py

@@ -0,0 +1,357 @@
+#!/usr/bin/python
+#
+# Copyright 2016 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.
+
+import ctypes
+import errno
+import os
+import socket
+import struct
+import unittest
+
+from bpf import *  # pylint: disable=wildcard-import
+import csocket
+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, 4, 0)
+HAVE_EBPF_ACCOUNTING = net_test.LINUX_VERSION >= (4, 9, 0)
+KEY_SIZE = 8
+VALUE_SIZE = 4
+TOTAL_ENTRIES = 20
+# Offset to store the map key in stack register REG10
+key_offset = -8
+# Offset to store the map value in stack register REG10
+value_offset = -16
+
+# Debug usage only.
+def PrintMapInfo(map_fd):
+  # A random key that the map does not contain.
+  key = 10086
+  while 1:
+    try:
+      nextKey = GetNextKey(map_fd, key).value
+      value = LookupMap(map_fd, nextKey)
+      print repr(nextKey) + " : " + repr(value.value)
+      key = nextKey
+    except:
+      print "no value"
+      break
+
+
+# A dummy loopback function that causes a socket to send traffic to itself.
+def SocketUDPLoopBack(packet_count, version, prog_fd):
+  family = {4: socket.AF_INET, 6: socket.AF_INET6}[version]
+  sock = socket.socket(family, socket.SOCK_DGRAM, 0)
+  if prog_fd is not None:
+    BpfProgAttachSocket(sock.fileno(), prog_fd)
+  net_test.SetNonBlocking(sock)
+  addr = {4: "127.0.0.1", 6: "::1"}[version]
+  sock.bind((addr, 0))
+  addr = sock.getsockname()
+  sockaddr = csocket.Sockaddr(addr)
+  for i in xrange(packet_count):
+    sock.sendto("foo", addr)
+    data, retaddr = csocket.Recvfrom(sock, 4096, 0)
+    assert "foo" == data
+    assert sockaddr == retaddr
+  return sock
+
+
+# The main code block for eBPF packet counting program. It takes a preloaded
+# key from BPF_REG_0 and use it to look up the bpf map, if the element does not
+# exist in the map yet, the program will update the map with a new <key, 1>
+# pair. Otherwise it will jump to next code block to handle it.
+# REG0: regiter storing return value from helper function and the final return
+# value of eBPF program.
+# REG1 - REG5: temporary register used for storing values and load parameters
+# into eBPF helper function. After calling helper function, the value for these
+# registers will be reset.
+# REG6 - REG9: registers store values that will not be cleared when calling
+# eBPF helper function.
+# REG10: A stack stores values need to be accessed by the address. Program can
+# retrieve the address of a value by specifying the position of the value in
+# the stack.
+def BpfFuncCountPacketInit(map_fd):
+  key_pos = BPF_REG_7
+  insPackCountStart = [
+      # Get a preloaded key from BPF_REG_0 and store it at BPF_REG_7
+      BpfMov64Reg(key_pos, BPF_REG_10),
+      BpfAlu64Imm(BPF_ADD, key_pos, key_offset),
+      # Load map fd and look up the key in the map
+      BpfLoadMapFd(map_fd, BPF_REG_1),
+      BpfMov64Reg(BPF_REG_2, key_pos),
+      BpfFuncCall(BPF_FUNC_map_lookup_elem),
+      # if the map element already exist, jump out of this
+      # code block and let next part to handle it
+      BpfJumpImm(BPF_AND, BPF_REG_0, 0, 10),
+      BpfLoadMapFd(map_fd, BPF_REG_1),
+      BpfMov64Reg(BPF_REG_2, key_pos),
+      # Initial a new <key, value> pair with value equal to 1 and update to map
+      BpfStMem(BPF_W, BPF_REG_10, value_offset, 1),
+      BpfMov64Reg(BPF_REG_3, BPF_REG_10),
+      BpfAlu64Imm(BPF_ADD, BPF_REG_3, value_offset),
+      BpfMov64Imm(BPF_REG_4, 0),
+      BpfFuncCall(BPF_FUNC_map_update_elem)
+  ]
+  return insPackCountStart
+
+
+INS_BPF_EXIT_BLOCK = [
+    BpfMov64Imm(BPF_REG_0, 0),
+    BpfExitInsn()
+]
+
+# Bpf instruction for cgroup bpf filter to accept a packet and exit.
+INS_CGROUP_ACCEPT = [
+    # Set return value to 1 and exit.
+    BpfMov64Imm(BPF_REG_0, 1),
+    BpfExitInsn()
+]
+
+# Bpf instruction for socket bpf filter to accept a packet and exit.
+INS_SK_FILTER_ACCEPT = [
+    # Precondition: BPF_REG_6 = sk_buff context
+    # Load the packet length from BPF_REG_6 and store it in BPF_REG_0 as the
+    # return value.
+    BpfLdxMem(BPF_W, BPF_REG_0, BPF_REG_6, 0),
+    BpfExitInsn()
+]
+
+# Update a existing map element with +1.
+INS_PACK_COUNT_UPDATE = [
+    # Precondition: BPF_REG_0 = Value retrieved from BPF maps
+    # Add one to the corresponding eBPF value field for a specific eBPF key.
+    BpfMov64Reg(BPF_REG_2, BPF_REG_0),
+    BpfMov64Imm(BPF_REG_1, 1),
+    BpfRawInsn(BPF_STX | BPF_XADD | BPF_W, BPF_REG_2, BPF_REG_1, 0, 0),
+]
+
+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")
+class BpfTest(net_test.NetworkTest):
+
+  def setUp(self):
+    self.map_fd = -1
+    self.prog_fd = -1
+    self.sock = None
+
+  def tearDown(self):
+    if self.prog_fd >= 0:
+      os.close(self.prog_fd)
+    if self.map_fd >= 0:
+      os.close(self.map_fd)
+    if self.sock:
+      self.sock.close()
+
+  def testCreateMap(self):
+    key, value = 1, 1
+    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
+                            TOTAL_ENTRIES)
+    UpdateMap(self.map_fd, key, value)
+    self.assertEquals(value, LookupMap(self.map_fd, key).value)
+    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
+    count = 0
+    while 1:
+      if count == TOTAL_ENTRIES - 1:
+        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.assertEquals(value, LookupMap(self.map_fd, key).value)
+        count += 1
+
+  def testProgLoad(self):
+    # Move skb to BPF_REG_6 for further usage
+    instructions = [
+        BpfMov64Reg(BPF_REG_6, BPF_REG_1)
+    ]
+    instructions += INS_SK_FILTER_ACCEPT
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions)
+    SocketUDPLoopBack(1, 4, self.prog_fd)
+    SocketUDPLoopBack(1, 6, self.prog_fd)
+
+  def testPacketBlock(self):
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, INS_BPF_EXIT_BLOCK)
+    self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 4, self.prog_fd)
+    self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 6, self.prog_fd)
+
+  def testPacketCount(self):
+    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
+                            TOTAL_ENTRIES)
+    key = 0xf0f0
+    # Set up instruction block with key loaded at BPF_REG_0.
+    instructions = [
+        BpfMov64Reg(BPF_REG_6, BPF_REG_1),
+        BpfMov64Imm(BPF_REG_0, key)
+    ]
+    # Concatenate the generic packet count bpf program to it.
+    instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd)
+                     + INS_SK_FILTER_ACCEPT + INS_PACK_COUNT_UPDATE
+                     + INS_SK_FILTER_ACCEPT)
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions)
+    packet_count = 10
+    SocketUDPLoopBack(packet_count, 4, self.prog_fd)
+    SocketUDPLoopBack(packet_count, 6, self.prog_fd)
+    self.assertEquals(packet_count * 2, LookupMap(self.map_fd, key).value)
+
+  @unittest.skipUnless(HAVE_EBPF_ACCOUNTING,
+                       "BPF helper function is not fully supported")
+  def testGetSocketCookie(self):
+    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
+                            TOTAL_ENTRIES)
+    # Move skb to REG6 for further usage, call helper function to get socket
+    # cookie of current skb and return the cookie at REG0 for next code block
+    instructions = [
+        BpfMov64Reg(BPF_REG_6, BPF_REG_1),
+        BpfFuncCall(BPF_FUNC_get_socket_cookie)
+    ]
+    instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd)
+                     + INS_SK_FILTER_ACCEPT + INS_PACK_COUNT_UPDATE
+                     + INS_SK_FILTER_ACCEPT)
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions)
+    packet_count = 10
+    def PacketCountByCookie(version):
+      self.sock = SocketUDPLoopBack(packet_count, version, self.prog_fd)
+      cookie = sock_diag.SockDiag.GetSocketCookie(self.sock)
+      self.assertEquals(packet_count, LookupMap(self.map_fd, cookie).value)
+      self.sock.close()
+    PacketCountByCookie(4)
+    PacketCountByCookie(6)
+
+  @unittest.skipUnless(HAVE_EBPF_ACCOUNTING,
+                       "BPF helper function is not fully supported")
+  def testGetSocketUid(self):
+    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
+                            TOTAL_ENTRIES)
+    # Set up the instruction with uid at BPF_REG_0.
+    instructions = [
+        BpfMov64Reg(BPF_REG_6, BPF_REG_1),
+        BpfFuncCall(BPF_FUNC_get_socket_uid)
+    ]
+    # Concatenate the generic packet count bpf program to it.
+    instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd)
+                     + INS_SK_FILTER_ACCEPT + INS_PACK_COUNT_UPDATE
+                     + INS_SK_FILTER_ACCEPT)
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions)
+    packet_count = 10
+    uid = 12345
+    with net_test.RunAsUid(uid):
+      self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, uid)
+      SocketUDPLoopBack(packet_count, 4, self.prog_fd)
+      self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value)
+      DeleteMap(self.map_fd, uid);
+      SocketUDPLoopBack(packet_count, 6, self.prog_fd)
+      self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value)
+
+@unittest.skipUnless(HAVE_EBPF_ACCOUNTING,
+                     "Cgroup BPF is not fully supported")
+class BpfCgroupTest(net_test.NetworkTest):
+
+  @classmethod
+  def setUpClass(cls):
+    if not os.path.isdir("/tmp"):
+      os.mkdir('/tmp')
+    os.system('mount -t cgroup2 cg_bpf /tmp')
+    cls._cg_fd = os.open('/tmp', os.O_DIRECTORY | os.O_RDONLY)
+
+  @classmethod
+  def tearDownClass(cls):
+    os.close(cls._cg_fd)
+    os.system('umount cg_bpf')
+
+  def setUp(self):
+    self.prog_fd = -1
+    self.map_fd = -1
+
+  def tearDown(self):
+    if self.prog_fd >= 0:
+      os.close(self.prog_fd)
+    if self.map_fd >= 0:
+      os.close(self.map_fd)
+    try:
+      BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_EGRESS)
+    except socket.error:
+      pass
+    try:
+      BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS)
+    except socket.error:
+      pass
+
+  def testCgroupBpfAttach(self):
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, INS_BPF_EXIT_BLOCK)
+    BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_INGRESS)
+    BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS)
+
+  def testCgroupIngress(self):
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, INS_BPF_EXIT_BLOCK)
+    BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_INGRESS)
+    self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 4, None)
+    self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 6, None)
+    BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS)
+    SocketUDPLoopBack(1, 4, None)
+    SocketUDPLoopBack(1, 6, None)
+
+  def testCgroupEgress(self):
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, INS_BPF_EXIT_BLOCK)
+    BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_EGRESS)
+    self.assertRaisesErrno(errno.EPERM, SocketUDPLoopBack, 1, 4, None)
+    self.assertRaisesErrno(errno.EPERM, SocketUDPLoopBack, 1, 6, None)
+    BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_EGRESS)
+    SocketUDPLoopBack( 1, 4, None)
+    SocketUDPLoopBack( 1, 6, None)
+
+  def testCgroupBpfUid(self):
+    self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE,
+                            TOTAL_ENTRIES)
+    # Similar to the program used in testGetSocketUid.
+    instructions = [
+        BpfMov64Reg(BPF_REG_6, BPF_REG_1),
+        BpfFuncCall(BPF_FUNC_get_socket_uid)
+    ]
+    instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd)
+                     + INS_CGROUP_ACCEPT + INS_PACK_COUNT_UPDATE + INS_CGROUP_ACCEPT)
+    self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, instructions)
+    BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_INGRESS)
+    uid = os.getuid()
+    packet_count = 20
+    SocketUDPLoopBack(packet_count, 4, None)
+    self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value)
+    DeleteMap(self.map_fd, uid);
+    SocketUDPLoopBack(packet_count, 6, None)
+    self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value)
+    BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS)
+
+if __name__ == "__main__":
+  unittest.main()