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.

Author: rabbitstack

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.ImageFilename)
+		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,7 +826,7 @@ 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) AddFileMapping(kevt *kevent.Kevent) error               { return nil }
 func (s *NoopPsSnapshotter) RemoveFileMapping(pid uint32, address va.Address) error { return nil }
@@ -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/ps/snapshotter.go

@@ -39,11 +39,11 @@ 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
 	// AddFileMapping adds a new data memory-mapped file to this process state.
 	AddFileMapping(*kevent.Kevent) error
 	// RemoveFileMapping removes memory-mapped file at the given base address.
-	RemoveFileMapping(pid uint32, address va.Address) error
+	RemoveFileMapping(pid uint32, addr va.Address) error
 	// WriteFromKcap appends a new process state to the snapshotter from the captured kernel event.
 	WriteFromKcap(kevt *kevent.Kevent) error
 	// Remove deletes process's state from the snapshotter.

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

@@ -241,14 +241,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/types/types_windows.go

@@ -407,17 +407,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
 		}
@@ -434,6 +434,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 {