Restore Windows privileges on shutdown, refuse approval on a locked session, and stop sessions reading a recycled capture buffer

client/vnc/server/capture_windows.go

@@ -278,6 +278,11 @@ type DesktopCapturer struct {
 // buffered to size 1 so the worker never blocks on a sender that's gone.
 type captureReq struct {
 	reply chan captureReply
+	// into, when non-nil, is copied into on the worker goroutine before the
+	// reply is sent. The worker recycles its output buffers, so copying there
+	// is what keeps a session from reading one that a later capture has
+	// already started overwriting.
+	into *image.RGBA
 }
 
 type captureReply struct {
@@ -385,6 +390,30 @@ func (c *DesktopCapturer) Capture() (*image.RGBA, error) {
 	}
 }
 
+// CaptureInto fills dst with a freshly captured frame.
+//
+// The worker owns a small ring of output buffers and recycles them, so the
+// pointer Capture returns stops being stable a couple of captures later. This
+// path does the copy on the worker instead, while it still owns the buffer, so
+// a session reading at its own pace can never be overtaken. It also skips the
+// staleness cache: that cache exists to share one DXGI round-trip between
+// sessions asking at the same moment, and copying from it would read the same
+// recycled buffer this is avoiding.
+func (c *DesktopCapturer) CaptureInto(dst *image.RGBA) error {
+	reply := make(chan captureReply, 1)
+	select {
+	case c.reqCh <- captureReq{reply: reply, into: dst}:
+	case <-c.done:
+		return fmt.Errorf("capturer closed")
+	}
+	select {
+	case r := <-reply:
+		return r.err
+	case <-c.done:
+		return fmt.Errorf("capturer closed")
+	}
+}
+
 // waitForClient blocks until a client connects or the capturer is closed.
 func (c *DesktopCapturer) waitForClient() bool {
 	if c.clients.Load() > 0 {
@@ -498,6 +527,17 @@ func (w *captureWorker) serveRequest(req captureReq) {
 	} else {
 		w.c.cursorState.store(snap)
 	}
+
+	if req.into != nil {
+		if req.into.Rect != img.Rect {
+			req.reply <- captureReply{err: fmt.Errorf("dst size mismatch: dst=%v capturer=%v", req.into.Rect, img.Rect)}
+			return
+		}
+		copy(req.into.Pix, img.Pix)
+		req.reply <- captureReply{img: req.into}
+		return
+	}
+
 	req.reply <- captureReply{img: img}
 }
 
@@ -595,3 +635,8 @@ func (w *captureWorker) closeCapturer() {
 		w.cap = nil
 	}
 }
+
+var (
+	_ ScreenCapturer = (*DesktopCapturer)(nil)
+	_ captureIntoer  = (*DesktopCapturer)(nil)
+)

client/vnc/server/console_user_windows.go

@@ -1,8 +1,36 @@
 package server
 
-// interactiveUserError returns nil when there is a logged-in user session
-// on the box. At the lock/login screen WTSQueryUserName is empty, which
-// means there is nobody to display an approval prompt to.
+import (
+	"unsafe"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// wtsSessionInfoEx is WTSSessionInfoEx, the WTSQuerySessionInformation info
+// class that reports a session's lock state.
+const wtsSessionInfoEx = 25
+
+// WTS_SESSIONSTATE_* from wtsapi32.h. Windows 7 and Server 2008 R2 shipped
+// these two inverted; every release the client supports reports lock as 0.
+const (
+	wtsSessionStateLock   int32 = 0
+	wtsSessionStateUnlock int32 = 1
+)
+
+// wtsInfoEx mirrors WTSINFOEXW as far as the field we read. The union that
+// follows Level holds LARGE_INTEGER members, so it is 8-byte aligned and Level
+// is followed by four bytes of padding.
+type wtsInfoEx struct {
+	Level        uint32
+	_            uint32
+	SessionID    uint32
+	SessionState uint32
+	SessionFlags int32
+}
+
+// interactiveUserError returns nil when there is a logged-in user session on
+// the box who could actually see an approval prompt. It is the guard that keeps
+// the prompt failing closed when there is nobody to consent.
 func interactiveUserError() error {
 	sid := getActiveSessionID()
 	if sid == 0 {
@@ -11,5 +39,55 @@ func interactiveUserError() error {
 	if !wtsSessionHasUser(sid) {
 		return errNoConsoleUser
 	}
+
+	// A logged-in session can still be locked, and WTSQueryUserName keeps
+	// returning the user's name throughout: the earlier check only rules out
+	// the login screen, where no session exists yet. While the workstation is
+	// locked, Windows shows the Winlogon secure desktop and the prompt renders
+	// on Default, so nobody can see or answer it — the request would sit there
+	// until it timed out. Refuse up front instead.
+	if locked, known := consoleSessionLocked(sid); known && locked {
+		return errNoConsoleUser
+	}
 	return nil
 }
+
+// consoleSessionLocked reports whether sessionID is locked. known is false when
+// the lock state cannot be determined, in which case the caller keeps its
+// previous behaviour and lets the prompt through: an unanswered prompt still
+// times out into a denial, so guessing "locked" here would only cost sessions
+// on hosts whose lock state is unreadable.
+func consoleSessionLocked(sessionID uint32) (locked, known bool) {
+	var buf uintptr
+	var bytesReturned uint32
+	r, _, _ := procWTSQuerySessionInformation.Call(
+		0, // WTS_CURRENT_SERVER_HANDLE
+		uintptr(sessionID),
+		uintptr(wtsSessionInfoEx),
+		uintptr(unsafe.Pointer(&buf)),
+		uintptr(unsafe.Pointer(&bytesReturned)),
+	)
+	if r == 0 || buf == 0 {
+		return false, false
+	}
+	defer func() { _, _, _ = procWTSFreeMemory.Call(buf) }()
+
+	if uintptr(bytesReturned) < unsafe.Sizeof(wtsInfoEx{}) {
+		log.Debugf("WTSSessionInfoEx returned %d bytes, too short to read the lock state", bytesReturned)
+		return false, false
+	}
+	info := (*wtsInfoEx)(unsafe.Pointer(buf))
+	if info.Level != 1 {
+		return false, false
+	}
+
+	switch info.SessionFlags {
+	case wtsSessionStateLock:
+		return true, true
+	case wtsSessionStateUnlock:
+		return false, true
+	default:
+		// WTS_SESSIONSTATE_UNKNOWN, reported while a session is still settling.
+		return false, false
+	}
+}

client/vnc/server/server_windows.go

@@ -265,8 +265,52 @@ func runSASListenerLoop(ctx context.Context, ev windows.Handle) {
 	}
 }
 
-// enablePrivilege enables a named privilege on the current process token.
-func enablePrivilege(name string) error {
+// priorPrivileges holds the state of the privileges platformInit enabled, so
+// platformShutdown can restore them. Package scope rather than a Server field
+// because these live on the process token, not on any one server instance.
+var (
+	privMu          sync.Mutex
+	priorPrivileges []windows.Tokenprivileges
+)
+
+// enablePrivilege enables a named privilege on the current process token and
+// returns the state it held beforehand, so the caller can put it back. A result
+// with PrivilegeCount == 0 means there was nothing to restore.
+func enablePrivilege(name string) (windows.Tokenprivileges, error) {
+	var prev windows.Tokenprivileges
+
+	var token windows.Token
+	if err := windows.OpenProcessToken(windows.CurrentProcess(),
+		windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &token); err != nil {
+		return prev, err
+	}
+	defer token.Close()
+
+	var luid windows.LUID
+	namePtr, err := windows.UTF16PtrFromString(name)
+	if err != nil {
+		return prev, fmt.Errorf("UTF16 privilege name: %w", err)
+	}
+	if err := windows.LookupPrivilegeValue(nil, namePtr, &luid); err != nil {
+		return prev, err
+	}
+	tp := windows.Tokenprivileges{PrivilegeCount: 1}
+	tp.Privileges[0].Luid = luid
+	tp.Privileges[0].Attributes = windows.SE_PRIVILEGE_ENABLED
+
+	var retLen uint32
+	if err := windows.AdjustTokenPrivileges(token, false, &tp,
+		uint32(unsafe.Sizeof(prev)), &prev, &retLen); err != nil {
+		return windows.Tokenprivileges{}, err
+	}
+	return prev, nil
+}
+
+// restorePrivilege puts back a privilege state captured by enablePrivilege.
+func restorePrivilege(prev windows.Tokenprivileges) error {
+	if prev.PrivilegeCount == 0 {
+		return nil
+	}
 	var token windows.Token
 	if err := windows.OpenProcessToken(windows.CurrentProcess(),
 		windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &token); err != nil {
@@ -274,18 +318,7 @@ func enablePrivilege(name string) error {
 	}
 	defer token.Close()
 
-	var luid windows.LUID
-	namePtr, err := windows.UTF16PtrFromString(name)
-	if err != nil {
-		return fmt.Errorf("UTF16 privilege name: %w", err)
-	}
-	if err := windows.LookupPrivilegeValue(nil, namePtr, &luid); err != nil {
-		return err
-	}
-	tp := windows.Tokenprivileges{PrivilegeCount: 1}
-	tp.Privileges[0].Luid = luid
-	tp.Privileges[0].Attributes = windows.SE_PRIVILEGE_ENABLED
-	return windows.AdjustTokenPrivileges(token, false, &tp, 0, nil, nil)
+	return windows.AdjustTokenPrivileges(token, false, &prev, 0, nil, nil)
 }
 
 func (s *Server) platformSessionManager() virtualSessionManager {
@@ -295,16 +328,40 @@ func (s *Server) platformSessionManager() virtualSessionManager {
 // platformShutdown restores any machine state mutated by platformInit.
 func (s *Server) platformShutdown() {
 	disableSoftwareSAS()
+
+	// Hand back the privileges platformInit took. The daemon token outlives
+	// the VNC server, so leaving SeTcb and SeAssignPrimaryToken enabled would
+	// keep two of the most powerful privileges on the box switched on for the
+	// rest of the process after VNC is turned off again.
+	privMu.Lock()
+	prev := priorPrivileges
+	priorPrivileges = nil
+	privMu.Unlock()
+
+	for _, p := range prev {
+		if err := restorePrivilege(p); err != nil {
+			log.Debugf("restore privilege: %v", err)
+		}
+	}
 }
 
 // platformInit starts the SAS listener and enables privileges needed for
 // Session 0 operations (agent spawning, SendSAS).
 func (s *Server) platformInit() {
+	var prior []windows.Tokenprivileges
 	for _, priv := range []string{"SeTcbPrivilege", "SeAssignPrimaryTokenPrivilege"} {
-		if err := enablePrivilege(priv); err != nil {
+		prev, err := enablePrivilege(priv)
+		if err != nil {
 			log.Debugf("enable %s: %v", priv, err)
+			continue
 		}
+		prior = append(prior, prev)
 	}
+
+	privMu.Lock()
+	priorPrivileges = prior
+	privMu.Unlock()
+
 	startSASListener(s.ctx)
 }