From 14ad5651194e5b558bd8ff7eacb614742e9ac8d9 Mon Sep 17 00:00:00 2001
From: Seoyoung Park <se0y0ung@unist.ac.kr>
Date: Wed, 26 May 2021 20:50:33 -0700
Subject: [PATCH] Add kprobeEvOpts EventNameOption in Module for unique
 eventName per Module

---
 elf/module.go             | 38 ++++++++++++++++++++++++++++++++++----
 elf/module_unsupported.go |  4 ++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/elf/module.go b/elf/module.go
index 28b7047c..4ff667d2 100644
--- a/elf/module.go
+++ b/elf/module.go
@@ -138,6 +138,8 @@ type Module struct {
 	schedPrograms      map[string]*SchedProgram
 	xdpPrograms        map[string]*XDPProgram
 
+	kprobeEvOpts EventNameOption
+
 	compatProbe bool // try to be automatically convert function names depending on kernel versions (SyS_ and __x64_sys_)
 }
 
@@ -150,6 +152,10 @@ type Kprobe struct {
 	efd   int
 }
 
+// EventNameOption is a function type which take eventName
+// then return the new eventName string
+type EventNameOption func(eventName string) string
+
 type Uprobe struct {
 	Name  string
 	insns *C.struct_bpf_insn
@@ -331,6 +337,30 @@ func (b *Module) EnableOptionCompatProbe() {
 	b.compatProbe = true
 }
 
+// TrySetKprobeEventNameOpts Try to set kprobe EventNameOption function.
+// If Module.kprobeEvOpts already set, return error.
+func (b *Module) TrySetKprobeEventNameOpts(evOpts EventNameOption) error {
+	if b.kprobeEvOpts != nil {
+		return fmt.Errorf("kprobe eventname option is already set. this might corrupt")
+	}
+	// To prevent any error, detour safeEventName function to remove illegal characters.
+	b.kprobeEvOpts = func(eventName string) string {
+		return safeEventName(evOpts(eventName))
+	}
+	return nil
+}
+
+func (b *Module) kprobeEventName(originalEventName string) string {
+	var eventName string
+	if b.kprobeEvOpts != nil {
+		eventName = b.kprobeEvOpts(originalEventName)
+	} else {
+		eventName = originalEventName
+	}
+	return eventName
+}
+
+
 // EnableKprobe enables a kprobe/kretprobe identified by secName.
 // For kretprobes, you can configure the maximum number of instances
 // of the function that can be probed simultaneously with maxactive.
@@ -356,7 +386,7 @@ func (b *Module) EnableKprobe(secName string, maxactive int) error {
 		probeType = "p"
 		funcName = strings.TrimPrefix(secName, "kprobe/")
 	}
-	eventName := probeType + funcName
+	eventName := b.kprobeEventName(probeType + funcName)
 
 	kprobeId, err := writeKprobeEvent(probeType, eventName, funcName, maxactiveStr)
 	// fallback without maxactive
@@ -745,10 +775,10 @@ func (b *Module) closeProbes() error {
 		var err error
 		if isKretprobe {
 			funcName = strings.TrimPrefix(name, "kretprobe/")
-			err = disableKprobe("r" + funcName)
+			err = disableKprobe(b.kprobeEventName("r" + funcName))
 		} else {
 			funcName = strings.TrimPrefix(name, "kprobe/")
-			err = disableKprobe("p" + funcName)
+			err = disableKprobe(b.kprobeEventName("p" + funcName))
 		}
 		if err != nil {
 			return fmt.Errorf("error clearing probe: %v", err)
@@ -923,4 +953,4 @@ func (b *Module) CloseExt(options map[string]CloseOptions) error {
 		return err
 	}
 	return nil
-}
+}
\ No newline at end of file
diff --git a/elf/module_unsupported.go b/elf/module_unsupported.go
index 975b68b2..ecbffc42 100644
--- a/elf/module_unsupported.go
+++ b/elf/module_unsupported.go
@@ -68,6 +68,10 @@ func (b *Module) EnableTracepoint(secName string) error {
 	return errNotSupported
 }
 
+func (b *Module) TrySetKprobeEventNameOpts(evOpts EventNameOption) error {
+	return errNotSupported
+}
+
 func (b *Module) IterMaps() <-chan *Map {
 	return nil
 }