internal/resource/resource_windows_test.go - third_party/google-guest-agent - Git at Google

 //  Copyright 2024 Google LLC
 //
 //  Licensed under the Apache License, Version 2.0 (the "License");
 //  you may not use this file except in compliance with the License.
 //  You may obtain a copy of the License at
 //
 //     https://www.apache.org/licenses/LICENSE-2.0
 //
 //  Unless required by applicable law or agreed to in writing, software
 //  distributed under the License is distributed on an "AS IS" BASIS,
 //  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 //  See the License for the specific language governing permissions and
 //  limitations under the License.

 //go:build windows

 package resource

 import (
 	"context"
 	"fmt"
 	"os"
 	"syscall"
 	"testing"
 	"unsafe"

 	"github.com/GoogleCloudPlatform/google-guest-agent/internal/run"
 	"golang.org/x/sys/windows"
 )

 // constraintTestOpts contains options for setting up the test. This is primarily used for
 // error injection.
 type constraintTestOpts struct {
 	// overrideAll indicates whether to override all functions with mock variants.
 	overrideAll bool
 	// createJobObjErr indicates whether to return an error when creating job object.
 	createJobObjErr bool
 	// overrideCreateJob indicates whether to override the createJob function.
 	overrideCreateJob bool
 	// openProcessErr indicates whether to return an error when opening process.
 	openProcessErr bool
 	// overrideOpenProcess indicates whether to override the openProcess function.
 	overrideOpenProcess bool
 	// assignProccessErr indicates whether to return an error when assigning process to job object.
 	assignProccessErr bool
 	// overrideAssignProcess indicates whether to override the assignProcessToJobObj function.
 	overrideAssignProcess bool
 	// setInfoErr indicates whether to return an error when setting information job object.
 	setInfoErr bool
 	// overrideSetInfo indicates whether to override the setInformationJobObj function.
 	overrideSetInfo bool
 }

 func constraintTestSetup(t *testing.T, opts constraintTestOpts) *windowsClient {
 	t.Helper()
 	client := &windowsClient{
 		createJobObject:          windows.CreateJobObject,
 		getProcessHandle:         windows.OpenProcess,
 		assignProcessToJobObject: windows.AssignProcessToJobObject,
 		setInformationJobObject:  windows.SetInformationJobObject,
 		handles:                  map[string]windows.Handle{},
 	}

 	if opts.overrideAll || opts.overrideCreateJob {
 		client.createJobObject = func(_ *windows.SecurityAttributes, _ *uint16) (windows.Handle, error) {
 			if opts.createJobObjErr {
 				return 0, fmt.Errorf("mock fail createJobObj")
 			}
 			return 1, nil
 		}
 	}

 	if opts.overrideAll || opts.overrideOpenProcess {
 		client.getProcessHandle = func(_ uint32, _ bool, _ uint32) (windows.Handle, error) {
 			if opts.openProcessErr {
 				return 0, fmt.Errorf("mock fail openProcess")
 			}
 			return 1, nil
 		}
 	}

 	if opts.overrideAll || opts.overrideAssignProcess {
 		client.assignProcessToJobObject = func(_ windows.Handle, _ windows.Handle) error {
 			if opts.assignProccessErr {
 				return fmt.Errorf("mock fail assignProcessToJobObj")
 			}
 			return nil
 		}
 	}

 	if opts.overrideAll || opts.overrideSetInfo {
 		client.setInformationJobObject = func(_ windows.Handle, _ uint32, _ uintptr, _ uint32) (ret int, err error) {
 			if opts.setInfoErr {
 				return 0, fmt.Errorf("mock fail setInformationJobObj")
 			}
 			return 1, nil
 		}
 	}

 	t.Cleanup(func() {
 		// Make sure to close all the handles.
 		for _, handle := range client.handles {
 			windows.CloseHandle(handle)
 		}
 	})
 	return client
 }

 // TestApplyResourceConstraintsError tests the error handling of applyResourceConstraints().
 func TestApplyResourceConstraintsError(t *testing.T) {
 	test := []struct {
 		// name is the name of the test.
 		name string
 		// pages is the number of pages for setting workingsetsize.
 		pages int
 		// opts is the options for setting up the test.
 		opts constraintTestOpts
 	}{
 		{
 			name: "create-job-error",
 			opts: constraintTestOpts{
 				overrideAll:     true,
 				createJobObjErr: true,
 			},
 		},
 		{
 			name: "get-process-error",
 			opts: constraintTestOpts{
 				overrideAll:    true,
 				openProcessErr: true,
 			},
 		},
 		{
 			name: "assign-process-error",
 			opts: constraintTestOpts{
 				overrideAll:       true,
 				assignProccessErr: true,
 			},
 		},
 		{
 			name:  "memory-too-small",
 			pages: 1,
 			opts: constraintTestOpts{
 				overrideAll: true,
 			},
 		},
 		{
 			name: "set-info-error",
 			opts: constraintTestOpts{
 				overrideAll: true,
 				setInfoErr:  true,
 			},
 		},
 	}

 	for _, test := range test {
 		t.Run(test.name, func(t *testing.T) {
 			c := constraintTestSetup(t, test.opts)

 			var maxMemory int64
 			if test.pages > 0 {
 				maxMemory = int64(test.pages * os.Getpagesize())
 			} else {
 				maxMemory = int64(21 * os.Getpagesize())
 			}
 			err := c.Apply(Constraint{Name: "testPlugin_A", MaxMemoryUsage: maxMemory, MaxCPUUsage: 10})
 			if err == nil {
 				t.Errorf("Apply() returned no error, error expected")
 			}
 		})
 	}
 }

 // TestApplyResourceConstraints tests if resource constraints are successfully applied to the
 // created JobObject.
 func TestApplyResourceConstraints(t *testing.T) {
 	c := constraintTestSetup(t, constraintTestOpts{})
 	testName := "testPlugin_A"

 	// Start a process for adding to the JobObject.
 	options := run.Options{
 		Name:       "powershell",
 		Args:       []string{"-NoProfile", "Start-Sleep -s 30"},
 		ExecMode:   run.ExecModeAsync,
 		OutputType: run.OutputNone,
 	}
 	res, err := run.WithContext(context.Background(), options)
 	if err != nil {
 		t.Fatalf("failed to start up test program: %v", err)
 	}
 	t.Cleanup(func() {
 		// Kill the test process if it's still running.
 		process, _ := os.FindProcess(res.Pid)
 		if process != nil {
 			process.Kill()
 		}
 	})

 	testConstraint := Constraint{
 		PID:            res.Pid,
 		Name:           testName,
 		MaxMemoryUsage: int64(21 * os.Getpagesize()),
 		MaxCPUUsage:    50,
 	}
 	if err := c.Apply(testConstraint); err != nil {
 		t.Fatalf("Apply(%+v) returned unexpected error: %v", testConstraint, err)
 	}

 	// Get the handle.
 	handle, found := c.handles[testName]
 	if !found {
 		t.Fatalf("failed to find job object handle for %s", testName)
 	}
 	defer windows.CloseHandle(handle)

 	// Check that the memory limits are set correctly.
 	var memoryLimitInfo windows.JOBOBJECT_BASIC_LIMIT_INFORMATION
 	if err := windows.QueryInformationJobObject(
 		handle,
 		windows.JobObjectBasicLimitInformation,
 		uintptr(unsafe.Pointer(&memoryLimitInfo)),
 		uint32(unsafe.Sizeof(memoryLimitInfo)),
 		nil,
 	); err != nil {
 		t.Fatalf("failed to query basic limit information: %v", err)
 	}

 	// Since the value pointed to by the pointer reside in the Plugin struct, this should be safe.
 	actual := int64(memoryLimitInfo.MaximumWorkingSetSize)
 	if actual != testConstraint.MaxMemoryUsage {
 		t.Errorf("Apply(%+v) set MaximumWorkingSetSize %d, expected %d", testConstraint, actual, testConstraint.MaxMemoryUsage)
 	}

 	// Check that the CPU limits are set correctly.
 	var cpuLimitInfo jobObjectCPURateControlInfo
 	if err := windows.QueryInformationJobObject(
 		handle,
 		windows.JobObjectCpuRateControlInformation,
 		uintptr(unsafe.Pointer(&cpuLimitInfo)),
 		uint32(unsafe.Sizeof(cpuLimitInfo)),
 		nil,
 	); err != nil {
 		t.Fatalf("failed to query CPU rate control information: %v", err)
 	}
 	expectedCPULimit := int(testConstraint.MaxCPUUsage * 100)
 	if int(cpuLimitInfo.Value) != expectedCPULimit {
 		t.Errorf("Apply(%+v) set CPU rate control value %d, expected %d", testConstraint, cpuLimitInfo.Value, expectedCPULimit)
 	}
 }

 func TestRemoveConstraint(t *testing.T) {
 	c := constraintTestSetup(t, constraintTestOpts{
 		overrideOpenProcess:   true,
 		overrideAssignProcess: true,
 		overrideSetInfo:       true,
 	})
 	testName := "testPlugin_B"
 	testConstraint := Constraint{
 		PID:            4,
 		Name:           testName,
 		MaxMemoryUsage: int64(21 * os.Getpagesize()),
 		MaxCPUUsage:    50,
 	}
 	if err := c.Apply(testConstraint); err != nil {
 		t.Fatalf("Apply(ctx, %+v) returned unexpected error: %v", testConstraint, err)
 	}

 	// Attempt to close it. It should not error.
 	if err := c.RemoveConstraint(context.Background(), testName); err != nil {
 		t.Fatalf("RemoveConstraint(ctx, %s) returned unexpected error: %v", testName, err)
 	}

 	// Check if the handle was successfully deleted from the map.
 	_, found := c.handles[testName]
 	if found {
 		t.Fatalf("RemoveConstraint(ctx, %s) did not remove job object handle", testName)
 	}

 	// Attempt to find the JobObject. It should no longer exist.
 	modKernel32 := windows.NewLazyDLL("kernel32.dll")
 	procOpenJobObject := modKernel32.NewProc("OpenJobObjectW")
 	name, err := windows.UTF16PtrFromString(testName)
 	if err != nil {
 		t.Fatalf("failed to parse job object name: %v", err)
 	}
 	h, _, _ := syscall.SyscallN(procOpenJobObject.Addr(), uintptr(0x1F001F), 0, uintptr(unsafe.Pointer(name)))
 	if h != 0 {
 		t.Fatalf("JobObject still exists")
 	}
 }
	// Copyright 2024 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	//go:build windows

	package resource

	import (
	"context"
	"fmt"
	"os"
	"syscall"
	"testing"
	"unsafe"

	"github.com/GoogleCloudPlatform/google-guest-agent/internal/run"
	"golang.org/x/sys/windows"
	)

	// constraintTestOpts contains options for setting up the test. This is primarily used for
	// error injection.
	type constraintTestOpts struct {
	// overrideAll indicates whether to override all functions with mock variants.
	overrideAll bool
	// createJobObjErr indicates whether to return an error when creating job object.
	createJobObjErr bool
	// overrideCreateJob indicates whether to override the createJob function.
	overrideCreateJob bool
	// openProcessErr indicates whether to return an error when opening process.
	openProcessErr bool
	// overrideOpenProcess indicates whether to override the openProcess function.
	overrideOpenProcess bool
	// assignProccessErr indicates whether to return an error when assigning process to job object.
	assignProccessErr bool
	// overrideAssignProcess indicates whether to override the assignProcessToJobObj function.
	overrideAssignProcess bool
	// setInfoErr indicates whether to return an error when setting information job object.
	setInfoErr bool
	// overrideSetInfo indicates whether to override the setInformationJobObj function.
	overrideSetInfo bool
	}

	func constraintTestSetup(t testing.T, opts constraintTestOpts) windowsClient {
	t.Helper()
	client := &windowsClient{
	createJobObject: windows.CreateJobObject,
	getProcessHandle: windows.OpenProcess,
	assignProcessToJobObject: windows.AssignProcessToJobObject,
	setInformationJobObject: windows.SetInformationJobObject,
	handles: map[string]windows.Handle{},
	}

	if opts.overrideAll \|\| opts.overrideCreateJob {
	client.createJobObject = func(_ windows.SecurityAttributes, _ uint16) (windows.Handle, error) {
	if opts.createJobObjErr {
	return 0, fmt.Errorf("mock fail createJobObj")
	}
	return 1, nil
	}
	}

	if opts.overrideAll \|\| opts.overrideOpenProcess {
	client.getProcessHandle = func(_ uint32, _ bool, _ uint32) (windows.Handle, error) {
	if opts.openProcessErr {
	return 0, fmt.Errorf("mock fail openProcess")
	}
	return 1, nil
	}
	}

	if opts.overrideAll \|\| opts.overrideAssignProcess {
	client.assignProcessToJobObject = func(_ windows.Handle, _ windows.Handle) error {
	if opts.assignProccessErr {
	return fmt.Errorf("mock fail assignProcessToJobObj")
	}
	return nil
	}
	}

	if opts.overrideAll \|\| opts.overrideSetInfo {
	client.setInformationJobObject = func(_ windows.Handle, _ uint32, _ uintptr, _ uint32) (ret int, err error) {
	if opts.setInfoErr {
	return 0, fmt.Errorf("mock fail setInformationJobObj")
	}
	return 1, nil
	}
	}

	t.Cleanup(func() {
	// Make sure to close all the handles.
	for _, handle := range client.handles {
	windows.CloseHandle(handle)
	}
	})
	return client
	}

	// TestApplyResourceConstraintsError tests the error handling of applyResourceConstraints().
	func TestApplyResourceConstraintsError(t *testing.T) {
	test := []struct {
	// name is the name of the test.
	name string
	// pages is the number of pages for setting workingsetsize.
	pages int
	// opts is the options for setting up the test.
	opts constraintTestOpts
	}{
	{
	name: "create-job-error",
	opts: constraintTestOpts{
	overrideAll: true,
	createJobObjErr: true,
	},
	},
	{
	name: "get-process-error",
	opts: constraintTestOpts{
	overrideAll: true,
	openProcessErr: true,
	},
	},
	{
	name: "assign-process-error",
	opts: constraintTestOpts{
	overrideAll: true,
	assignProccessErr: true,
	},
	},
	{
	name: "memory-too-small",
	pages: 1,
	opts: constraintTestOpts{
	overrideAll: true,
	},
	},
	{
	name: "set-info-error",
	opts: constraintTestOpts{
	overrideAll: true,
	setInfoErr: true,
	},
	},
	}

	for _, test := range test {
	t.Run(test.name, func(t *testing.T) {
	c := constraintTestSetup(t, test.opts)

	var maxMemory int64
	if test.pages > 0 {
	maxMemory = int64(test.pages * os.Getpagesize())
	} else {
	maxMemory = int64(21 * os.Getpagesize())
	}
	err := c.Apply(Constraint{Name: "testPlugin_A", MaxMemoryUsage: maxMemory, MaxCPUUsage: 10})
	if err == nil {
	t.Errorf("Apply() returned no error, error expected")
	}
	})
	}
	}

	// TestApplyResourceConstraints tests if resource constraints are successfully applied to the
	// created JobObject.
	func TestApplyResourceConstraints(t *testing.T) {
	c := constraintTestSetup(t, constraintTestOpts{})
	testName := "testPlugin_A"

	// Start a process for adding to the JobObject.
	options := run.Options{
	Name: "powershell",
	Args: []string{"-NoProfile", "Start-Sleep -s 30"},
	ExecMode: run.ExecModeAsync,
	OutputType: run.OutputNone,
	}
	res, err := run.WithContext(context.Background(), options)
	if err != nil {
	t.Fatalf("failed to start up test program: %v", err)
	}
	t.Cleanup(func() {
	// Kill the test process if it's still running.
	process, _ := os.FindProcess(res.Pid)
	if process != nil {
	process.Kill()
	}
	})

	testConstraint := Constraint{
	PID: res.Pid,
	Name: testName,
	MaxMemoryUsage: int64(21 * os.Getpagesize()),
	MaxCPUUsage: 50,
	}
	if err := c.Apply(testConstraint); err != nil {
	t.Fatalf("Apply(%+v) returned unexpected error: %v", testConstraint, err)
	}

	// Get the handle.
	handle, found := c.handles[testName]
	if !found {
	t.Fatalf("failed to find job object handle for %s", testName)
	}
	defer windows.CloseHandle(handle)

	// Check that the memory limits are set correctly.
	var memoryLimitInfo windows.JOBOBJECT_BASIC_LIMIT_INFORMATION
	if err := windows.QueryInformationJobObject(
	handle,
	windows.JobObjectBasicLimitInformation,
	uintptr(unsafe.Pointer(&memoryLimitInfo)),
	uint32(unsafe.Sizeof(memoryLimitInfo)),
	nil,
	); err != nil {
	t.Fatalf("failed to query basic limit information: %v", err)
	}

	// Since the value pointed to by the pointer reside in the Plugin struct, this should be safe.
	actual := int64(memoryLimitInfo.MaximumWorkingSetSize)
	if actual != testConstraint.MaxMemoryUsage {
	t.Errorf("Apply(%+v) set MaximumWorkingSetSize %d, expected %d", testConstraint, actual, testConstraint.MaxMemoryUsage)
	}

	// Check that the CPU limits are set correctly.
	var cpuLimitInfo jobObjectCPURateControlInfo
	if err := windows.QueryInformationJobObject(
	handle,
	windows.JobObjectCpuRateControlInformation,
	uintptr(unsafe.Pointer(&cpuLimitInfo)),
	uint32(unsafe.Sizeof(cpuLimitInfo)),
	nil,
	); err != nil {
	t.Fatalf("failed to query CPU rate control information: %v", err)
	}
	expectedCPULimit := int(testConstraint.MaxCPUUsage * 100)
	if int(cpuLimitInfo.Value) != expectedCPULimit {
	t.Errorf("Apply(%+v) set CPU rate control value %d, expected %d", testConstraint, cpuLimitInfo.Value, expectedCPULimit)
	}
	}

	func TestRemoveConstraint(t *testing.T) {
	c := constraintTestSetup(t, constraintTestOpts{
	overrideOpenProcess: true,
	overrideAssignProcess: true,
	overrideSetInfo: true,
	})
	testName := "testPlugin_B"
	testConstraint := Constraint{
	PID: 4,
	Name: testName,
	MaxMemoryUsage: int64(21 * os.Getpagesize()),
	MaxCPUUsage: 50,
	}
	if err := c.Apply(testConstraint); err != nil {
	t.Fatalf("Apply(ctx, %+v) returned unexpected error: %v", testConstraint, err)
	}

	// Attempt to close it. It should not error.
	if err := c.RemoveConstraint(context.Background(), testName); err != nil {
	t.Fatalf("RemoveConstraint(ctx, %s) returned unexpected error: %v", testName, err)
	}

	// Check if the handle was successfully deleted from the map.
	_, found := c.handles[testName]
	if found {
	t.Fatalf("RemoveConstraint(ctx, %s) did not remove job object handle", testName)
	}

	// Attempt to find the JobObject. It should no longer exist.
	modKernel32 := windows.NewLazyDLL("kernel32.dll")
	procOpenJobObject := modKernel32.NewProc("OpenJobObjectW")
	name, err := windows.UTF16PtrFromString(testName)
	if err != nil {
	t.Fatalf("failed to parse job object name: %v", err)
	}
	h, _, _ := syscall.SyscallN(procOpenJobObject.Addr(), uintptr(0x1F001F), 0, uintptr(unsafe.Pointer(name)))
	if h != 0 {
	t.Fatalf("JobObject still exists")
	}
	}