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
17import ctypes
18
19import csocket
20import cstruct
21import net_test
22import socket
23
24# TODO: figure out how to make this arch-dependent if we run these tests
25# on non-X86
26__NR_bpf = 321
27LOG_LEVEL = 1
28LOG_SIZE = 65536
29
30# BPF syscall commands constants.
31BPF_MAP_CREATE = 0
32BPF_MAP_LOOKUP_ELEM = 1
33BPF_MAP_UPDATE_ELEM = 2
34BPF_MAP_DELETE_ELEM = 3
35BPF_MAP_GET_NEXT_KEY = 4
36BPF_PROG_LOAD = 5
37BPF_OBJ_PIN = 6
38BPF_OBJ_GET = 7
39SO_ATTACH_BPF = 50
40
41# BPF map type constant.
42BPF_MAP_TYPE_UNSPEC = 0
43BPF_MAP_TYPE_HASH = 1
44BPF_MAP_TYPE_ARRAY = 2
45BPF_MAP_TYPE_PROG_ARRAY = 3
46BPF_MAP_TYPE_PERF_EVENT_ARRAY = 4
47
48# BPF program type constant.
49BPF_PROG_TYPE_UNSPEC = 0
50BPF_PROG_TYPE_SOCKET_FILTER = 1
51BPF_PROG_TYPE_KPROBE = 2
52BPF_PROG_TYPE_SCHED_CLS = 3
53BPF_PROG_TYPE_SCHED_ACT = 4
54
55# BPF register constant
56BPF_REG_0 = 0
57BPF_REG_1 = 1
58BPF_REG_2 = 2
59BPF_REG_3 = 3
60BPF_REG_4 = 4
61BPF_REG_5 = 5
62BPF_REG_6 = 6
63BPF_REG_7 = 7
64BPF_REG_8 = 8
65BPF_REG_9 = 9
66BPF_REG_10 = 10
67
68# BPF instruction constants
69BPF_PSEUDO_MAP_FD = 1
70BPF_LD = 0x00
71BPF_LDX = 0x01
72BPF_ST = 0x02
73BPF_STX = 0x03
74BPF_ALU = 0x04
75BPF_JMP = 0x05
76BPF_RET = 0x06
77BPF_MISC = 0x07
78BPF_W = 0x00
79BPF_H = 0x08
80BPF_B = 0x10
81BPF_IMM = 0x00
82BPF_ABS = 0x20
83BPF_IND = 0x40
84BPF_MEM = 0x60
85BPF_LEN = 0x80
86BPF_MSH = 0xa0
87BPF_ADD = 0x00
88BPF_SUB = 0x10
89BPF_MUL = 0x20
90BPF_DIV = 0x30
91BPF_OR = 0x40
92BPF_AND = 0x50
93BPF_LSH = 0x60
94BPF_RSH = 0x70
95BPF_NEG = 0x80
96BPF_MOD = 0x90
97BPF_XOR = 0xa0
98BPF_JA = 0x00
99BPF_JEQ = 0x10
100BPF_JGT = 0x20
101BPF_JGE = 0x30
102BPF_JSET = 0x40
103BPF_K = 0x00
104BPF_X = 0x08
105BPF_ALU64 = 0x07
106BPF_DW = 0x18
107BPF_XADD = 0xc0
108BPF_MOV = 0xb0
109
110BPF_ARSH = 0xc0
111BPF_END = 0xd0
112BPF_TO_LE = 0x00
113BPF_TO_BE = 0x08
114
115BPF_JNE = 0x50
116BPF_JSGT = 0x60
117
118BPF_JSGE = 0x70
119BPF_CALL = 0x80
120BPF_EXIT = 0x90
121
122# BPF helper function constants
123BPF_FUNC_unspec = 0
124BPF_FUNC_map_lookup_elem = 1
125BPF_FUNC_map_update_elem = 2
126BPF_FUNC_map_delete_elem = 3
127
128# BPF attr struct
129BpfAttrCreate = cstruct.Struct("bpf_attr_create", "=IIII",
130                               "map_type key_size value_size max_entries")
131BpfAttrOps = cstruct.Struct("bpf_attr_ops", "=QQQQ",
132                            "map_fd key_ptr value_ptr flags")
133BpfAttrProgLoad = cstruct.Struct(
134    "bpf_attr_prog_load", "=IIQQIIQI", "prog_type insn_cnt insns"
135    " license log_level log_size log_buf kern_version")
136BpfInsn = cstruct.Struct("bpf_insn", "=BBhi", "code dst_src_reg off imm")
137
138libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
139HAVE_EBPF_SUPPORT = net_test.LINUX_VERSION >= (4, 4, 0)
140
141
142# BPF program syscalls
143def CreateMap(map_type, key_size, value_size, max_entries):
144  attr = BpfAttrCreate((map_type, key_size, value_size, max_entries))
145  ret = libc.syscall(__NR_bpf, BPF_MAP_CREATE, attr.CPointer(), len(attr))
146  csocket.MaybeRaiseSocketError(ret)
147  return ret
148
149
150def UpdateMap(map_fd, key, value, flags=0):
151  c_value = ctypes.c_uint32(value)
152  c_key = ctypes.c_uint32(key)
153  value_ptr = ctypes.addressof(c_value)
154  key_ptr = ctypes.addressof(c_key)
155  attr = BpfAttrOps((map_fd, key_ptr, value_ptr, flags))
156  ret = libc.syscall(__NR_bpf, BPF_MAP_UPDATE_ELEM,
157                     attr.CPointer(), len(attr))
158  csocket.MaybeRaiseSocketError(ret)
159
160
161def LookupMap(map_fd, key):
162  c_value = ctypes.c_uint32(0)
163  c_key = ctypes.c_uint32(key)
164  attr = BpfAttrOps(
165      (map_fd, ctypes.addressof(c_key), ctypes.addressof(c_value), 0))
166  ret = libc.syscall(__NR_bpf, BPF_MAP_LOOKUP_ELEM,
167                     attr.CPointer(), len(attr))
168  csocket.MaybeRaiseSocketError(ret)
169  return c_value
170
171
172def GetNextKey(map_fd, key):
173  c_key = ctypes.c_uint32(key)
174  c_next_key = ctypes.c_uint32(0)
175  attr = BpfAttrOps(
176      (map_fd, ctypes.addressof(c_key), ctypes.addressof(c_next_key), 0))
177  ret = libc.syscall(__NR_bpf, BPF_MAP_GET_NEXT_KEY,
178                     attr.CPointer(), len(attr))
179  csocket.MaybeRaiseSocketError(ret)
180  return c_next_key
181
182
183def DeleteMap(map_fd, key):
184  c_key = ctypes.c_uint32(key)
185  attr = BpfAttrOps((map_fd, ctypes.addressof(c_key), 0, 0))
186  ret = libc.syscall(__NR_bpf, BPF_MAP_DELETE_ELEM,
187                     attr.CPointer(), len(attr))
188  csocket.MaybeRaiseSocketError(ret)
189
190
191def BpfProgLoad(prog_type, insn_ptr, prog_len, insn_len):
192  gpl_license = ctypes.create_string_buffer(b"GPL")
193  log_buf = ctypes.create_string_buffer(b"", LOG_SIZE)
194  attr = BpfAttrProgLoad(
195      (prog_type, prog_len / insn_len, insn_ptr, ctypes.addressof(gpl_license),
196       LOG_LEVEL, LOG_SIZE, ctypes.addressof(log_buf), 0))
197  ret = libc.syscall(__NR_bpf, BPF_PROG_LOAD, attr.CPointer(), len(attr))
198  csocket.MaybeRaiseSocketError(ret)
199  return ret
200
201
202def BpfProgAttach(sock_fd, prog_fd):
203  prog_ptr = ctypes.c_uint32(prog_fd)
204  ret = libc.setsockopt(sock_fd, socket.SOL_SOCKET, SO_ATTACH_BPF,
205                        ctypes.addressof(prog_ptr), ctypes.sizeof(prog_ptr))
206  csocket.MaybeRaiseSocketError(ret)
207
208
209# BPF program command constructors
210def BpfMov64Reg(dst, src):
211  code = BPF_ALU64 | BPF_MOV | BPF_X
212  dst_src = src << 4 | dst
213  ret = BpfInsn((code, dst_src, 0, 0))
214  return ret.Pack()
215
216
217def BpfLdxMem(size, dst, src, off):
218  code = BPF_LDX | (size & 0x18) | BPF_MEM
219  dst_src = src << 4 | dst
220  ret = BpfInsn((code, dst_src, off, 0))
221  return ret.Pack()
222
223
224def BpfStxMem(size, dst, src, off):
225  code = BPF_STX | (size & 0x18) | BPF_MEM
226  dst_src = src << 4 | dst
227  ret = BpfInsn((code, dst_src, off, 0))
228  return ret.Pack()
229
230
231def BpfStMem(size, dst, off, imm):
232  code = BPF_ST | (size & 0x18) | BPF_MEM
233  dst_src = dst
234  ret = BpfInsn((code, dst_src, off, imm))
235  return ret.Pack()
236
237
238def BpfAlu64Imm(op, dst, imm):
239  code = BPF_ALU64 | (op & 0xf0) | BPF_K
240  dst_src = dst
241  ret = BpfInsn((code, dst_src, 0, imm))
242  return ret.Pack()
243
244
245def BpfJumpImm(op, dst, imm, off):
246  code = BPF_JMP | (op & 0xf0) | BPF_K
247  dst_src = dst
248  ret = BpfInsn((code, dst_src, off, imm))
249  return ret.Pack()
250
251
252def BpfRawInsn(code, dst, src, off, imm):
253  ret = BpfInsn((code, (src << 4 | dst), off, imm))
254  return ret.Pack()
255
256
257def BpfMov64Imm(dst, imm):
258  code = BPF_ALU64 | BPF_MOV | BPF_K
259  dst_src = dst
260  ret = BpfInsn((code, dst_src, 0, imm))
261  return ret.Pack()
262
263
264def BpfExitInsn():
265  code = BPF_JMP | BPF_EXIT
266  ret = BpfInsn((code, 0, 0, 0))
267  return ret.Pack()
268
269
270def BpfLoadMapFd(map_fd, dst):
271  code = BPF_LD | BPF_DW | BPF_IMM
272  dst_src = BPF_PSEUDO_MAP_FD << 4 | dst
273  insn1 = BpfInsn((code, dst_src, 0, map_fd))
274  insn2 = BpfInsn((0, 0, 0, map_fd >> 32))
275  return insn1.Pack() + insn2.Pack()
276
277
278def BpfFuncLookupMap():
279  code = BPF_JMP | BPF_CALL
280  dst_src = 0
281  ret = BpfInsn((code, dst_src, 0, BPF_FUNC_map_lookup_elem))
282  return ret.Pack()
283
284
285def BpfFuncUpdateMap():
286  code = BPF_JMP | BPF_CALL
287  dst_src = 0
288  ret = BpfInsn((code, dst_src, 0, BPF_FUNC_map_update_elem))
289  return ret.Pack()
290
291
292def BpfFuncDeleteMap():
293  code = BPF_JMP | BPF_CALL
294  dst_src = 0
295  ret = BpfInsn((code, dst_src, 0, BPF_FUNC_map_delete_elem))
296  return ret.Pack()
297