chore(ps): Append/remove module by base address

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)

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
 			},

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")

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,

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.

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)
 }
 

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
 }

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)
 }
 

pkg/ps/types/types_windows.go

@@ -415,17 +415,17 @@ 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
 	}
 	ps.Modules = append(ps.Modules, mod)
 }
 
 // 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 {