Browse Source

windows: add token environment functions

This brings the x/sys package into parity with the capabilities provided
indirectly in CL 176619, and adds a helper to make it useful.

Change-Id: I81f2d205bbb0c2b2c530b2bd991c1e6ff30cc94e
Reviewed-on: https://go-review.googlesource.com/c/sys/+/176620
Run-TryBot: Jason Donenfeld <Jason@zx2c4.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Alex Brainman <alex.brainman@gmail.com>
Jason A. Donenfeld 6 years ago
parent
commit
06a5c49444
3 changed files with 68 additions and 2 deletions
  1. 33 1
      windows/env_windows.go
  2. 2 0
      windows/syscall_windows.go
  3. 33 1
      windows/zsyscall_windows.go

+ 33 - 1
windows/env_windows.go

@@ -6,7 +6,11 @@
 
 package windows
 
-import "syscall"
+import (
+	"syscall"
+	"unicode/utf16"
+	"unsafe"
+)
 
 func Getenv(key string) (value string, found bool) {
 	return syscall.Getenv(key)
@@ -24,6 +28,34 @@ func Environ() []string {
 	return syscall.Environ()
 }
 
+// Returns a default environment associated with the token, rather than the current
+// process. If inheritExisting is true, then this environment also inherits the
+// environment of the current process.
+func (token Token) Environ(inheritExisting bool) (env []string, err error) {
+	var block *uint16
+	err = CreateEnvironmentBlock(&block, token, inheritExisting)
+	if err != nil {
+		return nil, err
+	}
+	defer DestroyEnvironmentBlock(block)
+	blockp := uintptr(unsafe.Pointer(block))
+	for {
+		entry := (*[(1 << 30) - 1]uint16)(unsafe.Pointer(blockp))[:]
+		for i, v := range entry {
+			if v == 0 {
+				entry = entry[:i]
+				break
+			}
+		}
+		if len(entry) == 0 {
+			break
+		}
+		env = append(env, string(utf16.Decode(entry)))
+		blockp += 2 * (uintptr(len(entry)) + 1)
+	}
+	return env, nil
+}
+
 func Unsetenv(key string) error {
 	return syscall.Unsetenv(key)
 }

+ 2 - 0
windows/syscall_windows.go

@@ -190,6 +190,8 @@ func NewCallbackCDecl(fn interface{}) uintptr {
 //sys	FreeEnvironmentStrings(envs *uint16) (err error) = kernel32.FreeEnvironmentStringsW
 //sys	GetEnvironmentVariable(name *uint16, buffer *uint16, size uint32) (n uint32, err error) = kernel32.GetEnvironmentVariableW
 //sys	SetEnvironmentVariable(name *uint16, value *uint16) (err error) = kernel32.SetEnvironmentVariableW
+//sys	CreateEnvironmentBlock(block **uint16, token Token, inheritExisting bool) (err error) = userenv.CreateEnvironmentBlock
+//sys	DestroyEnvironmentBlock(block *uint16) (err error) = userenv.DestroyEnvironmentBlock
 //sys	SetFileTime(handle Handle, ctime *Filetime, atime *Filetime, wtime *Filetime) (err error)
 //sys	GetFileAttributes(name *uint16) (attrs uint32, err error) [failretval==INVALID_FILE_ATTRIBUTES] = kernel32.GetFileAttributesW
 //sys	SetFileAttributes(name *uint16, attrs uint32) (err error) = kernel32.SetFileAttributesW

+ 33 - 1
windows/zsyscall_windows.go

@@ -37,6 +37,7 @@ func errnoErr(e syscall.Errno) error {
 var (
 	modadvapi32 = NewLazySystemDLL("advapi32.dll")
 	modkernel32 = NewLazySystemDLL("kernel32.dll")
+	moduserenv  = NewLazySystemDLL("userenv.dll")
 	modshell32  = NewLazySystemDLL("shell32.dll")
 	modmswsock  = NewLazySystemDLL("mswsock.dll")
 	modcrypt32  = NewLazySystemDLL("crypt32.dll")
@@ -45,7 +46,6 @@ var (
 	modiphlpapi = NewLazySystemDLL("iphlpapi.dll")
 	modsecur32  = NewLazySystemDLL("secur32.dll")
 	modnetapi32 = NewLazySystemDLL("netapi32.dll")
-	moduserenv  = NewLazySystemDLL("userenv.dll")
 	modwtsapi32 = NewLazySystemDLL("wtsapi32.dll")
 
 	procRegisterEventSourceW               = modadvapi32.NewProc("RegisterEventSourceW")
@@ -128,6 +128,8 @@ var (
 	procFreeEnvironmentStringsW            = modkernel32.NewProc("FreeEnvironmentStringsW")
 	procGetEnvironmentVariableW            = modkernel32.NewProc("GetEnvironmentVariableW")
 	procSetEnvironmentVariableW            = modkernel32.NewProc("SetEnvironmentVariableW")
+	procCreateEnvironmentBlock             = moduserenv.NewProc("CreateEnvironmentBlock")
+	procDestroyEnvironmentBlock            = moduserenv.NewProc("DestroyEnvironmentBlock")
 	procSetFileTime                        = modkernel32.NewProc("SetFileTime")
 	procGetFileAttributesW                 = modkernel32.NewProc("GetFileAttributesW")
 	procSetFileAttributesW                 = modkernel32.NewProc("SetFileAttributesW")
@@ -1300,6 +1302,36 @@ func SetEnvironmentVariable(name *uint16, value *uint16) (err error) {
 	return
 }
 
+func CreateEnvironmentBlock(block **uint16, token Token, inheritExisting bool) (err error) {
+	var _p0 uint32
+	if inheritExisting {
+		_p0 = 1
+	} else {
+		_p0 = 0
+	}
+	r1, _, e1 := syscall.Syscall(procCreateEnvironmentBlock.Addr(), 3, uintptr(unsafe.Pointer(block)), uintptr(token), uintptr(_p0))
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func DestroyEnvironmentBlock(block *uint16) (err error) {
+	r1, _, e1 := syscall.Syscall(procDestroyEnvironmentBlock.Addr(), 1, uintptr(unsafe.Pointer(block)), 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
 func SetFileTime(handle Handle, ctime *Filetime, atime *Filetime, wtime *Filetime) (err error) {
 	r1, _, e1 := syscall.Syscall6(procSetFileTime.Addr(), 4, uintptr(handle), uintptr(unsafe.Pointer(ctime)), uintptr(unsafe.Pointer(atime)), uintptr(unsafe.Pointer(wtime)), 0, 0)
 	if r1 == 0 {