add program bypass support (#184)

* add program bypass support

* be more strict about start offset

* sanity check on registers used before inject

manager.go

@@ -11,7 +11,9 @@ import (
 	"time"
 
 	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/asm"
 	"github.com/cilium/ebpf/btf"
+	"github.com/cilium/ebpf/features"
 	"golang.org/x/sys/unix"
 )
 
@@ -126,6 +128,9 @@ type Options struct {
 
 	// KernelModuleBTFLoadFunc is a function to provide custom loading of BTF for kernel modules on-demand as programs are loaded
 	KernelModuleBTFLoadFunc func(kmodName string) (*btf.Spec, error)
+
+	// BypassEnabled controls whether program bypass is enabled for this manager
+	BypassEnabled bool
 }
 
 // InstructionPatcherFunc - A function that patches the instructions of a program
@@ -139,6 +144,8 @@ type Manager struct {
 	netlinkSocketCache *netlinkSocketCache
 	state              state
 	stateLock          sync.RWMutex
+	bypassIndexes      map[string]uint32
+	maxBypassIndex     uint32
 
 	// Probes - List of probes handled by the manager
 	Probes []*Probe
@@ -510,6 +517,14 @@ func (m *Manager) InitWithOptions(elf io.ReaderAt, options Options) error {
 			i++
 		}
 	}
+
+	// must run before map exclusion in case bypass is disabled
+	bypassMap, err := m.setupBypass()
+	if err != nil {
+		m.stateLock.Unlock()
+		return err
+	}
+
 	// Remove excluded maps
 	for _, excludeMapName := range m.options.ExcludedMaps {
 		delete(m.collectionSpec.Maps, excludeMapName)
@@ -530,6 +545,18 @@ func (m *Manager) InitWithOptions(elf io.ReaderAt, options Options) error {
 
 	// Configure activated probes
 	m.activateProbes()
+
+	// populate bypass indexes on Probe objects
+	// this must run after matchSpecs due to CopyProgram handling
+	if bypassMap != nil {
+		for _, mProbe := range m.Probes {
+			if idx, ok := m.bypassIndexes[mProbe.GetEBPFFuncName()]; ok {
+				mProbe.bypassIndex = idx
+				mProbe.bypassMap = bypassMap
+			}
+		}
+	}
+
 	m.state = initialized
 	m.stateLock.Unlock()
 	resetManager := func(m *Manager) {
@@ -618,6 +645,128 @@ func (m *Manager) InitWithOptions(elf io.ReaderAt, options Options) error {
 	return nil
 }
 
+func (m *Manager) setupBypass() (*Map, error) {
+	_, hasBypassMapSpec := m.collectionSpec.Maps[bypassMapName]
+	if !hasBypassMapSpec {
+		return nil, nil
+	}
+	if !m.options.BypassEnabled {
+		m.options.ExcludedMaps = append(m.options.ExcludedMaps, bypassMapName)
+		return nil, nil
+	}
+	// start with 1, so we know if programs even have a valid index set
+	m.maxBypassIndex = 1
+
+	const stackOffset = -8
+	// place a limit on how far we will inject from the start of a program
+	// otherwise we aren't sure what register we need to save/restore, and it could inflate the number of instructions.
+	const maxInstructionOffsetFromProgramStart = 1
+	// setup bypass constants for all programs
+	m.bypassIndexes = make(map[string]uint32, len(m.collectionSpec.Programs))
+	for name, p := range m.collectionSpec.Programs {
+		for i := 0; i < len(p.Instructions); i++ {
+			ins := p.Instructions[i]
+			if ins.Reference() != bypassOptInReference {
+				continue
+			}
+			// return error here to ensure we only error on programs that do have a bypass reference
+			if i > maxInstructionOffsetFromProgramStart {
+				return nil, fmt.Errorf("unable to inject bypass instructions into program %s: bypass reference occurs too late in program", name)
+			}
+			if i > 0 && p.Instructions[i-1].Src != asm.R1 {
+				return nil, fmt.Errorf("unable to inject bypass instructions into program %s: register other than r1 used before injection point", name)
+			}
+
+			m.bypassIndexes[name] = m.maxBypassIndex
+			newInsns := append([]asm.Instruction{
+				asm.Mov.Reg(asm.R6, asm.R1),
+				// save bypass index to stack
+				asm.StoreImm(asm.RFP, stackOffset, int64(m.maxBypassIndex), asm.Word),
+				// store pointer to bypass index
+				asm.Mov.Reg(asm.R2, asm.RFP),
+				asm.Add.Imm(asm.R2, stackOffset),
+				// load map reference
+				asm.LoadMapPtr(asm.R1, 0).WithReference(bypassMapName),
+				// bpf_map_lookup_elem
+				asm.FnMapLookupElem.Call(),
+				// if ret == 0, jump to `return 0`
+				{
+					OpCode:   asm.JEq.Op(asm.ImmSource),
+					Dst:      asm.R0,
+					Offset:   3, // jump TO return
+					Constant: int64(0),
+				},
+				// pointer indirection of result from map lookup
+				asm.LoadMem(asm.R1, asm.R0, 0, asm.Word),
+				// if bypass NOT enabled, jump over return
+				{
+					OpCode:   asm.JEq.Op(asm.ImmSource),
+					Dst:      asm.R1,
+					Offset:   2, // jump over return on next instruction
+					Constant: int64(0),
+				},
+				asm.Return(),
+				// zero out used stack slot
+				asm.StoreImm(asm.RFP, stackOffset, 0, asm.Word),
+				asm.Mov.Reg(asm.R1, asm.R6),
+			}, p.Instructions[i+1:]...)
+			// necessary to keep kernel happy about source information for start of program
+			newInsns[0] = newInsns[0].WithSource(ins.Source())
+			p.Instructions = append(p.Instructions[:i], newInsns...)
+			m.maxBypassIndex += 1
+			break
+		}
+	}
+	// no programs modified
+	if m.maxBypassIndex == 1 {
+		m.options.ExcludedMaps = append(m.options.ExcludedMaps, bypassMapName)
+		return nil, nil
+	}
+
+	hasPerCPU := false
+	if err := features.HaveMapType(ebpf.PerCPUArray); err == nil {
+		hasPerCPU = true
+	}
+
+	bypassMap := &Map{Name: bypassMapName}
+	m.Maps = append(m.Maps, bypassMap)
+
+	if m.options.MapSpecEditors == nil {
+		m.options.MapSpecEditors = make(map[string]MapSpecEditor)
+	}
+	m.options.MapSpecEditors[bypassMapName] = MapSpecEditor{
+		MaxEntries: m.maxBypassIndex + 1,
+		EditorFlag: EditMaxEntries,
+	}
+
+	if !hasPerCPU {
+		// use scalar value for bypass/enable
+		bypassValue = 1
+		enableValue = 0
+		return bypassMap, nil
+	}
+
+	// upgrade map type to per-cpu, if available
+	specEditor := m.options.MapSpecEditors[bypassMapName]
+	specEditor.Type = ebpf.PerCPUArray
+	specEditor.EditorFlag |= EditType
+	m.options.MapSpecEditors[bypassMapName] = specEditor
+
+	// allocate per-cpu slices used for bypass/enable
+	cpus, err := ebpf.PossibleCPU()
+	if err != nil {
+		return nil, err
+	}
+	if bypassValue == nil {
+		bypassValue = makeAndSet(cpus, uint32(1))
+	}
+	if enableValue == nil {
+		enableValue = makeAndSet(cpus, uint32(0))
+	}
+
+	return bypassMap, nil
+}
+
 // Start - Attach eBPF programs, start perf ring readers and apply maps and tail calls routing.
 func (m *Manager) Start() error {
 	m.stateLock.Lock()
@@ -716,6 +865,10 @@ func (m *Manager) Pause() error {
 	if m.state <= initialized {
 		return ErrManagerNotStarted
 	}
+	if !m.options.BypassEnabled {
+		return nil
+	}
+
 	for _, probe := range m.Probes {
 		if err := probe.Pause(); err != nil {
 			return err
@@ -734,6 +887,10 @@ func (m *Manager) Resume() error {
 	if m.state <= initialized {
 		return ErrManagerNotStarted
 	}
+	if !m.options.BypassEnabled {
+		return nil
+	}
+
 	for _, probe := range m.Probes {
 		if err := probe.Resume(); err != nil {
 			return err
@@ -908,6 +1065,19 @@ func (m *Manager) AddHook(UID string, newProbe *Probe) error {
 	newProbe.program = clonedProg
 	newProbe.programSpec = progSpec
 
+	var bypassMap *Map
+	for _, mp := range m.Maps {
+		if mp.Name == bypassMapName {
+			bypassMap = mp
+			break
+		}
+	}
+	bypassIndex, ok := m.bypassIndexes[newProbe.EBPFFuncName]
+	if ok && bypassMap != nil {
+		newProbe.bypassIndex = bypassIndex
+		newProbe.bypassMap = bypassMap
+	}
+
 	// init program
 	if err = newProbe.init(m); err != nil {
 		// clean up

perf_event.go

@@ -85,14 +85,6 @@ func (pe *perfEventLink) Close() error {
 	return pe.fd.Close()
 }
 
-func (pe *perfEventLink) Pause() error {
-	return ioctlPerfEventDisable(pe.fd)
-}
-
-func (pe *perfEventLink) Resume() error {
-	return ioctlPerfEventEnable(pe.fd)
-}
-
 func attachPerfEvent(pe *perfEventLink, prog *ebpf.Program) error {
 	if err := ioctlPerfEventSetBPF(pe.fd, prog.FD()); err != nil {
 		return fmt.Errorf("set perf event bpf: %w", err)
@@ -211,7 +203,3 @@ func ioctlPerfEventSetBPF(perfEventOpenFD *fd, progFD int) error {
 func ioctlPerfEventEnable(perfEventOpenFD *fd) error {
 	return unix.IoctlSetInt(int(perfEventOpenFD.raw), unix.PERF_EVENT_IOC_ENABLE, 0)
 }
-
-func ioctlPerfEventDisable(perfEventOpenFD *fd) error {
-	return unix.IoctlSetInt(int(perfEventOpenFD.raw), unix.PERF_EVENT_IOC_DISABLE, 0)
-}

probe.go

@@ -24,6 +24,13 @@ const (
 	AttachWithProbeEvents
 )
 
+const bypassMapName = "program_bypassed"
+const bypassOptInReference = "bypass_program"
+
+// values used for map update to bypass/enable programs
+var bypassValue any
+var enableValue any
+
 // Probe - Main eBPF probe wrapper. This structure is used to store the required data to attach a loaded eBPF
 // program to its hook point.
 type Probe struct {
@@ -45,6 +52,8 @@ type Probe struct {
 	tcFilter                netlink.BpfFilter
 	tcClsActQdisc           netlink.Qdisc
 	progLink                io.Closer
+	bypassIndex             uint32
+	bypassMap               *Map
 
 	// lastError - stores the last error that the probe encountered, it is used to surface a more useful error message
 	// when one of the validators (see Options.ActivatedProbes) fails.
@@ -583,16 +592,22 @@ func (p *Probe) pause() error {
 	}
 
 	v, ok := p.progLink.(pauser)
-	if !ok {
-		return fmt.Errorf("pause not supported for program type %s", p.programSpec.Type)
+	if ok {
+		if err := v.Pause(); err != nil {
+			p.lastError = err
+			return fmt.Errorf("error pausing probe %s: %w", p.ProbeIdentificationPair, err)
+		}
+		p.state = paused
+		return nil
 	}
 
-	if err := v.Pause(); err != nil {
-		p.lastError = err
-		return fmt.Errorf("error pausing probe %s: %w", p.ProbeIdentificationPair, err)
+	if p.bypassIndex > 0 && p.bypassMap != nil {
+		if err := p.bypassMap.array.Update(p.bypassIndex, bypassValue, ebpf.UpdateExist); err != nil {
+			return err
+		}
+		p.state = paused
+		return nil
 	}
-
-	p.state = paused
 	return nil
 }
 
@@ -604,16 +619,22 @@ func (p *Probe) resume() error {
 	}
 
 	v, ok := p.progLink.(pauser)
-	if !ok {
-		return fmt.Errorf("resume not supported for program type %s", p.programSpec.Type)
+	if ok {
+		if err := v.Resume(); err != nil {
+			p.lastError = err
+			return fmt.Errorf("error resuming probe %s: %w", p.ProbeIdentificationPair, err)
+		}
+		p.state = running
+		return nil
 	}
 
-	if err := v.Resume(); err != nil {
-		p.lastError = err
-		return fmt.Errorf("error resuming probe %s: %w", p.ProbeIdentificationPair, err)
+	if p.bypassIndex > 0 && p.bypassMap != nil {
+		if err := p.bypassMap.array.Update(p.bypassIndex, enableValue, ebpf.UpdateExist); err != nil {
+			return err
+		}
+		p.state = running
+		return nil
 	}
-
-	p.state = running
 	return nil
 }
 

utils.go

@@ -51,3 +51,12 @@ func cleanupProgramSpec(spec *ebpf.ProgramSpec) {
 		spec.Instructions = nil
 	}
 }
+
+// create slice of length n and fill with fillVal
+func makeAndSet[E any](n int, fillVal E) []E {
+	s := make([]E, n)
+	for i := range s {
+		s[i] = fillVal
+	}
+	return s
+}