windows: add Windows directory accessors

These are useful for the same reason that the already existing
GetSystemDirectory is.

Change-Id: I3041ce6cbeb66a4f8a5960fbaf39381c8c9c80d6
Reviewed-on: https://go-review.googlesource.com/c/sys/+/191837
Run-TryBot: Jason A. Donenfeld <Jason@zx2c4.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
diff --git a/windows/security_windows.go b/windows/security_windows.go
index 61b4964..7dfe201 100644
--- a/windows/security_windows.go
+++ b/windows/security_windows.go
@@ -644,6 +644,8 @@
 //sys	DuplicateTokenEx(existingToken Token, desiredAccess uint32, tokenAttributes *SecurityAttributes, impersonationLevel uint32, tokenType uint32, newToken *Token) (err error) = advapi32.DuplicateTokenEx
 //sys	GetUserProfileDirectory(t Token, dir *uint16, dirLen *uint32) (err error) = userenv.GetUserProfileDirectoryW
 //sys	getSystemDirectory(dir *uint16, dirLen uint32) (len uint32, err error) = kernel32.GetSystemDirectoryW
+//sys	getWindowsDirectory(dir *uint16, dirLen uint32) (len uint32, err error) = kernel32.GetWindowsDirectoryW
+//sys	getSystemWindowsDirectory(dir *uint16, dirLen uint32) (len uint32, err error) = kernel32.GetSystemWindowsDirectoryW
 
 // An access token contains the security information for a logon session.
 // The system creates an access token when a user logs on, and every
@@ -785,8 +787,8 @@
 	return linkedToken, nil
 }
 
-// GetSystemDirectory retrieves path to current location of the system
-// directory, which is typically, though not always, C:\Windows\System32.
+// GetSystemDirectory retrieves the path to current location of the system
+// directory, which is typically, though not always, `C:\Windows\System32`.
 func GetSystemDirectory() (string, error) {
 	n := uint32(MAX_PATH)
 	for {
@@ -802,6 +804,42 @@
 	}
 }
 
+// GetWindowsDirectory retrieves the path to current location of the Windows
+// directory, which is typically, though not always, `C:\Windows`. This may
+// be a private user directory in the case that the application is running
+// under a terminal server.
+func GetWindowsDirectory() (string, error) {
+	n := uint32(MAX_PATH)
+	for {
+		b := make([]uint16, n)
+		l, e := getWindowsDirectory(&b[0], n)
+		if e != nil {
+			return "", e
+		}
+		if l <= n {
+			return UTF16ToString(b[:l]), nil
+		}
+		n = l
+	}
+}
+
+// GetSystemWindowsDirectory retrieves the path to current location of the
+// Windows directory, which is typically, though not always, `C:\Windows`.
+func GetSystemWindowsDirectory() (string, error) {
+	n := uint32(MAX_PATH)
+	for {
+		b := make([]uint16, n)
+		l, e := getSystemWindowsDirectory(&b[0], n)
+		if e != nil {
+			return "", e
+		}
+		if l <= n {
+			return UTF16ToString(b[:l]), nil
+		}
+		n = l
+	}
+}
+
 // IsMember reports whether the access token t is a member of the provided SID.
 func (t Token) IsMember(sid *SID) (bool, error) {
 	var b int32
diff --git a/windows/syscall_test.go b/windows/syscall_test.go
index f09c6dd..7e28834 100644
--- a/windows/syscall_test.go
+++ b/windows/syscall_test.go
@@ -62,3 +62,20 @@
 		t.Fatalf("System directory does not end in system32: %s", d)
 	}
 }
+
+func TestGetWindowsDirectory(t *testing.T) {
+	d1, err := windows.GetWindowsDirectory()
+	if err != nil {
+		t.Fatalf("Failed to get Windows directory: %s", err)
+	}
+	d2, err := windows.GetSystemWindowsDirectory()
+	if err != nil {
+		t.Fatalf("Failed to get system Windows directory: %s", err)
+	}
+	if !strings.HasSuffix(strings.ToLower(d1), `\windows`) {
+		t.Fatalf("Windows directory does not end in windows: %s", d1)
+	}
+	if !strings.HasSuffix(strings.ToLower(d2), `\windows`) {
+		t.Fatalf("System Windows directory does not end in windows: %s", d2)
+	}
+}
diff --git a/windows/zsyscall_windows.go b/windows/zsyscall_windows.go
index 0f600db..1c275cb 100644
--- a/windows/zsyscall_windows.go
+++ b/windows/zsyscall_windows.go
@@ -305,6 +305,8 @@
 	procDuplicateTokenEx                   = modadvapi32.NewProc("DuplicateTokenEx")
 	procGetUserProfileDirectoryW           = moduserenv.NewProc("GetUserProfileDirectoryW")
 	procGetSystemDirectoryW                = modkernel32.NewProc("GetSystemDirectoryW")
+	procGetWindowsDirectoryW               = modkernel32.NewProc("GetWindowsDirectoryW")
+	procGetSystemWindowsDirectoryW         = modkernel32.NewProc("GetSystemWindowsDirectoryW")
 	procWTSQueryUserToken                  = modwtsapi32.NewProc("WTSQueryUserToken")
 	procWTSEnumerateSessionsW              = modwtsapi32.NewProc("WTSEnumerateSessionsW")
 	procWTSFreeMemory                      = modwtsapi32.NewProc("WTSFreeMemory")
@@ -3332,6 +3334,32 @@
 	return
 }
 
+func getWindowsDirectory(dir *uint16, dirLen uint32) (len uint32, err error) {
+	r0, _, e1 := syscall.Syscall(procGetWindowsDirectoryW.Addr(), 2, uintptr(unsafe.Pointer(dir)), uintptr(dirLen), 0)
+	len = uint32(r0)
+	if len == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func getSystemWindowsDirectory(dir *uint16, dirLen uint32) (len uint32, err error) {
+	r0, _, e1 := syscall.Syscall(procGetSystemWindowsDirectoryW.Addr(), 2, uintptr(unsafe.Pointer(dir)), uintptr(dirLen), 0)
+	len = uint32(r0)
+	if len == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
 func WTSQueryUserToken(session uint32, token *Token) (err error) {
 	r1, _, e1 := syscall.Syscall(procWTSQueryUserToken.Addr(), 2, uintptr(session), uintptr(unsafe.Pointer(token)), 0)
 	if r1 == 0 {