From e081f3e5b99a9c3bc959a91b22ac248ee27cb229 Mon Sep 17 00:00:00 2001 From: Kathryn Baldauf Date: Mon, 21 Dec 2020 22:26:30 -0800 Subject: [PATCH] Create tool to get/set job object resource limits Signed-off-by: Kathryn Baldauf --- cmd/jobobject-util/main.go | 205 +++++++++++++++++++ internal/jobobject/jobobject.go | 220 ++++++++++---------- internal/jobobject/limits.go | 305 ++++++++++++++++++++++++++++ internal/safefile/safeopen.go | 2 +- internal/winapi/filesystem.go | 2 +- internal/winapi/jobobject.go | 30 +++ internal/winapi/utils.go | 20 +- internal/winapi/zsyscall_windows.go | 78 ++++--- internal/winobjdir/object_dir.go | 2 +- 9 files changed, 714 insertions(+), 150 deletions(-) create mode 100644 cmd/jobobject-util/main.go create mode 100644 internal/jobobject/limits.go diff --git a/cmd/jobobject-util/main.go b/cmd/jobobject-util/main.go new file mode 100644 index 0000000000..34b18a31f0 --- /dev/null +++ b/cmd/jobobject-util/main.go @@ -0,0 +1,205 @@ +package main + +import ( + "context" + "errors" + "fmt" + "os" + "strconv" + + "github.com/Microsoft/hcsshim/internal/jobobject" + "github.com/urfave/cli" +) + +const ( + cpuLimitFlag = "cpu-limit" + cpuWeightFlag = "cpu-weight" + memoryLimitFlag = "memory-limit" + affinityFlag = "cpu-affinity" + useNTVariantFlag = "use-nt" + + usage = `jobobject-util is a command line tool for getting and setting job object limits` +) + +func main() { + app := cli.NewApp() + app.Name = "jobobject-util" + app.Commands = []cli.Command{ + getJobObjectLimitsCommand, + setJobObjectLimitsCommand, + } + app.Usage = usage + + if err := app.Run(os.Args); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} + +var getJobObjectLimitsCommand = cli.Command{ + Name: "get", + Usage: "gets the job object's resource limits", + ArgsUsage: "get [flags] ", + Flags: []cli.Flag{ + cli.BoolFlag{ + Name: useNTVariantFlag, + Usage: `Optional: indicates if the command should use the NT variant of job object Open/Create calls. `, + }, + cli.BoolFlag{ + Name: cpuLimitFlag, + Usage: "Optional: get job object's CPU limit", + }, + cli.BoolFlag{ + Name: cpuWeightFlag, + Usage: "Optional: get job object's CPU weight.", + }, + cli.BoolFlag{ + Name: memoryLimitFlag, + Usage: "Optional: get job object's memory limit in bytes.", + }, + cli.BoolFlag{ + Name: affinityFlag, + Usage: "Optional: get job object's CPU affinity as a bitmask.", + }, + }, + Action: func(cli *cli.Context) error { + ctx := context.Background() + name := cli.Args().First() + if name == "" { + return errors.New("`get` command must specify a target job object name") + } + options := &jobobject.Options{ + Name: name, + Notifications: false, + UseNTVariant: cli.Bool(useNTVariantFlag), + } + job, err := jobobject.Open(ctx, options) + if err != nil { + return err + } + defer job.Close() + + output := "" + + // Only allow one processor related flag since limit and weight are + // mutually exclusive for a job object + if cli.IsSet(cpuLimitFlag) && cli.IsSet(cpuWeightFlag) { + return errors.New("cpu limit and weight are mutually exclusive") + } + if cli.IsSet(cpuLimitFlag) { + cpuRate, err := job.GetCPULimit(jobobject.RateBased) + if err != nil { + return err + } + output += fmt.Sprintf("%s: %d\n", cpuLimitFlag, cpuRate) + } else if cli.IsSet(cpuWeightFlag) { + cpuWeight, err := job.GetCPULimit(jobobject.WeightBased) + if err != nil { + return err + } + output += fmt.Sprintf("%s: %d\n", cpuWeightFlag, cpuWeight) + } + + if cli.IsSet(memoryLimitFlag) { + jobObjMemLimit, err := job.GetMemoryLimit() + if err != nil { + return err + } + output += fmt.Sprintf("%s: %d\n", memoryLimitFlag, jobObjMemLimit) + } + + if cli.IsSet(affinityFlag) { + affinity, err := job.GetCPUAffinity() + if err != nil { + return err + } + affinityString := strconv.FormatUint(affinity, 2) + output += fmt.Sprintf("%s: %s\n", affinityFlag, affinityString) + } + fmt.Fprintln(os.Stdout, output) + return nil + }, +} + +var setJobObjectLimitsCommand = cli.Command{ + Name: "set", + Usage: "tool used to set resource limits on job objects", + ArgsUsage: "set [flags] ", + Flags: []cli.Flag{ + cli.BoolFlag{ + Name: useNTVariantFlag, + Usage: `Optional: indicates if the command should use the NT variant of job object Open/Create calls. `, + }, + cli.Uint64Flag{ + Name: cpuLimitFlag, + Usage: "Optional: set job object's CPU limit", + }, + cli.Uint64Flag{ + Name: cpuWeightFlag, + Usage: "Optional: set job object's CPU weight.", + }, + cli.Uint64Flag{ + Name: memoryLimitFlag, + Usage: "Optional: set job object's memory limit in bytes.", + }, + cli.StringFlag{ + Name: affinityFlag, + Usage: "Optional: set job object's CPU affinity given a bitmask", + }, + }, + Action: func(cli *cli.Context) error { + ctx := context.Background() + name := cli.Args().First() + if name == "" { + return errors.New("`set` command must specify a job object name") + } + + options := &jobobject.Options{ + Name: name, + Notifications: false, + UseNTVariant: cli.Bool(useNTVariantFlag), + } + job, err := jobobject.Open(ctx, options) + if err != nil { + return err + } + defer job.Close() + + // Only allow one processor related flag since limit and weight are + // mutually exclusive for a job object + if cli.IsSet(cpuLimitFlag) && cli.IsSet(cpuWeightFlag) { + return errors.New("cpu limit and weight are mutually exclusive") + } + if cli.IsSet(cpuLimitFlag) { + cpuRate := uint32(cli.Uint64(cpuLimitFlag)) + if err := job.SetCPULimit(jobobject.RateBased, cpuRate); err != nil { + return err + } + } else if cli.IsSet(cpuWeightFlag) { + cpuWeight := uint32(cli.Uint64(cpuWeightFlag)) + if err := job.SetCPULimit(jobobject.WeightBased, cpuWeight); err != nil { + return err + } + } + + if cli.IsSet(memoryLimitFlag) { + memLimitInBytes := cli.Uint64(memoryLimitFlag) + if err := job.SetMemoryLimit(memLimitInBytes); err != nil { + return err + } + } + + if cli.IsSet(affinityFlag) { + affinityString := cli.String(affinityFlag) + affinity, err := strconv.ParseUint(affinityString, 2, 64) + if err != nil { + return err + } + if err := job.SetCPUAffinity(affinity); err != nil { + return err + } + } + + return nil + }, +} diff --git a/internal/jobobject/jobobject.go b/internal/jobobject/jobobject.go index 143e6b3126..9137c1825a 100644 --- a/internal/jobobject/jobobject.go +++ b/internal/jobobject/jobobject.go @@ -58,31 +58,52 @@ var ( ErrNotRegistered = errors.New("job is not registered to receive notifications") ) +type Options struct { + // `Name` specifies the name of the job object if a named job object is desired. + Name string + // `Notifications` specifies if the job will be registered to receive notifications. + // Defaults to false. + Notifications bool + // `UseNTVariant` specifies if we should use the `Nt` variant of Open/CreateJobObject. + // Defaults to false. + UseNTVariant bool +} + // Create creates a job object. // -// `name` specifies the name of the job object if a named job object is desired. If name -// is an empty string, the job will not be assigned a name. +// If name is an empty string, the job will not be assigned a name. +// +// If notifications are not enabled `PollNotifications` will return immediately with error `errNotRegistered`. // -// `notifications` specifies if the job will be registered to receive notifications. -// If this is false, `PollNotifications` will return immediately with error `errNotRegistered`. +// If `options` is nil, use default option values. // // Returns a JobObject structure and an error if there is one. -func Create(ctx context.Context, name string, notifications bool) (_ *JobObject, err error) { - var ( - jobName *uint16 - mq *queue.MessageQueue - ) +func Create(ctx context.Context, options *Options) (_ *JobObject, err error) { + var jobName *winapi.UnicodeString - if name != "" { - jobName, err = windows.UTF16PtrFromString(name) + if options != nil && options.Name != "" { + jobName, err = winapi.NewUnicodeString(options.Name) if err != nil { return nil, err } } - jobHandle, err := windows.CreateJobObject(nil, jobName) - if err != nil { - return nil, err + var jobHandle windows.Handle + if options != nil && options.UseNTVariant { + oa := winapi.ObjectAttributes{ + Length: unsafe.Sizeof(winapi.ObjectAttributes{}), + ObjectName: jobName, + Attributes: 0, + } + status := winapi.NtCreateJobObject(&jobHandle, winapi.JOB_OBJECT_ALL_ACCESS, &oa) + if status != 0 { + return nil, winapi.RtlNtStatusToDosError(status) + } + } else { + jobHandle, err = windows.CreateJobObject(nil, jobName.Buffer) + if err != nil { + return nil, err + } } defer func() { @@ -91,140 +112,109 @@ func Create(ctx context.Context, name string, notifications bool) (_ *JobObject, } }() + job := &JobObject{ + handle: jobHandle, + } + // If the IOCP we'll be using to receive messages for all jobs hasn't been // created, create it and start polling. - if notifications { - ioInitOnce.Do(func() { - h, err := windows.CreateIoCompletionPort(windows.InvalidHandle, 0, 0, 0xffffffff) - if err != nil { - initIOErr = err - return - } - ioCompletionPort = h - go pollIOCP(ctx, h) - }) - - if initIOErr != nil { - return nil, initIOErr - } - - mq = queue.NewMessageQueue() - jobMap.Store(uintptr(jobHandle), mq) - if err = attachIOCP(jobHandle, ioCompletionPort); err != nil { - jobMap.Delete(uintptr(jobHandle)) + if options != nil && options.Notifications { + mq, err := setupNotifications(ctx, job) + if err != nil { return nil, err } + job.mq = mq } - return &JobObject{ - handle: jobHandle, - mq: mq, - }, nil + return job, nil } -// SetResourceLimits sets resource limits on the job object (cpu, memory, storage). -func (job *JobObject) SetResourceLimits(limits *JobLimits) error { - // Go through and check what limits were specified and apply them to the job. - if limits.MemoryLimitInBytes != 0 { - if err := job.SetMemoryLimit(limits.MemoryLimitInBytes); err != nil { - return errors.Wrap(err, "failed to set job object memory limit") - } +// Open opens an existing job object with name provided in `options`. If no name is provided +// return an error since we need to know what job object to open. +// +// If notifications are not enabled `PollNotifications` will return immediately with error `errNotRegistered`. +// +// Returns a JobObject structure and an error if there is one. +func Open(ctx context.Context, options *Options) (_ *JobObject, err error) { + if options != nil && options.Name == "" { + return nil, errors.New("no job object name specified to open") + } + unicodeJobName, err := winapi.NewUnicodeString(options.Name) + if err != nil { + return nil, err } - if limits.CPULimit != 0 { - if err := job.SetCPULimit(RateBased, limits.CPULimit); err != nil { - return errors.Wrap(err, "failed to set job object cpu limit") + var jobHandle windows.Handle + if options != nil && options.UseNTVariant { + oa := winapi.ObjectAttributes{ + Length: unsafe.Sizeof(winapi.ObjectAttributes{}), + ObjectName: unicodeJobName, + Attributes: 0, } - } else if limits.CPUWeight != 0 { - if err := job.SetCPULimit(WeightBased, limits.CPUWeight); err != nil { - return errors.Wrap(err, "failed to set job object cpu limit") + status := winapi.NtOpenJobObject(&jobHandle, winapi.JOB_OBJECT_ALL_ACCESS, &oa) + if status != 0 { + return nil, winapi.RtlNtStatusToDosError(status) } - } - - if limits.MaxBandwidth != 0 || limits.MaxIOPS != 0 { - if err := job.SetIOLimit(limits.MaxBandwidth, limits.MaxIOPS); err != nil { - return errors.Wrap(err, "failed to set io limit on job object") + } else { + jobHandle, err = winapi.OpenJobObject(winapi.JOB_OBJECT_ALL_ACCESS, false, unicodeJobName.Buffer) + if err != nil { + return nil, err } } - return nil -} -// SetCPULimit sets the CPU limit specified on the job object. -func (job *JobObject) SetCPULimit(rateControlType CPURateControlType, rateControlValue uint32) error { - job.handleLock.RLock() - defer job.handleLock.RUnlock() + defer func() { + if err != nil { + windows.Close(jobHandle) + } + }() - if job.handle == 0 { - return ErrAlreadyClosed + job := &JobObject{ + handle: jobHandle, } - var cpuInfo winapi.JOBOBJECT_CPU_RATE_CONTROL_INFORMATION - switch rateControlType { - case WeightBased: - if rateControlValue < CPUWeightMin || rateControlValue > CPUWeightMax { - return fmt.Errorf("processor weight value of `%d` is invalid", rateControlValue) - } - cpuInfo.ControlFlags = winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | winapi.JOB_OBJECT_CPU_RATE_CONTROL_WEIGHT_BASED - cpuInfo.Value = rateControlValue - case RateBased: - if rateControlValue < CPULimitMin || rateControlValue > CPULimitMax { - return fmt.Errorf("processor rate of `%d` is invalid", rateControlValue) + // If the IOCP we'll be using to receive messages for all jobs hasn't been + // created, create it and start polling. + if options != nil && options.Notifications { + mq, err := setupNotifications(ctx, job) + if err != nil { + return nil, err } - cpuInfo.ControlFlags = winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | winapi.JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP - cpuInfo.Value = rateControlValue - default: - return errors.New("invalid job object cpu rate control type") + job.mq = mq } - _, err := windows.SetInformationJobObject(job.handle, windows.JobObjectCpuRateControlInformation, uintptr(unsafe.Pointer(&cpuInfo)), uint32(unsafe.Sizeof(cpuInfo))) - if err != nil { - return fmt.Errorf("failed to set cpu limit info on job object: %s", err) - } - return nil + return job, nil } -// SetMemoryLimit sets the memory limit specified on the job object. -func (job *JobObject) SetMemoryLimit(memoryLimitInBytes uint64) error { - job.handleLock.RLock() - defer job.handleLock.RUnlock() - - if job.handle == 0 { - return ErrAlreadyClosed - } +// helper function to setup notifications for creating/opening a job object +func setupNotifications(ctx context.Context, job *JobObject) (*queue.MessageQueue, error) { + ioInitOnce.Do(func() { + h, err := windows.CreateIoCompletionPort(windows.InvalidHandle, 0, 0, 0xffffffff) + if err != nil { + initIOErr = err + return + } + ioCompletionPort = h + go pollIOCP(ctx, h) + }) - var eliInfo windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION - eliInfo.JobMemoryLimit = uintptr(memoryLimitInBytes) - eliInfo.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_JOB_MEMORY - _, err := windows.SetInformationJobObject(job.handle, windows.JobObjectExtendedLimitInformation, uintptr(unsafe.Pointer(&eliInfo)), uint32(unsafe.Sizeof(eliInfo))) - if err != nil { - return fmt.Errorf("failed to set extended limit info on job object: %s", err) + if initIOErr != nil { + return nil, initIOErr } - return nil -} -// SetIOLimit sets the IO limits specified on the job object. -func (job *JobObject) SetIOLimit(maxBandwidth, maxIOPS int64) error { job.handleLock.RLock() defer job.handleLock.RUnlock() if job.handle == 0 { - return ErrAlreadyClosed + return nil, ErrAlreadyClosed } - ioInfo := winapi.JOBOBJECT_IO_RATE_CONTROL_INFORMATION{ - ControlFlags: winapi.JOB_OBJECT_IO_RATE_CONTROL_ENABLE, - } - if maxBandwidth != 0 { - ioInfo.MaxBandwidth = maxBandwidth - } - if maxIOPS != 0 { - ioInfo.MaxIops = maxIOPS - } - _, err := winapi.SetIoRateControlInformationJobObject(job.handle, &ioInfo) - if err != nil { - return fmt.Errorf("failed to set IO limit info on job object: %s", err) + mq := queue.NewMessageQueue() + jobMap.Store(uintptr(job.handle), mq) + if err := attachIOCP(job.handle, ioCompletionPort); err != nil { + jobMap.Delete(uintptr(job.handle)) + return nil, err } - return nil + return mq, nil } // PollNotification will poll for a job object notification. This call should only be called once diff --git a/internal/jobobject/limits.go b/internal/jobobject/limits.go new file mode 100644 index 0000000000..990af100cf --- /dev/null +++ b/internal/jobobject/limits.go @@ -0,0 +1,305 @@ +package jobobject + +import ( + "fmt" + "unsafe" + + "github.com/Microsoft/hcsshim/internal/winapi" + "github.com/pkg/errors" + "golang.org/x/sys/windows" +) + +const ( + processorWeightMax float64 = 10000 + memoryLimitMax uint64 = 0xffffffffffffffff +) + +func isFlagSet(flag uint32, controlFlags uint32) bool { + return (flag & controlFlags) == flag +} + +// SetResourceLimits sets resource limits on the job object (cpu, memory, storage). +func (job *JobObject) SetResourceLimits(limits *JobLimits) error { + // Go through and check what limits were specified and apply them to the job. + if limits.MemoryLimitInBytes != 0 { + if err := job.SetMemoryLimit(limits.MemoryLimitInBytes); err != nil { + return errors.Wrap(err, "failed to set job object memory limit") + } + } + + if limits.CPULimit != 0 { + if err := job.SetCPULimit(RateBased, limits.CPULimit); err != nil { + return errors.Wrap(err, "failed to set job object cpu limit") + } + } else if limits.CPUWeight != 0 { + if err := job.SetCPULimit(WeightBased, limits.CPUWeight); err != nil { + return errors.Wrap(err, "failed to set job object cpu limit") + } + } + + if limits.MaxBandwidth != 0 || limits.MaxIOPS != 0 { + if err := job.SetIOLimit(limits.MaxBandwidth, limits.MaxIOPS); err != nil { + return errors.Wrap(err, "failed to set io limit on job object") + } + } + return nil +} + +// SetMemoryLimit sets the memory limit of the job object based on the given `memoryLimitInBytes`. +func (job *JobObject) SetMemoryLimit(memoryLimitInBytes uint64) error { + if memoryLimitInBytes >= memoryLimitMax { + return errors.New("memory limit specified exceeds the max size") + } + + info, err := job.getExtendedInformation() + if err != nil { + return err + } + + info.JobMemoryLimit = uintptr(memoryLimitInBytes) + info.BasicLimitInformation.LimitFlags |= windows.JOB_OBJECT_LIMIT_JOB_MEMORY + return job.setExtendedInformation(info) +} + +// GetMemoryLimit gets the memory limit in bytes of the job object. +func (job *JobObject) GetMemoryLimit() (uint64, error) { + info, err := job.getExtendedInformation() + if err != nil { + return 0, err + } + return uint64(info.JobMemoryLimit), nil +} + +// SetCPULimit sets the CPU limit depending on the specified `CPURateControlType` to +// `rateControlValue` for the job object. +func (job *JobObject) SetCPULimit(rateControlType CPURateControlType, rateControlValue uint32) error { + cpuInfo, err := job.getCPURateControlInformation() + if err != nil { + return err + } + switch rateControlType { + case WeightBased: + if rateControlValue < CPUWeightMin || rateControlValue > CPUWeightMax { + return fmt.Errorf("processor weight value of `%d` is invalid", rateControlValue) + } + cpuInfo.ControlFlags |= winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | winapi.JOB_OBJECT_CPU_RATE_CONTROL_WEIGHT_BASED + cpuInfo.Value = rateControlValue + case RateBased: + if rateControlValue < CPULimitMin || rateControlValue > CPULimitMax { + return fmt.Errorf("processor rate of `%d` is invalid", rateControlValue) + } + cpuInfo.ControlFlags |= winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | winapi.JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP + cpuInfo.Value = rateControlValue + default: + return errors.New("invalid job object cpu rate control type") + } + return job.setCPURateControlInfo(cpuInfo) +} + +// GetCPULimit gets the cpu limits for the job object. +// `rateControlType` is used to indicate what type of cpu limit to query for. +func (job *JobObject) GetCPULimit(rateControlType CPURateControlType) (uint32, error) { + info, err := job.getCPURateControlInformation() + if err != nil { + return 0, err + } + + if !isFlagSet(winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE, info.ControlFlags) { + return 0, errors.New("the job does not have cpu rate control enabled") + } + + switch rateControlType { + case WeightBased: + if !isFlagSet(winapi.JOB_OBJECT_CPU_RATE_CONTROL_WEIGHT_BASED, info.ControlFlags) { + return 0, errors.New("cannot get cpu weight for job object without cpu weight option set") + } + case RateBased: + if !isFlagSet(winapi.JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP, info.ControlFlags) { + return 0, errors.New("cannot get cpu rate hard cap for job object without cpu rate hard cap option set") + } + default: + return 0, errors.New("invalid job object cpu rate control type") + } + return info.Value, nil +} + +// SetCPUAffinity sets the processor affinity for the job object. +// The affinity is passed in as a bitmask. +func (job *JobObject) SetCPUAffinity(affinityBitMask uint64) error { + info, err := job.getExtendedInformation() + if err != nil { + return err + } + info.BasicLimitInformation.LimitFlags |= uint32(windows.JOB_OBJECT_LIMIT_AFFINITY) + info.BasicLimitInformation.Affinity = uintptr(affinityBitMask) + return job.setExtendedInformation(info) +} + +// GetCPUAffinity gets the processor affinity for the job object. +// The returned affinity is a bitmask. +func (job *JobObject) GetCPUAffinity() (uint64, error) { + info, err := job.getExtendedInformation() + if err != nil { + return 0, err + } + return uint64(info.BasicLimitInformation.Affinity), nil +} + +// SetIOLimit sets the IO limits specified on the job object. +func (job *JobObject) SetIOLimit(maxBandwidth, maxIOPS int64) error { + ioInfo, err := job.getIOLimit() + if err != nil { + return err + } + ioInfo.ControlFlags |= winapi.JOB_OBJECT_IO_RATE_CONTROL_ENABLE + if maxBandwidth != 0 { + ioInfo.MaxBandwidth = maxBandwidth + } + if maxIOPS != 0 { + ioInfo.MaxIops = maxIOPS + } + return job.setIORateControlInfo(ioInfo) +} + +// GetIOMaxBandwidthLimit gets the max bandwidth for the job object. +func (job *JobObject) GetIOMaxBandwidthLimit() (int64, error) { + info, err := job.getIOLimit() + if err != nil { + return 0, err + } + return info.MaxBandwidth, nil +} + +// GetIOMaxIopsLimit gets the max iops for the job object. +func (job *JobObject) GetIOMaxIopsLimit() (int64, error) { + info, err := job.getIOLimit() + if err != nil { + return 0, err + } + return info.MaxIops, nil +} + +// Helper function for getting a job object's extended information. +func (job *JobObject) getExtendedInformation() (*windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION, error) { + job.handleLock.RLock() + defer job.handleLock.RUnlock() + + if job.handle == 0 { + return nil, ErrAlreadyClosed + } + + info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{} + if err := winapi.QueryInformationJobObject( + job.handle, + windows.JobObjectExtendedLimitInformation, + uintptr(unsafe.Pointer(&info)), + uint32(unsafe.Sizeof(info)), + nil, + ); err != nil { + return nil, errors.Wrapf(err, "query %v returned error", info) + } + return &info, nil +} + +// Helper function for getting a job object's CPU rate control information. +func (job *JobObject) getCPURateControlInformation() (*winapi.JOBOBJECT_CPU_RATE_CONTROL_INFORMATION, error) { + job.handleLock.RLock() + defer job.handleLock.RUnlock() + + if job.handle == 0 { + return nil, ErrAlreadyClosed + } + + info := winapi.JOBOBJECT_CPU_RATE_CONTROL_INFORMATION{} + if err := winapi.QueryInformationJobObject( + job.handle, + windows.JobObjectCpuRateControlInformation, + uintptr(unsafe.Pointer(&info)), + uint32(unsafe.Sizeof(info)), + nil, + ); err != nil { + return nil, errors.Wrapf(err, "query %v returned error", info) + } + return &info, nil +} + +// Helper function for setting a job object's extended information. +func (job *JobObject) setExtendedInformation(info *windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION) error { + job.handleLock.RLock() + defer job.handleLock.RUnlock() + + if job.handle == 0 { + return ErrAlreadyClosed + } + + if _, err := windows.SetInformationJobObject( + job.handle, + windows.JobObjectExtendedLimitInformation, + uintptr(unsafe.Pointer(info)), + uint32(unsafe.Sizeof(*info)), + ); err != nil { + return errors.Wrapf(err, "failed to set Extended info %v on job object", info) + } + return nil +} + +// Helper function for querying job handle for IO limit information. +func (job *JobObject) getIOLimit() (*winapi.JOBOBJECT_IO_RATE_CONTROL_INFORMATION, error) { + job.handleLock.RLock() + defer job.handleLock.RUnlock() + + if job.handle == 0 { + return nil, ErrAlreadyClosed + } + + ioInfo := &winapi.JOBOBJECT_IO_RATE_CONTROL_INFORMATION{} + var blockCount uint32 = 1 + + if _, err := winapi.QueryIoRateControlInformationJobObject( + job.handle, + nil, + &ioInfo, + &blockCount, + ); err != nil { + return nil, errors.Wrapf(err, "query %v returned error", ioInfo) + } + + if !isFlagSet(winapi.JOB_OBJECT_IO_RATE_CONTROL_ENABLE, ioInfo.ControlFlags) { + return nil, fmt.Errorf("query %v cannot get IO limits for job object without IO rate control option set", ioInfo) + } + return ioInfo, nil +} + +// Helper function for setting a job object's IO rate control information. +func (job *JobObject) setIORateControlInfo(ioInfo *winapi.JOBOBJECT_IO_RATE_CONTROL_INFORMATION) error { + job.handleLock.RLock() + defer job.handleLock.RUnlock() + + if job.handle == 0 { + return ErrAlreadyClosed + } + + if _, err := winapi.SetIoRateControlInformationJobObject(job.handle, ioInfo); err != nil { + return errors.Wrapf(err, "failed to set IO limit info %v on job object", ioInfo) + } + return nil +} + +// Helper function for setting a job object's CPU rate control information. +func (job *JobObject) setCPURateControlInfo(cpuInfo *winapi.JOBOBJECT_CPU_RATE_CONTROL_INFORMATION) error { + job.handleLock.RLock() + defer job.handleLock.RUnlock() + + if job.handle == 0 { + return ErrAlreadyClosed + } + if _, err := windows.SetInformationJobObject( + job.handle, + windows.JobObjectCpuRateControlInformation, + uintptr(unsafe.Pointer(cpuInfo)), + uint32(unsafe.Sizeof(cpuInfo)), + ); err != nil { + return errors.Wrapf(err, "failed to set cpu limit info %v on job object", cpuInfo) + } + return nil +} diff --git a/internal/safefile/safeopen.go b/internal/safefile/safeopen.go index d484c212cd..2c243b948a 100644 --- a/internal/safefile/safeopen.go +++ b/internal/safefile/safeopen.go @@ -76,7 +76,7 @@ func openRelativeInternal(path string, root *os.File, accessMask uint32, shareFl } oa.Length = unsafe.Sizeof(oa) - oa.ObjectName = uintptr(unsafe.Pointer(pathUnicode)) + oa.ObjectName = pathUnicode oa.RootDirectory = uintptr(root.Fd()) oa.Attributes = winapi.OBJ_DONT_REPARSE status := winapi.NtCreateFile( diff --git a/internal/winapi/filesystem.go b/internal/winapi/filesystem.go index 490576b942..7ce52afd5e 100644 --- a/internal/winapi/filesystem.go +++ b/internal/winapi/filesystem.go @@ -79,7 +79,7 @@ type IOStatusBlock struct { type ObjectAttributes struct { Length uintptr RootDirectory uintptr - ObjectName uintptr + ObjectName *UnicodeString Attributes uintptr SecurityDescriptor uintptr SecurityQoS uintptr diff --git a/internal/winapi/jobobject.go b/internal/winapi/jobobject.go index 6cb411cc42..ba12b1ad92 100644 --- a/internal/winapi/jobobject.go +++ b/internal/winapi/jobobject.go @@ -21,6 +21,11 @@ const ( JOB_OBJECT_MSG_NOTIFICATION_LIMIT uint32 = 11 ) +// Access rights for creating or opening job objects. +// +// https://docs.microsoft.com/en-us/windows/win32/procthread/job-object-security-and-access-rights +const JOB_OBJECT_ALL_ACCESS = 0x1F001F + // IO limit flags // // https://docs.microsoft.com/en-us/windows/win32/api/jobapi2/ns-jobapi2-jobobject_io_rate_control_information @@ -183,3 +188,28 @@ type JOBOBJECT_ASSOCIATE_COMPLETION_PORT struct { // ); // //sys SetIoRateControlInformationJobObject(jobHandle windows.Handle, ioRateControlInfo *JOBOBJECT_IO_RATE_CONTROL_INFORMATION) (ret uint32, err error) = kernel32.SetIoRateControlInformationJobObject + +// DWORD QueryIoRateControlInformationJobObject( +// HANDLE hJob, +// PCWSTR VolumeName, +// JOBOBJECT_IO_RATE_CONTROL_INFORMATION **InfoBlocks, +// ULONG *InfoBlockCount +// ); +//sys QueryIoRateControlInformationJobObject(jobHandle windows.Handle, volumeName *uint16, ioRateControlInfo **JOBOBJECT_IO_RATE_CONTROL_INFORMATION, infoBlockCount *uint32) (ret uint32, err error) = kernel32.QueryIoRateControlInformationJobObject + +// NTSTATUS +// NtOpenJobObject ( +// _Out_ PHANDLE JobHandle, +// _In_ ACCESS_MASK DesiredAccess, +// _In_ POBJECT_ATTRIBUTES ObjectAttributes +// ); +//sys NtOpenJobObject(jobHandle *windows.Handle, desiredAccess uint32, objAttributes *ObjectAttributes) (status uint32) = ntdll.NtOpenJobObject + +// NTSTATUS +// NTAPI +// NtCreateJobObject ( +// _Out_ PHANDLE JobHandle, +// _In_ ACCESS_MASK DesiredAccess, +// _In_opt_ POBJECT_ATTRIBUTES ObjectAttributes +// ); +//sys NtCreateJobObject(jobHandle *windows.Handle, desiredAccess uint32, objAttributes *ObjectAttributes) (status uint32) = ntdll.NtCreateJobObject diff --git a/internal/winapi/utils.go b/internal/winapi/utils.go index f3055d4175..5d3f5a7b7f 100644 --- a/internal/winapi/utils.go +++ b/internal/winapi/utils.go @@ -3,8 +3,9 @@ package winapi import ( "errors" "syscall" - "unicode/utf16" "unsafe" + + "golang.org/x/sys/windows" ) type UnicodeString struct { @@ -26,17 +27,22 @@ func (uni UnicodeString) String() string { // NewUnicodeString allocates a new UnicodeString and copies `s` into // the buffer of the new UnicodeString. func NewUnicodeString(s string) (*UnicodeString, error) { - ws := utf16.Encode(([]rune)(s)) - if len(ws) > 32767 { + // Get length of original `s` to use in the UnicodeString since the `buf` + // created later will have an additional trailing null character + length := len(s) + if length > 32767 { return nil, syscall.ENAMETOOLONG } + buf, err := windows.UTF16FromString(s) + if err != nil { + return nil, err + } uni := &UnicodeString{ - Length: uint16(len(ws) * 2), - MaximumLength: uint16(len(ws) * 2), - Buffer: &make([]uint16, len(ws))[0], + Length: uint16(length * 2), + MaximumLength: uint16(length * 2), + Buffer: &buf[0], } - copy((*[32768]uint16)(unsafe.Pointer(uni.Buffer))[:], ws) return uni, nil } diff --git a/internal/winapi/zsyscall_windows.go b/internal/winapi/zsyscall_windows.go index 0a990951d7..3a54c1fa1b 100644 --- a/internal/winapi/zsyscall_windows.go +++ b/internal/winapi/zsyscall_windows.go @@ -39,31 +39,34 @@ func errnoErr(e syscall.Errno) error { var ( modiphlpapi = windows.NewLazySystemDLL("iphlpapi.dll") modkernel32 = windows.NewLazySystemDLL("kernel32.dll") + modntdll = windows.NewLazySystemDLL("ntdll.dll") modadvapi32 = windows.NewLazySystemDLL("advapi32.dll") modcfgmgr32 = windows.NewLazySystemDLL("cfgmgr32.dll") - modntdll = windows.NewLazySystemDLL("ntdll.dll") - procSetJobCompartmentId = modiphlpapi.NewProc("SetJobCompartmentId") - procIsProcessInJob = modkernel32.NewProc("IsProcessInJob") - procQueryInformationJobObject = modkernel32.NewProc("QueryInformationJobObject") - procOpenJobObjectW = modkernel32.NewProc("OpenJobObjectW") - procSetIoRateControlInformationJobObject = modkernel32.NewProc("SetIoRateControlInformationJobObject") - procGetQueuedCompletionStatus = modkernel32.NewProc("GetQueuedCompletionStatus") - procSearchPathW = modkernel32.NewProc("SearchPathW") - procLogonUserW = modadvapi32.NewProc("LogonUserW") - procRtlMoveMemory = modkernel32.NewProc("RtlMoveMemory") - procLocalAlloc = modkernel32.NewProc("LocalAlloc") - procLocalFree = modkernel32.NewProc("LocalFree") - procGetActiveProcessorCount = modkernel32.NewProc("GetActiveProcessorCount") - procCM_Get_Device_ID_List_SizeA = modcfgmgr32.NewProc("CM_Get_Device_ID_List_SizeA") - procCM_Get_Device_ID_ListA = modcfgmgr32.NewProc("CM_Get_Device_ID_ListA") - procCM_Locate_DevNodeW = modcfgmgr32.NewProc("CM_Locate_DevNodeW") - procCM_Get_DevNode_PropertyW = modcfgmgr32.NewProc("CM_Get_DevNode_PropertyW") - procNtCreateFile = modntdll.NewProc("NtCreateFile") - procNtSetInformationFile = modntdll.NewProc("NtSetInformationFile") - procNtOpenDirectoryObject = modntdll.NewProc("NtOpenDirectoryObject") - procNtQueryDirectoryObject = modntdll.NewProc("NtQueryDirectoryObject") - procRtlNtStatusToDosError = modntdll.NewProc("RtlNtStatusToDosError") + procSetJobCompartmentId = modiphlpapi.NewProc("SetJobCompartmentId") + procGetQueuedCompletionStatus = modkernel32.NewProc("GetQueuedCompletionStatus") + procIsProcessInJob = modkernel32.NewProc("IsProcessInJob") + procQueryInformationJobObject = modkernel32.NewProc("QueryInformationJobObject") + procOpenJobObjectW = modkernel32.NewProc("OpenJobObjectW") + procSetIoRateControlInformationJobObject = modkernel32.NewProc("SetIoRateControlInformationJobObject") + procQueryIoRateControlInformationJobObject = modkernel32.NewProc("QueryIoRateControlInformationJobObject") + procNtOpenJobObject = modntdll.NewProc("NtOpenJobObject") + procNtCreateJobObject = modntdll.NewProc("NtCreateJobObject") + procSearchPathW = modkernel32.NewProc("SearchPathW") + procLogonUserW = modadvapi32.NewProc("LogonUserW") + procRtlMoveMemory = modkernel32.NewProc("RtlMoveMemory") + procLocalAlloc = modkernel32.NewProc("LocalAlloc") + procLocalFree = modkernel32.NewProc("LocalFree") + procGetActiveProcessorCount = modkernel32.NewProc("GetActiveProcessorCount") + procCM_Get_Device_ID_List_SizeA = modcfgmgr32.NewProc("CM_Get_Device_ID_List_SizeA") + procCM_Get_Device_ID_ListA = modcfgmgr32.NewProc("CM_Get_Device_ID_ListA") + procCM_Locate_DevNodeW = modcfgmgr32.NewProc("CM_Locate_DevNodeW") + procCM_Get_DevNode_PropertyW = modcfgmgr32.NewProc("CM_Get_DevNode_PropertyW") + procNtCreateFile = modntdll.NewProc("NtCreateFile") + procNtSetInformationFile = modntdll.NewProc("NtSetInformationFile") + procNtOpenDirectoryObject = modntdll.NewProc("NtOpenDirectoryObject") + procNtQueryDirectoryObject = modntdll.NewProc("NtQueryDirectoryObject") + procRtlNtStatusToDosError = modntdll.NewProc("RtlNtStatusToDosError") ) func SetJobCompartmentId(handle windows.Handle, compartmentId uint32) (win32Err error) { @@ -74,6 +77,18 @@ func SetJobCompartmentId(handle windows.Handle, compartmentId uint32) (win32Err return } +func GetQueuedCompletionStatus(cphandle windows.Handle, qty *uint32, key *uintptr, overlapped **windows.Overlapped, timeout uint32) (err error) { + r1, _, e1 := syscall.Syscall6(procGetQueuedCompletionStatus.Addr(), 5, uintptr(cphandle), uintptr(unsafe.Pointer(qty)), uintptr(unsafe.Pointer(key)), uintptr(unsafe.Pointer(overlapped)), uintptr(timeout), 0) + if r1 == 0 { + if e1 != 0 { + err = errnoErr(e1) + } else { + err = syscall.EINVAL + } + } + return +} + func IsProcessInJob(procHandle windows.Handle, jobHandle windows.Handle, result *bool) (err error) { r1, _, e1 := syscall.Syscall(procIsProcessInJob.Addr(), 3, uintptr(procHandle), uintptr(jobHandle), uintptr(unsafe.Pointer(result))) if r1 == 0 { @@ -130,9 +145,10 @@ func SetIoRateControlInformationJobObject(jobHandle windows.Handle, ioRateContro return } -func GetQueuedCompletionStatus(cphandle windows.Handle, qty *uint32, key *uintptr, overlapped **windows.Overlapped, timeout uint32) (err error) { - r1, _, e1 := syscall.Syscall6(procGetQueuedCompletionStatus.Addr(), 5, uintptr(cphandle), uintptr(unsafe.Pointer(qty)), uintptr(unsafe.Pointer(key)), uintptr(unsafe.Pointer(overlapped)), uintptr(timeout), 0) - if r1 == 0 { +func QueryIoRateControlInformationJobObject(jobHandle windows.Handle, volumeName *uint16, ioRateControlInfo **JOBOBJECT_IO_RATE_CONTROL_INFORMATION, infoBlockCount *uint32) (ret uint32, err error) { + r0, _, e1 := syscall.Syscall6(procQueryIoRateControlInformationJobObject.Addr(), 4, uintptr(jobHandle), uintptr(unsafe.Pointer(volumeName)), uintptr(unsafe.Pointer(ioRateControlInfo)), uintptr(unsafe.Pointer(infoBlockCount)), 0, 0) + ret = uint32(r0) + if ret == 0 { if e1 != 0 { err = errnoErr(e1) } else { @@ -142,6 +158,18 @@ func GetQueuedCompletionStatus(cphandle windows.Handle, qty *uint32, key *uintpt return } +func NtOpenJobObject(jobHandle *windows.Handle, desiredAccess uint32, objAttributes *ObjectAttributes) (status uint32) { + r0, _, _ := syscall.Syscall(procNtOpenJobObject.Addr(), 3, uintptr(unsafe.Pointer(jobHandle)), uintptr(desiredAccess), uintptr(unsafe.Pointer(objAttributes))) + status = uint32(r0) + return +} + +func NtCreateJobObject(jobHandle *windows.Handle, desiredAccess uint32, objAttributes *ObjectAttributes) (status uint32) { + r0, _, _ := syscall.Syscall(procNtCreateJobObject.Addr(), 3, uintptr(unsafe.Pointer(jobHandle)), uintptr(desiredAccess), uintptr(unsafe.Pointer(objAttributes))) + status = uint32(r0) + return +} + func SearchPath(lpPath *uint16, lpFileName *uint16, lpExtension *uint16, nBufferLength uint32, lpBuffer *uint16, lpFilePath **uint16) (size uint32, err error) { r0, _, e1 := syscall.Syscall6(procSearchPathW.Addr(), 6, uintptr(unsafe.Pointer(lpPath)), uintptr(unsafe.Pointer(lpFileName)), uintptr(unsafe.Pointer(lpExtension)), uintptr(nBufferLength), uintptr(unsafe.Pointer(lpBuffer)), uintptr(unsafe.Pointer(lpFilePath))) size = uint32(r0) diff --git a/internal/winobjdir/object_dir.go b/internal/winobjdir/object_dir.go index 6d448c896b..83baf2551d 100644 --- a/internal/winobjdir/object_dir.go +++ b/internal/winobjdir/object_dir.go @@ -32,7 +32,7 @@ func EnumerateNTObjectDirectory(ntObjDirPath string) ([]string, error) { } oa.Length = unsafe.Sizeof(oa) - oa.ObjectName = uintptr(unsafe.Pointer(pathUnicode)) + oa.ObjectName = pathUnicode // open `ntObjDirPath` directory status := winapi.NtOpenDirectoryObject(