elevate: extend to support list of privileges to retainsr/elevate

Signed-off-by: Simon Rozman <simon@rozman.si>

3 files changed, 11 insertions, 10 deletions

diff --git a/elevate/privileges.go b/elevate/privileges.go
index e41507eb..c99dbd79 100644
--- a/elevate/privileges.go
+++ b/elevate/privileges.go
@@ -13,13 +13,15 @@ import (
 	"golang.org/x/sys/windows"
 )
 
-func DropAllPrivileges(retainDriverLoading bool) error {
-	var luid windows.LUID
-	if retainDriverLoading {
-		err := windows.LookupPrivilegeValue(nil, windows.StringToUTF16Ptr("SeLoadDriverPrivilege"), &luid)
+func DropAllPrivileges(retainPrivileges []string) error {
+	retainLUIDs := make(map[windows.LUID]bool)
+	for _, luidName := range retainPrivileges {
+		var luid windows.LUID
+		err := windows.LookupPrivilegeValue(nil, windows.StringToUTF16Ptr(luidName), &luid)
 		if err != nil {
 			return err
 		}
+		retainLUIDs[luid] = true
 	}
 	var processToken windows.Token
 	err := windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_READ|windows.TOKEN_WRITE, &processToken)
@@ -43,9 +45,8 @@ func DropAllPrivileges(retainDriverLoading bool) error {
 		return errors.New("GetTokenInformation returned incomplete data")
 	}
 	tokenPrivileges := (*windows.Tokenprivileges)(unsafe.Pointer(&buffer[0]))
-	for i := uint32(0); i < tokenPrivileges.PrivilegeCount; i++ {
-		item := (*windows.LUIDAndAttributes)(unsafe.Pointer(uintptr(unsafe.Pointer(&tokenPrivileges.Privileges[0])) + unsafe.Sizeof(tokenPrivileges.Privileges[0])*uintptr(i)))
-		if retainDriverLoading && item.Luid == luid {
+	for _, item := range tokenPrivileges.AllPrivileges() {
+		if retainLUIDs[item.Luid] {
 			continue
 		}
 		item.Attributes = windows.SE_PRIVILEGE_REMOVED
@@ -100,4 +101,4 @@ func SetDefaultObjectDacl() error {
 	}
 	//TODO: sacl?
 	return nil
-}
\ No newline at end of file
+}
diff --git a/main.go b/main.go
index b7a84010..4d88ed56 100644
--- a/main.go
+++ b/main.go
@@ -221,7 +221,7 @@ func main() {
 		if len(os.Args) != 6 {
 			usage()
 		}
-		err := elevate.DropAllPrivileges(false)
+		err := elevate.DropAllPrivileges([]string{})
 		if err != nil {
 			fatal(err)
 		}
diff --git a/tunnel/service.go b/tunnel/service.go
index 8b81f56d..6c97a3af 100644
--- a/tunnel/service.go
+++ b/tunnel/service.go
@@ -179,7 +179,7 @@ func (service *tunnelService) Execute(args []string, r <-chan svc.ChangeRequest,
 	}
 
 	log.Println("Dropping privileges")
-	err = elevate.DropAllPrivileges(true)
+	err = elevate.DropAllPrivileges([]string{"SeLoadDriverPrivilege"})
 	if err != nil {
 		serviceError = services.ErrorDropPrivileges
 		return