#!/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 subprocess import tempfile 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_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_ACCOUNTING, "BPF helper function is 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 CheckAllMapEntry(self, nonexistent_key, totalEntries, value): count = 0 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.assertGreaterEqual(key, 0) self.assertEquals(value, LookupMap(self.map_fd, key).value) count += 1 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) value = 1024 key = 1 self.assertRaisesErrno(errno.EPERM, UpdateMap, self.map_fd, key, value) self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, key) @unittest.skipUnless(HAVE_EBPF_ACCOUNTING, "BPF helper function is not fully supported") def testWrOnlyMap(self): self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, TOTAL_ENTRIES, map_flags=BPF_F_WRONLY) value = 1024 key = 1 UpdateMap(self.map_fd, key, value) self.assertRaisesErrno(errno.EPERM, LookupMap, self.map_fd, key) 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): cls._cg_dir = tempfile.mkdtemp(prefix="cg_bpf-") cmd = "mount -t cgroup2 cg_bpf %s" % cls._cg_dir try: subprocess.check_call(cmd.split()) except subprocess.CalledProcessError: # If an exception is thrown in setUpClass, the test fails and # tearDownClass is not called. os.rmdir(cls._cg_dir) raise cls._cg_fd = os.open(cls._cg_dir, os.O_DIRECTORY | os.O_RDONLY) @classmethod def tearDownClass(cls): os.close(cls._cg_fd) subprocess.call(('umount %s' % cls._cg_dir).split()) os.rmdir(cls._cg_dir) 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()