From c3f1bd4c0af7615455c987b5709908e75e239733 Mon Sep 17 00:00:00 2001
From: Andrew Lytvynov <awly@tailscale.com>
Date: Tue, 14 Nov 2023 18:21:03 -0700
Subject: [PATCH] clientupdate: fix auto-update on Windows over RDP (#10242)

`winutil.WTSGetActiveConsoleSessionId` only works for physical desktop
logins and does not return the session ID for RDP logins. We need to
`windows.WTSEnumerateSessions` and find the active session.

Fixes https://github.com/tailscale/corp/issues/15772

Signed-off-by: Andrew Lytvynov <awly@tailscale.com>
---
 clientupdate/clientupdate_windows.go | 41 +++++++++++++++++++++++-----
 cmd/tailscale/depaware.txt           |  2 +-
 cmd/tailscaled/depaware.txt          |  2 +-
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/clientupdate/clientupdate_windows.go b/clientupdate/clientupdate_windows.go
index 299343e02..ef21cf05d 100644
--- a/clientupdate/clientupdate_windows.go
+++ b/clientupdate/clientupdate_windows.go
@@ -7,13 +7,15 @@
 package clientupdate
 
 import (
+	"errors"
+	"fmt"
 	"os/exec"
 	"os/user"
 	"path/filepath"
 	"syscall"
+	"unsafe"
 
 	"golang.org/x/sys/windows"
-	"tailscale.com/util/winutil"
 	"tailscale.com/util/winutil/authenticode"
 )
 
@@ -39,13 +41,14 @@ func launchTailscaleAsGUIUser(exePath string) error {
 
 	var token windows.Token
 	if u, err := user.Current(); err == nil && u.Name == "SYSTEM" {
-		sessionID := winutil.WTSGetActiveConsoleSessionId()
-		if sessionID != 0xFFFFFFFF {
-			if err := windows.WTSQueryUserToken(sessionID, &token); err != nil {
-				return err
-			}
-			defer token.Close()
+		sessionID, err := wtsGetActiveSessionID()
+		if err != nil {
+			return fmt.Errorf("wtsGetActiveSessionID(): %w", err)
 		}
+		if err := windows.WTSQueryUserToken(sessionID, &token); err != nil {
+			return fmt.Errorf("WTSQueryUserToken (0x%x): %w", sessionID, err)
+		}
+		defer token.Close()
 	}
 
 	cmd := exec.Command(exePath)
@@ -55,3 +58,27 @@ func launchTailscaleAsGUIUser(exePath string) error {
 	}
 	return cmd.Start()
 }
+
+func wtsGetActiveSessionID() (uint32, error) {
+	var (
+		sessionInfo *windows.WTS_SESSION_INFO
+		count       uint32 = 0
+	)
+
+	const WTS_CURRENT_SERVER_HANDLE = 0
+	if err := windows.WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &sessionInfo, &count); err != nil {
+		return 0, fmt.Errorf("WTSEnumerateSessions: %w", err)
+	}
+	defer windows.WTSFreeMemory(uintptr(unsafe.Pointer(sessionInfo)))
+
+	current := unsafe.Pointer(sessionInfo)
+	for i := uint32(0); i < count; i++ {
+		session := (*windows.WTS_SESSION_INFO)(current)
+		if session.State == windows.WTSActive {
+			return session.SessionID, nil
+		}
+		current = unsafe.Add(current, unsafe.Sizeof(windows.WTS_SESSION_INFO{}))
+	}
+
+	return 0, errors.New("no active desktop sessions found")
+}
diff --git a/cmd/tailscale/depaware.txt b/cmd/tailscale/depaware.txt
index 7225470f2..32131bdb2 100644
--- a/cmd/tailscale/depaware.txt
+++ b/cmd/tailscale/depaware.txt
@@ -76,7 +76,7 @@ tailscale.com/cmd/tailscale dependencies: (generated by github.com/tailscale/dep
         tailscale.com/client/tailscale                               from tailscale.com/cmd/tailscale/cli+
         tailscale.com/client/tailscale/apitype                       from tailscale.com/cmd/tailscale/cli+
         tailscale.com/client/web                                     from tailscale.com/cmd/tailscale/cli
-        tailscale.com/clientupdate                                   from tailscale.com/cmd/tailscale/cli
+     💣 tailscale.com/clientupdate                                   from tailscale.com/cmd/tailscale/cli
         tailscale.com/clientupdate/distsign                          from tailscale.com/clientupdate
         tailscale.com/cmd/tailscale/cli                              from tailscale.com/cmd/tailscale
         tailscale.com/control/controlbase                            from tailscale.com/control/controlhttp
diff --git a/cmd/tailscaled/depaware.txt b/cmd/tailscaled/depaware.txt
index 0a9afacb3..d7b43fa96 100644
--- a/cmd/tailscaled/depaware.txt
+++ b/cmd/tailscaled/depaware.txt
@@ -228,7 +228,7 @@ tailscale.com/cmd/tailscaled dependencies: (generated by github.com/tailscale/de
         tailscale.com/client/tailscale                               from tailscale.com/derp+
         tailscale.com/client/tailscale/apitype                       from tailscale.com/ipn/ipnlocal+
         tailscale.com/client/web                                     from tailscale.com/ipn/ipnlocal
-        tailscale.com/clientupdate                                   from tailscale.com/ipn/ipnlocal+
+     💣 tailscale.com/clientupdate                                   from tailscale.com/ipn/ipnlocal+
         tailscale.com/clientupdate/distsign                          from tailscale.com/clientupdate
         tailscale.com/cmd/tailscaled/childproc                       from tailscale.com/ssh/tailssh+
         tailscale.com/control/controlbase                            from tailscale.com/control/controlclient+