netbird/client/internal/engine_vnc.go

package internal

import (
	"context"
	"errors"
	"fmt"
	"net/netip"
	"os"
	"path/filepath"

	log "github.com/sirupsen/logrus"

	firewallManager "github.com/netbirdio/netbird/client/firewall/manager"
	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
	sshauth "github.com/netbirdio/netbird/client/ssh/auth"
	vncserver "github.com/netbirdio/netbird/client/vnc/server"
	mgmProto "github.com/netbirdio/netbird/shared/management/proto"
	sshuserhash "github.com/netbirdio/netbird/shared/sshauth"
)

const envVNCForceRecording = "NB_VNC_FORCE_RECORDING"

const (
	vncExternalPort uint16 = 5900
	vncInternalPort uint16 = 25900
)

type vncServer interface {
	Start(ctx context.Context, addr netip.AddrPort, network netip.Prefix) error
	Stop() error
}

func (e *Engine) setupVNCPortRedirection() error {
	if e.firewall == nil || e.wgInterface == nil {
		return nil
	}

	localAddr := e.wgInterface.Address().IP
	if !localAddr.IsValid() {
		return errors.New("invalid local NetBird address")
	}

	if err := e.firewall.AddInboundDNAT(localAddr, firewallManager.ProtocolTCP, vncExternalPort, vncInternalPort); err != nil {
		return fmt.Errorf("add VNC port redirection: %w", err)
	}
	log.Infof("VNC port redirection: %s:%d -> %s:%d", localAddr, vncExternalPort, localAddr, vncInternalPort)

	return nil
}

func (e *Engine) cleanupVNCPortRedirection() error {
	if e.firewall == nil || e.wgInterface == nil {
		return nil
	}

	localAddr := e.wgInterface.Address().IP
	if !localAddr.IsValid() {
		return errors.New("invalid local NetBird address")
	}

	if err := e.firewall.RemoveInboundDNAT(localAddr, firewallManager.ProtocolTCP, vncExternalPort, vncInternalPort); err != nil {
		return fmt.Errorf("remove VNC port redirection: %w", err)
	}

	return nil
}

// updateVNC handles starting/stopping the VNC server based on the config flag.
// sshConf provides the JWT identity provider config (shared with SSH).
func (e *Engine) updateVNC(sshConf *mgmProto.SSHConfig) error {
	if !e.config.ServerVNCAllowed {
		if e.vncSrv != nil {
			log.Info("VNC server disabled, stopping")
		}
		return e.stopVNCServer()
	}

	if e.config.BlockInbound {
		log.Info("VNC server disabled because inbound connections are blocked")
		return e.stopVNCServer()
	}

	if e.vncSrv != nil {
		// Update JWT config on existing server in case management sent new config.
		e.updateVNCServerJWT(sshConf)
		return nil
	}

	return e.startVNCServer(sshConf)
}

func (e *Engine) startVNCServer(sshConf *mgmProto.SSHConfig) error {
	if e.wgInterface == nil {
		return errors.New("wg interface not initialized")
	}

	capturer, injector := newPlatformVNC()
	if capturer == nil || injector == nil {
		log.Debug("VNC server not supported on this platform")
		return nil
	}

	netbirdIP := e.wgInterface.Address().IP

	srv := vncserver.New(capturer, injector, "")
	if vncNeedsServiceMode() {
		log.Info("VNC: running in Session 0, enabling service mode (agent proxy)")
		srv.SetServiceMode(true)
	}

	// Configure VNC authentication.
	if e.config.DisableVNCAuth != nil && *e.config.DisableVNCAuth {
		log.Info("VNC: authentication disabled by config")
		srv.SetDisableAuth(true)
	} else if protoJWT := sshConf.GetJwtConfig(); protoJWT != nil {
		audiences := protoJWT.GetAudiences()
		if len(audiences) == 0 && protoJWT.GetAudience() != "" {
			audiences = []string{protoJWT.GetAudience()}
		}
		srv.SetJWTConfig(&vncserver.JWTConfig{
			Issuer:       protoJWT.GetIssuer(),
			Audiences:    audiences,
			KeysLocation: protoJWT.GetKeysLocation(),
			MaxTokenAge:  protoJWT.GetMaxTokenAge(),
		})
		log.Debugf("VNC: JWT authentication configured (issuer=%s)", protoJWT.GetIssuer())
	}

	e.configureVNCRecording(srv, sshConf)

	if netstackNet := e.wgInterface.GetNet(); netstackNet != nil {
		srv.SetNetstackNet(netstackNet)
	}

	listenAddr := netip.AddrPortFrom(netbirdIP, vncInternalPort)
	network := e.wgInterface.Address().Network
	if err := srv.Start(e.ctx, listenAddr, network); err != nil {
		return fmt.Errorf("start VNC server: %w", err)
	}

	e.vncSrv = srv

	if registrar, ok := e.firewall.(interface {
		RegisterNetstackService(protocol nftypes.Protocol, port uint16)
	}); ok {
		registrar.RegisterNetstackService(nftypes.TCP, vncInternalPort)
		log.Debugf("registered VNC service for TCP:%d", vncInternalPort)
	}

	if err := e.setupVNCPortRedirection(); err != nil {
		log.Warnf("setup VNC port redirection: %v", err)
	}

	log.Info("VNC server enabled")
	return nil
}

// configureVNCRecording enables session recording on the VNC server from the
// management-supplied settings. The env var NB_VNC_FORCE_RECORDING overrides
// the API for local development: when set, recording is always enabled and
// writes into that directory. Otherwise recordings go next to the state file
// under vnc-recordings/.
func (e *Engine) configureVNCRecording(srv *vncserver.Server, sshConf *mgmProto.SSHConfig) {
	recDir := os.Getenv(envVNCForceRecording)
	apiEnabled := sshConf.GetEnableRecording()

	if recDir == "" && !apiEnabled {
		log.Debugf("VNC recording disabled (env=%q, api=%v)", recDir, apiEnabled)
		return
	}

	if recDir == "" {
		base := e.defaultRecordingBase()
		if base == "" {
			log.Warn("VNC recording requested by management but no state directory is available")
			return
		}
		recDir = filepath.Join(base, "vnc-recordings")
	} else {
		recDir = filepath.Join(recDir, "vnc")
	}

	srv.SetRecordingDir(recDir)
	log.Infof("VNC recording enabled (dir=%s, source=%s)", recDir, recordingSource(apiEnabled))

	encKey := string(sshConf.GetRecordingEncryptionKey())
	if encKey == "" {
		encKey = os.Getenv("NB_VNC_RECORDING_ENCRYPTION_KEY")
	}
	if encKey != "" {
		srv.SetRecordingEncryptionKey(encKey)
		log.Info("VNC recording encryption enabled")
	}
}

func (e *Engine) defaultRecordingBase() string {
	if e.stateManager == nil {
		return ""
	}
	p := e.stateManager.FilePath()
	if p == "" {
		return ""
	}
	return filepath.Dir(p)
}

func recordingSource(api bool) string {
	if api {
		return "management"
	}
	return "env"
}

// updateVNCServerJWT configures the JWT validation for the VNC server using
// the same JWT config as SSH (same identity provider).
func (e *Engine) updateVNCServerJWT(sshConf *mgmProto.SSHConfig) {
	if e.vncSrv == nil {
		return
	}

	vncSrv, ok := e.vncSrv.(*vncserver.Server)
	if !ok {
		return
	}

	if e.config.DisableVNCAuth != nil && *e.config.DisableVNCAuth {
		vncSrv.SetDisableAuth(true)
		return
	}

	protoJWT := sshConf.GetJwtConfig()
	if protoJWT == nil {
		return
	}

	audiences := protoJWT.GetAudiences()
	if len(audiences) == 0 && protoJWT.GetAudience() != "" {
		audiences = []string{protoJWT.GetAudience()}
	}

	vncSrv.SetJWTConfig(&vncserver.JWTConfig{
		Issuer:       protoJWT.GetIssuer(),
		Audiences:    audiences,
		KeysLocation: protoJWT.GetKeysLocation(),
		MaxTokenAge:  protoJWT.GetMaxTokenAge(),
	})
}

// updateVNCServerAuth updates VNC fine-grained access control from management.
func (e *Engine) updateVNCServerAuth(vncAuth *mgmProto.VNCAuth) {
	if vncAuth == nil || e.vncSrv == nil {
		return
	}

	vncSrv, ok := e.vncSrv.(*vncserver.Server)
	if !ok {
		return
	}

	protoUsers := vncAuth.GetAuthorizedUsers()
	authorizedUsers := make([]sshuserhash.UserIDHash, len(protoUsers))
	for i, hash := range protoUsers {
		if len(hash) != 16 {
			log.Warnf("invalid VNC auth hash length %d, expected 16", len(hash))
			return
		}
		authorizedUsers[i] = sshuserhash.UserIDHash(hash)
	}

	machineUsers := make(map[string][]uint32)
	for osUser, indexes := range vncAuth.GetMachineUsers() {
		machineUsers[osUser] = indexes.GetIndexes()
	}

	vncSrv.UpdateVNCAuth(&sshauth.Config{
		UserIDClaim:     vncAuth.GetUserIDClaim(),
		AuthorizedUsers: authorizedUsers,
		MachineUsers:    machineUsers,
	})
}

// GetVNCServerStatus returns whether the VNC server is running.
func (e *Engine) GetVNCServerStatus() bool {
	return e.vncSrv != nil
}

func (e *Engine) stopVNCServer() error {
	if e.vncSrv == nil {
		return nil
	}

	if err := e.cleanupVNCPortRedirection(); err != nil {
		log.Warnf("cleanup VNC port redirection: %v", err)
	}

	if registrar, ok := e.firewall.(interface {
		UnregisterNetstackService(protocol nftypes.Protocol, port uint16)
	}); ok {
		registrar.UnregisterNetstackService(nftypes.TCP, vncInternalPort)
	}

	log.Info("stopping VNC server")
	err := e.vncSrv.Stop()
	e.vncSrv = nil
	if err != nil {
		return fmt.Errorf("stop VNC server: %w", err)
	}
	return nil
}