From 409fc489cd23371e4e56fdf96b6aade79a90ff50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guillermo=20Juli=C3=A1n?= Date: Thu, 26 Dec 2024 12:57:56 +0100 Subject: [PATCH] [EBPF] Add BeforeStop/AfterStop method to modifiers (#32453) --- pkg/ebpf/helper_call_patcher.go | 6 +- pkg/ebpf/manager.go | 114 +++++++++++++++++++++++++------- pkg/ebpf/manager_test.go | 68 ++++++++++++++++--- pkg/ebpf/printk_patcher.go | 8 +-- 4 files changed, 155 insertions(+), 41 deletions(-) diff --git a/pkg/ebpf/helper_call_patcher.go b/pkg/ebpf/helper_call_patcher.go index b274d752270a9..380bbab529aeb 100644 --- a/pkg/ebpf/helper_call_patcher.go +++ b/pkg/ebpf/helper_call_patcher.go @@ -40,7 +40,7 @@ var replaceIns = asm.Mov.Imm(asm.R0, 0) // conditionally select eBPF helpers. This should be regarded as a last resort // when the aforementioned options don't apply (prebuilt artifacts, for // example). -func NewHelperCallRemover(helpers ...asm.BuiltinFunc) Modifier { +func NewHelperCallRemover(helpers ...asm.BuiltinFunc) ModifierBeforeInit { return &helperCallRemover{ helpers: helpers, } @@ -81,10 +81,6 @@ func (h *helperCallRemover) BeforeInit(m *manager.Manager, _ names.ModuleName, _ return nil } -func (h *helperCallRemover) AfterInit(*manager.Manager, names.ModuleName, *manager.Options) error { - return nil -} - func (h *helperCallRemover) String() string { return fmt.Sprintf("HelperCallRemover[%+v]", h.helpers) } diff --git a/pkg/ebpf/manager.go b/pkg/ebpf/manager.go index 06e790609a9ec..0cf37a7e1891a 100644 --- a/pkg/ebpf/manager.go +++ b/pkg/ebpf/manager.go @@ -8,6 +8,7 @@ package ebpf import ( + "errors" "fmt" "io" "sync" @@ -51,27 +52,76 @@ func NewManagerWithDefault(mgr *manager.Manager, name string, modifiers ...Modif return NewManager(mgr, name, append(defaultModifiers, modifiers...)...) } -// Modifier is an interface that can be implemented by a package to -// add functionality to the ebpf.Manager. It exposes a name to identify the modifier, -// two functions that will be called before and after the ebpf.Manager.InitWithOptions -// call, and a function that will be called when the manager is stopped. -// Note regarding internal state of the modifier: if the modifier is added to the list of modifiers -// enabled by default (pkg/ebpf/ebpf.go:registerDefaultModifiers), all managers with those default modifiers -// will share the same instance of the modifier. On the other hand, if the modifier is added to a specific -// manager, it can have its own instance of the modifier, unless the caller explicitly uses the same modifier -// instance with different managers. In other words, if the modifier is to have any internal state specific to -// each manager, it should not be added to the list of default modifiers, and developers using it -// should be aware of this behavior. +// Modifier is an interface that can be implemented by a package to add +// functionality to the ebpf.Manager. It exposes a name to identify the +// modifier, and then any of the functions Before/AfterInit, Before/AfterStart, +// Before/AfterStop, that will be called at the corresponding stage of the +// manager lifecycle. To avoid code churn and implementing unnecessary +// functions, the Modifier interface is split into sub-interfaces, each with a +// single function. This way, the developer can implement only the functions +// they need, and the manager will call them at the right time. Note regarding +// internal state of the modifier: if the modifier is added to the list of +// modifiers enabled by default (see NewManagerWithDefault above), all managers +// with those default modifiers will share the same instance of the modifier. On +// the other hand, if the modifier is added to a specific manager, it can have +// its own instance of the modifier, unless the caller explicitly uses the same +// modifier instance with different managers. In other words, if the modifier is +// to have any internal state specific to each manager, it should not be added +// to the list of default modifiers, and developers using it should be aware of +// this behavior. type Modifier interface { fmt.Stringer +} + +// ModifierBeforeInit is a sub-interface of Modifier that exposes a BeforeInit method +type ModifierBeforeInit interface { + Modifier + // BeforeInit is called before the ebpf.Manager.InitWithOptions call - // names.ModuleName refers to the name associated with Manager instance. + // names.ModuleName refers to the name associated with Manager instance. An + // error returned from this function will stop the initialization process. BeforeInit(*manager.Manager, names.ModuleName, *manager.Options) error +} + +// ModifierAfterInit is a sub-interface of Modifier that exposes an AfterInit method +type ModifierAfterInit interface { + Modifier // AfterInit is called after the ebpf.Manager.InitWithOptions call AfterInit(*manager.Manager, names.ModuleName, *manager.Options) error } +// ModifierBeforeStop is a sub-interface of Modifier that exposes a BeforeStop method +type ModifierBeforeStop interface { + Modifier + + // BeforeStop is called before the ebpf.Manager.Stop call. An error returned + // from this function will not prevent the manager from stopping, but it will + // be logged. + BeforeStop(*manager.Manager, names.ModuleName, manager.MapCleanupType) error +} + +// ModifierAfterStop is a sub-interface of Modifier that exposes an AfterStop method +type ModifierAfterStop interface { + Modifier + + // AfterStop is called after the ebpf.Manager.Stop call. An error returned + // from this function will be logged. + AfterStop(*manager.Manager, names.ModuleName, manager.MapCleanupType) error +} + +func runModifiersOfType[K Modifier](modifiers []Modifier, funcName string, runner func(K) error) error { + var errs error + for _, mod := range modifiers { + if as, ok := mod.(K); ok { + if err := runner(as); err != nil { + errs = errors.Join(errs, fmt.Errorf("error running %s manager modifier %s: %w", mod, funcName, err)) + } + } + } + return errs +} + // InitWithOptions is a wrapper around ebpf-manager.Manager.InitWithOptions func (m *Manager) InitWithOptions(bytecode io.ReaderAt, opts *manager.Options) error { // we must load the ELF file before initialization, @@ -81,22 +131,40 @@ func (m *Manager) InitWithOptions(bytecode io.ReaderAt, opts *manager.Options) e return fmt.Errorf("failed to load elf from reader: %w", err) } - for _, mod := range m.EnabledModifiers { - log.Tracef("Running %s manager modifier BeforeInit", mod) - if err := mod.BeforeInit(m.Manager, m.Name, opts); err != nil { - return fmt.Errorf("error running %s manager modifier: %w", mod, err) - } + err := runModifiersOfType(m.EnabledModifiers, "BeforeInit", func(mod ModifierBeforeInit) error { + return mod.BeforeInit(m.Manager, m.Name, opts) + }) + if err != nil { + return err } if err := m.Manager.InitWithOptions(nil, *opts); err != nil { return err } - for _, mod := range m.EnabledModifiers { - log.Tracef("Running %s manager modifier AfterInit", mod) - if err := mod.AfterInit(m.Manager, m.Name, opts); err != nil { - return fmt.Errorf("error running %s manager modifier: %w", mod, err) - } + return runModifiersOfType(m.EnabledModifiers, "AfterInit", func(mod ModifierAfterInit) error { + return mod.AfterInit(m.Manager, m.Name, opts) + }) +} + +// Stop is a wrapper around ebpf-manager.Manager.Stop +func (m *Manager) Stop(cleanupType manager.MapCleanupType) error { + var errs error + + err := runModifiersOfType(m.EnabledModifiers, "BeforeStop", func(mod ModifierBeforeStop) error { + return mod.BeforeStop(m.Manager, m.Name, cleanupType) + }) + if err != nil { + errs = errors.Join(errs, err) + } + + if err := m.Manager.Stop(cleanupType); err != nil { + errs = errors.Join(errs, fmt.Errorf("failed to stop manager %w", err)) } - return nil + + err = runModifiersOfType(m.EnabledModifiers, "AfterStop", func(mod ModifierAfterStop) error { + return mod.AfterStop(m.Manager, m.Name, cleanupType) + }) + + return errors.Join(errs, err) } diff --git a/pkg/ebpf/manager_test.go b/pkg/ebpf/manager_test.go index 32018e808fb0e..0a03659dadfa1 100644 --- a/pkg/ebpf/manager_test.go +++ b/pkg/ebpf/manager_test.go @@ -9,29 +9,44 @@ package ebpf import ( "testing" - "github.com/DataDog/datadog-agent/pkg/ebpf/names" manager "github.com/DataDog/ebpf-manager" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + "github.com/DataDog/datadog-agent/pkg/ebpf/bytecode" + "github.com/DataDog/datadog-agent/pkg/ebpf/names" ) -// PrintkPatcherModifier adds an InstructionPatcher to the manager that removes the newline character from log_debug calls if needed type dummyModifier struct { + mock.Mock } const dummyModifierName = "DummyModifier" func (t *dummyModifier) String() string { + // Do not mock this method for simplicity, to avoid having to define it always return dummyModifierName } -// BeforeInit adds the patchPrintkNewline function to the manager -func (t *dummyModifier) BeforeInit(_ *manager.Manager, _ names.ModuleName, _ *manager.Options) error { - return nil +func (t *dummyModifier) BeforeInit(m *manager.Manager, name names.ModuleName, opts *manager.Options) error { + args := t.Called(m, name, opts) + return args.Error(0) +} + +func (t *dummyModifier) AfterInit(m *manager.Manager, name names.ModuleName, opts *manager.Options) error { + args := t.Called(m, name, opts) + return args.Error(0) } -// AfterInit is a no-op for this modifier -func (t *dummyModifier) AfterInit(_ *manager.Manager, _ names.ModuleName, _ *manager.Options) error { - return nil +func (t *dummyModifier) BeforeStop(m *manager.Manager, name names.ModuleName, cleanupType manager.MapCleanupType) error { + args := t.Called(m, name, cleanupType) + return args.Error(0) +} + +func (t *dummyModifier) AfterStop(m *manager.Manager, name names.ModuleName, cleanupType manager.MapCleanupType) error { + args := t.Called(m, name, cleanupType) + return args.Error(0) } func TestNewManagerWithDefault(t *testing.T) { @@ -69,3 +84,40 @@ func TestNewManagerWithDefault(t *testing.T) { }) } } + +func TestManagerInitWithOptions(t *testing.T) { + modifier := &dummyModifier{} + modifier.On("BeforeInit", mock.Anything, mock.Anything, mock.Anything).Return(nil) + modifier.On("AfterInit", mock.Anything, mock.Anything, mock.Anything).Return(nil) + + mgr := NewManager(&manager.Manager{}, "test", modifier) + require.NotNil(t, mgr) + + // Load a simple eBPF program to test the modifiers + cfg := NewConfig() + require.NotNil(t, cfg) + + buf, err := bytecode.GetReader(cfg.BPFDir, "logdebug-test.o") + require.NoError(t, err) + t.Cleanup(func() { _ = buf.Close }) + + opts := manager.Options{} + err = mgr.InitWithOptions(buf, &opts) + require.NoError(t, err) + + modifier.AssertExpectations(t) +} + +func TestManagerStop(t *testing.T) { + modifier := &dummyModifier{} + modifier.On("BeforeStop", mock.Anything, mock.Anything, mock.Anything).Return(nil) + modifier.On("AfterStop", mock.Anything, mock.Anything, mock.Anything).Return(nil) + + mgr := NewManager(&manager.Manager{}, "test", modifier) + require.NotNil(t, mgr) + + // The Stop call will fail because the manager is not initialized, but the modifiers should still be called + _ = mgr.Stop(manager.CleanAll) + + modifier.AssertExpectations(t) +} diff --git a/pkg/ebpf/printk_patcher.go b/pkg/ebpf/printk_patcher.go index afda0f4869c6a..b417c8b011415 100644 --- a/pkg/ebpf/printk_patcher.go +++ b/pkg/ebpf/printk_patcher.go @@ -213,6 +213,9 @@ func patchPrintkInstructions(p *ebpf.ProgramSpec) (int, error) { type PrintkPatcherModifier struct { } +// ensure PrintkPatcherModifier implements the ModifierBeforeInit interface +var _ ModifierBeforeInit = &PrintkPatcherModifier{} + func (t *PrintkPatcherModifier) String() string { return "PrintkPatcherModifier" } @@ -222,8 +225,3 @@ func (t *PrintkPatcherModifier) BeforeInit(m *manager.Manager, _ names.ModuleNam m.InstructionPatchers = append(m.InstructionPatchers, patchPrintkNewline) return nil } - -// AfterInit is a no-op for this modifier -func (t *PrintkPatcherModifier) AfterInit(_ *manager.Manager, _ names.ModuleName, _ *manager.Options) error { - return nil -}