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 {