1# Copyright 2016 Sasha Goldshtein 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15from __future__ import print_function 16import ctypes as ct 17import sys 18from .libbcc import lib, _USDT_CB, _USDT_PROBE_CB, \ 19 bcc_usdt_location, bcc_usdt_argument, \ 20 BCC_USDT_ARGUMENT_FLAGS 21 22class USDTException(Exception): 23 pass 24 25class USDTProbeArgument(object): 26 def __init__(self, argument): 27 self.signed = argument.size < 0 28 self.size = abs(argument.size) 29 self.valid = argument.valid 30 if self.valid & BCC_USDT_ARGUMENT_FLAGS.CONSTANT != 0: 31 self.constant = argument.constant 32 if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0: 33 self.deref_offset = argument.deref_offset 34 if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT != 0: 35 self.deref_ident = argument.deref_ident 36 if self.valid & BCC_USDT_ARGUMENT_FLAGS.BASE_REGISTER_NAME != 0: 37 self.base_register_name = argument.base_register_name 38 if self.valid & BCC_USDT_ARGUMENT_FLAGS.INDEX_REGISTER_NAME != 0: 39 self.index_register_name = argument.index_register_name 40 if self.valid & BCC_USDT_ARGUMENT_FLAGS.SCALE != 0: 41 self.scale = argument.scale 42 43 def _size_prefix(self): 44 return "%d %s bytes" % \ 45 (self.size, "signed " if self.signed else "unsigned") 46 47 def _format(self): 48 # This mimics the logic in cc/usdt_args.cc that gives meaning to the 49 # various argument settings. A change there will require a change here. 50 if self.valid & BCC_USDT_ARGUMENT_FLAGS.CONSTANT != 0: 51 return "%d" % self.constant 52 if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET == 0: 53 return "%s" % self.base_register_name.decode() 54 if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0 and \ 55 self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT == 0: 56 if self.valid & BCC_USDT_ARGUMENT_FLAGS.INDEX_REGISTER_NAME != 0: 57 index_offset = " + %s" % self.index_register_name.decode() 58 if self.valid & BCC_USDT_ARGUMENT_FLAGS.SCALE != 0: 59 index_offset += " * %d" % self.scale 60 else: 61 index_offset = "" 62 sign = '+' if self.deref_offset >= 0 else '-' 63 return "*(%s %s %d%s)" % (self.base_register_name.decode(), 64 sign, abs(self.deref_offset), index_offset) 65 if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0 and \ 66 self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT != 0 and \ 67 self.valid & BCC_USDT_ARGUMENT_FLAGS.BASE_REGISTER_NAME != 0 and \ 68 self.base_register_name == "ip": 69 sign = '+' if self.deref_offset >= 0 else '-' 70 return "*(&%s %s %d)" % (self.deref_ident.decode(), 71 sign, abs(self.deref_offset)) 72 # If we got here, this is an unrecognized case. Doesn't mean it's 73 # necessarily bad, so just provide the raw data. It just means that 74 # other tools won't be able to work with this argument. 75 return "unrecognized argument format, flags %d" % self.valid 76 77 def __str__(self): 78 return "%s @ %s" % (self._size_prefix(), self._format()) 79 80class USDTProbeLocation(object): 81 def __init__(self, probe, index, location): 82 self.probe = probe 83 self.index = index 84 self.num_arguments = probe.num_arguments 85 self.address = location.address 86 self.bin_path = location.bin_path 87 88 def __str__(self): 89 return "%s 0x%x" % (self.bin_path.decode(), self.address) 90 91 def get_argument(self, index): 92 arg = bcc_usdt_argument() 93 res = lib.bcc_usdt_get_argument(self.probe.context, self.probe.provider, 94 self.probe.name, 95 self.index, index, ct.byref(arg)) 96 if res != 0: 97 raise USDTException( 98 "error retrieving probe argument %d location %d" % 99 (index, self.index)) 100 return USDTProbeArgument(arg) 101 102class USDTProbe(object): 103 def __init__(self, context, probe): 104 self.context = context 105 self.provider = probe.provider 106 self.name = probe.name 107 self.bin_path = probe.bin_path 108 self.semaphore = probe.semaphore 109 self.num_locations = probe.num_locations 110 self.num_arguments = probe.num_arguments 111 112 def __str__(self): 113 return "%s:%s [sema 0x%x]" % \ 114 (self.provider.decode(), self.name.decode(), self.semaphore) 115 116 def short_name(self): 117 return "%s:%s" % (self.provider.decode(), self.name.decode()) 118 119 def get_location(self, index): 120 loc = bcc_usdt_location() 121 res = lib.bcc_usdt_get_location(self.context, self.provider, self.name, 122 index, ct.byref(loc)) 123 if res != 0: 124 raise USDTException("error retrieving probe location %d" % index) 125 return USDTProbeLocation(self, index, loc) 126 127class USDT(object): 128 def __init__(self, pid=None, path=None): 129 if pid and pid != -1: 130 self.pid = pid 131 if path: 132 self.context = lib.bcc_usdt_new_frompid(pid, path.encode('ascii')) 133 else: 134 self.context = lib.bcc_usdt_new_frompid(pid, ct.c_char_p(0)) 135 if self.context == None: 136 raise USDTException("USDT failed to instrument PID %d" % pid) 137 elif path: 138 self.path = path 139 self.context = lib.bcc_usdt_new_frompath(path.encode('ascii')) 140 if self.context == None: 141 raise USDTException("USDT failed to instrument path %s" % path) 142 else: 143 raise USDTException( 144 "either a pid or a binary path must be specified") 145 146 def __del__(self): 147 lib.bcc_usdt_close(self.context) 148 149 def enable_probe(self, probe, fn_name): 150 probe_parts = probe.split(":", 1) 151 if len(probe_parts) == 1: 152 ret = lib.bcc_usdt_enable_probe( 153 self.context, probe.encode('ascii'), fn_name.encode('ascii')) 154 else: 155 (provider_name, probe_name) = probe_parts 156 ret = lib.bcc_usdt_enable_fully_specified_probe( 157 self.context, provider_name.encode('ascii'), probe_name.encode('ascii'), 158 fn_name.encode('ascii')) 159 160 if ret != 0: 161 raise USDTException( 162"""Failed to enable USDT probe '%s': 163the specified pid might not contain the given language's runtime, 164or the runtime was not built with the required USDT probes. Look 165for a configure flag similar to --with-dtrace or --enable-dtrace. 166To check which probes are present in the process, use the tplist tool. 167""" % probe) 168 169 def enable_probe_or_bail(self, probe, fn_name): 170 try: 171 self.enable_probe(probe, fn_name) 172 except USDTException as e: 173 print(e, file=sys.stderr) 174 sys.exit(1) 175 176 def get_context(self): 177 return self.context 178 179 def get_text(self): 180 ctx_array = (ct.c_void_p * 1)() 181 ctx_array[0] = ct.c_void_p(self.context) 182 return lib.bcc_usdt_genargs(ctx_array, 1).decode() 183 184 def get_probe_arg_ctype(self, probe_name, arg_index): 185 probe_parts = probe_name.split(":", 1) 186 if len(probe_parts) == 1: 187 return lib.bcc_usdt_get_probe_argctype( 188 self.context, probe_name.encode('ascii'), arg_index).decode() 189 else: 190 (provider_name, probe) = probe_parts 191 return lib.bcc_usdt_get_fully_specified_probe_argctype( 192 self.context, provider_name.encode('ascii'), probe.encode('ascii'), arg_index).decode() 193 194 def enumerate_probes(self): 195 probes = [] 196 def _add_probe(probe): 197 probes.append(USDTProbe(self.context, probe.contents)) 198 199 lib.bcc_usdt_foreach(self.context, _USDT_CB(_add_probe)) 200 return probes 201 202 # This is called by the BPF module's __init__ when it realizes that there 203 # is a USDT context and probes need to be attached. 204 def attach_uprobes(self, bpf, attach_usdt_ignore_pid): 205 probes = self.enumerate_active_probes() 206 for (binpath, fn_name, addr, pid) in probes: 207 if attach_usdt_ignore_pid: 208 pid = -1 209 bpf.attach_uprobe(name=binpath, fn_name=fn_name, 210 addr=addr, pid=pid) 211 212 def enumerate_active_probes(self): 213 probes = [] 214 def _add_probe(binpath, fn_name, addr, pid): 215 probes.append((binpath, fn_name, addr, pid)) 216 217 lib.bcc_usdt_foreach_uprobe(self.context, _USDT_PROBE_CB(_add_probe)) 218 return probes 219