From 226ff32320da7b90d0b5bc2365f4e359c466fb78 Mon Sep 17 00:00:00 2001
From: Awn Umar <awn@spacetime.dev>
Date: Wed, 10 Jun 2020 10:21:08 +0000
Subject: [PATCH] windows: add process working size system calls

Fixes golang/go#39422
Related to https://github.com/awnumar/memcall/issues/3

Change-Id: Idf3eec42c3077b39fe033091eea6d62b6a9d8d32
GitHub-Last-Rev: 7f57085a5f8ca2f0b5d3a6fa041d5fff03e7878d
GitHub-Pull-Request: golang/sys#72
Reviewed-on: https://go-review.googlesource.com/c/sys/+/236680
Run-TryBot: Alex Brainman <alex.brainman@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Alex Brainman <alex.brainman@gmail.com>
---
 windows/memory_windows.go       |  5 +++++
 windows/syscall_windows.go      |  2 ++
 windows/syscall_windows_test.go | 17 +++++++++++++++++
 windows/zsyscall_windows.go     | 19 +++++++++++++++++++
 4 files changed, 43 insertions(+)

diff --git a/windows/memory_windows.go b/windows/memory_windows.go
index f80a4204..e409d76f 100644
--- a/windows/memory_windows.go
+++ b/windows/memory_windows.go
@@ -23,4 +23,9 @@ const (
 	PAGE_EXECUTE_READ      = 0x20
 	PAGE_EXECUTE_READWRITE = 0x40
 	PAGE_EXECUTE_WRITECOPY = 0x80
+
+	QUOTA_LIMITS_HARDWS_MIN_DISABLE = 0x00000002
+	QUOTA_LIMITS_HARDWS_MIN_ENABLE  = 0x00000001
+	QUOTA_LIMITS_HARDWS_MAX_DISABLE = 0x00000008
+	QUOTA_LIMITS_HARDWS_MAX_ENABLE  = 0x00000004
 )
diff --git a/windows/syscall_windows.go b/windows/syscall_windows.go
index 12c0544c..62cf70e9 100644
--- a/windows/syscall_windows.go
+++ b/windows/syscall_windows.go
@@ -308,6 +308,8 @@ func NewCallbackCDecl(fn interface{}) uintptr {
 //sys	GetProcessId(process Handle) (id uint32, err error)
 //sys	OpenThread(desiredAccess uint32, inheritHandle bool, threadId uint32) (handle Handle, err error)
 //sys	SetProcessPriorityBoost(process Handle, disable bool) (err error) = kernel32.SetProcessPriorityBoost
+//sys	GetProcessWorkingSetSizeEx(hProcess Handle, lpMinimumWorkingSetSize *uintptr, lpMaximumWorkingSetSize *uintptr, flags *uint32)
+//sys	SetProcessWorkingSetSizeEx(hProcess Handle, dwMinimumWorkingSetSize uintptr, dwMaximumWorkingSetSize uintptr, flags uint32) (err error)
 
 // Volume Management Functions
 //sys	DefineDosDevice(flags uint32, deviceName *uint16, targetPath *uint16) (err error) = DefineDosDeviceW
diff --git a/windows/syscall_windows_test.go b/windows/syscall_windows_test.go
index 86d0d7b1..0ee8d2c7 100644
--- a/windows/syscall_windows_test.go
+++ b/windows/syscall_windows_test.go
@@ -404,3 +404,20 @@ func TestGetPreferredUILanguages(t *testing.T) {
 		}
 	}
 }
+
+func TestProcessWorkingSetSizeEx(t *testing.T) {
+	// Grab a handle to the current process
+	hProcess := windows.CurrentProcess()
+
+	// Allocate memory to store the result of the query
+	var minimumWorkingSetSize, maximumWorkingSetSize uintptr
+
+	// Make the system-call
+	var flag uint32
+	windows.GetProcessWorkingSetSizeEx(hProcess, &minimumWorkingSetSize, &maximumWorkingSetSize, &flag)
+
+	// Set the new limits to the current ones
+	if err := windows.SetProcessWorkingSetSizeEx(hProcess, minimumWorkingSetSize, maximumWorkingSetSize, flag); err != nil {
+		t.Error(err)
+	}
+}
diff --git a/windows/zsyscall_windows.go b/windows/zsyscall_windows.go
index 2aa4fa64..8a562fee 100644
--- a/windows/zsyscall_windows.go
+++ b/windows/zsyscall_windows.go
@@ -217,6 +217,8 @@ var (
 	procGetProcessId                                         = modkernel32.NewProc("GetProcessId")
 	procOpenThread                                           = modkernel32.NewProc("OpenThread")
 	procSetProcessPriorityBoost                              = modkernel32.NewProc("SetProcessPriorityBoost")
+	procGetProcessWorkingSetSizeEx                           = modkernel32.NewProc("GetProcessWorkingSetSizeEx")
+	procSetProcessWorkingSetSizeEx                           = modkernel32.NewProc("SetProcessWorkingSetSizeEx")
 	procDefineDosDeviceW                                     = modkernel32.NewProc("DefineDosDeviceW")
 	procDeleteVolumeMountPointW                              = modkernel32.NewProc("DeleteVolumeMountPointW")
 	procFindFirstVolumeW                                     = modkernel32.NewProc("FindFirstVolumeW")
@@ -2414,6 +2416,23 @@ func SetProcessPriorityBoost(process Handle, disable bool) (err error) {
 	return
 }
 
+func GetProcessWorkingSetSizeEx(hProcess Handle, lpMinimumWorkingSetSize *uintptr, lpMaximumWorkingSetSize *uintptr, flags *uint32) {
+	syscall.Syscall6(procGetProcessWorkingSetSizeEx.Addr(), 4, uintptr(hProcess), uintptr(unsafe.Pointer(lpMinimumWorkingSetSize)), uintptr(unsafe.Pointer(lpMaximumWorkingSetSize)), uintptr(unsafe.Pointer(flags)), 0, 0)
+	return
+}
+
+func SetProcessWorkingSetSizeEx(hProcess Handle, dwMinimumWorkingSetSize uintptr, dwMaximumWorkingSetSize uintptr, flags uint32) (err error) {
+	r1, _, e1 := syscall.Syscall6(procSetProcessWorkingSetSizeEx.Addr(), 4, uintptr(hProcess), uintptr(dwMinimumWorkingSetSize), uintptr(dwMaximumWorkingSetSize), uintptr(flags), 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
 func DefineDosDevice(flags uint32, deviceName *uint16, targetPath *uint16) (err error) {
 	r1, _, e1 := syscall.Syscall(procDefineDosDeviceW.Addr(), 3, uintptr(flags), uintptr(unsafe.Pointer(deviceName)), uintptr(unsafe.Pointer(targetPath)))
 	if r1 == 0 {