From f093f234243516f04332770b671902547547971e Mon Sep 17 00:00:00 2001 From: Harsh Rawat Date: Sun, 1 Mar 2026 01:37:29 +0530 Subject: [PATCH] add vmmem process lookup functionality + windows API abstraction The original implementation used a brute-force approach to locate the Hyper-V memory management process. It relied on enumerating every single Process ID (PID) on the system and attempting to open each one to query its image path. This meant the system was executing expensive OS-level calls and generating ignored access-denied errors hundreds of times for unrelated processes, just to check if the name matched "vmmem". Ref- https://github.com/microsoft/hcsshim/blob/178d662f94b76b24fda758dfc41a7d9e6a2ee614/internal/uvm/stats.go#L76 The new enhancement significantly optimizes this discovery process by using windows API `CreateToolhelp32Snapshot`, which captures process names upfront without opening the processes. Now, the code iterates through the snapshot and pre-filters by the executable name. It only executes the expensive `OpenProcess` and security token checks (`LookupAccount`) if the process is already confirmed to be named `vmmem` or `vmmem.exe`. Furthermore, the updated code introduces an interface to wrap the Windows OS calls. This dependency injection allows the system calls to be easily mocked, making the function fully unit-testable without requiring a live Hyper-V virtual machine. Signed-off-by: Harsh Rawat --- internal/vm/vmutils/vmmem.go | 84 +++++++ internal/vm/vmutils/vmmem_test.go | 306 ++++++++++++++++++++++++++ internal/windows/doc.go | 76 +++++++ internal/windows/interface.go | 37 ++++ internal/windows/mock/mock_windows.go | 173 +++++++++++++++ internal/windows/windows.go | 55 +++++ test/go.sum | 2 + 7 files changed, 733 insertions(+) create mode 100644 internal/vm/vmutils/vmmem.go create mode 100644 internal/vm/vmutils/vmmem_test.go create mode 100644 internal/windows/doc.go create mode 100644 internal/windows/interface.go create mode 100644 internal/windows/mock/mock_windows.go create mode 100644 internal/windows/windows.go diff --git a/internal/vm/vmutils/vmmem.go b/internal/vm/vmutils/vmmem.go new file mode 100644 index 0000000000..facb231c46 --- /dev/null +++ b/internal/vm/vmutils/vmmem.go @@ -0,0 +1,84 @@ +//go:build windows + +package vmutils + +import ( + "context" + "errors" + "fmt" + "strings" + "unsafe" + + "github.com/Microsoft/hcsshim/internal/log" + iwin "github.com/Microsoft/hcsshim/internal/windows" + + "github.com/Microsoft/go-winio/pkg/guid" + "golang.org/x/sys/windows" +) + +const ( + // vmmemProcessName is the name of the Hyper-V memory management process. + vmmemProcessName = "vmmem" + // vmmemProcessNameExt is the name of the process with .exe extension. + vmmemProcessNameExt = "vmmem.exe" + // ntVirtualMachineDomain is the domain name for Hyper-V virtual machine security principals. + ntVirtualMachineDomain = "NT VIRTUAL MACHINE" +) + +// LookupVMMEM locates the vmmem process for a VM given the VM ID. +// It enumerates processes using Toolhelp32 to filter by name, then validates +// the token using LookupAccount to match the "NT VIRTUAL MACHINE\" identity. +func LookupVMMEM(ctx context.Context, vmID guid.GUID, win iwin.API) (windows.Handle, error) { + vmIDStr := strings.ToUpper(vmID.String()) + log.G(ctx).WithField("vmID", vmIDStr).Debug("looking up vmmem via LookupAccount") + + // 1. Take a snapshot of all processes to grab names without opening them. + snapshot, err := win.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0) + if err != nil { + return 0, fmt.Errorf("failed to create process snapshot: %w", err) + } + defer func(win iwin.API, h windows.Handle) { + _ = win.CloseHandle(h) + }(win, snapshot) + + var pe32 windows.ProcessEntry32 + pe32.Size = uint32(unsafe.Sizeof(pe32)) + + err = win.Process32First(snapshot, &pe32) + for err == nil { + exeName := windows.UTF16ToString(pe32.ExeFile[:]) + + // 2. Only target processes named vmmem or vmmem.exe. + if strings.EqualFold(exeName, vmmemProcessName) || strings.EqualFold(exeName, vmmemProcessNameExt) { + + // 3. Open the process to inspect its security token. + pHandle, err := win.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pe32.ProcessID) + if err == nil { + var t windows.Token + if err := win.OpenProcessToken(pHandle, windows.TOKEN_QUERY, &t); err == nil { + tUser, err := win.GetTokenUser(t) + + // 4. Use the OS API to resolve the SID to account and domain strings. + if err == nil { + account, domain, _, err := win.LookupAccount(tUser.User.Sid, "") + if err == nil { + // 5. Compare against the expected Hyper-V UVM identity. + if strings.EqualFold(domain, ntVirtualMachineDomain) && strings.EqualFold(account, vmIDStr) { + _ = win.CloseToken(t) + log.G(ctx).WithField("pid", pe32.ProcessID).Debug("found vmmem match") + return pHandle, nil + } + } + } + _ = win.CloseToken(t) + } + // Close the process handle if it's not the exact VM we are looking for. + _ = win.CloseHandle(pHandle) + } + } + // Move to the next process in the snapshot. + err = win.Process32Next(snapshot, &pe32) + } + + return 0, errors.New("failed to find matching vmmem process") +} diff --git a/internal/vm/vmutils/vmmem_test.go b/internal/vm/vmutils/vmmem_test.go new file mode 100644 index 0000000000..7308468cf5 --- /dev/null +++ b/internal/vm/vmutils/vmmem_test.go @@ -0,0 +1,306 @@ +//go:build windows + +package vmutils + +import ( + "context" + "errors" + "strings" + "testing" + "unsafe" + + "github.com/Microsoft/hcsshim/internal/windows/mock" + + "github.com/Microsoft/go-winio/pkg/guid" + "go.uber.org/mock/gomock" + "golang.org/x/sys/windows" +) + +const ( + testSnapshot windows.Handle = 1000 + testProcessHandle windows.Handle = 2000 + testToken windows.Token = 3000 + testPID uint32 = 1234 +) + +var ( + testVMID = guid.GUID{ + Data1: 0x12345678, + Data2: 0x1234, + Data3: 0x5678, + Data4: [8]byte{0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0}, + } + testVMIDStr = testVMID.String() + mockSID = &windows.SID{} + mockTokenUser = &windows.Tokenuser{ + User: windows.SIDAndAttributes{Sid: mockSID}, + } + errNoMoreProcesses = errors.New("no more processes") +) + +func TestLookupVMMEM(t *testing.T) { + tests := []struct { + name string + setupMock func(*mockHelper) + expectError bool + errorContains string + }{ + { + name: "successful lookup - vmmem.exe", + setupMock: func(h *mockHelper) { + h.expectSuccessfulMatch("vmmem.exe") + }, + }, + { + name: "successful lookup - vmmem without extension", + setupMock: func(h *mockHelper) { + h.expectSuccessfulMatch("vmmem") + }, + }, + { + name: "failed to create snapshot", + setupMock: func(h *mockHelper) { + h.m.EXPECT(). + CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)). + Return(windows.Handle(0), errors.New("access denied")) + }, + expectError: true, + errorContains: "failed to create process snapshot", + }, + { + name: "no vmmem process found", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.expectProcess32(testPID, "explorer.exe") + h.m.EXPECT().Process32Next(testSnapshot, gomock.Any()).Return(errNoMoreProcesses) + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + { + name: "vmmem found but wrong VM ID", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.expectProcess32(testPID, "vmmem.exe") + h.expectOpenProcess(testPID, testProcessHandle, nil) + h.expectOpenProcessToken(nil) + h.expectGetTokenUser(mockTokenUser, nil) + h.expectLookupAccount("DIFFERENT-VM-ID", "NT VIRTUAL MACHINE", nil) + h.expectCloseToken() + h.expectCloseProcessHandle() + h.expectNoMoreProcesses() + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + { + name: "vmmem found but wrong domain", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.expectProcess32(testPID, "vmmem.exe") + h.expectOpenProcess(testPID, testProcessHandle, nil) + h.expectOpenProcessToken(nil) + h.expectGetTokenUser(mockTokenUser, nil) + h.expectLookupAccount(testVMIDStr, "WORKGROUP", nil) + h.expectCloseToken() + h.expectCloseProcessHandle() + h.expectNoMoreProcesses() + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + { + name: "OpenProcess fails", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.expectProcess32(testPID, "vmmem.exe") + h.expectOpenProcess(testPID, 0, errors.New("access denied")) + h.expectNoMoreProcesses() + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + { + name: "OpenProcessToken fails", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.expectProcess32(testPID, "vmmem.exe") + h.expectOpenProcess(testPID, testProcessHandle, nil) + h.expectOpenProcessToken(errors.New("token access denied")) + h.expectCloseProcessHandle() + h.expectNoMoreProcesses() + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + { + name: "GetTokenUser fails", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.expectProcess32(testPID, "vmmem.exe") + h.expectOpenProcess(testPID, testProcessHandle, nil) + h.expectOpenProcessToken(nil) + h.expectGetTokenUser(nil, errors.New("failed to get token user")) + h.expectCloseToken() + h.expectCloseProcessHandle() + h.expectNoMoreProcesses() + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + { + name: "LookupAccount fails", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.expectProcess32(testPID, "vmmem.exe") + h.expectOpenProcess(testPID, testProcessHandle, nil) + h.expectOpenProcessToken(nil) + h.expectGetTokenUser(mockTokenUser, nil) + h.expectLookupAccount("", "", errors.New("lookup failed")) + h.expectCloseToken() + h.expectCloseProcessHandle() + h.expectNoMoreProcesses() + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + { + name: "Process32First fails", + setupMock: func(h *mockHelper) { + h.expectSnapshot() + h.m.EXPECT(). + Process32First(testSnapshot, gomock.Any()). + Return(errors.New("failed to get first process")) + h.expectCloseSnapshot() + }, + expectError: true, + errorContains: "failed to find matching vmmem process", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockAPI := mock.NewMockAPI(ctrl) + helper := &mockHelper{m: mockAPI} + tt.setupMock(helper) + + handle, err := LookupVMMEM(context.Background(), testVMID, mockAPI) + + if tt.expectError { + if err == nil { + t.Errorf("expected error but got none") + } else if tt.errorContains != "" && !strings.Contains(err.Error(), tt.errorContains) { + t.Errorf("expected error to contain %q, but got: %v", tt.errorContains, err) + } + if handle != 0 { + t.Errorf("expected handle to be 0 on error, got: %v", handle) + } + } else { + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if handle == 0 { + t.Errorf("expected non-zero handle on success") + } + } + }) + } +} + +// makeProcessEntry creates a ProcessEntry32 with the given name and PID. +func makeProcessEntry(pid uint32, exeName string) windows.ProcessEntry32 { + var pe windows.ProcessEntry32 + pe.Size = uint32(unsafe.Sizeof(pe)) + pe.ProcessID = pid + utf16Name, _ := windows.UTF16FromString(exeName) + copy(pe.ExeFile[:], utf16Name) + return pe +} + +// mockHelper provides common mock setup operations to reduce duplication. +type mockHelper struct { + m *mock.MockAPI +} + +func (h *mockHelper) expectSnapshot() { + h.m.EXPECT(). + CreateToolhelp32Snapshot(uint32(windows.TH32CS_SNAPPROCESS), uint32(0)). + Return(testSnapshot, nil) +} + +func (h *mockHelper) expectCloseSnapshot() { + h.m.EXPECT().CloseHandle(testSnapshot).Return(nil) +} + +func (h *mockHelper) expectProcess32(pid uint32, name string) { + h.m.EXPECT(). + Process32First(testSnapshot, gomock.Any()). + DoAndReturn(func(_ windows.Handle, pe *windows.ProcessEntry32) error { + *pe = makeProcessEntry(pid, name) + return nil + }) +} + +func (h *mockHelper) expectNoMoreProcesses() { + h.m.EXPECT(). + Process32Next(testSnapshot, gomock.Any()). + Return(errNoMoreProcesses).AnyTimes() +} + +func (h *mockHelper) expectOpenProcess(pid uint32, handle windows.Handle, err error) { + h.m.EXPECT(). + OpenProcess(uint32(windows.PROCESS_QUERY_LIMITED_INFORMATION), false, pid). + Return(handle, err) +} + +func (h *mockHelper) expectCloseProcessHandle() { + h.m.EXPECT().CloseHandle(testProcessHandle).Return(nil) +} + +func (h *mockHelper) expectOpenProcessToken(err error) { + call := h.m.EXPECT(). + OpenProcessToken(testProcessHandle, uint32(windows.TOKEN_QUERY), gomock.Any()) + if err != nil { + call.Return(err) + } else { + call.DoAndReturn(func(_ windows.Handle, _ uint32, token *windows.Token) error { + *token = testToken + return nil + }) + } +} + +func (h *mockHelper) expectGetTokenUser(user *windows.Tokenuser, err error) { + h.m.EXPECT().GetTokenUser(testToken).Return(user, err) +} + +func (h *mockHelper) expectLookupAccount(account, domain string, err error) { + h.m.EXPECT().LookupAccount(mockSID, "").Return(account, domain, uint32(0), err) +} + +func (h *mockHelper) expectCloseToken() { + h.m.EXPECT().CloseToken(testToken).Return(nil) +} + +// expectSuccessfulMatch sets up all mocks for a successful vmmem match. +func (h *mockHelper) expectSuccessfulMatch(processName string) { + h.expectSnapshot() + h.expectProcess32(testPID, processName) + h.expectOpenProcess(testPID, testProcessHandle, nil) + h.expectOpenProcessToken(nil) + h.expectGetTokenUser(mockTokenUser, nil) + h.expectLookupAccount(testVMIDStr, "NT VIRTUAL MACHINE", nil) + h.expectCloseToken() + h.expectNoMoreProcesses() + h.expectCloseSnapshot() +} diff --git a/internal/windows/doc.go b/internal/windows/doc.go new file mode 100644 index 0000000000..1c7dc006d1 --- /dev/null +++ b/internal/windows/doc.go @@ -0,0 +1,76 @@ +//go:build windows + +/* +Package windows provides an abstraction layer over Windows API calls to enable better testability. + +# Purpose + +The API interface abstracts Windows system calls, allowing them to be mocked in tests. +This is particularly useful for testing code that interacts with processes, tokens, and security +identifiers without requiring actual Windows system resources. + +# Usage + +## Production Code + +In production code, use the NewWindowsAPI() function to get the real implementation: + + import ( + "github.com/Microsoft/hcsshim/internal/windows" + ) + + func MyFunction(ctx context.Context) error { + winAPI := &windows.WinAPI{} + // Use winAPI instead of calling windows package directly + snapshot, err := winAPI.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0) + // ... + } + +## Testing + +In tests, create a mock implementation of the API interface. Here's a simple example: + + // Function to test + func GetProcessSnapshot(api windows.API) (windows.Handle, error) { + snapshot, err := api.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0) + if err != nil { + return 0, err + } + return snapshot, nil + } + + // Mock implementation for testing + type mockAPI struct { + snapshot windows.Handle + err error + } + + func (m *mockAPI) CreateToolhelp32Snapshot(flags uint32, processID uint32) (windows.Handle, error) { + return m.snapshot, m.err + } + + func (m *mockAPI) CloseHandle(h windows.Handle) error { + return nil + } + + // ... implement other API methods as needed ... + + // Test using the mock + func TestGetProcessSnapshot(t *testing.T) { + mockWinAPI := &mockAPI{ + snapshot: windows.Handle(12345), + err: nil, + } + + snapshot, err := GetProcessSnapshot(mockWinAPI) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if snapshot != windows.Handle(12345) { + t.Fatalf("expected snapshot 12345, got %v", snapshot) + } + } + +All method names match the underlying Windows API calls for easy understanding and reference. +*/ +package windows diff --git a/internal/windows/interface.go b/internal/windows/interface.go new file mode 100644 index 0000000000..4dda52d7f6 --- /dev/null +++ b/internal/windows/interface.go @@ -0,0 +1,37 @@ +//go:build windows + +package windows + +import ( + "golang.org/x/sys/windows" +) + +// API abstracts Windows API calls for testing purposes. +type API interface { + // CreateToolhelp32Snapshot takes a snapshot of the specified processes. + CreateToolhelp32Snapshot(flags uint32, processID uint32) (windows.Handle, error) + + // CloseHandle closes an open object handle. + CloseHandle(h windows.Handle) error + + // Process32First retrieves information about the first process in a snapshot. + Process32First(snapshot windows.Handle, pe *windows.ProcessEntry32) error + + // Process32Next retrieves information about the next process in a snapshot. + Process32Next(snapshot windows.Handle, pe *windows.ProcessEntry32) error + + // OpenProcess opens an existing local process object. + OpenProcess(desiredAccess uint32, inheritHandle bool, processID uint32) (windows.Handle, error) + + // OpenProcessToken opens the access token associated with a process. + OpenProcessToken(process windows.Handle, desiredAccess uint32, token *windows.Token) error + + // GetTokenUser retrieves the user account of the token. + GetTokenUser(token windows.Token) (*windows.Tokenuser, error) + + // LookupAccount retrieves the name of the account for a SID and the name of the first domain on which this SID is found. + LookupAccount(sid *windows.SID, system string) (account string, domain string, accType uint32, err error) + + // CloseToken closes a token handle. + CloseToken(token windows.Token) error +} diff --git a/internal/windows/mock/mock_windows.go b/internal/windows/mock/mock_windows.go new file mode 100644 index 0000000000..b4cd1f16c1 --- /dev/null +++ b/internal/windows/mock/mock_windows.go @@ -0,0 +1,173 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: interface.go +// +// Generated by this command: +// +// mockgen -source=interface.go -package=mock -destination=mock\mock_windows.go +// + +// Package mock is a generated GoMock package. +package mock + +import ( + reflect "reflect" + + gomock "go.uber.org/mock/gomock" + windows "golang.org/x/sys/windows" +) + +// MockAPI is a mock of API interface. +type MockAPI struct { + ctrl *gomock.Controller + recorder *MockAPIMockRecorder + isgomock struct{} +} + +// MockAPIMockRecorder is the mock recorder for MockAPI. +type MockAPIMockRecorder struct { + mock *MockAPI +} + +// NewMockAPI creates a new mock instance. +func NewMockAPI(ctrl *gomock.Controller) *MockAPI { + mock := &MockAPI{ctrl: ctrl} + mock.recorder = &MockAPIMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockAPI) EXPECT() *MockAPIMockRecorder { + return m.recorder +} + +// CloseHandle mocks base method. +func (m *MockAPI) CloseHandle(h windows.Handle) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CloseHandle", h) + ret0, _ := ret[0].(error) + return ret0 +} + +// CloseHandle indicates an expected call of CloseHandle. +func (mr *MockAPIMockRecorder) CloseHandle(h any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CloseHandle", reflect.TypeOf((*MockAPI)(nil).CloseHandle), h) +} + +// CloseToken mocks base method. +func (m *MockAPI) CloseToken(token windows.Token) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CloseToken", token) + ret0, _ := ret[0].(error) + return ret0 +} + +// CloseToken indicates an expected call of CloseToken. +func (mr *MockAPIMockRecorder) CloseToken(token any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CloseToken", reflect.TypeOf((*MockAPI)(nil).CloseToken), token) +} + +// CreateToolhelp32Snapshot mocks base method. +func (m *MockAPI) CreateToolhelp32Snapshot(flags, processID uint32) (windows.Handle, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CreateToolhelp32Snapshot", flags, processID) + ret0, _ := ret[0].(windows.Handle) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// CreateToolhelp32Snapshot indicates an expected call of CreateToolhelp32Snapshot. +func (mr *MockAPIMockRecorder) CreateToolhelp32Snapshot(flags, processID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateToolhelp32Snapshot", reflect.TypeOf((*MockAPI)(nil).CreateToolhelp32Snapshot), flags, processID) +} + +// GetTokenUser mocks base method. +func (m *MockAPI) GetTokenUser(token windows.Token) (*windows.Tokenuser, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetTokenUser", token) + ret0, _ := ret[0].(*windows.Tokenuser) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetTokenUser indicates an expected call of GetTokenUser. +func (mr *MockAPIMockRecorder) GetTokenUser(token any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTokenUser", reflect.TypeOf((*MockAPI)(nil).GetTokenUser), token) +} + +// LookupAccount mocks base method. +func (m *MockAPI) LookupAccount(sid *windows.SID, system string) (string, string, uint32, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupAccount", sid, system) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(string) + ret2, _ := ret[2].(uint32) + ret3, _ := ret[3].(error) + return ret0, ret1, ret2, ret3 +} + +// LookupAccount indicates an expected call of LookupAccount. +func (mr *MockAPIMockRecorder) LookupAccount(sid, system any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupAccount", reflect.TypeOf((*MockAPI)(nil).LookupAccount), sid, system) +} + +// OpenProcess mocks base method. +func (m *MockAPI) OpenProcess(desiredAccess uint32, inheritHandle bool, processID uint32) (windows.Handle, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OpenProcess", desiredAccess, inheritHandle, processID) + ret0, _ := ret[0].(windows.Handle) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OpenProcess indicates an expected call of OpenProcess. +func (mr *MockAPIMockRecorder) OpenProcess(desiredAccess, inheritHandle, processID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenProcess", reflect.TypeOf((*MockAPI)(nil).OpenProcess), desiredAccess, inheritHandle, processID) +} + +// OpenProcessToken mocks base method. +func (m *MockAPI) OpenProcessToken(process windows.Handle, desiredAccess uint32, token *windows.Token) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OpenProcessToken", process, desiredAccess, token) + ret0, _ := ret[0].(error) + return ret0 +} + +// OpenProcessToken indicates an expected call of OpenProcessToken. +func (mr *MockAPIMockRecorder) OpenProcessToken(process, desiredAccess, token any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OpenProcessToken", reflect.TypeOf((*MockAPI)(nil).OpenProcessToken), process, desiredAccess, token) +} + +// Process32First mocks base method. +func (m *MockAPI) Process32First(snapshot windows.Handle, pe *windows.ProcessEntry32) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Process32First", snapshot, pe) + ret0, _ := ret[0].(error) + return ret0 +} + +// Process32First indicates an expected call of Process32First. +func (mr *MockAPIMockRecorder) Process32First(snapshot, pe any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Process32First", reflect.TypeOf((*MockAPI)(nil).Process32First), snapshot, pe) +} + +// Process32Next mocks base method. +func (m *MockAPI) Process32Next(snapshot windows.Handle, pe *windows.ProcessEntry32) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Process32Next", snapshot, pe) + ret0, _ := ret[0].(error) + return ret0 +} + +// Process32Next indicates an expected call of Process32Next. +func (mr *MockAPIMockRecorder) Process32Next(snapshot, pe any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Process32Next", reflect.TypeOf((*MockAPI)(nil).Process32Next), snapshot, pe) +} diff --git a/internal/windows/windows.go b/internal/windows/windows.go new file mode 100644 index 0000000000..53940065cf --- /dev/null +++ b/internal/windows/windows.go @@ -0,0 +1,55 @@ +//go:build windows + +package windows + +import ( + "golang.org/x/sys/windows" +) + +// WinAPI is the real implementation of windows.API that delegates to the actual Windows API. +type WinAPI struct{} + +// CreateToolhelp32Snapshot takes a snapshot of the specified processes. +func (w *WinAPI) CreateToolhelp32Snapshot(flags uint32, processID uint32) (windows.Handle, error) { + return windows.CreateToolhelp32Snapshot(flags, processID) +} + +// CloseHandle closes an open object handle. +func (w *WinAPI) CloseHandle(h windows.Handle) error { + return windows.CloseHandle(h) +} + +// Process32First retrieves information about the first process in a snapshot. +func (w *WinAPI) Process32First(snapshot windows.Handle, pe *windows.ProcessEntry32) error { + return windows.Process32First(snapshot, pe) +} + +// Process32Next retrieves information about the next process in a snapshot. +func (w *WinAPI) Process32Next(snapshot windows.Handle, pe *windows.ProcessEntry32) error { + return windows.Process32Next(snapshot, pe) +} + +// OpenProcess opens an existing local process object. +func (w *WinAPI) OpenProcess(desiredAccess uint32, inheritHandle bool, processID uint32) (windows.Handle, error) { + return windows.OpenProcess(desiredAccess, inheritHandle, processID) +} + +// OpenProcessToken opens the access token associated with a process. +func (w *WinAPI) OpenProcessToken(process windows.Handle, desiredAccess uint32, token *windows.Token) error { + return windows.OpenProcessToken(process, desiredAccess, token) +} + +// GetTokenUser retrieves the user account of the token. +func (w *WinAPI) GetTokenUser(token windows.Token) (*windows.Tokenuser, error) { + return token.GetTokenUser() +} + +// LookupAccount retrieves the name of the account for a SID and the name of the first domain on which this SID is found. +func (w *WinAPI) LookupAccount(sid *windows.SID, system string) (account string, domain string, accType uint32, err error) { + return sid.LookupAccount(system) +} + +// CloseToken closes a token handle. +func (w *WinAPI) CloseToken(token windows.Token) error { + return token.Close() +} diff --git a/test/go.sum b/test/go.sum index 11b84610fc..e3eb3b89d1 100644 --- a/test/go.sum +++ b/test/go.sum @@ -345,6 +345,8 @@ go.opentelemetry.io/proto/otlp v1.5.0 h1:xJvq7gMzB31/d406fB8U5CBdyQGw4P399D1aQWU go.opentelemetry.io/proto/otlp v1.5.0/go.mod h1:keN8WnHxOy8PG0rQZjJJ5A2ebUoafqWp0eVQ4yIXvJ4= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= +go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=