Merge pull request #6 from Microsoft/privilege

Add functions for running with privileges

privilege.go

@@ -0,0 +1,147 @@
+package winio
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"runtime"
+	"syscall"
+	"unicode/utf16"
+)
+
+//sys adjustTokenPrivileges(token syscall.Handle, releaseAll bool, input *byte, outputSize uint32, output *byte, requiredSize *uint32) (err error) = 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 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
+//sys lookupPrivilegeDisplayName(systemName string, name *uint16, buffer *uint16, size *uint32, languageId *uint32) (err error) = advapi32.LookupPrivilegeDisplayNameW
+
+const (
+	SE_PRIVILEGE_ENABLED = 2
+
+	SeBackupPrivilege  = "SeBackupPrivilege"
+	SeRestorePrivilege = "SeRestorePrivilege"
+)
+
+const (
+	securityAnonymous = iota
+	securityIdentification
+	securityImpersonation
+	securityDelegation
+)
+
+type PrivilegeError struct {
+	privileges []uint64
+}
+
+func (e *PrivilegeError) Error() string {
+	s := ""
+	if len(e.privileges) > 1 {
+		s = "Could not enable privileges "
+	} else {
+		s = "Could not enable privilege "
+	}
+	for i, p := range e.privileges {
+		if i != 0 {
+			s += ", "
+		}
+		s += `"`
+		s += getPrivilegeName(p)
+		s += `"`
+	}
+	return s
+}
+
+func RunWithPrivilege(name string, fn func() error) error {
+	return RunWithPrivileges([]string{name}, fn)
+}
+
+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)
+	}
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+	token, err := newThreadToken()
+	if err != nil {
+		return err
+	}
+	defer releaseThreadToken(token)
+	err = adjustPrivileges(token, privileges)
+	if err != nil {
+		return err
+	}
+	return fn()
+}
+
+func adjustPrivileges(token syscall.Handle, privileges []uint64) error {
+	var b bytes.Buffer
+	binary.Write(&b, binary.LittleEndian, uint32(len(privileges)))
+	for _, p := range privileges {
+		binary.Write(&b, binary.LittleEndian, p)
+		binary.Write(&b, binary.LittleEndian, uint32(SE_PRIVILEGE_ENABLED))
+	}
+	prevState := make([]byte, b.Len())
+	reqSize := uint32(0)
+	if err := adjustTokenPrivileges(token, false, &b.Bytes()[0], uint32(len(prevState)), &prevState[0], &reqSize); err != nil {
+		return err
+	}
+	if int(binary.LittleEndian.Uint32(prevState[0:4])) < len(privileges) {
+		return &PrivilegeError{privileges}
+	}
+	return nil
+}
+
+func getPrivilegeName(luid uint64) string {
+	var nameBuffer [256]uint16
+	bufSize := uint32(len(nameBuffer))
+	err := lookupPrivilegeName("", &luid, &nameBuffer[0], &bufSize)
+	if err != nil {
+		return fmt.Sprintf("<unknown privilege %d>", luid)
+	}
+
+	var displayNameBuffer [256]uint16
+	displayBufSize := uint32(len(displayNameBuffer))
+	var langId uint32
+	err = lookupPrivilegeDisplayName("", &nameBuffer[0], &displayNameBuffer[0], &displayBufSize, &langId)
+	if err != nil {
+		return fmt.Sprintf("<unknown privilege %s>", utf16.Decode(nameBuffer[:bufSize]))
+	}
+
+	return string(utf16.Decode(displayNameBuffer[:displayBufSize]))
+}
+
+func newThreadToken() (syscall.Handle, error) {
+	err := impersonateSelf(securityImpersonation)
+	if err != nil {
+		panic(err)
+		return 0, err
+	}
+
+	var token syscall.Handle
+	err = openThreadToken(getCurrentThread(), syscall.TOKEN_ADJUST_PRIVILEGES|syscall.TOKEN_QUERY, false, &token)
+	if err != nil {
+		rerr := revertToSelf()
+		if rerr != nil {
+			panic(rerr)
+		}
+		return 0, err
+	}
+	return token, nil
+}
+
+func releaseThreadToken(h syscall.Handle) {
+	err := revertToSelf()
+	if err != nil {
+		panic(err)
+	}
+	syscall.Close(h)
+}

privileges_test.go

@@ -0,0 +1,17 @@
+package winio
+
+import "testing"
+
+func TestRunWithUnavailablePrivilege(t *testing.T) {
+	err := RunWithPrivilege("SeCreateTokenPrivilege", func() error { return nil })
+	if _, ok := err.(*PrivilegeError); err == nil || !ok {
+		t.Fatal("expected PrivilegeError")
+	}
+}
+
+func TestRunWithPrivileges(t *testing.T) {
+	err := RunWithPrivilege("SeShutdownPrivilege", func() error { return nil })
+	if err != nil {
+		t.Fatal(err)
+	}
+}

syscall.go

@@ -1,3 +1,3 @@
 package winio
 
-//go:generate go run $GOROOT/src/syscall/mksyscall_windows.go -output zsyscall.go file.go pipe.go sd.go fileinfo.go
+//go:generate go run $GOROOT/src/syscall/mksyscall_windows.go -output zsyscall.go file.go pipe.go sd.go fileinfo.go privilege.go

zsyscall.go

@@ -26,6 +26,14 @@ var (
 	procGetSecurityDescriptorLength                          = modadvapi32.NewProc("GetSecurityDescriptorLength")
 	procGetFileInformationByHandleEx                         = modkernel32.NewProc("GetFileInformationByHandleEx")
 	procSetFileInformationByHandle                           = modkernel32.NewProc("SetFileInformationByHandle")
+	procAdjustTokenPrivileges                                = modadvapi32.NewProc("AdjustTokenPrivileges")
+	procImpersonateSelf                                      = modadvapi32.NewProc("ImpersonateSelf")
+	procRevertToSelf                                         = modadvapi32.NewProc("RevertToSelf")
+	procOpenThreadToken                                      = modadvapi32.NewProc("OpenThreadToken")
+	procGetCurrentThread                                     = modkernel32.NewProc("GetCurrentThread")
+	procLookupPrivilegeValueW                                = modadvapi32.NewProc("LookupPrivilegeValueW")
+	procLookupPrivilegeNameW                                 = modadvapi32.NewProc("LookupPrivilegeNameW")
+	procLookupPrivilegeDisplayNameW                          = modadvapi32.NewProc("LookupPrivilegeDisplayNameW")
 )
 
 func cancelIoEx(file syscall.Handle, o *syscall.Overlapped) (err error) {
@@ -242,3 +250,137 @@ 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) (err error) {
+	var _p0 uint32
+	if releaseAll {
+		_p0 = 1
+	} else {
+		_p0 = 0
+	}
+	r1, _, e1 := syscall.Syscall6(procAdjustTokenPrivileges.Addr(), 6, uintptr(token), uintptr(_p0), uintptr(unsafe.Pointer(input)), uintptr(outputSize), uintptr(unsafe.Pointer(output)), uintptr(unsafe.Pointer(requiredSize)))
+	if r1 == 0 {
+		if e1 != 0 {
+			err = error(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func impersonateSelf(level uint32) (err error) {
+	r1, _, e1 := syscall.Syscall(procImpersonateSelf.Addr(), 1, uintptr(level), 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = error(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func revertToSelf() (err error) {
+	r1, _, e1 := syscall.Syscall(procRevertToSelf.Addr(), 0, 0, 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = error(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func openThreadToken(thread syscall.Handle, accessMask uint32, openAsSelf bool, token *syscall.Handle) (err error) {
+	var _p0 uint32
+	if openAsSelf {
+		_p0 = 1
+	} else {
+		_p0 = 0
+	}
+	r1, _, e1 := syscall.Syscall6(procOpenThreadToken.Addr(), 4, uintptr(thread), uintptr(accessMask), uintptr(_p0), uintptr(unsafe.Pointer(token)), 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = error(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func getCurrentThread() (h syscall.Handle) {
+	r0, _, _ := syscall.Syscall(procGetCurrentThread.Addr(), 0, 0, 0, 0)
+	h = syscall.Handle(r0)
+	return
+}
+
+func lookupPrivilegeValue(systemName string, name string, luid *uint64) (err error) {
+	var _p0 *uint16
+	_p0, err = syscall.UTF16PtrFromString(systemName)
+	if err != nil {
+		return
+	}
+	var _p1 *uint16
+	_p1, err = syscall.UTF16PtrFromString(name)
+	if err != nil {
+		return
+	}
+	return _lookupPrivilegeValue(_p0, _p1, luid)
+}
+
+func _lookupPrivilegeValue(systemName *uint16, name *uint16, luid *uint64) (err error) {
+	r1, _, e1 := syscall.Syscall(procLookupPrivilegeValueW.Addr(), 3, uintptr(unsafe.Pointer(systemName)), uintptr(unsafe.Pointer(name)), uintptr(unsafe.Pointer(luid)))
+	if r1 == 0 {
+		if e1 != 0 {
+			err = error(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func lookupPrivilegeName(systemName string, luid *uint64, buffer *uint16, size *uint32) (err error) {
+	var _p0 *uint16
+	_p0, err = syscall.UTF16PtrFromString(systemName)
+	if err != nil {
+		return
+	}
+	return _lookupPrivilegeName(_p0, luid, buffer, size)
+}
+
+func _lookupPrivilegeName(systemName *uint16, luid *uint64, buffer *uint16, size *uint32) (err error) {
+	r1, _, e1 := syscall.Syscall6(procLookupPrivilegeNameW.Addr(), 4, uintptr(unsafe.Pointer(systemName)), uintptr(unsafe.Pointer(luid)), uintptr(unsafe.Pointer(buffer)), uintptr(unsafe.Pointer(size)), 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = error(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func lookupPrivilegeDisplayName(systemName string, name *uint16, buffer *uint16, size *uint32, languageId *uint32) (err error) {
+	var _p0 *uint16
+	_p0, err = syscall.UTF16PtrFromString(systemName)
+	if err != nil {
+		return
+	}
+	return _lookupPrivilegeDisplayName(_p0, name, buffer, size, languageId)
+}
+
+func _lookupPrivilegeDisplayName(systemName *uint16, name *uint16, buffer *uint16, size *uint32, languageId *uint32) (err error) {
+	r1, _, e1 := syscall.Syscall6(procLookupPrivilegeDisplayNameW.Addr(), 5, uintptr(unsafe.Pointer(systemName)), uintptr(unsafe.Pointer(name)), uintptr(unsafe.Pointer(buffer)), uintptr(unsafe.Pointer(size)), uintptr(unsafe.Pointer(languageId)), 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = error(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}