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