From f1a78c8aaf47eb9911b6d354a7d6529bb4940f3c Mon Sep 17 00:00:00 2001 From: John Starks Date: Thu, 12 May 2016 21:31:45 -0700 Subject: [PATCH] Add function for setting privileges globally Also optimize privilege setting with the same privilege name multiple times. --- privilege.go | 79 +++++++++++++++++++++++++++++++++++++++------------- zsyscall.go | 12 +++++--- 2 files changed, 68 insertions(+), 23 deletions(-) diff --git a/privilege.go b/privilege.go index 81f9af7..3d59412 100644 --- a/privilege.go +++ b/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 - for _, name := range names { - p := uint64(0) - err := lookupPrivilegeValue("", name, &p) - if err != nil { - return err - } - privileges = append(privileges, p) + privileges, err := mapPrivileges(names) + if err != nil { + return err } 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 - err = lookupPrivilegeDisplayName("", &nameBuffer[0], &displayNameBuffer[0], &displayBufSize, &langId) + var langID uint32 + err = lookupPrivilegeDisplayName("", &nameBuffer[0], &displayNameBuffer[0], &displayBufSize, &langID) if err != nil { - return fmt.Sprintf("", utf16.Decode(nameBuffer[:bufSize])) + return fmt.Sprintf("", 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() } diff --git a/zsyscall.go b/zsyscall.go index 74b6e97..6d047d3 100644 --- a/zsyscall.go +++ b/zsyscall.go @@ -2,8 +2,12 @@ package winio -import "unsafe" -import "syscall" +import ( + "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