xref: /aosp_15_r20/kernel/tests/net/test/bpf.py (revision 2f2c4c7ab4226c71756b9c31670392fdd6887c4f)
1#!/usr/bin/python3
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"""kernel net test library for bpf testing."""
18
19import ctypes
20import errno
21import os
22import resource
23import socket
24import sys
25
26import csocket
27import cstruct
28import net_test
29
30# __NR_bpf syscall numbers for various architectures.
31# NOTE: If python inherited COMPAT_UTS_MACHINE, uname's 'machine' field will
32# return the 32-bit architecture name, even if python itself is 64-bit. To work
33# around this problem and pick the right syscall nr, we can additionally check
34# the bitness of the python interpreter. Assume that the 64-bit architectures
35# are not running with COMPAT_UTS_MACHINE and must be 64-bit at all times.
36__NR_bpf = {  # pylint: disable=invalid-name
37    "aarch64-32bit": 386,
38    "aarch64-64bit": 280,
39    "armv7l-32bit": 386,
40    "armv8l-32bit": 386,
41    "armv8l-64bit": 280,
42    "i686-32bit": 357,
43    "i686-64bit": 321,
44    "x86_64-32bit": 357,
45    "x86_64-64bit": 321,
46    "riscv64-64bit": 280,
47}[os.uname()[4] + "-" + ("64" if sys.maxsize > 0x7FFFFFFF else "32") + "bit"]
48
49# After ACK merge of 5.10.168 is when support for this was backported from
50# upstream Linux 5.14 and was merged into ACK android{12,13}-5.10 branches.
51# Require support to be backported to any 5.10+ kernel.
52HAVE_SO_NETNS_COOKIE = net_test.LINUX_VERSION >= (5, 10, 0)
53
54# Note: This is *not* correct for parisc & sparc architectures
55SO_NETNS_COOKIE = 71
56
57LOG_LEVEL = 1
58LOG_SIZE = 65536
59
60# BPF syscall commands constants.
61BPF_MAP_CREATE = 0
62BPF_MAP_LOOKUP_ELEM = 1
63BPF_MAP_UPDATE_ELEM = 2
64BPF_MAP_DELETE_ELEM = 3
65BPF_MAP_GET_NEXT_KEY = 4
66BPF_PROG_LOAD = 5
67BPF_OBJ_PIN = 6
68BPF_OBJ_GET = 7
69BPF_PROG_ATTACH = 8
70BPF_PROG_DETACH = 9
71BPF_PROG_TEST_RUN = 10
72BPF_PROG_GET_NEXT_ID = 11
73BPF_MAP_GET_NEXT_ID = 12
74BPF_PROG_GET_FD_BY_ID = 13
75BPF_MAP_GET_FD_BY_ID = 14
76BPF_OBJ_GET_INFO_BY_FD = 15
77BPF_PROG_QUERY = 16
78
79# setsockopt SOL_SOCKET constants
80SO_ATTACH_BPF = 50
81
82# BPF map type constant.
83BPF_MAP_TYPE_UNSPEC = 0
84BPF_MAP_TYPE_HASH = 1
85BPF_MAP_TYPE_ARRAY = 2
86BPF_MAP_TYPE_PROG_ARRAY = 3
87BPF_MAP_TYPE_PERF_EVENT_ARRAY = 4
88
89# BPF program type constant.
90BPF_PROG_TYPE_UNSPEC = 0
91BPF_PROG_TYPE_SOCKET_FILTER = 1
92BPF_PROG_TYPE_KPROBE = 2
93BPF_PROG_TYPE_SCHED_CLS = 3
94BPF_PROG_TYPE_SCHED_ACT = 4
95BPF_PROG_TYPE_TRACEPOINT = 5
96BPF_PROG_TYPE_XDP = 6
97BPF_PROG_TYPE_PERF_EVENT = 7
98BPF_PROG_TYPE_CGROUP_SKB = 8
99BPF_PROG_TYPE_CGROUP_SOCK = 9
100
101# BPF program attach type.
102BPF_CGROUP_INET_INGRESS = 0
103BPF_CGROUP_INET_EGRESS = 1
104BPF_CGROUP_INET_SOCK_CREATE = 2
105
106# BPF register constant
107BPF_REG_0 = 0
108BPF_REG_1 = 1
109BPF_REG_2 = 2
110BPF_REG_3 = 3
111BPF_REG_4 = 4
112BPF_REG_5 = 5
113BPF_REG_6 = 6
114BPF_REG_7 = 7
115BPF_REG_8 = 8
116BPF_REG_9 = 9
117BPF_REG_10 = 10
118
119# BPF instruction constants
120BPF_PSEUDO_MAP_FD = 1
121BPF_LD = 0x00
122BPF_LDX = 0x01
123BPF_ST = 0x02
124BPF_STX = 0x03
125BPF_ALU = 0x04
126BPF_JMP = 0x05
127BPF_RET = 0x06
128BPF_MISC = 0x07
129BPF_W = 0x00
130BPF_H = 0x08
131BPF_B = 0x10
132BPF_IMM = 0x00
133BPF_ABS = 0x20
134BPF_IND = 0x40
135BPF_MEM = 0x60
136BPF_LEN = 0x80
137BPF_MSH = 0xa0
138BPF_ADD = 0x00
139BPF_SUB = 0x10
140BPF_MUL = 0x20
141BPF_DIV = 0x30
142BPF_OR = 0x40
143BPF_AND = 0x50
144BPF_LSH = 0x60
145BPF_RSH = 0x70
146BPF_NEG = 0x80
147BPF_MOD = 0x90
148BPF_XOR = 0xa0
149BPF_JA = 0x00
150BPF_JEQ = 0x10
151BPF_JGT = 0x20
152BPF_JGE = 0x30
153BPF_JSET = 0x40
154BPF_K = 0x00
155BPF_X = 0x08
156BPF_ALU64 = 0x07
157BPF_DW = 0x18
158BPF_XADD = 0xc0
159BPF_MOV = 0xb0
160
161BPF_ARSH = 0xc0
162BPF_END = 0xd0
163BPF_TO_LE = 0x00
164BPF_TO_BE = 0x08
165
166BPF_JNE = 0x50
167BPF_JSGT = 0x60
168
169BPF_JSGE = 0x70
170BPF_CALL = 0x80
171BPF_EXIT = 0x90
172
173# BPF helper function constants
174# pylint: disable=invalid-name
175BPF_FUNC_unspec = 0
176BPF_FUNC_map_lookup_elem = 1
177BPF_FUNC_map_update_elem = 2
178BPF_FUNC_map_delete_elem = 3
179BPF_FUNC_ktime_get_ns = 5
180BPF_FUNC_get_current_uid_gid = 15
181BPF_FUNC_skb_change_head = 43
182BPF_FUNC_get_socket_cookie = 46
183BPF_FUNC_get_socket_uid = 47
184BPF_FUNC_ktime_get_boot_ns = 125
185# pylint: enable=invalid-name
186
187BPF_F_RDONLY = 1 << 3
188BPF_F_WRONLY = 1 << 4
189
190#  These object below belongs to the same kernel union and the types below
191#  (e.g., bpf_attr_create) aren't kernel struct names but just different
192#  variants of the union.
193# pylint: disable=invalid-name
194BpfAttrCreate = cstruct.Struct(
195    "bpf_attr_create", "=IIIII",
196    "map_type key_size value_size max_entries, map_flags")
197BpfAttrOps = cstruct.Struct(
198    "bpf_attr_ops", "=QQQQ",
199    "map_fd key_ptr value_ptr flags")
200BpfAttrProgLoad = cstruct.Struct(
201    "bpf_attr_prog_load", "=IIQQIIQI", "prog_type insn_cnt insns"
202    " license log_level log_size log_buf kern_version")
203BpfAttrProgAttach = cstruct.Struct(
204    "bpf_attr_prog_attach", "=III", "target_fd attach_bpf_fd attach_type")
205BpfAttrGetFdById = cstruct.Struct(
206    "bpf_attr_get_fd_by_id", "=III", "id next_id open_flags")
207BpfAttrProgQuery = cstruct.Struct(
208    "bpf_attr_prog_query", "=IIIIQIQ", "target_fd attach_type query_flags attach_flags prog_ids_ptr prog_cnt prog_attach_flags")
209BpfInsn = cstruct.Struct("bpf_insn", "=BBhi", "code dst_src_reg off imm")
210# pylint: enable=invalid-name
211
212libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
213
214# set memlock resource 1 GiB
215resource.setrlimit(resource.RLIMIT_MEMLOCK, (1073741824, 1073741824))
216
217
218# BPF program syscalls
219def BpfSyscall(op, attr):
220  ret = libc.syscall(__NR_bpf, op, csocket.VoidPointer(attr), len(attr))
221  csocket.MaybeRaiseSocketError(ret)
222  return ret
223
224
225def CreateMap(map_type, key_size, value_size, max_entries, map_flags=0):
226  attr = BpfAttrCreate((map_type, key_size, value_size, max_entries, map_flags))
227  return BpfSyscall(BPF_MAP_CREATE, attr)
228
229
230def UpdateMap(map_fd, key, value, flags=0):
231  c_value = ctypes.c_uint32(value)
232  c_key = ctypes.c_uint32(key)
233  value_ptr = ctypes.addressof(c_value)
234  key_ptr = ctypes.addressof(c_key)
235  attr = BpfAttrOps((map_fd, key_ptr, value_ptr, flags))
236  BpfSyscall(BPF_MAP_UPDATE_ELEM, attr)
237
238
239def LookupMap(map_fd, key):
240  c_value = ctypes.c_uint32(0)
241  c_key = ctypes.c_uint32(key)
242  attr = BpfAttrOps(
243      (map_fd, ctypes.addressof(c_key), ctypes.addressof(c_value), 0))
244  BpfSyscall(BPF_MAP_LOOKUP_ELEM, attr)
245  return c_value
246
247
248def GetNextKey(map_fd, key):
249  """Get the next key in the map after the specified key."""
250  if key is not None:
251    c_key = ctypes.c_uint32(key)
252    c_next_key = ctypes.c_uint32(0)
253    key_ptr = ctypes.addressof(c_key)
254  else:
255    key_ptr = 0
256  c_next_key = ctypes.c_uint32(0)
257  attr = BpfAttrOps(
258      (map_fd, key_ptr, ctypes.addressof(c_next_key), 0))
259  BpfSyscall(BPF_MAP_GET_NEXT_KEY, attr)
260  return c_next_key
261
262
263def GetFirstKey(map_fd):
264  return GetNextKey(map_fd, None)
265
266
267def DeleteMap(map_fd, key):
268  c_key = ctypes.c_uint32(key)
269  attr = BpfAttrOps((map_fd, ctypes.addressof(c_key), 0, 0))
270  BpfSyscall(BPF_MAP_DELETE_ELEM, attr)
271
272
273def BpfProgLoad(prog_type, instructions, prog_license=b"GPL"):
274  bpf_prog = b"".join(instructions)
275  insn_buff = ctypes.create_string_buffer(bpf_prog)
276  gpl_license = ctypes.create_string_buffer(prog_license)
277  log_buf = ctypes.create_string_buffer(b"", LOG_SIZE)
278  attr = BpfAttrProgLoad((prog_type, len(insn_buff) // len(BpfInsn),
279                          ctypes.addressof(insn_buff),
280                          ctypes.addressof(gpl_license), LOG_LEVEL,
281                          LOG_SIZE, ctypes.addressof(log_buf), 0))
282  return BpfSyscall(BPF_PROG_LOAD, attr)
283
284
285# Attach a socket eBPF filter to a target socket
286def BpfProgAttachSocket(sock_fd, prog_fd):
287  uint_fd = ctypes.c_uint32(prog_fd)
288  ret = libc.setsockopt(sock_fd, socket.SOL_SOCKET, SO_ATTACH_BPF,
289                        ctypes.pointer(uint_fd), ctypes.sizeof(uint_fd))
290  csocket.MaybeRaiseSocketError(ret)
291
292
293# Attach a eBPF filter to a cgroup
294def BpfProgAttach(prog_fd, target_fd, prog_type):
295  attr = BpfAttrProgAttach((target_fd, prog_fd, prog_type))
296  return BpfSyscall(BPF_PROG_ATTACH, attr)
297
298
299# Detach a eBPF filter from a cgroup
300def BpfProgDetach(target_fd, prog_type):
301  attr = BpfAttrProgAttach((target_fd, 0, prog_type))
302  try:
303    return BpfSyscall(BPF_PROG_DETACH, attr)
304  except socket.error as e:
305    if e.errno != errno.ENOENT:
306      raise
307
308
309# Convert a BPF program ID into an open file descriptor
310def BpfProgGetFdById(prog_id):
311  if prog_id is None:
312    return None
313  attr = BpfAttrGetFdById((prog_id, 0, 0))
314  return BpfSyscall(BPF_PROG_GET_FD_BY_ID, attr)
315
316
317# Convert a BPF map ID into an open file descriptor
318def BpfMapGetFdById(map_id):
319  if map_id is None:
320    return None
321  attr = BpfAttrGetFdById((map_id, 0, 0))
322  return BpfSyscall(BPF_MAP_GET_FD_BY_ID, attr)
323
324
325# Return BPF program id attached to a given cgroup & attach point
326# Note: as written this only supports a *single* program per attach point
327def BpfProgQuery(target_fd, attach_type, query_flags, attach_flags):
328  prog_id = ctypes.c_uint32(-1)
329  minus_one = prog_id.value   # but unsigned, so really 4294967295
330  attr = BpfAttrProgQuery((target_fd, attach_type, query_flags, attach_flags, ctypes.addressof(prog_id), 1, 0))
331  if BpfSyscall(BPF_PROG_QUERY, attr) == 0:
332    # to see kernel updates we have to convert back from the buffer that actually went to the kernel...
333    attr._Parse(attr._buffer)
334    assert attr.prog_cnt >= 0, "prog_cnt is %s" % attr.prog_cnt
335    assert attr.prog_cnt <= 1, "prog_cnt is %s" % attr.prog_cnt  # we don't support more atm
336    if attr.prog_cnt == 0:
337      assert prog_id.value == minus_one, "prog_id is %s" % prog_id
338      return None
339    else:
340      assert prog_id.value != minus_one, "prog_id is %s" % prog_id
341      return prog_id.value
342  else:
343    return None
344
345
346# BPF program command constructors
347def BpfMov64Reg(dst, src):
348  code = BPF_ALU64 | BPF_MOV | BPF_X
349  dst_src = src << 4 | dst
350  ret = BpfInsn((code, dst_src, 0, 0))
351  return ret.Pack()
352
353
354def BpfLdxMem(size, dst, src, off):
355  code = BPF_LDX | (size & 0x18) | BPF_MEM
356  dst_src = src << 4 | dst
357  ret = BpfInsn((code, dst_src, off, 0))
358  return ret.Pack()
359
360
361def BpfStxMem(size, dst, src, off):
362  code = BPF_STX | (size & 0x18) | BPF_MEM
363  dst_src = src << 4 | dst
364  ret = BpfInsn((code, dst_src, off, 0))
365  return ret.Pack()
366
367
368def BpfStMem(size, dst, off, imm):
369  code = BPF_ST | (size & 0x18) | BPF_MEM
370  dst_src = dst
371  ret = BpfInsn((code, dst_src, off, imm))
372  return ret.Pack()
373
374
375def BpfAlu64Imm(op, dst, imm):
376  code = BPF_ALU64 | (op & 0xf0) | BPF_K
377  dst_src = dst
378  ret = BpfInsn((code, dst_src, 0, imm))
379  return ret.Pack()
380
381
382def BpfJumpImm(op, dst, imm, off):
383  code = BPF_JMP | (op & 0xf0) | BPF_K
384  dst_src = dst
385  ret = BpfInsn((code, dst_src, off, imm))
386  return ret.Pack()
387
388
389def BpfRawInsn(code, dst, src, off, imm):
390  ret = BpfInsn((code, (src << 4 | dst), off, imm))
391  return ret.Pack()
392
393
394def BpfMov64Imm(dst, imm):
395  code = BPF_ALU64 | BPF_MOV | BPF_K
396  dst_src = dst
397  ret = BpfInsn((code, dst_src, 0, imm))
398  return ret.Pack()
399
400
401def BpfExitInsn():
402  code = BPF_JMP | BPF_EXIT
403  ret = BpfInsn((code, 0, 0, 0))
404  return ret.Pack()
405
406
407def BpfLoadMapFd(map_fd, dst):
408  code = BPF_LD | BPF_DW | BPF_IMM
409  dst_src = BPF_PSEUDO_MAP_FD << 4 | dst
410  insn1 = BpfInsn((code, dst_src, 0, map_fd))
411  insn2 = BpfInsn((0, 0, 0, map_fd >> 32))
412  return insn1.Pack() + insn2.Pack()
413
414
415def BpfFuncCall(func):
416  code = BPF_JMP | BPF_CALL
417  dst_src = 0
418  ret = BpfInsn((code, dst_src, 0, func))
419  return ret.Pack()
420