Track active VNC sessions in status and address CodeRabbit findings

This commit is contained in:
Viktor Liu
2026-05-16 17:06:19 +02:00
parent 9f0aa1ce26
commit 62cf9e873b
12 changed files with 708 additions and 471 deletions

View File

@@ -24,6 +24,7 @@ const (
type vncServer interface {
Start(ctx context.Context, addr netip.AddrPort, network netip.Prefix) error
Stop() error
ActiveSessions() []vncserver.ActiveSessionInfo
}
func (e *Engine) setupVNCPortRedirection() error {
@@ -208,9 +209,13 @@ func (e *Engine) updateVNCServerAuth(vncAuth *mgmProto.VNCAuth) {
})
}
// GetVNCServerStatus returns whether the VNC server is running.
func (e *Engine) GetVNCServerStatus() bool {
return e.vncSrv != nil
// GetVNCServerStatus returns whether the VNC server is running and the list
// of active VNC sessions.
func (e *Engine) GetVNCServerStatus() (enabled bool, sessions []vncserver.ActiveSessionInfo) {
if e.vncSrv == nil {
return false, nil
}
return true, e.vncSrv.ActiveSessions()
}
func (e *Engine) stopVNCServer() error {

View File

@@ -21,8 +21,7 @@ func newConsoleVNC() (vncserver.ScreenCapturer, vncserver.InputInjector, error)
}
inj, err := vncserver.NewUInputInjector(w, h)
if err != nil {
poller.Close()
return nil, nil, fmt.Errorf("uinput init: %w", err)
return poller, &vncserver.StubInputInjector{}, nil
}
return poller, inj, nil
}

File diff suppressed because it is too large Load Diff

View File

@@ -398,9 +398,18 @@ message SSHServerState {
repeated SSHSessionInfo sessions = 2;
}
// VNCSessionInfo contains information about an active VNC session
message VNCSessionInfo {
string remoteAddress = 1;
string mode = 2;
string username = 3;
string jwtUsername = 4;
}
// VNCServerState contains the latest state of the VNC server
message VNCServerState {
bool enabled = 1;
repeated VNCSessionInfo sessions = 2;
}
// FullStatus contains the full state held by the Status instance

View File

@@ -1192,8 +1192,19 @@ func (s *Server) getVNCServerState() *proto.VNCServerState {
return nil
}
enabled, sessions := engine.GetVNCServerStatus()
pbSessions := make([]*proto.VNCSessionInfo, 0, len(sessions))
for _, sess := range sessions {
pbSessions = append(pbSessions, &proto.VNCSessionInfo{
RemoteAddress: sess.RemoteAddress,
Mode: sess.Mode,
Username: sess.Username,
JwtUsername: sess.JWTUsername,
})
}
return &proto.VNCServerState{
Enabled: engine.GetVNCServerStatus(),
Enabled: enabled,
Sessions: pbSessions,
}
}

View File

@@ -131,8 +131,16 @@ type SSHServerStateOutput struct {
Sessions []SSHSessionOutput `json:"sessions" yaml:"sessions"`
}
type VNCSessionOutput struct {
RemoteAddress string `json:"remoteAddress" yaml:"remoteAddress"`
Mode string `json:"mode" yaml:"mode"`
Username string `json:"username,omitempty" yaml:"username,omitempty"`
JWTUsername string `json:"jwtUsername,omitempty" yaml:"jwtUsername,omitempty"`
}
type VNCServerStateOutput struct {
Enabled bool `json:"enabled" yaml:"enabled"`
Enabled bool `json:"enabled" yaml:"enabled"`
Sessions []VNCSessionOutput `json:"sessions" yaml:"sessions"`
}
type OutputOverview struct {
@@ -178,9 +186,7 @@ func ConvertToStatusOutputOverview(pbFullStatus *proto.FullStatus, opts ConvertO
relayOverview := mapRelays(pbFullStatus.GetRelays())
sshServerOverview := mapSSHServer(pbFullStatus.GetSshServerState())
vncServerOverview := VNCServerStateOutput{
Enabled: pbFullStatus.GetVncServerState().GetEnabled(),
}
vncServerOverview := mapVNCServer(pbFullStatus.GetVncServerState())
peersOverview := mapPeers(pbFullStatus.GetPeers(), opts.StatusFilter, opts.PrefixNamesFilter, opts.PrefixNamesFilterMap, opts.IPsFilter, opts.ConnectionTypeFilter)
overview := OutputOverview{
@@ -280,6 +286,25 @@ func mapSSHServer(sshServerState *proto.SSHServerState) SSHServerStateOutput {
}
}
func mapVNCServer(state *proto.VNCServerState) VNCServerStateOutput {
if state == nil {
return VNCServerStateOutput{Sessions: []VNCSessionOutput{}}
}
sessions := make([]VNCSessionOutput, 0, len(state.GetSessions()))
for _, sess := range state.GetSessions() {
sessions = append(sessions, VNCSessionOutput{
RemoteAddress: sess.GetRemoteAddress(),
Mode: sess.GetMode(),
Username: sess.GetUsername(),
JWTUsername: sess.GetJwtUsername(),
})
}
return VNCServerStateOutput{
Enabled: state.GetEnabled(),
Sessions: sessions,
}
}
func mapPeers(
peers []*proto.PeerState,
statusFilter string,
@@ -544,7 +569,30 @@ func (o *OutputOverview) GeneralSummary(showURL bool, showRelays bool, showNameS
vncServerStatus := "Disabled"
if o.VNCServerState.Enabled {
vncServerStatus = "Enabled"
vncSessionCount := len(o.VNCServerState.Sessions)
if vncSessionCount > 0 {
sessionWord := "session"
if vncSessionCount > 1 {
sessionWord = "sessions"
}
vncServerStatus = fmt.Sprintf("Enabled (%d active %s)", vncSessionCount, sessionWord)
} else {
vncServerStatus = "Enabled"
}
if showSSHSessions && vncSessionCount > 0 {
for _, sess := range o.VNCServerState.Sessions {
var line string
if sess.JWTUsername != "" {
line = fmt.Sprintf("[%s@%s -> %s] mode=%s",
sess.JWTUsername, sess.RemoteAddress, sess.Username, sess.Mode)
} else {
line = fmt.Sprintf("[%s] mode=%s user=%s",
sess.RemoteAddress, sess.Mode, sess.Username)
}
vncServerStatus += "\n " + line
}
}
}
peersCountString := fmt.Sprintf("%d/%d Connected", o.Peers.Connected, o.Peers.Total)
@@ -1011,4 +1059,12 @@ func anonymizeOverview(a *anonymize.Anonymizer, overview *OutputOverview) {
}
overview.SSHServerState.Sessions[i].Command = a.AnonymizeString(session.Command)
}
for i, sess := range overview.VNCServerState.Sessions {
if host, port, err := net.SplitHostPort(sess.RemoteAddress); err == nil {
overview.VNCServerState.Sessions[i].RemoteAddress = fmt.Sprintf("%s:%s", a.AnonymizeIPString(host), port)
} else {
overview.VNCServerState.Sessions[i].RemoteAddress = a.AnonymizeIPString(sess.RemoteAddress)
}
}
}

View File

@@ -240,6 +240,10 @@ var overview = OutputOverview{
Enabled: false,
Sessions: []SSHSessionOutput{},
},
VNCServerState: VNCServerStateOutput{
Enabled: false,
Sessions: []VNCSessionOutput{},
},
}
func TestConversionFromFullStatusToOutputOverview(t *testing.T) {
@@ -406,7 +410,8 @@ func TestParsingToJSON(t *testing.T) {
"sessions":[]
},
"vncServer":{
"enabled":false
"enabled":false,
"sessions":[]
}
}`
// @formatter:on
@@ -518,6 +523,7 @@ sshServer:
sessions: []
vncServer:
enabled: false
sessions: []
`
assert.Equal(t, expectedYAML, yaml)

View File

@@ -285,21 +285,25 @@ func (c *CGCapturer) Capture() (*image.RGBA, error) {
case bytesPerPixel == 4 && ds == 2:
convertBGRAToRGBADownscale2(img.Pix, img.Stride, src, bytesPerRow, outW, outH)
default:
convertBGRAToRGBAGeneric(img.Pix, img.Stride, src, bytesPerRow, outW, outH, bytesPerPixel, ds)
convertBGRAToRGBAGeneric(img.Pix, img.Stride, src, bytesPerRow, bgraDownscaleParams{outW: outW, outH: outH, bytesPerPixel: bytesPerPixel, ds: ds})
}
return img, nil
}
type bgraDownscaleParams struct {
outW, outH, bytesPerPixel, ds int
}
// convertBGRAToRGBAGeneric is the slow per-pixel fallback for non-4-bytes
// or non-1/2 downscale formats. Always available regardless of the source
// format quirks the fast paths optimize for.
func convertBGRAToRGBAGeneric(dst []byte, dstStride int, src []byte, srcStride, outW, outH, bytesPerPixel, ds int) {
for row := 0; row < outH; row++ {
srcOff := row * ds * srcStride
func convertBGRAToRGBAGeneric(dst []byte, dstStride int, src []byte, srcStride int, p bgraDownscaleParams) {
for row := 0; row < p.outH; row++ {
srcOff := row * p.ds * srcStride
dstOff := row * dstStride
for col := 0; col < outW; col++ {
si := srcOff + col*ds*bytesPerPixel
for col := 0; col < p.outW; col++ {
si := srcOff + col*p.ds*p.bytesPerPixel
di := dstOff + col*4
dst[di+0] = src[si+2]
dst[di+1] = src[si+1]

View File

@@ -321,7 +321,7 @@ func (c *DesktopCapturer) Width() int {
c.mu.Lock()
w := c.w
c.mu.Unlock()
if w == 0 {
if w == 0 && c.clients.Load() > 0 {
_, _ = c.Capture()
c.mu.Lock()
w = c.w
@@ -331,12 +331,13 @@ func (c *DesktopCapturer) Width() int {
}
// Height returns the current screen height, triggering a capture if the
// worker hasn't initialised yet (see Width).
// worker hasn't initialised yet (see Width). Returns 0 while no client is
// connected so callers don't deadlock against a parked worker.
func (c *DesktopCapturer) Height() int {
c.mu.Lock()
h := c.h
c.mu.Unlock()
if h == 0 {
if h == 0 && c.clients.Load() > 0 {
_, _ = c.Capture()
c.mu.Lock()
h = c.h

View File

@@ -35,9 +35,10 @@ const (
wheelDelta = 120
keyeventfKeyUp = 0x0002
keyeventfUnicode = 0x0004
keyeventfScanCode = 0x0008
keyeventfExtendedKey = 0x0001
keyeventfKeyUp = 0x0002
keyeventfUnicode = 0x0004
keyeventfScanCode = 0x0008
)
// winlogonDesktopName is the name of the Windows secure desktop that hosts the
@@ -234,7 +235,7 @@ func (w *WindowsInputInjector) doInjectKey(keysym uint32, down bool) {
flags |= keyeventfKeyUp
}
if extended {
flags |= keyeventfScanCode
flags |= keyeventfExtendedKey
}
sendKeyInput(vk, 0, flags)
}

View File

@@ -364,8 +364,8 @@ type rectCoalescer struct {
curY int
}
func newRectCoalescer(cap int) *rectCoalescer {
return &rectCoalescer{out: make([][4]int, 0, cap)}
func newRectCoalescer(capacity int) *rectCoalescer {
return &rectCoalescer{out: make([][4]int, 0, capacity)}
}
// consume processes one rect from the (row-ordered) input.

View File

@@ -147,6 +147,18 @@ type Server struct {
authorizer *sshauth.Authorizer
netstackNet *netstack.Net
agentToken []byte // raw token bytes for agent-mode auth
sessionsMu sync.Mutex
sessionSeq uint64
sessions map[uint64]ActiveSessionInfo
}
// ActiveSessionInfo describes a currently connected VNC client.
type ActiveSessionInfo struct {
RemoteAddress string
Mode string
Username string
JWTUsername string
}
// vncSession provides capturer and injector for a virtual display session.
@@ -174,9 +186,36 @@ func New(capturer ScreenCapturer, injector InputInjector, password string) *Serv
password: password,
authorizer: sshauth.NewAuthorizer(),
log: log.WithField("component", "vnc-server"),
sessions: make(map[uint64]ActiveSessionInfo),
}
}
// ActiveSessions returns a snapshot of currently connected VNC clients.
func (s *Server) ActiveSessions() []ActiveSessionInfo {
s.sessionsMu.Lock()
defer s.sessionsMu.Unlock()
out := make([]ActiveSessionInfo, 0, len(s.sessions))
for _, info := range s.sessions {
out = append(out, info)
}
return out
}
func (s *Server) addSession(info ActiveSessionInfo) uint64 {
s.sessionsMu.Lock()
defer s.sessionsMu.Unlock()
s.sessionSeq++
id := s.sessionSeq
s.sessions[id] = info
return id
}
func (s *Server) removeSession(id uint64) {
s.sessionsMu.Lock()
defer s.sessionsMu.Unlock()
delete(s.sessions, id)
}
// SetServiceMode enables proxy-to-agent mode for Windows service operation.
func (s *Server) SetServiceMode(enabled bool) {
s.serviceMode = enabled
@@ -408,7 +447,7 @@ func (s *Server) handleConnection(conn net.Conn) {
conn.Close()
return
}
connLog, ok := s.authorizeJWT(conn, header, connLog)
connLog, jwtUserID, ok := s.authorizeJWT(conn, header, connLog)
if !ok {
return
}
@@ -419,6 +458,14 @@ func (s *Server) handleConnection(conn net.Conn) {
}
defer sessionCleanup()
sessionID := s.addSession(ActiveSessionInfo{
RemoteAddress: conn.RemoteAddr().String(),
Mode: modeString(header.mode),
Username: header.username,
JWTUsername: jwtUserID,
})
defer s.removeSession(sessionID)
if err := s.validateCapturer(capturer); err != nil {
rejectConnection(conn, codeMessage(RejectCodeCapturerError, fmt.Sprintf("screen capturer: %v", err)))
connLog.Warnf("capturer not ready: %v", err)
@@ -686,23 +733,24 @@ func (s *Server) verifyAgentToken(conn net.Conn, connLog *log.Entry) bool {
}
// authorizeJWT performs JWT validation when auth is enabled. Returns the
// enriched log entry and ok=false if the connection was rejected.
func (s *Server) authorizeJWT(conn net.Conn, header *connectionHeader, connLog *log.Entry) (*log.Entry, bool) {
// enriched log entry, jwt user ID (empty when auth disabled), and ok=false
// if the connection was rejected.
func (s *Server) authorizeJWT(conn net.Conn, header *connectionHeader, connLog *log.Entry) (*log.Entry, string, bool) {
if s.disableAuth {
return connLog, true
return connLog, "", true
}
if s.jwtConfig == nil {
rejectConnection(conn, codeMessage(RejectCodeAuthConfig, "auth enabled but no identity provider configured"))
connLog.Warn("auth rejected: no identity provider configured")
return connLog, false
return connLog, "", false
}
jwtUserID, err := s.authenticateJWT(header)
if err != nil {
rejectConnection(conn, codeMessage(jwtErrorCode(err), err.Error()))
connLog.Warnf("auth rejected: %v", err)
return connLog, false
return connLog, "", false
}
return connLog.WithField("jwt_user", jwtUserID), true
return connLog.WithField("jwt_user", jwtUserID), jwtUserID, true
}
// acquireSessionResources returns the capturer/injector to use for this
@@ -752,3 +800,15 @@ func (s *Server) acquireAttachSession() ScreenCapturer {
func attachSessionCleanup() {
// Attach mode keeps the shared capturer; nothing to release per session.
}
// modeString returns a human-readable session mode name.
func modeString(m byte) string {
switch m {
case ModeAttach:
return "attach"
case ModeSession:
return "session"
default:
return "unknown"
}
}