1 #!/usr/bin/python 2 # 3 # Copyright 2016 The Android Open Source Project 4 # 5 # Licensed under the Apache License, Version 2.0 (the "License"); 6 # you may not use this file except in compliance with the License. 7 # You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 17 import ctypes 18 import errno 19 import os 20 import socket 21 import struct 22 import subprocess 23 import tempfile 24 import unittest 25 26 from bpf import * # pylint: disable=wildcard-import 27 import csocket 28 import net_test 29 import sock_diag 30 31 libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True) 32 HAVE_EBPF_ACCOUNTING = net_test.LINUX_VERSION >= (4, 9, 0) 33 KEY_SIZE = 8 34 VALUE_SIZE = 4 35 TOTAL_ENTRIES = 20 36 # Offset to store the map key in stack register REG10 37 key_offset = -8 38 # Offset to store the map value in stack register REG10 39 value_offset = -16 40 41 # Debug usage only. 42 def PrintMapInfo(map_fd): 43 # A random key that the map does not contain. 44 key = 10086 45 while 1: 46 try: 47 nextKey = GetNextKey(map_fd, key).value 48 value = LookupMap(map_fd, nextKey) 49 print repr(nextKey) + " : " + repr(value.value) 50 key = nextKey 51 except: 52 print "no value" 53 break 54 55 56 # A dummy loopback function that causes a socket to send traffic to itself. 57 def SocketUDPLoopBack(packet_count, version, prog_fd): 58 family = {4: socket.AF_INET, 6: socket.AF_INET6}[version] 59 sock = socket.socket(family, socket.SOCK_DGRAM, 0) 60 if prog_fd is not None: 61 BpfProgAttachSocket(sock.fileno(), prog_fd) 62 net_test.SetNonBlocking(sock) 63 addr = {4: "127.0.0.1", 6: "::1"}[version] 64 sock.bind((addr, 0)) 65 addr = sock.getsockname() 66 sockaddr = csocket.Sockaddr(addr) 67 for i in xrange(packet_count): 68 sock.sendto("foo", addr) 69 data, retaddr = csocket.Recvfrom(sock, 4096, 0) 70 assert "foo" == data 71 assert sockaddr == retaddr 72 return sock 73 74 75 # The main code block for eBPF packet counting program. It takes a preloaded 76 # key from BPF_REG_0 and use it to look up the bpf map, if the element does not 77 # exist in the map yet, the program will update the map with a new <key, 1> 78 # pair. Otherwise it will jump to next code block to handle it. 79 # REG0: regiter storing return value from helper function and the final return 80 # value of eBPF program. 81 # REG1 - REG5: temporary register used for storing values and load parameters 82 # into eBPF helper function. After calling helper function, the value for these 83 # registers will be reset. 84 # REG6 - REG9: registers store values that will not be cleared when calling 85 # eBPF helper function. 86 # REG10: A stack stores values need to be accessed by the address. Program can 87 # retrieve the address of a value by specifying the position of the value in 88 # the stack. 89 def BpfFuncCountPacketInit(map_fd): 90 key_pos = BPF_REG_7 91 insPackCountStart = [ 92 # Get a preloaded key from BPF_REG_0 and store it at BPF_REG_7 93 BpfMov64Reg(key_pos, BPF_REG_10), 94 BpfAlu64Imm(BPF_ADD, key_pos, key_offset), 95 # Load map fd and look up the key in the map 96 BpfLoadMapFd(map_fd, BPF_REG_1), 97 BpfMov64Reg(BPF_REG_2, key_pos), 98 BpfFuncCall(BPF_FUNC_map_lookup_elem), 99 # if the map element already exist, jump out of this 100 # code block and let next part to handle it 101 BpfJumpImm(BPF_AND, BPF_REG_0, 0, 10), 102 BpfLoadMapFd(map_fd, BPF_REG_1), 103 BpfMov64Reg(BPF_REG_2, key_pos), 104 # Initial a new <key, value> pair with value equal to 1 and update to map 105 BpfStMem(BPF_W, BPF_REG_10, value_offset, 1), 106 BpfMov64Reg(BPF_REG_3, BPF_REG_10), 107 BpfAlu64Imm(BPF_ADD, BPF_REG_3, value_offset), 108 BpfMov64Imm(BPF_REG_4, 0), 109 BpfFuncCall(BPF_FUNC_map_update_elem) 110 ] 111 return insPackCountStart 112 113 114 INS_BPF_EXIT_BLOCK = [ 115 BpfMov64Imm(BPF_REG_0, 0), 116 BpfExitInsn() 117 ] 118 119 # Bpf instruction for cgroup bpf filter to accept a packet and exit. 120 INS_CGROUP_ACCEPT = [ 121 # Set return value to 1 and exit. 122 BpfMov64Imm(BPF_REG_0, 1), 123 BpfExitInsn() 124 ] 125 126 # Bpf instruction for socket bpf filter to accept a packet and exit. 127 INS_SK_FILTER_ACCEPT = [ 128 # Precondition: BPF_REG_6 = sk_buff context 129 # Load the packet length from BPF_REG_6 and store it in BPF_REG_0 as the 130 # return value. 131 BpfLdxMem(BPF_W, BPF_REG_0, BPF_REG_6, 0), 132 BpfExitInsn() 133 ] 134 135 # Update a existing map element with +1. 136 INS_PACK_COUNT_UPDATE = [ 137 # Precondition: BPF_REG_0 = Value retrieved from BPF maps 138 # Add one to the corresponding eBPF value field for a specific eBPF key. 139 BpfMov64Reg(BPF_REG_2, BPF_REG_0), 140 BpfMov64Imm(BPF_REG_1, 1), 141 BpfRawInsn(BPF_STX | BPF_XADD | BPF_W, BPF_REG_2, BPF_REG_1, 0, 0), 142 ] 143 144 INS_BPF_PARAM_STORE = [ 145 BpfStxMem(BPF_DW, BPF_REG_10, BPF_REG_0, key_offset), 146 ] 147 148 @unittest.skipUnless(HAVE_EBPF_ACCOUNTING, 149 "BPF helper function is not fully supported") 150 class BpfTest(net_test.NetworkTest): 151 152 def setUp(self): 153 self.map_fd = -1 154 self.prog_fd = -1 155 self.sock = None 156 157 def tearDown(self): 158 if self.prog_fd >= 0: 159 os.close(self.prog_fd) 160 if self.map_fd >= 0: 161 os.close(self.map_fd) 162 if self.sock: 163 self.sock.close() 164 165 def testCreateMap(self): 166 key, value = 1, 1 167 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 168 TOTAL_ENTRIES) 169 UpdateMap(self.map_fd, key, value) 170 self.assertEquals(value, LookupMap(self.map_fd, key).value) 171 DeleteMap(self.map_fd, key) 172 self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, key) 173 174 def CheckAllMapEntry(self, nonexistent_key, totalEntries, value): 175 count = 0 176 key = nonexistent_key 177 while True: 178 if count == totalEntries: 179 self.assertRaisesErrno(errno.ENOENT, GetNextKey, self.map_fd, key) 180 break 181 else: 182 result = GetNextKey(self.map_fd, key) 183 key = result.value 184 self.assertGreaterEqual(key, 0) 185 self.assertEquals(value, LookupMap(self.map_fd, key).value) 186 count += 1 187 188 def testIterateMap(self): 189 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 190 TOTAL_ENTRIES) 191 value = 1024 192 for key in xrange(0, TOTAL_ENTRIES): 193 UpdateMap(self.map_fd, key, value) 194 for key in xrange(0, TOTAL_ENTRIES): 195 self.assertEquals(value, LookupMap(self.map_fd, key).value) 196 self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, 101) 197 nonexistent_key = -1 198 self.CheckAllMapEntry(nonexistent_key, TOTAL_ENTRIES, value) 199 200 def testFindFirstMapKey(self): 201 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 202 TOTAL_ENTRIES) 203 value = 1024 204 for key in xrange(0, TOTAL_ENTRIES): 205 UpdateMap(self.map_fd, key, value) 206 firstKey = GetFirstKey(self.map_fd) 207 key = firstKey.value 208 self.CheckAllMapEntry(key, TOTAL_ENTRIES - 1, value) 209 210 211 def testRdOnlyMap(self): 212 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 213 TOTAL_ENTRIES, map_flags=BPF_F_RDONLY) 214 value = 1024 215 key = 1 216 self.assertRaisesErrno(errno.EPERM, UpdateMap, self.map_fd, key, value) 217 self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, key) 218 219 @unittest.skipUnless(HAVE_EBPF_ACCOUNTING, 220 "BPF helper function is not fully supported") 221 def testWrOnlyMap(self): 222 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 223 TOTAL_ENTRIES, map_flags=BPF_F_WRONLY) 224 value = 1024 225 key = 1 226 UpdateMap(self.map_fd, key, value) 227 self.assertRaisesErrno(errno.EPERM, LookupMap, self.map_fd, key) 228 229 def testProgLoad(self): 230 # Move skb to BPF_REG_6 for further usage 231 instructions = [ 232 BpfMov64Reg(BPF_REG_6, BPF_REG_1) 233 ] 234 instructions += INS_SK_FILTER_ACCEPT 235 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions) 236 SocketUDPLoopBack(1, 4, self.prog_fd) 237 SocketUDPLoopBack(1, 6, self.prog_fd) 238 239 def testPacketBlock(self): 240 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, INS_BPF_EXIT_BLOCK) 241 self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 4, self.prog_fd) 242 self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 6, self.prog_fd) 243 244 def testPacketCount(self): 245 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 246 TOTAL_ENTRIES) 247 key = 0xf0f0 248 # Set up instruction block with key loaded at BPF_REG_0. 249 instructions = [ 250 BpfMov64Reg(BPF_REG_6, BPF_REG_1), 251 BpfMov64Imm(BPF_REG_0, key) 252 ] 253 # Concatenate the generic packet count bpf program to it. 254 instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd) 255 + INS_SK_FILTER_ACCEPT + INS_PACK_COUNT_UPDATE 256 + INS_SK_FILTER_ACCEPT) 257 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions) 258 packet_count = 10 259 SocketUDPLoopBack(packet_count, 4, self.prog_fd) 260 SocketUDPLoopBack(packet_count, 6, self.prog_fd) 261 self.assertEquals(packet_count * 2, LookupMap(self.map_fd, key).value) 262 263 @unittest.skipUnless(HAVE_EBPF_ACCOUNTING, 264 "BPF helper function is not fully supported") 265 def testGetSocketCookie(self): 266 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 267 TOTAL_ENTRIES) 268 # Move skb to REG6 for further usage, call helper function to get socket 269 # cookie of current skb and return the cookie at REG0 for next code block 270 instructions = [ 271 BpfMov64Reg(BPF_REG_6, BPF_REG_1), 272 BpfFuncCall(BPF_FUNC_get_socket_cookie) 273 ] 274 instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd) 275 + INS_SK_FILTER_ACCEPT + INS_PACK_COUNT_UPDATE 276 + INS_SK_FILTER_ACCEPT) 277 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions) 278 packet_count = 10 279 def PacketCountByCookie(version): 280 self.sock = SocketUDPLoopBack(packet_count, version, self.prog_fd) 281 cookie = sock_diag.SockDiag.GetSocketCookie(self.sock) 282 self.assertEquals(packet_count, LookupMap(self.map_fd, cookie).value) 283 self.sock.close() 284 PacketCountByCookie(4) 285 PacketCountByCookie(6) 286 287 @unittest.skipUnless(HAVE_EBPF_ACCOUNTING, 288 "BPF helper function is not fully supported") 289 def testGetSocketUid(self): 290 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 291 TOTAL_ENTRIES) 292 # Set up the instruction with uid at BPF_REG_0. 293 instructions = [ 294 BpfMov64Reg(BPF_REG_6, BPF_REG_1), 295 BpfFuncCall(BPF_FUNC_get_socket_uid) 296 ] 297 # Concatenate the generic packet count bpf program to it. 298 instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd) 299 + INS_SK_FILTER_ACCEPT + INS_PACK_COUNT_UPDATE 300 + INS_SK_FILTER_ACCEPT) 301 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_SOCKET_FILTER, instructions) 302 packet_count = 10 303 uid = 12345 304 with net_test.RunAsUid(uid): 305 self.assertRaisesErrno(errno.ENOENT, LookupMap, self.map_fd, uid) 306 SocketUDPLoopBack(packet_count, 4, self.prog_fd) 307 self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value) 308 DeleteMap(self.map_fd, uid); 309 SocketUDPLoopBack(packet_count, 6, self.prog_fd) 310 self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value) 311 312 @unittest.skipUnless(HAVE_EBPF_ACCOUNTING, 313 "Cgroup BPF is not fully supported") 314 class BpfCgroupTest(net_test.NetworkTest): 315 316 @classmethod 317 def setUpClass(cls): 318 cls._cg_dir = tempfile.mkdtemp(prefix="cg_bpf-") 319 cmd = "mount -t cgroup2 cg_bpf %s" % cls._cg_dir 320 try: 321 subprocess.check_call(cmd.split()) 322 except subprocess.CalledProcessError: 323 # If an exception is thrown in setUpClass, the test fails and 324 # tearDownClass is not called. 325 os.rmdir(cls._cg_dir) 326 raise 327 cls._cg_fd = os.open(cls._cg_dir, os.O_DIRECTORY | os.O_RDONLY) 328 329 @classmethod 330 def tearDownClass(cls): 331 os.close(cls._cg_fd) 332 subprocess.call(('umount %s' % cls._cg_dir).split()) 333 os.rmdir(cls._cg_dir) 334 335 def setUp(self): 336 self.prog_fd = -1 337 self.map_fd = -1 338 339 def tearDown(self): 340 if self.prog_fd >= 0: 341 os.close(self.prog_fd) 342 if self.map_fd >= 0: 343 os.close(self.map_fd) 344 try: 345 BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_EGRESS) 346 except socket.error: 347 pass 348 try: 349 BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS) 350 except socket.error: 351 pass 352 353 def testCgroupBpfAttach(self): 354 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, INS_BPF_EXIT_BLOCK) 355 BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_INGRESS) 356 BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS) 357 358 def testCgroupIngress(self): 359 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, INS_BPF_EXIT_BLOCK) 360 BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_INGRESS) 361 self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 4, None) 362 self.assertRaisesErrno(errno.EAGAIN, SocketUDPLoopBack, 1, 6, None) 363 BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS) 364 SocketUDPLoopBack(1, 4, None) 365 SocketUDPLoopBack(1, 6, None) 366 367 def testCgroupEgress(self): 368 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, INS_BPF_EXIT_BLOCK) 369 BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_EGRESS) 370 self.assertRaisesErrno(errno.EPERM, SocketUDPLoopBack, 1, 4, None) 371 self.assertRaisesErrno(errno.EPERM, SocketUDPLoopBack, 1, 6, None) 372 BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_EGRESS) 373 SocketUDPLoopBack( 1, 4, None) 374 SocketUDPLoopBack( 1, 6, None) 375 376 def testCgroupBpfUid(self): 377 self.map_fd = CreateMap(BPF_MAP_TYPE_HASH, KEY_SIZE, VALUE_SIZE, 378 TOTAL_ENTRIES) 379 # Similar to the program used in testGetSocketUid. 380 instructions = [ 381 BpfMov64Reg(BPF_REG_6, BPF_REG_1), 382 BpfFuncCall(BPF_FUNC_get_socket_uid) 383 ] 384 instructions += (INS_BPF_PARAM_STORE + BpfFuncCountPacketInit(self.map_fd) 385 + INS_CGROUP_ACCEPT + INS_PACK_COUNT_UPDATE + INS_CGROUP_ACCEPT) 386 self.prog_fd = BpfProgLoad(BPF_PROG_TYPE_CGROUP_SKB, instructions) 387 BpfProgAttach(self.prog_fd, self._cg_fd, BPF_CGROUP_INET_INGRESS) 388 uid = os.getuid() 389 packet_count = 20 390 SocketUDPLoopBack(packet_count, 4, None) 391 self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value) 392 DeleteMap(self.map_fd, uid); 393 SocketUDPLoopBack(packet_count, 6, None) 394 self.assertEquals(packet_count, LookupMap(self.map_fd, uid).value) 395 BpfProgDetach(self._cg_fd, BPF_CGROUP_INET_INGRESS) 396 397 if __name__ == "__main__": 398 unittest.main() 399