From af110e89bcf2deb44c12e03dbbb855fdc64e7edd Mon Sep 17 00:00:00 2001 From: rabbitstack Date: Wed, 11 Dec 2024 18:58:23 +0100 Subject: [PATCH] chore(ps): Append/remove module by base address Instead of appending/removing the process modules by its image name, we're now using the module base address to determine if the module was already added. The same logic is used to remove the existing module. --- internal/etw/processors/image_windows.go | 4 ++-- internal/etw/processors/image_windows_test.go | 3 ++- internal/etw/source_test.go | 6 +++--- pkg/kcap/reader_windows.go | 4 ++-- pkg/ps/snapshotter.go | 2 +- pkg/ps/snapshotter_mock.go | 4 ++-- pkg/ps/snapshotter_windows.go | 4 ++-- pkg/ps/snapshotter_windows_test.go | 3 ++- pkg/ps/types/types_windows.go | 16 +++++++++++++--- 9 files changed, 29 insertions(+), 17 deletions(-) diff --git a/internal/etw/processors/image_windows.go b/internal/etw/processors/image_windows.go index 325862b0e..0b14981c9 100644 --- a/internal/etw/processors/image_windows.go +++ b/internal/etw/processors/image_windows.go @@ -44,11 +44,11 @@ func (m *imageProcessor) ProcessEvent(e *kevent.Kevent) (*kevent.Kevent, bool, e } if e.IsUnloadImage() { pid := e.Kparams.MustGetPid() - mod := e.GetParamAsString(kparams.ImagePath) + addr := e.Kparams.TryGetAddress(kparams.ImageBase) if pid == 0 { pid = e.PID } - return e, false, m.psnap.RemoveModule(pid, mod) + return e, false, m.psnap.RemoveModule(pid, addr) } if e.IsLoadImage() || e.IsImageRundown() { return e, false, m.psnap.AddModule(e) diff --git a/internal/etw/processors/image_windows_test.go b/internal/etw/processors/image_windows_test.go index 9a4911d5e..2a90e8d06 100644 --- a/internal/etw/processors/image_windows_test.go +++ b/internal/etw/processors/image_windows_test.go @@ -24,6 +24,7 @@ import ( "github.com/rabbitstack/fibratus/pkg/kevent/ktypes" "github.com/rabbitstack/fibratus/pkg/ps" "github.com/rabbitstack/fibratus/pkg/util/signature" + "github.com/rabbitstack/fibratus/pkg/util/va" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" @@ -79,7 +80,7 @@ func TestImageProcessor(t *testing.T) { }, func() *ps.SnapshotterMock { psnap := new(ps.SnapshotterMock) - psnap.On("RemoveModule", uint32(676), "C:\\Windows\\system32\\kernel32.dll").Return(nil) + psnap.On("RemoveModule", uint32(676), va.Address(0xfffb313833a3)).Return(nil) psnap.On("FindModule", mock.Anything).Return(false, nil) return psnap }, diff --git a/internal/etw/source_test.go b/internal/etw/source_test.go index db58db85a..ea92d1376 100644 --- a/internal/etw/source_test.go +++ b/internal/etw/source_test.go @@ -826,10 +826,10 @@ func (s *NoopPsSnapshotter) AddThread(kevt *kevent.Kevent) error func (s *NoopPsSnapshotter) AddModule(kevt *kevent.Kevent) error { return nil } func (s *NoopPsSnapshotter) FindModule(addr va.Address) (bool, *pstypes.Module) { return false, nil } func (s *NoopPsSnapshotter) RemoveThread(pid uint32, tid uint32) error { return nil } -func (s *NoopPsSnapshotter) RemoveModule(pid uint32, mod string) error { return nil } +func (s *NoopPsSnapshotter) RemoveModule(pid uint32, addr va.Address) error { return nil } func (s *NoopPsSnapshotter) WriteFromKcap(kevt *kevent.Kevent) error { return nil } func (s *NoopPsSnapshotter) AddMmap(kevt *kevent.Kevent) error { return nil } -func (s *NoopPsSnapshotter) RemoveMmap(pid uint32, address va.Address) error { return nil } +func (s *NoopPsSnapshotter) RemoveMmap(pid uint32, addr va.Address) error { return nil } func TestCallstackEnrichment(t *testing.T) { hsnap := new(handle.SnapshotterMock) @@ -840,7 +840,7 @@ func TestCallstackEnrichment(t *testing.T) { // exercise callstack enrichment with a noop // process snapshotter. This will make the - // symbolizer to always fallback to Debug Help + // symbolizer to always fall back to Debug Help // API when resolving symbolic information nopsnap := new(NoopPsSnapshotter) log.Info("test callstack enrichment with noop ps snapshotter") diff --git a/pkg/kcap/reader_windows.go b/pkg/kcap/reader_windows.go index 3c2d95950..049f2cc63 100644 --- a/pkg/kcap/reader_windows.go +++ b/pkg/kcap/reader_windows.go @@ -191,8 +191,8 @@ func (r *reader) updateSnapshotters(kevt *kevent.Kevent) error { } case ktypes.UnloadImage: pid := kevt.Kparams.MustGetPid() - mod := kevt.GetParamAsString(kparams.ImagePath) - if err := r.psnapshotter.RemoveModule(pid, mod); err != nil { + addr := kevt.Kparams.TryGetAddress(kparams.ImageBase) + if err := r.psnapshotter.RemoveModule(pid, addr); err != nil { return err } case ktypes.CreateProcess, diff --git a/pkg/ps/snapshotter.go b/pkg/ps/snapshotter.go index 8135489a3..0dcb60f60 100644 --- a/pkg/ps/snapshotter.go +++ b/pkg/ps/snapshotter.go @@ -39,7 +39,7 @@ type Snapshotter interface { // RemoveThread removes the thread from the given process. RemoveThread(pid uint32, tid uint32) error // RemoveModule removes the module the given process. - RemoveModule(pid uint32, mod string) error + RemoveModule(pid uint32, addr va.Address) error // AddMmap adds a new memory mapping (data memory-mapped file, image, or pagefile) to this process state. AddMmap(*kevent.Kevent) error // RemoveMmap removes memory mapping at the given base address. diff --git a/pkg/ps/snapshotter_mock.go b/pkg/ps/snapshotter_mock.go index 4e36d44c6..ee4458efd 100644 --- a/pkg/ps/snapshotter_mock.go +++ b/pkg/ps/snapshotter_mock.go @@ -97,8 +97,8 @@ func (s *SnapshotterMock) RemoveThread(pid uint32, tid uint32) error { } // RemoveModule method -func (s *SnapshotterMock) RemoveModule(pid uint32, mod string) error { - args := s.Called(pid, mod) +func (s *SnapshotterMock) RemoveModule(pid uint32, addr va.Address) error { + args := s.Called(pid, addr) return args.Error(0) } diff --git a/pkg/ps/snapshotter_windows.go b/pkg/ps/snapshotter_windows.go index a6d6f9fdf..3ee582e65 100644 --- a/pkg/ps/snapshotter_windows.go +++ b/pkg/ps/snapshotter_windows.go @@ -240,14 +240,14 @@ func (s *snapshotter) RemoveThread(pid uint32, tid uint32) error { return nil } -func (s *snapshotter) RemoveModule(pid uint32, module string) error { +func (s *snapshotter) RemoveModule(pid uint32, addr va.Address) error { s.mu.Lock() defer s.mu.Unlock() proc, ok := s.procs[pid] if !ok { return nil } - proc.RemoveModule(module) + proc.RemoveModule(addr) moduleCount.Add(-1) return nil } diff --git a/pkg/ps/snapshotter_windows_test.go b/pkg/ps/snapshotter_windows_test.go index f6127e1b7..8a2a74955 100644 --- a/pkg/ps/snapshotter_windows_test.go +++ b/pkg/ps/snapshotter_windows_test.go @@ -453,6 +453,7 @@ func TestRemoveModule(t *testing.T) { Kparams: kevent.Kparams{ kparams.ProcessID: {Name: kparams.ProcessID, Type: kparams.PID, Value: uint32(os.Getpid())}, kparams.ImagePath: {Name: kparams.ImagePath, Type: kparams.UnicodeString, Value: "C:\\Windows\\System32\\notepad.exe"}, + kparams.ImageBase: {Name: kparams.ImageBase, Type: kparams.Address, Value: uint64(0xffff7656)}, }, } @@ -462,7 +463,7 @@ func TestRemoveModule(t *testing.T) { require.True(t, ok) require.NotNil(t, ps) require.Len(t, ps.Modules, 1) - require.NoError(t, psnap.RemoveModule(uint32(os.Getpid()), "C:\\Windows\\System32\\notepad.exe")) + require.NoError(t, psnap.RemoveModule(uint32(os.Getpid()), va.Address(0xffff7656))) require.Len(t, ps.Modules, 0) } diff --git a/pkg/ps/types/types_windows.go b/pkg/ps/types/types_windows.go index 619e58ba5..b16fe4bfd 100644 --- a/pkg/ps/types/types_windows.go +++ b/pkg/ps/types/types_windows.go @@ -415,7 +415,7 @@ func (ps *PS) RemoveHandle(handle windows.Handle) { // AddModule adds a new module to this process state. func (ps *PS) AddModule(mod Module) { - m := ps.FindModule(mod.Name) + m := ps.FindModuleByAddr(mod.BaseAddress) if m != nil { return } @@ -423,9 +423,9 @@ func (ps *PS) AddModule(mod Module) { } // RemoveModule removes a specified module from this process state. -func (ps *PS) RemoveModule(path string) { +func (ps *PS) RemoveModule(addr va.Address) { for i, mod := range ps.Modules { - if filepath.Base(mod.Name) == filepath.Base(path) { + if mod.BaseAddress == addr { ps.Modules = append(ps.Modules[:i], ps.Modules[i+1:]...) break } @@ -442,6 +442,16 @@ func (ps *PS) FindModule(path string) *Module { return nil } +// FindModuleByAddr finds the module by its base address. +func (ps *PS) FindModuleByAddr(addr va.Address) *Module { + for _, mod := range ps.Modules { + if mod.BaseAddress == addr { + return &mod + } + } + return nil +} + // FindModuleByVa finds the module name by // probing the range of the given virtual address. func (ps *PS) FindModuleByVa(addr va.Address) *Module {