<commit_before>#!/usr/bin/python
# @lint-avoid-python-3-compatibility-imports
#
# funccount Count functions, tracepoints, and USDT probes.
#           For Linux, uses BCC, eBPF.
#
# USAGE: funccount [-h] [-p PID] [-i INTERVAL] [-T] [-r] pattern
#
# The pattern is a string with optional '*' wildcards, similar to file
# globbing. If you'd prefer to use regular expressions, use the -r option.
#
# Copyright (c) 2015 <NAME>.
# Licensed under the Apache License, Version 2.0 (the "License")
#
# 09-Sep-2015   <NAME>       Created this.
# 18-Oct-2016   <NAME>    Generalized for uprobes, tracepoints, USDT.

from __future__ import print_function
from bcc import BPF, USDT
from time import sleep, strftime
import argparse
import os
import re
import signal
import sys
import traceback

debug = False

class Probe(object):
    def __init__(self, pattern, use_regex=False, pid=None):
        """Init a new probe.

        Init the probe from the pattern provided by the user. The supported
        patterns mimic the 'trace' and 'argdist' tools, but are simpler because
        we don't have to distinguish between probes and retprobes.

            func            -- probe a kernel function
            lib:func        -- probe a user-space function in the library 'lib'
            p::func         -- same thing as 'func'
            p:lib:func      -- same thing as 'lib:func'
            t:cat:event     -- probe a kernel tracepoint
            u:lib:probe     -- probe a USDT tracepoint
        """
        parts = pattern.split(':')
        if len(parts) == 1:
            parts = ["p", "", parts[0]]
        elif len(parts) == 2:
            parts = ["p", parts[0], parts[1]]
        elif len(parts) == 3:
            if parts[0] == "t":
                parts = ["t", "", "%s:%s" % tuple(parts[1:])]
            if parts[0] not in ["p", "t", "u"]:
                raise Exception("Type must be 'p', 't', or 'u', but got %s" %
                                parts[0])
        else:
            raise Exception("Too many ':'-separated components in pattern %s" %
                            pattern)

        (self.type, self.library, self.pattern) = parts
        if not use_regex:
            self.pattern = self.pattern.replace('*', '.*')
            self.pattern = '^' + self.pattern + '$'

        if (self.type == "p" and self.library) or self.type == "u":
            libpath = BPF.find_library(self.library)
            if libpath is None:
                # This might be an executable (e.g. 'bash')
                libpath = BPF.find_exe(self.library)
            if libpath is None or len(libpath) == 0:
                raise Exception("unable to find library %s" % self.library)
            self.library = libpath

        self.pid = pid
        self.matched = 0
        self.trace_functions = {}   # map location number to function name

    def is_kernel_probe(self):
        return self.type == "t" or (self.type == "p" and self.library == "")

    def attach(self):
        if self.type == "p" and not self.library:
            for index, function in self.trace_functions.items():
                self.bpf.attach_kprobe(
                        event=function,
                        fn_name="trace_count_%d" % index,
                        pid=self.pid or -1)
        elif self.type == "p" and self.library:
            for index, function in self.trace_functions.items():
                self.bpf.attach_uprobe(
                        name=self.library,
                        sym=function,
                        fn_name="trace_count_%d" % index,
                        pid=self.pid or -1)
        elif self.type == "t":
            for index, function in self.trace_functions.items():
                self.bpf.attach_tracepoint(
                        tp=function,
                        fn_name="trace_count_%d" % index,
                        pid=self.pid or -1)
        elif self.type == "u":
            pass    # Nothing to do -- attach already happened in `load`

        if self.matched == 0:
            raise Exception("No functions matched by pattern %s" %
                            self.pattern)

    def _add_function(self, template, probe_name):
        new_func = "trace_count_%d" % self.matched
        text = template.replace("PROBE_FUNCTION", new_func)
        text = text.replace("LOCATION", str(self.matched))
        self.trace_functions[self.matched] = probe_name
        self.matched += 1
        return text

    def _generate_functions(self, template):
        self.usdt = None
        text = ""
        if self.type == "p" and not self.library:
            for function in BPF.get_kprobe_functions(self.pattern):
                text += self._add_function(template, function)
        elif self.type == "p" and self.library:
            # uprobes are tricky because the same function may have multiple
            # addresses, and the same address may be mapped to multiple
            # functions. We aren't allowed to create more than one uprobe
            # per address, so track unique addresses and ignore functions that
            # map to an address that we've already seen. Also ignore functions
            # that may repeat multiple times with different addresses.
            addresses, functions = (set(), set())
            for function, address in BPF.get_user_functions_and_addresses(
                                        self.library, self.pattern):
                if address in addresses or function in functions:
                    continue
                addresses.add(address)
                functions.add(function)
                text += self._add_function(template, function)
        elif self.type == "t":
            for tracepoint in BPF.get_tracepoints(self.pattern):
                text += self._add_function(template, tracepoint)
        elif self.type == "u":
            self.usdt = USDT(path=self.library, pid=self.pid)
            for probe in self.usdt.enumerate_probes():
                if not self.pid and (probe.bin_path != self.library):
                    continue
                if re.match(self.pattern, probe.name):
                    new_func = "trace_count_%d" % self.matched
                    text += self._add_function(template, probe.name)
                    self.usdt.enable_probe(probe.name, new_func)
            if debug:<commit_msg>funccount: Verify probe max limit (#771)<commit_after>#!/usr/bin/python
# @lint-avoid-python-3-compatibility-imports
#
# funccount Count functions, tracepoints, and USDT probes.
#           For Linux, uses BCC, eBPF.
#
# USAGE: funccount [-h] [-p PID] [-i INTERVAL] [-T] [-r] pattern
#
# The pattern is a string with optional '*' wildcards, similar to file
# globbing. If you'd prefer to use regular expressions, use the -r option.
#
# Copyright (c) 2015 <NAME>.
# Licensed under the Apache License, Version 2.0 (the "License")
#
# 09-Sep-2015   <NAME>       Created this.
# 18-Oct-2016   <NAME>    Generalized for uprobes, tracepoints, USDT.

from __future__ import print_function
from bcc import BPF, USDT
from time import sleep, strftime
import argparse
import os
import re
import signal
import sys
import traceback

debug = False

def verify_limit(num):
    probe_limit = 1000
    if num > probe_limit:
        raise Exception("maximum of %d probes allowed, attempted %d" %
                        (probe_limit, num))

class Probe(object):
    def __init__(self, pattern, use_regex=False, pid=None):
        """Init a new probe.

        Init the probe from the pattern provided by the user. The supported
        patterns mimic the 'trace' and 'argdist' tools, but are simpler because
        we don't have to distinguish between probes and retprobes.

            func            -- probe a kernel function
            lib:func        -- probe a user-space function in the library 'lib'
            p::func         -- same thing as 'func'
            p:lib:func      -- same thing as 'lib:func'
            t:cat:event     -- probe a kernel tracepoint
            u:lib:probe     -- probe a USDT tracepoint
        """
        parts = pattern.split(':')
        if len(parts) == 1:
            parts = ["p", "", parts[0]]
        elif len(parts) == 2:
            parts = ["p", parts[0], parts[1]]
        elif len(parts) == 3:
            if parts[0] == "t":
                parts = ["t", "", "%s:%s" % tuple(parts[1:])]
            if parts[0] not in ["p", "t", "u"]:
                raise Exception("Type must be 'p', 't', or 'u', but got %s" %
                                parts[0])
        else:
            raise Exception("Too many ':'-separated components in pattern %s" %
                            pattern)

        (self.type, self.library, self.pattern) = parts
        if not use_regex:
            self.pattern = self.pattern.replace('*', '.*')
            self.pattern = '^' + self.pattern + '$'

        if (self.type == "p" and self.library) or self.type == "u":
            libpath = BPF.find_library(self.library)
            if libpath is None:
                # This might be an executable (e.g. 'bash')
                libpath = BPF.find_exe(self.library)
            if libpath is None or len(libpath) == 0:
                raise Exception("unable to find library %s" % self.library)
            self.library = libpath

        self.pid = pid
        self.matched = 0
        self.trace_functions = {}   # map location number to function name

    def is_kernel_probe(self):
        return self.type == "t" or (self.type == "p" and self.library == "")

    def attach(self):
        if self.type == "p" and not self.library:
            for index, function in self.trace_functions.items():
                self.bpf.attach_kprobe(
                        event=function,
                        fn_name="trace_count_%d" % index,
                        pid=self.pid or -1)
        elif self.type == "p" and self.library:
            for index, function in self.trace_functions.items():
                self.bpf.attach_uprobe(
                        name=self.library,
                        sym=function,
                        fn_name="trace_count_%d" % index,
                        pid=self.pid or -1)
        elif self.type == "t":
            for index, function in self.trace_functions.items():
                self.bpf.attach_tracepoint(
                        tp=function,
                        fn_name="trace_count_%d" % index,
                        pid=self.pid or -1)
        elif self.type == "u":
            pass    # Nothing to do -- attach already happened in `load`

        if self.matched == 0:
            raise Exception("No functions matched by pattern %s" %
                            self.pattern)

    def _add_function(self, template, probe_name):
        new_func = "trace_count_%d" % self.matched
        text = template.replace("PROBE_FUNCTION", new_func)
        text = text.replace("LOCATION", str(self.matched))
        self.trace_functions[self.matched] = probe_name
        self.matched += 1
        return text

    def _generate_functions(self, template):
        self.usdt = None
        text = ""
        if self.type == "p" and not self.library:
            functions = BPF.get_kprobe_functions(self.pattern)
            verify_limit(len(functions))
            for function in functions:
                text += self._add_function(template, function)
        elif self.type == "p" and self.library:
            # uprobes are tricky because the same function may have multiple
            # addresses, and the same address may be mapped to multiple
            # functions. We aren't allowed to create more than one uprobe
            # per address, so track unique addresses and ignore functions that
            # map to an address that we've already seen. Also ignore functions
            # that may repeat multiple times with different addresses.
            addresses, functions = (set(), set())
            functions_and_addresses = BPF.get_user_functions_and_addresses(
                                        self.library, self.pattern)
            verify_limit(len(functions_and_addresses))
            for function, address in functions_and_addresses:
                if address in addresses or function in functions:
                    continue
                addresses.add(address)
                functions.add(function)
                text += self._add_function(template, function)
        elif self.type == "t":
            tracepoints = BPF.get_tracepoints(self.pattern)
            verify_limit(len(tracepoints))
            for tracepoint in tracepoints:
                text += self._add_function(template, tracepoint)
        elif self.type == "u":
            self.usdt = USDT(path=self.library, pid=self.pid)
            matches = []
            for probe in self.usdt.enumerate_probes():
                if not self.pid and (probe.bin_path != self.library):
                    continue
                if re.match(self.pattern, probe.name):
                    matches.append(probe.name)
            verify_limit(len(matches))
            for match in matches:
                new_func = "trace_count_%d" % self.matched
                text += self._add_function(template, match)
                self.usdt.enable_probe(match, new_func)
            if debug:
{'fraction_non_alphanumeric': 0.10224457623301554, 'fraction_numerical': 0.005930485699271826, 'mean_word_length': 1.76734524304113, 'pattern_counts': {'":': 12, '<': 9, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7086555', 'n_tokens_mistral': 3705, 'n_tokens_neox': 3569, 'n_words': 1125}