Skip to content

Commit 9f2f130

Browse files
rabbitstackrabbitstack
committed
refactor(ps): Store memory mappings per process
The process state is refactored to store all the memory mappings created by the process. This effectively eliminates the need to keep the memory mappings in the fs processor hence improving code quality, readability, and performance.
1 parent 0799c5e commit 9f2f130

File tree

8 files changed

+138
-171
lines changed

8 files changed

+138
-171
lines changed

‎internal/etw/processors/fs_windows.go‎

Lines changed: 37 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,6 @@ var (
4949
type fsProcessor struct {
5050
// files stores the file metadata indexed by file object
5151
files map[uint64]*FileInfo
52-
// mmaps stores memory-mapped files by pid and file object
53-
mmaps map[uint32]map[uint64]*MmapInfo
5452

5553
hsnap handle.Snapshotter
5654
psnap ps.Snapshotter
@@ -75,13 +73,6 @@ type FileInfo struct {
7573
Type fs.FileType
7674
}
7775

78-
// MmapInfo stores information of the memory-mapped file.
79-
type MmapInfo struct {
80-
File string
81-
BaseAddr uint64
82-
Size uint64
83-
}
84-
8576
func newFsProcessor(
8677
hsnap handle.Snapshotter,
8778
psnap ps.Snapshotter,
@@ -91,7 +82,6 @@ func newFsProcessor(
9182
) Processor {
9283
f := &fsProcessor{
9384
files: make(map[uint64]*FileInfo),
94-
mmaps: make(map[uint32]map[uint64]*MmapInfo),
9585
irps: make(map[uint64]*kevent.Kevent),
9686
hsnap: hsnap,
9787
psnap: psnap,
@@ -139,35 +129,25 @@ func (f *fsProcessor) processEvent(e *kevent.Kevent) (*kevent.Kevent, error) {
139129
f.files[fileObject] = &FileInfo{Name: filepath, Type: fs.GetFileType(filepath, 0)}
140130
}
141131
case ktypes.MapFileRundown:
142-
// if the memory-mapped view refers to the image/data file
143-
// we store it in internal state for each process. The state
144-
// is consulted later when we process unmap events
145-
sec := e.Kparams.MustGetUint32(kparams.FileViewSectionType)
146-
isMapped := sec != va.SectionPagefile && sec != va.SectionPhysical
147-
if !isMapped {
148-
return e, nil
149-
}
150132
fileKey := e.Kparams.MustGetUint64(kparams.FileKey)
151-
viewBase := e.Kparams.MustGetUint64(kparams.FileViewBase)
152-
viewSize := e.Kparams.MustGetUint64(kparams.FileViewSize)
153-
f.initMmap(e.PID)
154133
fileinfo := f.files[fileKey]
134+
155135
if fileinfo != nil {
156136
totalMapRundownFiles.Add(1)
157-
f.mmaps[e.PID][fileKey] = &MmapInfo{File: fileinfo.Name, BaseAddr: viewBase, Size: viewSize}
137+
e.AppendParam(kparams.FilePath, kparams.Path, fileinfo.Name)
158138
} else {
159-
process, err := windows.OpenProcess(windows.PROCESS_QUERY_INFORMATION, false, e.PID)
160-
if err != nil {
161-
return e, nil
139+
// if the view of section is backed by the data/image file
140+
// try to get the mapped file name and append it to params
141+
sec := e.Kparams.MustGetUint32(kparams.FileViewSectionType)
142+
isMapped := sec != va.SectionPagefile && sec != va.SectionPhysical
143+
if isMapped {
144+
totalMapRundownFiles.Add(1)
145+
addr := e.Kparams.MustGetUint64(kparams.FileViewBase) + (e.Kparams.MustGetUint64(kparams.FileOffset))
146+
e.AppendParam(kparams.FilePath, kparams.Path, f.getMappedFile(e.PID, addr))
162147
}
163-
defer windows.Close(process)
164-
totalMapRundownFiles.Add(1)
165-
addr := e.Kparams.MustGetUint64(kparams.FileViewBase) + (e.Kparams.MustGetUint64(kparams.FileOffset))
166-
name := f.devMapper.Convert(sys.GetMappedFile(process, uintptr(addr)))
167-
f.mmaps[e.PID][fileKey] = &MmapInfo{File: name, BaseAddr: viewBase, Size: viewSize}
168148
}
169-
e.AppendParam(kparams.FilePath, kparams.Path, f.mmaps[e.PID][fileKey].File)
170-
return e, f.psnap.AddFileMapping(e)
149+
150+
return e, f.psnap.AddMmap(e)
171151
case ktypes.CreateFile:
172152
// we defer the processing of the CreateFile event until we get
173153
// the matching FileOpEnd event. This event contains the operation
@@ -255,24 +235,22 @@ func (f *fsProcessor) processEvent(e *kevent.Kevent) (*kevent.Kevent, error) {
255235
fileObject := e.Kparams.MustGetUint64(kparams.FileObject)
256236
delete(f.files, fileObject)
257237
case ktypes.UnmapViewFile:
258-
_ = f.psnap.RemoveFileMapping(e.PID, e.Kparams.TryGetAddress(kparams.FileViewBase))
259-
fileKey := e.Kparams.MustGetUint64(kparams.FileKey)
260-
if _, ok := f.mmaps[e.PID]; !ok {
261-
return e, nil
262-
}
263-
mmapinfo := f.mmaps[e.PID][fileKey]
264-
if mmapinfo != nil {
265-
e.AppendParam(kparams.FilePath, kparams.Path, mmapinfo.File)
238+
ok, proc := f.psnap.Find(e.PID)
239+
addr := e.Kparams.TryGetAddress(kparams.FileViewBase)
240+
if ok {
241+
mmap := proc.FindMmap(addr)
242+
if mmap != nil {
243+
e.AppendParam(kparams.FilePath, kparams.Path, mmap.File)
244+
}
266245
}
246+
267247
totalMapRundownFiles.Add(-1)
268-
delete(f.mmaps[e.PID], fileKey)
269-
if len(f.mmaps[e.PID]) == 0 {
270-
// process terminated, all files unmapped
271-
f.removeMmap(e.PID)
272-
}
248+
249+
return e, f.psnap.RemoveMmap(e.PID, addr)
273250
default:
274251
var fileObject uint64
275252
fileKey := e.Kparams.MustGetUint64(kparams.FileKey)
253+
276254
if !e.IsMapViewFile() {
277255
fileObject = e.Kparams.MustGetUint64(kparams.FileObject)
278256
}
@@ -286,28 +264,18 @@ func (f *fsProcessor) processEvent(e *kevent.Kevent) (*kevent.Kevent, error) {
286264
if fileinfo == nil && e.IsMapViewFile() {
287265
sec := e.Kparams.MustGetUint32(kparams.FileViewSectionType)
288266
isMapped := sec != va.SectionPagefile && sec != va.SectionPhysical
289-
if !isMapped {
290-
return e, nil
267+
if isMapped {
268+
totalMapRundownFiles.Add(1)
269+
addr := e.Kparams.MustGetUint64(kparams.FileViewBase) + (e.Kparams.MustGetUint64(kparams.FileOffset))
270+
e.AppendParam(kparams.FilePath, kparams.Path, f.getMappedFile(e.PID, addr))
291271
}
292-
process, err := windows.OpenProcess(windows.PROCESS_QUERY_INFORMATION, false, e.PID)
293-
if err != nil {
294-
return e, nil
295-
}
296-
defer windows.Close(process)
297-
viewBase := e.Kparams.MustGetUint64(kparams.FileViewBase)
298-
viewSize := e.Kparams.MustGetUint64(kparams.FileViewSize)
299-
addr := e.Kparams.MustGetUint64(kparams.FileViewBase) + (e.Kparams.MustGetUint64(kparams.FileOffset))
300-
name := f.devMapper.Convert(sys.GetMappedFile(process, uintptr(addr)))
301-
f.initMmap(e.PID)
302-
f.mmaps[e.PID][fileKey] = &MmapInfo{File: name, BaseAddr: viewBase, Size: viewSize}
303-
e.AppendParam(kparams.FilePath, kparams.Path, name)
304-
return e, f.psnap.AddFileMapping(e)
305272
}
306273

307274
// ignore object misses that are produced by CloseFile
308275
if fileinfo == nil && !e.IsCloseFile() {
309276
fileObjectMisses.Add(1)
310277
}
278+
311279
if e.IsDeleteFile() {
312280
delete(f.files, fileObject)
313281
}
@@ -317,14 +285,16 @@ func (f *fsProcessor) processEvent(e *kevent.Kevent) (*kevent.Kevent, error) {
317285
}
318286
return e, nil
319287
}
288+
320289
if fileinfo != nil {
321290
if fileinfo.Type != fs.Unknown {
322291
e.AppendEnum(kparams.FileType, uint32(fileinfo.Type), fs.FileTypes)
323292
}
324293
e.AppendParam(kparams.FilePath, kparams.Path, fileinfo.Name)
325294
}
295+
326296
if e.IsMapViewFile() {
327-
return e, f.psnap.AddFileMapping(e)
297+
return e, f.psnap.AddMmap(e)
328298
}
329299
}
330300
return e, nil
@@ -352,15 +322,13 @@ func (f *fsProcessor) findFile(fileKey, fileObject uint64) *FileInfo {
352322
return nil
353323
}
354324

355-
func (f *fsProcessor) initMmap(pid uint32) {
356-
m := f.mmaps[pid]
357-
if m == nil {
358-
f.mmaps[pid] = make(map[uint64]*MmapInfo)
325+
func (f *fsProcessor) getMappedFile(pid uint32, addr uint64) string {
326+
process, err := windows.OpenProcess(windows.PROCESS_QUERY_INFORMATION, false, pid)
327+
if err != nil {
328+
return ""
359329
}
360-
}
361-
362-
func (f *fsProcessor) removeMmap(pid uint32) {
363-
delete(f.mmaps, pid)
330+
defer windows.Close(process)
331+
return f.devMapper.Convert(sys.GetMappedFile(process, uintptr(addr)))
364332
}
365333

366334
func (f *fsProcessor) purge() {

‎internal/etw/processors/fs_windows_test.go‎

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"github.com/rabbitstack/fibratus/pkg/kevent/kparams"
2828
"github.com/rabbitstack/fibratus/pkg/kevent/ktypes"
2929
"github.com/rabbitstack/fibratus/pkg/ps"
30+
pstypes "github.com/rabbitstack/fibratus/pkg/ps/types"
3031
"github.com/rabbitstack/fibratus/pkg/util/va"
3132
"github.com/stretchr/testify/assert"
3233
"github.com/stretchr/testify/mock"
@@ -89,12 +90,11 @@ func TestFsProcessor(t *testing.T) {
8990
},
9091
func(e *kevent.Kevent, t *testing.T, hsnap *handle.SnapshotterMock, p Processor) {
9192
fsProcessor := p.(*fsProcessor)
92-
assert.Contains(t, fsProcessor.mmaps, uint32(10233))
93-
mapinfo := fsProcessor.mmaps[10233][124567380264]
94-
require.NotNil(t, mapinfo)
95-
assert.Equal(t, "C:\\Windows\\System32\\kernel32.dll", mapinfo.File)
96-
assert.Equal(t, uint64(3098), mapinfo.Size)
97-
assert.Equal(t, uint64(0xffff23433), mapinfo.BaseAddr)
93+
94+
assert.Equal(t, "C:\\Windows\\System32\\kernel32.dll", e.GetParamAsString(kparams.FilePath))
95+
96+
psnap := fsProcessor.psnap.(*ps.SnapshotterMock)
97+
psnap.AssertNumberOfCalls(t, "AddMmap", 1)
9898
},
9999
},
100100
{
@@ -201,19 +201,16 @@ func TestFsProcessor(t *testing.T) {
201201
kparams.FileViewSectionType: {Name: kparams.FileViewSectionType, Type: kparams.Enum, Value: uint32(va.SectionImage), Enum: kevent.ViewSectionTypes},
202202
},
203203
},
204-
func(p Processor) {
205-
fsProcessor := p.(*fsProcessor)
206-
fsProcessor.mmaps[10233] = make(map[uint64]*MmapInfo)
207-
fsProcessor.mmaps[10233][124567380264] = &MmapInfo{File: "C:\\Windows\\System32\\kernel32.dll"}
208-
},
204+
nil,
209205
func() *handle.SnapshotterMock {
210206
hsnap := new(handle.SnapshotterMock)
211207
return hsnap
212208
},
213209
func(e *kevent.Kevent, t *testing.T, hsnap *handle.SnapshotterMock, p Processor) {
214210
fsProcessor := p.(*fsProcessor)
215-
assert.True(t, e.Kparams.Contains(kparams.FilePath))
216-
assert.Nil(t, fsProcessor.mmaps[3098][124567380264])
211+
212+
psnap := fsProcessor.psnap.(*ps.SnapshotterMock)
213+
psnap.AssertNumberOfCalls(t, "RemoveMmap", 1)
217214
},
218215
},
219216
{
@@ -298,8 +295,13 @@ func TestFsProcessor(t *testing.T) {
298295
t.Run(tt.name, func(t *testing.T) {
299296
hsnap := tt.hsnap()
300297
psnap := new(ps.SnapshotterMock)
301-
psnap.On("AddFileMapping", mock.Anything).Return(nil)
302-
psnap.On("RemoveFileMapping", mock.Anything, mock.Anything).Return(nil)
298+
psnap.On("AddMmap", mock.Anything).Return(nil)
299+
psnap.On("RemoveMmap", mock.Anything, mock.Anything).Return(nil)
300+
psnap.On("Find", mock.Anything).Return(true, &pstypes.PS{
301+
Mmaps: []pstypes.Mmap{
302+
{File: "C:\\Windows\\System32\\kernel32.dll", BaseAddress: va.Address(0xffff23433), Size: 3098},
303+
},
304+
})
303305
p := newFsProcessor(hsnap, psnap, fs.NewDevMapper(), fs.NewDevPathResolver(), &config.Config{})
304306
if tt.setupProcessor != nil {
305307
tt.setupProcessor(p)

‎internal/etw/source_test.go‎

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,8 @@ func TestEventSourceStartTraces(t *testing.T) {
6868
psnap.On("Write", mock.Anything).Return(nil)
6969
psnap.On("AddThread", mock.Anything).Return(nil)
7070
psnap.On("AddModule", mock.Anything).Return(nil)
71-
psnap.On("AddFileMapping", mock.Anything).Return(nil)
72-
psnap.On("RemoveFileMapping", mock.Anything, mock.Anything).Return(nil)
71+
psnap.On("AddMmap", mock.Anything).Return(nil)
72+
psnap.On("RemoveMmap", mock.Anything, mock.Anything).Return(nil)
7373
psnap.On("RemoveThread", mock.Anything, mock.Anything).Return(nil)
7474
psnap.On("RemoveModule", mock.Anything, mock.Anything).Return(nil)
7575
psnap.On("FindModule", mock.Anything).Return(false, nil)
@@ -153,8 +153,8 @@ func TestEventSourceEnableFlagsDynamically(t *testing.T) {
153153
psnap.On("Write", mock.Anything).Return(nil)
154154
psnap.On("AddThread", mock.Anything).Return(nil)
155155
psnap.On("AddModule", mock.Anything).Return(nil)
156-
psnap.On("AddFileMapping", mock.Anything).Return(nil)
157-
psnap.On("RemoveFileMapping", mock.Anything, mock.Anything).Return(nil)
156+
psnap.On("AddMmap", mock.Anything).Return(nil)
157+
psnap.On("RemoveMmap", mock.Anything, mock.Anything).Return(nil)
158158
psnap.On("RemoveThread", mock.Anything, mock.Anything).Return(nil)
159159
psnap.On("RemoveModule", mock.Anything, mock.Anything).Return(nil)
160160
psnap.On("FindModule", mock.Anything).Return(false, nil)
@@ -236,8 +236,8 @@ func TestEventSourceEnableFlagsDynamicallyWithYaraEnabled(t *testing.T) {
236236
psnap.On("Write", mock.Anything).Return(nil)
237237
psnap.On("AddThread", mock.Anything).Return(nil)
238238
psnap.On("AddModule", mock.Anything).Return(nil)
239-
psnap.On("AddFileMapping", mock.Anything).Return(nil)
240-
psnap.On("RemoveFileMapping", mock.Anything, mock.Anything).Return(nil)
239+
psnap.On("AddMmap", mock.Anything).Return(nil)
240+
psnap.On("RemoveMmap", mock.Anything, mock.Anything).Return(nil)
241241
psnap.On("RemoveThread", mock.Anything, mock.Anything).Return(nil)
242242
psnap.On("RemoveModule", mock.Anything, mock.Anything).Return(nil)
243243
psnap.On("FindModule", mock.Anything).Return(false, nil)
@@ -311,8 +311,8 @@ func TestEventSourceRundownEvents(t *testing.T) {
311311
psnap.On("Write", mock.Anything).Return(nil)
312312
psnap.On("AddThread", mock.Anything).Return(nil)
313313
psnap.On("AddModule", mock.Anything).Return(nil)
314-
psnap.On("AddFileMapping", mock.Anything).Return(nil)
315-
psnap.On("RemoveFileMapping", mock.Anything, mock.Anything).Return(nil)
314+
psnap.On("AddMmap", mock.Anything).Return(nil)
315+
psnap.On("RemoveMmap", mock.Anything, mock.Anything).Return(nil)
316316
psnap.On("RemoveThread", mock.Anything, mock.Anything).Return(nil)
317317
psnap.On("RemoveModule", mock.Anything, mock.Anything).Return(nil)
318318
psnap.On("FindModule", mock.Anything).Return(false, nil)
@@ -721,11 +721,11 @@ func TestEventSourceAllEvents(t *testing.T) {
721721
psnap.On("Write", mock.Anything).Return(nil)
722722
psnap.On("AddThread", mock.Anything).Return(nil)
723723
psnap.On("AddModule", mock.Anything).Return(nil)
724-
psnap.On("AddFileMapping", mock.Anything).Return(nil)
724+
psnap.On("AddMmap", mock.Anything).Return(nil)
725725
psnap.On("RemoveThread", mock.Anything, mock.Anything).Return(nil)
726726
psnap.On("RemoveModule", mock.Anything, mock.Anything).Return(nil)
727727
psnap.On("FindModule", mock.Anything).Return(false, nil)
728-
psnap.On("RemoveFileMapping", mock.Anything, mock.Anything).Return(nil)
728+
psnap.On("RemoveMmap", mock.Anything, mock.Anything).Return(nil)
729729
psnap.On("FindAndPut", mock.Anything).Return(&pstypes.PS{})
730730
psnap.On("Find", mock.Anything).Return(true, &pstypes.PS{})
731731
psnap.On("Remove", mock.Anything).Return(nil)
@@ -814,22 +814,22 @@ type NoopPsSnapshotter struct{}
814814

815815
var fakeProc = &pstypes.PS{PID: 111111, Name: "fake.exe"}
816816

817-
func (s *NoopPsSnapshotter) Write(kevt *kevent.Kevent) error { return nil }
818-
func (s *NoopPsSnapshotter) Remove(kevt *kevent.Kevent) error { return nil }
819-
func (s *NoopPsSnapshotter) Find(pid uint32) (bool, *pstypes.PS) { return true, fakeProc }
820-
func (s *NoopPsSnapshotter) FindAndPut(pid uint32) *pstypes.PS { return fakeProc }
821-
func (s *NoopPsSnapshotter) Put(ps *pstypes.PS) {}
822-
func (s *NoopPsSnapshotter) Size() uint32 { return 1 }
823-
func (s *NoopPsSnapshotter) Close() error { return nil }
824-
func (s *NoopPsSnapshotter) GetSnapshot() []*pstypes.PS { return nil }
825-
func (s *NoopPsSnapshotter) AddThread(kevt *kevent.Kevent) error { return nil }
826-
func (s *NoopPsSnapshotter) AddModule(kevt *kevent.Kevent) error { return nil }
827-
func (s *NoopPsSnapshotter) FindModule(addr va.Address) (bool, *pstypes.Module) { return false, nil }
828-
func (s *NoopPsSnapshotter) RemoveThread(pid uint32, tid uint32) error { return nil }
829-
func (s *NoopPsSnapshotter) RemoveModule(pid uint32, mod string) error { return nil }
830-
func (s *NoopPsSnapshotter) WriteFromKcap(kevt *kevent.Kevent) error { return nil }
831-
func (s *NoopPsSnapshotter) AddFileMapping(kevt *kevent.Kevent) error { return nil }
832-
func (s *NoopPsSnapshotter) RemoveFileMapping(pid uint32, address va.Address) error { return nil }
817+
func (s *NoopPsSnapshotter) Write(kevt *kevent.Kevent) error { return nil }
818+
func (s *NoopPsSnapshotter) Remove(kevt *kevent.Kevent) error { return nil }
819+
func (s *NoopPsSnapshotter) Find(pid uint32) (bool, *pstypes.PS) { return true, fakeProc }
820+
func (s *NoopPsSnapshotter) FindAndPut(pid uint32) *pstypes.PS { return fakeProc }
821+
func (s *NoopPsSnapshotter) Put(ps *pstypes.PS) {}
822+
func (s *NoopPsSnapshotter) Size() uint32 { return 1 }
823+
func (s *NoopPsSnapshotter) Close() error { return nil }
824+
func (s *NoopPsSnapshotter) GetSnapshot() []*pstypes.PS { return nil }
825+
func (s *NoopPsSnapshotter) AddThread(kevt *kevent.Kevent) error { return nil }
826+
func (s *NoopPsSnapshotter) AddModule(kevt *kevent.Kevent) error { return nil }
827+
func (s *NoopPsSnapshotter) FindModule(addr va.Address) (bool, *pstypes.Module) { return false, nil }
828+
func (s *NoopPsSnapshotter) RemoveThread(pid uint32, tid uint32) error { return nil }
829+
func (s *NoopPsSnapshotter) RemoveModule(pid uint32, mod string) error { return nil }
830+
func (s *NoopPsSnapshotter) WriteFromKcap(kevt *kevent.Kevent) error { return nil }
831+
func (s *NoopPsSnapshotter) AddMmap(kevt *kevent.Kevent) error { return nil }
832+
func (s *NoopPsSnapshotter) RemoveMmap(pid uint32, address va.Address) error { return nil }
833833

834834
func TestCallstackEnrichment(t *testing.T) {
835835
hsnap := new(handle.SnapshotterMock)

‎pkg/ps/snapshotter.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,10 +40,10 @@ type Snapshotter interface {
4040
RemoveThread(pid uint32, tid uint32) error
4141
// RemoveModule removes the module the given process.
4242
RemoveModule(pid uint32, mod string) error
43-
// AddFileMapping adds a new data memory-mapped file to this process state.
44-
AddFileMapping(*kevent.Kevent) error
45-
// RemoveFileMapping removes memory-mapped file at the given base address.
46-
RemoveFileMapping(pid uint32, address va.Address) error
43+
// AddMmap adds a new memory mapping (data memory-mapped file, image, or pagefile) to this process state.
44+
AddMmap(*kevent.Kevent) error
45+
// RemoveMmap removes memory mapping at the given base address.
46+
RemoveMmap(pid uint32, address va.Address) error
4747
// WriteFromKcap appends a new process state to the snapshotter from the captured kernel event.
4848
WriteFromKcap(kevt *kevent.Kevent) error
4949
// Remove deletes process's state from the snapshotter.

‎pkg/ps/snapshotter_mock.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,13 +106,13 @@ func (s *SnapshotterMock) RemoveModule(pid uint32, mod string) error {
106106
func (s *SnapshotterMock) WriteFromKcap(kevt *kevent.Kevent) error { return nil }
107107

108108
// AddFileMapping method
109-
func (s *SnapshotterMock) AddFileMapping(kevt *kevent.Kevent) error {
109+
func (s *SnapshotterMock) AddMmap(kevt *kevent.Kevent) error {
110110
args := s.Called(kevt)
111111
return args.Error(0)
112112
}
113113

114114
// RemoveFileMapping method
115-
func (s *SnapshotterMock) RemoveFileMapping(pid uint32, address va.Address) error {
115+
func (s *SnapshotterMock) RemoveMmap(pid uint32, address va.Address) error {
116116
args := s.Called(pid, address)
117117
return args.Error(0)
118118
}

0 commit comments

Comments
 (0)