refactor(alertsenders): Launch systray server manually

rabbitstack · rabbitstack · commit 78475525d98b · 2024-11-30T23:34:44.000+01:00
The systray server is automatically launched when the
systray alert sender is enabled. The process is created
either via the regular CreateProcess API or CreateProcessAsUser
when Fibratus is running inside Windows Service.
This requires querying the user token of the currently logged-in
console session user. The TCB privilege must be present in the
user token.
diff --git a/build/msi/fibratus.wxs b/build/msi/fibratus.wxs
@@ -18,11 +18,6 @@
 
     <UI Id="UI">
       <ui:WixUI Id="WixUI_InstallDir" InstallDirectory="INSTALLDIR" />
-      <Publish Dialog="ExitDialog"
-              Control="Finish"
-              Event="DoAction"
-              Value="LaunchSystray"
-              Condition="NOT Installed"/>
     </UI>
 
     <!-- Custom banners -->
@@ -39,8 +34,6 @@
       <Files Include="!(bindpath.dir)**">
         <!-- Exclude from harvesting as its need fine-grained authoring in Windows Service -->
         <Exclude Files="!(bindpath.dir)Bin\fibratus.exe" />
-        <!-- Exclude from harvesting as it is used in custom action -->
-        <Exclude Files="!(bindpath.dir)Bin\fibratus-systray.exe" />
       </Files>
 
       <!-- Windows Service -->
@@ -57,14 +50,6 @@
         </ServiceInstall>
         <ServiceControl Id="Fibratus" Name="Fibratus" Start="install" Stop="both" Remove="uninstall" Wait="yes" />
       </Component>
-      <Component Directory="BINDIR" Id="Systray" Guid="F3C06EDD-C830-4FCD-BAFA-0D15C697EE76">
-        <File Id="Systray" Source="!(bindpath.dir)Bin\fibratus-systray.exe" KeyPath="yes" Checksum="yes"/>
-        <RegistryValue Root="HKMU" Action="write"
-                       Key="Software\Microsoft\Windows\CurrentVersion\Run"
-                       Name="Fibratus Systray"
-                       Value="[BINDIR]fibratus-systray.exe"
-                       Type="string" />
-      </Component>
     </ComponentGroup>
 
     <StandardDirectory Id="ProgramFiles64Folder">
@@ -85,10 +70,7 @@
     <Property Id="ConfigureServiceRecovery" Value="&quot;SC.EXE&quot; failureflag fibratus 1" />
     <CustomAction Id="ConfigureServiceRecovery" BinaryRef="Wix4UtilCA_$(sys.BUILDARCHSHORT)" DllEntry="WixQuietExec" Impersonate="no" Execute="deferred" Return="ignore" />
 
-    <!-- Systray -->
-    <Property Id="WixShellExecTarget" Value="[#Systray]"/>
-    <CustomAction Id="LaunchSystray" BinaryRef="Wix4UtilCA_$(sys.BUILDARCHSHORT)" DllEntry="WixShellExec" Impersonate="yes" />
-    <util:CloseApplication Id="CloseSystray" CloseMessage="yes" Target="fibratus-systray" RebootPrompt="no" />
+    <util:CloseApplication Id="CloseSystray" CloseMessage="yes" Target="fibratus-systray.exe" RebootPrompt="no" />
 
     <InstallExecuteSequence>
       <Custom Action="ConfigureServiceRecovery" After="InstallServices" Condition="NOT REMOVE" />
diff --git a/cmd/systray/main_windows.go b/cmd/systray/main_windows.go
@@ -181,26 +181,25 @@ func (s *Systray) loadIconFromResource(mod windows.Handle) (windows.Handle, erro
 
 func (s *Systray) handlePipeClient(conn net.Conn) {
 	defer conn.Close()
-	for {
-		buf, err := io.ReadAll(conn)
-		if err != nil {
-			if err != io.EOF {
-				logrus.Errorf("pipe read: %v", err)
-			}
-			break
-		}
-		var m Msg
-		err = json.Unmarshal(buf, &m)
-		if err != nil {
-			logrus.Error(err)
-			break
-		}
-		err = s.handleMessage(m)
-		if err != nil {
-			logrus.Error(err)
-			break
+	buf, err := io.ReadAll(conn)
+	if err != nil {
+		if err != io.EOF {
+			logrus.Errorf("pipe read: %v", err)
 		}
 	}
+	if len(buf) == 0 {
+		return
+	}
+	var m Msg
+	err = json.Unmarshal(buf, &m)
+	if err != nil {
+		logrus.Error(err)
+		return
+	}
+	err = s.handleMessage(m)
+	if err != nil {
+		logrus.Error(err)
+	}
 }
 
 func (s *Systray) handleMessage(m Msg) error {
@@ -209,6 +208,7 @@ func (s *Systray) handleMessage(m Msg) error {
 		var c systray.Config
 		err := m.decode(&c)
 		if err != nil {
+			logrus.Errorf("unable to decode systray server config: %v", err)
 			return err
 		}
 		s.config = c
@@ -231,7 +231,7 @@ func (s *Systray) handleMessage(m Msg) error {
 }
 
 func main() {
-	err := log.InitFromConfig(log.Config{Level: "info", LogStdout: true}, "fibratus-systray.log")
+	err := log.InitFromConfig(log.Config{Level: "info", LogStdout: true, Formatter: "text"}, "fibratus-systray.log")
 	if err != nil {
 		fmt.Printf("%v", err)
 		os.Exit(1)
@@ -258,9 +258,6 @@ func main() {
 		}
 	}()
 
-	// detach console
-	sys.FreeConsole()
-
 	// server loop
 	for {
 		conn, err := l.Accept()
diff --git a/pkg/alertsender/systray/systray.go b/pkg/alertsender/systray/systray.go
@@ -26,8 +26,12 @@ import (
 	"github.com/cenkalti/backoff/v4"
 	"github.com/rabbitstack/fibratus/pkg/alertsender"
 	"github.com/rabbitstack/fibratus/pkg/kevent"
+	"github.com/rabbitstack/fibratus/pkg/sys"
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/sys/windows"
+	"math"
 	"os"
+	"path/filepath"
 	"time"
 )
 
@@ -45,6 +49,7 @@ const systrayPipe = `\\.\pipe\fibratus-systray`
 // by the systray sender.
 type systray struct {
 	config Config
+	proc   windows.Handle // systray server process handle
 }
 
 // MsgType determines the type of the message sent
@@ -85,6 +90,78 @@ func makeSender(config alertsender.Config) (alertsender.Sender, error) {
 		return &systray{}, nil
 	}
 
+	// spin up systray server process
+	var si windows.StartupInfo
+	var pi windows.ProcessInformation
+
+	exe, err := os.Executable()
+	if err != nil {
+		return nil, err
+	}
+	cmdline := filepath.Join(filepath.Dir(exe), "fibratus-systray.exe")
+	argv, err := windows.UTF16PtrFromString(cmdline)
+	if err != nil {
+		return nil, err
+	}
+
+	log.Infof("starting systray server process %q", cmdline)
+
+	if sys.IsWindowsService() {
+		// if we're running inside Windows Service, the systray server
+		// process must be created in the console session and with the
+		// currently logged-in user access token
+		var token windows.Token
+
+		// enable TCB privilege to obtain the console session user token
+		sys.SetTcbPrivilege()
+
+		consoleSessionID := sys.WTSGetActiveConsoleSessionID()
+		if consoleSessionID == math.MaxUint32 && sys.IsInSandbox() {
+			// if we failed to obtain the console session ID and
+			// are running inside Windows Sandbox, use session 1
+			consoleSessionID = 1
+		}
+
+		if sys.WTSQueryUserToken(consoleSessionID, &token) {
+			log.Infof("obtained user token for console session ID %d", consoleSessionID)
+			err = windows.CreateProcessAsUser(
+				token,
+				nil,
+				argv,
+				nil,
+				nil,
+				false,
+				windows.CREATE_NO_WINDOW,
+				nil,
+				nil,
+				&si,
+				&pi,
+			)
+		} else {
+			err = fmt.Errorf("unable to obtain user token for console session ID %d", consoleSessionID)
+		}
+
+		// drop TCB privilege and close the token handle
+		sys.RemoveTcbPrivilege()
+		_ = windows.CloseHandle(windows.Handle(token))
+	} else {
+		err = windows.CreateProcess(
+			nil,
+			argv,
+			nil,
+			nil,
+			false,
+			windows.CREATE_NO_WINDOW,
+			nil,
+			nil,
+			&si,
+			&pi)
+	}
+
+	if err != nil {
+		log.Warnf("unable to start systray server process: %v", err)
+	}
+
 	s := &systray{config: c}
 	b := &backoff.ExponentialBackOff{
 		// first backoff timeout will be somewhere in the 100 - 300 ms range given the default multiplier
@@ -112,6 +189,10 @@ func makeSender(config alertsender.Config) (alertsender.Sender, error) {
 		break
 	}
 
+	log.Info("established connection to systray server")
+
+	s.proc = pi.Process
+
 	return s, s.send(&Msg{Type: Conf, Data: c})
 }
 
@@ -124,11 +205,17 @@ func (s *systray) Send(alert alertsender.Alert) error {
 func (*systray) Type() alertsender.Type { return alertsender.Systray }
 func (*systray) SupportsMarkdown() bool { return false }
 
-func (s *systray) Shutdown() error { return nil }
+func (s *systray) Shutdown() error {
+	if s.proc != 0 && sys.IsProcessRunning(s.proc) {
+		return windows.TerminateProcess(s.proc, 1)
+	}
+	return nil
+}
 
 func (s *systray) send(m *Msg) error {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
+
 	conn, err := winio.DialPipeContext(ctx, systrayPipe)
 	if err != nil {
 		return fmt.Errorf("unable to dial %s pipe: %v", systrayPipe, err)
@@ -151,6 +238,5 @@ func (s *systray) send(m *Msg) error {
 
 func pipeExists() bool {
 	_, err := os.Stat(systrayPipe)
-	log.Warnf("pipe not found: %v", err)
 	return err == nil
 }
diff --git a/pkg/sys/privilege.go b/pkg/sys/privilege.go
@@ -29,6 +29,8 @@ import (
 const (
 	// SeDebugPrivilege is the name of the privilege used to debug programs.
 	SeDebugPrivilege = "SeDebugPrivilege"
+	// SeTcbPrivilege privilege identifies its holder as part of the trusted computer base.
+	SeTcbPrivilege = "SeTcbPrivilege"
 )
 
 // Errors returned by AdjustTokenPrivileges.
@@ -41,6 +43,8 @@ const (
 const (
 	// PrivilegedEnabled enables the privilege.
 	PrivilegedEnabled uint32 = 0x00000002
+	// PrivilegeRemoved removes the privilege.
+	PrivilegeRemoved uint32 = 0x00000004
 )
 
 // mapPrivileges maps privilege names to LUID values.
@@ -57,10 +61,10 @@ func mapPrivileges(names []string) ([]windows.LUID, error) {
 	return privileges, nil
 }
 
-// EnableTokenPrivileges enables the specified privileges in the given
+// adjustTokenPrivileges enables/disables the specified privileges in the given
 // Token. The token must have TOKEN_ADJUST_PRIVILEGES access. If the token
 // does not already contain the privilege it cannot be enabled.
-func EnableTokenPrivileges(token windows.Token, privileges ...string) error {
+func adjustTokenPrivileges(token windows.Token, state uint32, privileges ...string) error {
 	privValues, err := mapPrivileges(privileges)
 	if err != nil {
 		return err
@@ -74,7 +78,7 @@ func EnableTokenPrivileges(token windows.Token, privileges ...string) error {
 		if err := binary.Write(&b, binary.LittleEndian, p); err != nil {
 			continue
 		}
-		if err := binary.Write(&b, binary.LittleEndian, PrivilegedEnabled); err != nil {
+		if err := binary.Write(&b, binary.LittleEndian, state); err != nil {
 			continue
 		}
 	}
@@ -90,9 +94,29 @@ func EnableTokenPrivileges(token windows.Token, privileges ...string) error {
 	return nil
 }
 
-// SetDebugPrivilege sets the debug privilege in the current running process.
+// SetDebugPrivilege sets the debug privilege in the current process token.
 func SetDebugPrivilege() {
+	enablePrivileges(SeDebugPrivilege)
+}
+
+// SetTcbPrivilege sets the TCB privilege in the current process token.
+func SetTcbPrivilege() {
+	enablePrivileges(SeTcbPrivilege)
+}
+
+// RemoveTcbPrivilege removes the TCB privilege from the access token of the current process.
+func RemoveTcbPrivilege() {
+	removePrivileges(SeTcbPrivilege)
+}
+
+func enablePrivileges(privs ...string) {
+	var token windows.Token
+	_ = windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &token)
+	_ = adjustTokenPrivileges(token, PrivilegedEnabled, privs...)
+}
+
+func removePrivileges(privs ...string) {
 	var token windows.Token
 	_ = windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &token)
-	_ = EnableTokenPrivileges(token, SeDebugPrivilege)
+	_ = adjustTokenPrivileges(token, PrivilegeRemoved, privs...)
 }
diff --git a/pkg/sys/syscall.go b/pkg/sys/syscall.go
@@ -44,6 +44,8 @@ package sys
 
 // Windows Terminal Server Functions
 //sys WTSQuerySessionInformationA(handle windows.Handle, sessionID uint32, klass uint8, buf **uint16, size *uint32) (err error) = wtsapi32.WTSQuerySessionInformationW
+//sys WTSGetActiveConsoleSessionID() (n uint32) = kernel32.WTSGetActiveConsoleSessionId
+//sys WTSQueryUserToken(sessionID uint32, token *windows.Token) (ok bool) = wtsapi32.WTSQueryUserToken
 
 // Windows Trust Functions
 //sys WinVerifyTrust(handle windows.Handle, action *windows.GUID, data *WintrustData) (ret uint32, err error) [failretval!=0] = wintrust.WinVerifyTrust
diff --git a/pkg/sys/util.go b/pkg/sys/util.go
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2021-2022 by Nedim Sabic Sabic
+ * https://www.fibratus.io
+ * All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package sys
+
+import "golang.org/x/sys/windows/registry"
+
+// IsInSandbox indicates if the process is running inside an
+// isolated environment such as Windows Containers or Windows
+// Sandbox.
+func IsInSandbox() bool {
+	key, err := registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM\CurrentControlSet\Control`, registry.QUERY_VALUE)
+	if err != nil {
+		return false
+	}
+	defer key.Close()
+	v, _, err := key.GetIntegerValue("ContainerType")
+
+	return err == nil && v > 0
+}
diff --git a/pkg/sys/zsyscall_windows.go b/pkg/sys/zsyscall_windows.go
