Merge pull request #22 from Microsoft/priv_opts

Add function for setting privileges globally

privilege.go

@@ -5,14 +5,17 @@ import (
 	"encoding/binary"
 	"fmt"
 	"runtime"
+	"sync"
 	"syscall"
 	"unicode/utf16"
+
+	"golang.org/x/sys/windows"
 )
 
-//sys adjustTokenPrivileges(token syscall.Handle, releaseAll bool, input *byte, outputSize uint32, output *byte, requiredSize *uint32) (success bool, err error) [true] = advapi32.AdjustTokenPrivileges
+//sys adjustTokenPrivileges(token windows.Token, releaseAll bool, input *byte, outputSize uint32, output *byte, requiredSize *uint32) (success bool, err error) [true] = advapi32.AdjustTokenPrivileges
 //sys impersonateSelf(level uint32) (err error) = advapi32.ImpersonateSelf
 //sys revertToSelf() (err error) = advapi32.RevertToSelf
-//sys openThreadToken(thread syscall.Handle, accessMask uint32, openAsSelf bool, token *syscall.Handle) (err error) = advapi32.OpenThreadToken
+//sys openThreadToken(thread syscall.Handle, accessMask uint32, openAsSelf bool, token *windows.Token) (err error) = advapi32.OpenThreadToken
 //sys getCurrentThread() (h syscall.Handle) = GetCurrentThread
 //sys lookupPrivilegeValue(systemName string, name string, luid *uint64) (err error) = advapi32.LookupPrivilegeValueW
 //sys lookupPrivilegeName(systemName string, luid *uint64, buffer *uint16, size *uint32) (err error) = advapi32.LookupPrivilegeNameW
@@ -34,6 +37,12 @@ const (
 	securityDelegation
 )
 
+var (
+	privNames     = make(map[string]uint64)
+	privNameMutex sync.Mutex
+)
+
+// PrivilegeError represents an error enabling privileges.
 type PrivilegeError struct {
 	privileges []uint64
 }
@@ -56,19 +65,16 @@ func (e *PrivilegeError) Error() string {
 	return s
 }
 
+// RunWithPrivilege enables a single privilege for a function call.
 func RunWithPrivilege(name string, fn func() error) error {
 	return RunWithPrivileges([]string{name}, fn)
 }
 
+// RunWithPrivileges enables privileges for a function call.
 func RunWithPrivileges(names []string, fn func() error) error {
-	var privileges []uint64
+	privileges, err := mapPrivileges(names)
-	for _, name := range names {
+	if err != nil {
-		p := uint64(0)
+		return err
-		err := lookupPrivilegeValue("", name, &p)
-		if err != nil {
-			return err
-		}
-		privileges = append(privileges, p)
 	}
 	runtime.LockOSThread()
 	defer runtime.UnlockOSThread()
@@ -84,7 +90,43 @@ func RunWithPrivileges(names []string, fn func() error) error {
 	return fn()
 }
 
-func adjustPrivileges(token syscall.Handle, privileges []uint64) error {
+func mapPrivileges(names []string) ([]uint64, error) {
+	var privileges []uint64
+	privNameMutex.Lock()
+	defer privNameMutex.Unlock()
+	for _, name := range names {
+		p, ok := privNames[name]
+		if !ok {
+			err := lookupPrivilegeValue("", name, &p)
+			if err != nil {
+				return nil, err
+			}
+			privNames[name] = p
+		}
+		privileges = append(privileges, p)
+	}
+	return privileges, nil
+}
+
+// EnableProcessPrivileges enables privileges globally for the process.
+func EnableProcessPrivileges(names []string) error {
+	privileges, err := mapPrivileges(names)
+	if err != nil {
+		return err
+	}
+
+	p, _ := windows.GetCurrentProcess()
+	var token windows.Token
+	err = windows.OpenProcessToken(p, windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &token)
+	if err != nil {
+		return err
+	}
+
+	defer token.Close()
+	return adjustPrivileges(token, privileges)
+}
+
+func adjustPrivileges(token windows.Token, privileges []uint64) error {
 	var b bytes.Buffer
 	binary.Write(&b, binary.LittleEndian, uint32(len(privileges)))
 	for _, p := range privileges {
@@ -113,23 +155,22 @@ func getPrivilegeName(luid uint64) string {
 
 	var displayNameBuffer [256]uint16
 	displayBufSize := uint32(len(displayNameBuffer))
-	var langId uint32
+	var langID uint32
-	err = lookupPrivilegeDisplayName("", &nameBuffer[0], &displayNameBuffer[0], &displayBufSize, &langId)
+	err = lookupPrivilegeDisplayName("", &nameBuffer[0], &displayNameBuffer[0], &displayBufSize, &langID)
 	if err != nil {
-		return fmt.Sprintf("<unknown privilege %s>", utf16.Decode(nameBuffer[:bufSize]))
+		return fmt.Sprintf("<unknown privilege %s>", string(utf16.Decode(nameBuffer[:bufSize])))
 	}
 
 	return string(utf16.Decode(displayNameBuffer[:displayBufSize]))
 }
 
-func newThreadToken() (syscall.Handle, error) {
+func newThreadToken() (windows.Token, error) {
 	err := impersonateSelf(securityImpersonation)
 	if err != nil {
-		panic(err)
 		return 0, err
 	}
 
-	var token syscall.Handle
+	var token windows.Token
 	err = openThreadToken(getCurrentThread(), syscall.TOKEN_ADJUST_PRIVILEGES|syscall.TOKEN_QUERY, false, &token)
 	if err != nil {
 		rerr := revertToSelf()
@@ -141,10 +182,10 @@ func newThreadToken() (syscall.Handle, error) {
 	return token, nil
 }
 
-func releaseThreadToken(h syscall.Handle) {
+func releaseThreadToken(h windows.Token) {
 	err := revertToSelf()
 	if err != nil {
 		panic(err)
 	}
-	syscall.Close(h)
+	h.Close()
 }

zsyscall.go

@@ -2,8 +2,12 @@
 
 package winio
 
-import "unsafe"
+import (
-import "syscall"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
 
 var _ unsafe.Pointer
 
@@ -300,7 +304,7 @@ func setFileInformationByHandle(h syscall.Handle, class uint32, buffer *byte, si
 	return
 }
 
-func adjustTokenPrivileges(token syscall.Handle, releaseAll bool, input *byte, outputSize uint32, output *byte, requiredSize *uint32) (success bool, err error) {
+func adjustTokenPrivileges(token windows.Token, releaseAll bool, input *byte, outputSize uint32, output *byte, requiredSize *uint32) (success bool, err error) {
 	var _p0 uint32
 	if releaseAll {
 		_p0 = 1
@@ -343,7 +347,7 @@ func revertToSelf() (err error) {
 	return
 }
 
-func openThreadToken(thread syscall.Handle, accessMask uint32, openAsSelf bool, token *syscall.Handle) (err error) {
+func openThreadToken(thread syscall.Handle, accessMask uint32, openAsSelf bool, token *windows.Token) (err error) {
 	var _p0 uint32
 	if openAsSelf {
 		_p0 = 1