add support for uprobes

cli_flags.go

@@ -69,6 +69,10 @@ var (
 		defaultOffCPUThreshold)
 	envVarsHelp = "Comma separated list of environment variables that will be reported with the" +
 		"captured profiling samples."
+	probeLinkHelper = "Attach a probe to a symbol of an executable. " +
+		"Expected format: /path/to/executable:symbol"
+	loadProbeHelper = "Load generic eBPF program that can be attached externally to " +
+		"various user or kernel space hooks."
 )
 
 // Package-scope variable, so that conditionally compiled other components can refer
@@ -127,6 +131,13 @@ func parseArgs() (*controller.Config, error) {
 
 	fs.StringVar(&args.IncludeEnvVars, "env-vars", defaultEnvVarsValue, envVarsHelp)
 
+	fs.Func("uprobe-link", probeLinkHelper, func(link string) error {
+		args.UProbeLinks = append(args.UProbeLinks, link)
+		return nil
+	})
+
+	fs.BoolVar(&args.LoadProbe, "load-probe", false, loadProbeHelper)
+
 	fs.Usage = func() {
 		fs.PrintDefaults()
 	}

internal/controller/config.go

@@ -32,6 +32,8 @@ type Config struct {
 	VerboseMode            bool
 	Version                bool
 	OffCPUThreshold        float64
+	UProbeLinks            []string
+	LoadProbe              bool
 
 	Reporter reporter.Reporter
 

internal/controller/controller.go

@@ -98,6 +98,8 @@ func (c *Controller) Start(ctx context.Context) error {
 		ProbabilisticThreshold: c.config.ProbabilisticThreshold,
 		OffCPUThreshold:        uint32(c.config.OffCPUThreshold * float64(math.MaxUint32)),
 		IncludeEnvVars:         envVars,
+		UProbeLinks:            c.config.UProbeLinks,
+		LoadProbe:              c.config.LoadProbe,
 	})
 	if err != nil {
 		return fmt.Errorf("failed to load eBPF tracer: %w", err)
@@ -125,6 +127,13 @@ func (c *Controller) Start(ctx context.Context) error {
 		log.Printf("Enabled off-cpu profiling with p=%f", c.config.OffCPUThreshold)
 	}
 
+	if len(c.config.UProbeLinks) > 0 {
+		if err := trc.AttachUProbes(c.config.UProbeLinks); err != nil {
+			return fmt.Errorf("failed to attach uprobes: %v", err)
+		}
+		log.Printf("Attached uprobes")
+	}
+
 	if c.config.ProbabilisticThreshold < tracer.ProbabilisticThresholdMax {
 		trc.StartProbabilisticProfiling(ctx)
 		log.Printf("Enabled probabilistic profiling")

reporter/base_reporter.go

@@ -76,8 +76,11 @@ func (b *baseReporter) ExecutableMetadata(args *ExecutableMetadataArgs) {
 }
 
 func (b *baseReporter) ReportTraceEvent(trace *libpf.Trace, meta *samples.TraceEventMeta) error {
-	if meta.Origin != support.TraceOriginSampling && meta.Origin != support.TraceOriginOffCPU {
-		// At the moment only on-CPU and off-CPU traces are reported.
+	switch meta.Origin {
+	case support.TraceOriginSampling:
+	case support.TraceOriginOffCPU:
+	case support.TraceOriginUProbe:
+	default:
 		return fmt.Errorf("skip reporting trace for %d origin: %w", meta.Origin,
 			errUnknownOrigin)
 	}

reporter/internal/pdata/generate.go

@@ -64,6 +64,7 @@ func (p *Pdata) Generate(tree samples.TraceEventsTree,
 		for _, origin := range []libpf.Origin{
 			support.TraceOriginSampling,
 			support.TraceOriginOffCPU,
+			support.TraceOriginUProbe,
 		} {
 			if len(originToEvents[origin]) == 0 {
 				// Do not append empty profiles.
@@ -125,6 +126,9 @@ func (p *Pdata) setProfile(
 	case support.TraceOriginOffCPU:
 		st.SetTypeStrindex(stringSet.Add("events"))
 		st.SetUnitStrindex(stringSet.Add("nanoseconds"))
+	case support.TraceOriginUProbe:
+		st.SetTypeStrindex(stringSet.Add("events"))
+		st.SetUnitStrindex(stringSet.Add("count"))
 	default:
 		// Should never happen
 		return fmt.Errorf("generating profile for unsupported origin %d", origin)
@@ -149,6 +153,8 @@ func (p *Pdata) setProfile(
 			sample.Value().Append(1)
 		case support.TraceOriginOffCPU:
 			sample.Value().Append(traceInfo.OffTimes...)
+		case support.TraceOriginUProbe:
+			sample.Value().Append(1)
 		}
 
 		// Walk every frame of the trace.

support/ebpf/off_cpu.ebpf.c

@@ -51,10 +51,10 @@ int tracepoint__sched_switch(UNUSED void *ctx)
   return 0;
 }
 
-// dummy is never loaded or called. It just makes sure kprobe_progs is
+// kprobe__dummy is never loaded or called. It just makes sure kprobe_progs is
 // referenced and make the compiler and linker happy.
 SEC("kprobe/dummy")
-int dummy(struct pt_regs *ctx)
+int kprobe__dummy(struct pt_regs *ctx)
 {
   bpf_tail_call(ctx, &kprobe_progs, 0);
   return 0;

support/ebpf/types.h

@@ -344,6 +344,7 @@ typedef enum TraceOrigin {
   TRACE_UNKNOWN,
   TRACE_SAMPLING,
   TRACE_OFF_CPU,
+  TRACE_UPROBE,
 } TraceOrigin;
 
 // MAX_FRAME_UNWINDS defines the maximum number of frames per

support/ebpf/uprobe.ebpf.c

@@ -0,0 +1,20 @@
+#include "bpfdefs.h"
+#include "tracemgmt.h"
+#include "types.h"
+
+// uprobe__generic serves as entry point for uprobe based profiling.
+SEC("uprobe/generic")
+int uprobe__generic(void *ctx)
+{
+  u64 pid_tgid = bpf_get_current_pid_tgid();
+  u32 pid      = pid_tgid >> 32;
+  u32 tid      = pid_tgid & 0xFFFFFFFF;
+
+  if (pid == 0 || tid == 0) {
+    return 0;
+  }
+
+  u64 ts = bpf_ktime_get_ns();
+
+  return collect_trace(ctx, TRACE_UPROBE, pid, tid, ts, 0);
+}

support/types_def.go

@@ -98,6 +98,7 @@ const (
 	TraceOriginUnknown  = C.TRACE_UNKNOWN
 	TraceOriginSampling = C.TRACE_SAMPLING
 	TraceOriginOffCPU   = C.TRACE_OFF_CPU
+	TraceOriginUProbe   = C.TRACE_UPROBE
 )
 
 type ApmSpanID C.ApmSpanID

tracer/tracer.go

@@ -141,6 +141,12 @@ type Config struct {
 	// IncludeEnvVars holds a list of environment variables that should be captured and reported
 	// from processes
 	IncludeEnvVars libpf.Set[string]
+	// UProbes holds a list of executable:symbol elements to which
+	// a uprobe will be attached.
+	UProbeLinks []string
+	// LoadProbe inidicates whether the generic eBPF program should be loaded
+	// without being attached to something.
+	LoadProbe bool
 }
 
 // hookPoint specifies the group and name of the hooked point in the kernel.
@@ -389,13 +395,47 @@ func initializeMapsAndPrograms(kmod *kallsyms.Module, cfg *Config) (
 		return nil, nil, fmt.Errorf("failed to load perf eBPF programs: %v", err)
 	}
 
+	if cfg.OffCPUThreshold > 0 || len(cfg.UProbeLinks) > 0 || cfg.LoadProbe {
+		// Load the tail call destinations if any kind of event profiling is enabled.
+		if err = loadProbeUnwinders(coll, ebpfProgs, ebpfMaps["kprobe_progs"], tailCallProgs,
+			cfg.BPFVerifierLogLevel, ebpfMaps["perf_progs"].FD()); err != nil {
+			return nil, nil, fmt.Errorf("failed to load kprobe eBPF programs: %v", err)
+		}
+	}
+
 	if cfg.OffCPUThreshold > 0 {
-		if err = loadKProbeUnwinders(coll, ebpfProgs, ebpfMaps["kprobe_progs"], tailCallProgs,
+		offCPUProgs := []progLoaderHelper{
+			{
+				name:             "finish_task_switch",
+				noTailCallTarget: true,
+				enable:           true,
+			},
+			{
+				name:             "tracepoint__sched_switch",
+				noTailCallTarget: true,
+				enable:           true,
+			},
+		}
+		if err = loadProbeUnwinders(coll, ebpfProgs, ebpfMaps["kprobe_progs"], offCPUProgs,
 			cfg.BPFVerifierLogLevel, ebpfMaps["perf_progs"].FD()); err != nil {
 			return nil, nil, fmt.Errorf("failed to load kprobe eBPF programs: %v", err)
 		}
 	}
 
+	if len(cfg.UProbeLinks) > 0 || cfg.LoadProbe {
+		uprobeProgs := []progLoaderHelper{
+			{
+				name:             "uprobe__generic",
+				noTailCallTarget: true,
+				enable:           true,
+			},
+		}
+		if err = loadProbeUnwinders(coll, ebpfProgs, ebpfMaps["kprobe_progs"], uprobeProgs,
+			cfg.BPFVerifierLogLevel, ebpfMaps["perf_progs"].FD()); err != nil {
+			return nil, nil, fmt.Errorf("failed to load uprobe eBPF programs: %v", err)
+		}
+	}
+
 	if err = loadSystemConfig(coll, ebpfMaps, kmod, cfg.IncludeTracers,
 		cfg.OffCPUThreshold, cfg.FilterErrorFrames); err != nil {
 		return nil, nil, fmt.Errorf("failed to load system config: %v", err)
@@ -558,31 +598,16 @@ func progArrayReferences(perfTailCallMapFD int, insns asm.Instructions) []int {
 	return insNos
 }
 
-// loadKProbeUnwinders reuses large parts of loadPerfUnwinders. By default all eBPF programs
-// are written as perf event eBPF programs. loadKProbeUnwinders dynamically rewrites the
-// specification of these programs to kprobe eBPF programs and adjusts tail call maps.
-func loadKProbeUnwinders(coll *cebpf.CollectionSpec, ebpfProgs map[string]*cebpf.Program,
-	tailcallMap *cebpf.Map, tailCallProgs []progLoaderHelper,
+// loadProbeUnwinders reuses large parts of loadPerfUnwinders. By default all eBPF programs
+// are written as perf event eBPF programs. loadProbeUnwinders dynamically rewrites the
+// specification of these programs to xProbe eBPF programs and adjusts tail call maps.
+func loadProbeUnwinders(coll *cebpf.CollectionSpec, ebpfProgs map[string]*cebpf.Program,
+	tailcallMap *cebpf.Map, progs []progLoaderHelper,
 	bpfVerifierLogLevel uint32, perfTailCallMapFD int) error {
 	programOptions := cebpf.ProgramOptions{
 		LogLevel: cebpf.LogLevel(bpfVerifierLogLevel),
 	}
 
-	progs := make([]progLoaderHelper, len(tailCallProgs)+2)
-	copy(progs, tailCallProgs)
-	progs = append(progs,
-		progLoaderHelper{
-			name:             "finish_task_switch",
-			noTailCallTarget: true,
-			enable:           true,
-		},
-		progLoaderHelper{
-			name:             "tracepoint__sched_switch",
-			noTailCallTarget: true,
-			enable:           true,
-		},
-	)
-
 	for _, unwindProg := range progs {
 		if !unwindProg.enable {
 			continue
@@ -873,7 +898,11 @@ func (t *Tracer) loadBpfTrace(raw []byte, cpu int) *host.Trace {
 		EnvVars:          procMeta.EnvVariables,
 	}
 
-	if trace.Origin != support.TraceOriginSampling && trace.Origin != support.TraceOriginOffCPU {
+	switch trace.Origin {
+	case support.TraceOriginSampling:
+	case support.TraceOriginOffCPU:
+	case support.TraceOriginUProbe:
+	default:
 		log.Warnf("Skip handling trace from unexpected %d origin", trace.Origin)
 		return nil
 	}
@@ -1120,6 +1149,27 @@ func (t *Tracer) StartOffCPUProfiling() error {
 	return nil
 }
 
+func (t *Tracer) AttachUProbes(uprobes []string) error {
+	uProbeProg, ok := t.ebpfProgs["uprobe__generic"]
+	if !ok {
+		return errors.New("uprobe__generic is not available")
+	}
+	for _, uprobeStr := range uprobes {
+		split := strings.SplitN(uprobeStr, ":", 2)
+
+		exec, err := link.OpenExecutable(split[0])
+		if err != nil {
+			return err
+		}
+		uprobeLink, err := exec.Uprobe(split[1], uProbeProg, nil)
+		if err != nil {
+			return err
+		}
+		t.hooks[hookPoint{group: "uprobe", name: uprobeStr}] = uprobeLink
+	}
+	return nil
+}
+
 // TraceProcessor gets the trace processor.
 func (t *Tracer) TraceProcessor() tracehandler.TraceProcessor {
 	return t.processManager