Skip to content

Commit 3e63b84

Browse files
shreyanshjain7174shreyanshjain7174
authored
Merge pull request microsoft#2628 from shreyanshjain7174/vmutils/add-test-coverage
vmutils: add unit tests for allProcessEntries & LookupVMMEM
2 parents 393a2b6 + 3a1ba4d commit 3e63b84

1 file changed

Lines changed: 302 additions & 0 deletions

File tree

‎internal/vm/vmutils/vmmem_test.go‎

Lines changed: 302 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,12 @@ const (
2121
testProcessHandle windows.Handle = 2000
2222
testToken windows.Token = 3000
2323
testPID uint32 = 1234
24+
25+
// Second set of handles for multi-process test scenarios where two
26+
// vmmem entries appear in the same snapshot.
27+
testProcessHandle2 windows.Handle = 4000
28+
testToken2 windows.Token = 5000
29+
testPID2 uint32 = 5678
2430
)
2531

2632
var (
@@ -39,6 +45,8 @@ var (
3945
)
4046

4147
func TestLookupVMMEM(t *testing.T) {
48+
t.Parallel()
49+
4250
tests := []struct {
4351
name string
4452
setupMock func(*mockHelper)
@@ -183,10 +191,68 @@ func TestLookupVMMEM(t *testing.T) {
183191
expectError: true,
184192
errorContains: "failed to find matching vmmem process",
185193
},
194+
195+
// --- Multi-process and case-sensitivity edge cases ---
196+
197+
{
198+
// Two vmmem processes in the snapshot. The first belongs to a
199+
// different VM; the function should close its handle and keep
200+
// scanning until it finds the one matching our VM ID.
201+
name: "skips wrong VM vmmem, matches second",
202+
setupMock: func(h *mockHelper) {
203+
h.expectSkipThenMatch(testVMIDStr)
204+
},
205+
},
206+
{
207+
// Multiple vmmem processes in the snapshot but none of them
208+
// belong to the VM we're looking for.
209+
name: "multiple vmmem processes, none match",
210+
setupMock: func(h *mockHelper) {
211+
h.expectMultipleVmmemNoneMatch()
212+
},
213+
expectError: true,
214+
errorContains: "failed to find matching vmmem process",
215+
},
216+
{
217+
// The process name uses unusual casing ("VMMEM.EXE"). The code
218+
// does a case-insensitive comparison via strings.EqualFold, so
219+
// this should still match.
220+
name: "case-insensitive process name match",
221+
setupMock: func(h *mockHelper) {
222+
h.expectSnapshot()
223+
h.expectProcess32(testPID, "VMMEM.EXE")
224+
h.expectOpenProcess(testPID, testProcessHandle, nil)
225+
h.expectOpenProcessToken(nil)
226+
h.expectGetTokenUser(mockTokenUser, nil)
227+
h.expectLookupAccount(testVMIDStr, "NT VIRTUAL MACHINE", nil)
228+
h.expectCloseToken()
229+
h.expectNoMoreProcesses()
230+
h.expectCloseSnapshot()
231+
},
232+
},
233+
{
234+
// LookupAccount returns the VM ID in lowercase. The GUID
235+
// comparison uses strings.EqualFold, so casing shouldn't matter.
236+
name: "case-insensitive VM ID match",
237+
setupMock: func(h *mockHelper) {
238+
h.expectSnapshot()
239+
h.expectProcess32(testPID, "vmmem")
240+
h.expectOpenProcess(testPID, testProcessHandle, nil)
241+
h.expectOpenProcessToken(nil)
242+
h.expectGetTokenUser(mockTokenUser, nil)
243+
// Return the VM ID in lowercase — EqualFold should still match.
244+
h.expectLookupAccount(strings.ToLower(testVMIDStr), "nt virtual machine", nil)
245+
h.expectCloseToken()
246+
h.expectNoMoreProcesses()
247+
h.expectCloseSnapshot()
248+
},
249+
},
186250
}
187251

188252
for _, tt := range tests {
189253
t.Run(tt.name, func(t *testing.T) {
254+
t.Parallel()
255+
190256
ctrl := gomock.NewController(t)
191257
defer ctrl.Finish()
192258

@@ -217,6 +283,168 @@ func TestLookupVMMEM(t *testing.T) {
217283
}
218284
}
219285

286+
// TestAllProcessEntries verifies the allProcessEntries iterator independently
287+
// from LookupVMMEM to ensure correct ordering, early termination cleanup,
288+
// and faithful forwarding of ProcessEntry32 fields.
289+
func TestAllProcessEntries(t *testing.T) {
290+
t.Parallel()
291+
type wantEntry struct {
292+
pid uint32
293+
name string
294+
}
295+
296+
tests := []struct {
297+
name string
298+
setupMock func(*mockHelper)
299+
breakAfter int // 0 = consume all entries; >0 = break after N
300+
want []wantEntry
301+
}{
302+
{
303+
// Snapshot contains exactly one process; iterator should yield it once
304+
// and stop when Process32Next signals the end.
305+
name: "single process entry",
306+
setupMock: func(h *mockHelper) {
307+
gomock.InOrder(
308+
h.m.EXPECT().
309+
CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)).
310+
Return(testSnapshot, nil),
311+
h.m.EXPECT().
312+
Process32First(testSnapshot, gomock.Any()).
313+
DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
314+
*pe = makeProcessEntry(100, "explorer.exe")
315+
return nil
316+
}),
317+
h.m.EXPECT().
318+
Process32Next(testSnapshot, gomock.Any()).
319+
Return(errNoMoreProcesses),
320+
h.m.EXPECT().CloseHandle(testSnapshot).Return(nil),
321+
)
322+
},
323+
want: []wantEntry{{100, "explorer.exe"}},
324+
},
325+
{
326+
// Three distinct processes in the snapshot; all must be yielded in
327+
// the exact order the OS returns them (First, Next, Next).
328+
name: "multiple process entries yielded in order",
329+
setupMock: func(h *mockHelper) {
330+
gomock.InOrder(
331+
h.m.EXPECT().
332+
CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)).
333+
Return(testSnapshot, nil),
334+
h.m.EXPECT().
335+
Process32First(testSnapshot, gomock.Any()).
336+
DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
337+
*pe = makeProcessEntry(10, "init.exe")
338+
return nil
339+
}),
340+
h.m.EXPECT().
341+
Process32Next(testSnapshot, gomock.Any()).
342+
DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
343+
*pe = makeProcessEntry(20, "svchost.exe")
344+
return nil
345+
}),
346+
h.m.EXPECT().
347+
Process32Next(testSnapshot, gomock.Any()).
348+
DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
349+
*pe = makeProcessEntry(30, "explorer.exe")
350+
return nil
351+
}),
352+
h.m.EXPECT().
353+
Process32Next(testSnapshot, gomock.Any()).
354+
Return(errNoMoreProcesses),
355+
h.m.EXPECT().CloseHandle(testSnapshot).Return(nil),
356+
)
357+
},
358+
want: []wantEntry{
359+
{10, "init.exe"},
360+
{20, "svchost.exe"},
361+
{30, "explorer.exe"},
362+
},
363+
},
364+
{
365+
// Consumer breaks out of the range loop after the first entry.
366+
// The snapshot handle must still be closed (gomock enforces this).
367+
name: "consumer breaks early — snapshot still closed",
368+
setupMock: func(h *mockHelper) {
369+
gomock.InOrder(
370+
h.m.EXPECT().
371+
CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)).
372+
Return(testSnapshot, nil),
373+
h.m.EXPECT().
374+
Process32First(testSnapshot, gomock.Any()).
375+
DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
376+
*pe = makeProcessEntry(10, "init.exe")
377+
return nil
378+
}),
379+
// No Process32Next expected — consumer breaks before it's called.
380+
h.m.EXPECT().CloseHandle(testSnapshot).Return(nil),
381+
)
382+
},
383+
breakAfter: 1,
384+
want: []wantEntry{{10, "init.exe"}},
385+
},
386+
{
387+
// Verify that the yielded ProcessEntry32 carries the correct PID
388+
// and ExeFile name end-to-end through the iterator.
389+
name: "validates ProcessEntry32 fields",
390+
setupMock: func(h *mockHelper) {
391+
gomock.InOrder(
392+
h.m.EXPECT().
393+
CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)).
394+
Return(testSnapshot, nil),
395+
h.m.EXPECT().
396+
Process32First(testSnapshot, gomock.Any()).
397+
DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
398+
*pe = makeProcessEntry(testPID, "vmmem.exe")
399+
return nil
400+
}),
401+
h.m.EXPECT().
402+
Process32Next(testSnapshot, gomock.Any()).
403+
Return(errNoMoreProcesses),
404+
h.m.EXPECT().CloseHandle(testSnapshot).Return(nil),
405+
)
406+
},
407+
want: []wantEntry{{testPID, "vmmem.exe"}},
408+
},
409+
}
410+
411+
for _, tt := range tests {
412+
t.Run(tt.name, func(t *testing.T) {
413+
t.Parallel()
414+
415+
ctrl := gomock.NewController(t)
416+
defer ctrl.Finish()
417+
418+
mockAPI := mock.NewMockAPI(ctrl)
419+
helper := &mockHelper{m: mockAPI}
420+
tt.setupMock(helper)
421+
422+
var got []wantEntry
423+
i := 0
424+
for pe := range allProcessEntries(context.Background(), mockAPI) {
425+
name := windows.UTF16ToString(pe.ExeFile[:])
426+
got = append(got, wantEntry{pe.ProcessID, name})
427+
i++
428+
if tt.breakAfter > 0 && i >= tt.breakAfter {
429+
break
430+
}
431+
}
432+
433+
if len(got) != len(tt.want) {
434+
t.Fatalf("got %d entries, want %d", len(got), len(tt.want))
435+
}
436+
for idx, w := range tt.want {
437+
if got[idx].pid != w.pid {
438+
t.Errorf("entry[%d] PID: got %d, want %d", idx, got[idx].pid, w.pid)
439+
}
440+
if got[idx].name != w.name {
441+
t.Errorf("entry[%d] ExeFile: got %q, want %q", idx, got[idx].name, w.name)
442+
}
443+
}
444+
})
445+
}
446+
}
447+
220448
// makeProcessEntry creates a ProcessEntry32 with the given name and PID.
221449
func makeProcessEntry(pid uint32, exeName string) windows.ProcessEntry32 {
222450
var pe windows.ProcessEntry32
@@ -304,3 +532,77 @@ func (h *mockHelper) expectSuccessfulMatch(processName string) {
304532
h.expectNoMoreProcesses()
305533
h.expectCloseSnapshot()
306534
}
535+
536+
// expectSkipThenMatch sets up a snapshot with two vmmem processes: the first
537+
// belongs to a different VM and should be skipped, while the second matches
538+
// the given vmIDStr.
539+
func (h *mockHelper) expectSkipThenMatch(vmIDStr string) {
540+
gomock.InOrder(
541+
h.m.EXPECT().CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)).Return(testSnapshot, nil),
542+
// First vmmem: wrong VM.
543+
h.m.EXPECT().Process32First(testSnapshot, gomock.Any()).DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
544+
*pe = makeProcessEntry(testPID, "vmmem")
545+
return nil
546+
}),
547+
h.m.EXPECT().OpenProcess(uint32(windows.PROCESS_QUERY_LIMITED_INFORMATION), false, testPID).Return(testProcessHandle, nil),
548+
h.m.EXPECT().OpenProcessToken(testProcessHandle, uint32(windows.TOKEN_QUERY), gomock.Any()).DoAndReturn(
549+
func(_ windows.Handle, _ uint32, token *windows.Token) error { *token = testToken; return nil },
550+
),
551+
h.m.EXPECT().GetTokenUser(testToken).Return(mockTokenUser, nil),
552+
h.m.EXPECT().LookupAccount(mockSID, "").Return("OTHER-GUID", "NT VIRTUAL MACHINE", uint32(0), nil),
553+
h.m.EXPECT().CloseToken(testToken).Return(nil),
554+
h.m.EXPECT().CloseHandle(testProcessHandle).Return(nil),
555+
// Second vmmem: correct VM.
556+
h.m.EXPECT().Process32Next(testSnapshot, gomock.Any()).DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
557+
*pe = makeProcessEntry(testPID2, "vmmem.exe")
558+
return nil
559+
}),
560+
h.m.EXPECT().OpenProcess(uint32(windows.PROCESS_QUERY_LIMITED_INFORMATION), false, testPID2).Return(testProcessHandle2, nil),
561+
h.m.EXPECT().OpenProcessToken(testProcessHandle2, uint32(windows.TOKEN_QUERY), gomock.Any()).DoAndReturn(
562+
func(_ windows.Handle, _ uint32, token *windows.Token) error { *token = testToken2; return nil },
563+
),
564+
h.m.EXPECT().GetTokenUser(testToken2).Return(mockTokenUser, nil),
565+
h.m.EXPECT().LookupAccount(mockSID, "").Return(vmIDStr, "NT VIRTUAL MACHINE", uint32(0), nil),
566+
h.m.EXPECT().CloseToken(testToken2).Return(nil),
567+
h.m.EXPECT().Process32Next(testSnapshot, gomock.Any()).Return(errNoMoreProcesses).AnyTimes(),
568+
h.m.EXPECT().CloseHandle(testSnapshot).Return(nil),
569+
)
570+
}
571+
572+
// expectMultipleVmmemNoneMatch sets up a snapshot with two vmmem processes,
573+
// both belonging to different VMs. Verifies that all handles are properly
574+
// closed when no match is found.
575+
func (h *mockHelper) expectMultipleVmmemNoneMatch() {
576+
gomock.InOrder(
577+
h.m.EXPECT().CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)).Return(testSnapshot, nil),
578+
// First vmmem: wrong VM.
579+
h.m.EXPECT().Process32First(testSnapshot, gomock.Any()).DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
580+
*pe = makeProcessEntry(testPID, "vmmem")
581+
return nil
582+
}),
583+
h.m.EXPECT().OpenProcess(uint32(windows.PROCESS_QUERY_LIMITED_INFORMATION), false, testPID).Return(testProcessHandle, nil),
584+
h.m.EXPECT().OpenProcessToken(testProcessHandle, uint32(windows.TOKEN_QUERY), gomock.Any()).DoAndReturn(
585+
func(_ windows.Handle, _ uint32, token *windows.Token) error { *token = testToken; return nil },
586+
),
587+
h.m.EXPECT().GetTokenUser(testToken).Return(mockTokenUser, nil),
588+
h.m.EXPECT().LookupAccount(mockSID, "").Return("WRONG-GUID-1", "NT VIRTUAL MACHINE", uint32(0), nil),
589+
h.m.EXPECT().CloseToken(testToken).Return(nil),
590+
h.m.EXPECT().CloseHandle(testProcessHandle).Return(nil),
591+
// Second vmmem: also wrong VM.
592+
h.m.EXPECT().Process32Next(testSnapshot, gomock.Any()).DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error {
593+
*pe = makeProcessEntry(testPID2, "vmmem.exe")
594+
return nil
595+
}),
596+
h.m.EXPECT().OpenProcess(uint32(windows.PROCESS_QUERY_LIMITED_INFORMATION), false, testPID2).Return(testProcessHandle2, nil),
597+
h.m.EXPECT().OpenProcessToken(testProcessHandle2, uint32(windows.TOKEN_QUERY), gomock.Any()).DoAndReturn(
598+
func(_ windows.Handle, _ uint32, token *windows.Token) error { *token = testToken2; return nil },
599+
),
600+
h.m.EXPECT().GetTokenUser(testToken2).Return(mockTokenUser, nil),
601+
h.m.EXPECT().LookupAccount(mockSID, "").Return("WRONG-GUID-2", "NT VIRTUAL MACHINE", uint32(0), nil),
602+
h.m.EXPECT().CloseToken(testToken2).Return(nil),
603+
h.m.EXPECT().CloseHandle(testProcessHandle2).Return(nil),
604+
// No more processes.
605+
h.m.EXPECT().Process32Next(testSnapshot, gomock.Any()).Return(errNoMoreProcesses),
606+
h.m.EXPECT().CloseHandle(testSnapshot).Return(nil),
607+
)
608+
}

0 commit comments

Comments
 (0)