rdp-proxy/backend/internal/modules/sessionbroker/dataplane.go

package sessionbroker

import (
	"crypto/rsa"
	"fmt"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"

	"github.com/example/remote-access-platform/backend/internal/platform/secrets"
	sessioncontracts "github.com/example/remote-access-platform/backend/pkg/contracts/session"
)

const (
	directWorkerTLSTrustModeSmokeInsecure = "smoke_insecure"
	directWorkerTLSTrustModePublicCA      = "public_ca"
	directWorkerTLSTrustModePlatformCA    = "platform_ca"
)

type DataPlaneTokenClaims struct {
	SessionID       string    `json:"session_id"`
	AttachmentID    string    `json:"attachment_id"`
	UserID          string    `json:"user_id"`
	OrganizationID  string    `json:"organization_id"`
	ClusterID       string    `json:"cluster_id,omitempty"`
	WorkerID        string    `json:"worker_id"`
	ResourceID      string    `json:"resource_id"`
	AllowedChannels []string  `json:"allowed_channels"`
	ExpiresAtValue  time.Time `json:"expires_at"`
	jwt.RegisteredClaims
}

func (s *Service) buildDataPlaneOffer(session RemoteSession, attachment SessionAttachment) (*sessioncontracts.DataPlaneOffer, error) {
	if s.cfg.DataPlane.TokenTTL <= 0 || s.cfg.DataPlane.TokenPrivateKeyPEM == "" {
		return nil, nil
	}

	now := s.now().UTC()
	expiresAt := now.Add(s.cfg.DataPlane.TokenTTL)
	allowedChannels := dataPlaneAllowedChannelsFromSession(session)
	jti := uuid.NewString()
	claims := DataPlaneTokenClaims{
		SessionID:       session.ID,
		AttachmentID:    attachment.ID,
		UserID:          attachment.UserID,
		OrganizationID:  session.OrganizationID,
		WorkerID:        session.WorkerID,
		ResourceID:      session.ResourceID,
		AllowedChannels: allowedChannels,
		ExpiresAtValue:  expiresAt,
		RegisteredClaims: jwt.RegisteredClaims{
			ID:        jti,
			Issuer:    s.cfg.Auth.Issuer,
			Subject:   attachment.UserID,
			Audience:  jwt.ClaimStrings{"rap-data-plane", "worker:" + session.WorkerID},
			IssuedAt:  jwt.NewNumericDate(now),
			NotBefore: jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(expiresAt),
		},
	}
	token, err := signDataPlaneToken(claims, s.cfg.DataPlane.TokenPrivateKeyPEM)
	if err != nil {
		return nil, err
	}

	candidates := s.buildDataPlaneCandidates(session)
	preferred := sessioncontracts.DataPlaneCandidateBackendGateway
	if len(candidates) > 0 {
		preferred = candidates[0].Type
	}

	return &sessioncontracts.DataPlaneOffer{
		Preferred:  preferred,
		Token:      token,
		ExpiresAt:  expiresAt,
		Candidates: candidates,
	}, nil
}

func (s *Service) buildDataPlaneCandidates(session RemoteSession) []sessioncontracts.DataPlaneCandidate {
	var candidates []sessioncontracts.DataPlaneCandidate
	if directURL := s.directWorkerWSSURL(session.WorkerID); directURL != "" && s.canAdvertiseDirectWorkerWSS() {
		metadata := map[string]any(nil)
		if s.cfg.DataPlane.DirectWorkerJSONRuntime {
			metadata = map[string]any{
				"runtime_transport": "json_v1",
				"traffic_ready":     true,
			}
			s.addDirectWorkerTLSTrustMetadata(metadata)
			if s.cfg.DataPlane.DirectWorkerBinaryRender {
				metadata["render_transport"] = "binary_v1"
				metadata["binary_render"] = true
				metadata["supported_color_modes"] = []string{"full_color", "grayscale"}
				metadata["default_color_mode"] = "full_color"
			}
		}
		candidates = append(candidates, sessioncontracts.DataPlaneCandidate{
			Type:     sessioncontracts.DataPlaneCandidateDirectWorkerWSS,
			URL:      directURL,
			WorkerID: session.WorkerID,
			Priority: 10,
			Metadata: metadata,
		})
	}
	if s.cfg.DataPlane.BackendGatewayURL != "" {
		candidates = append(candidates, sessioncontracts.DataPlaneCandidate{
			Type:     sessioncontracts.DataPlaneCandidateBackendGateway,
			URL:      s.cfg.DataPlane.BackendGatewayURL,
			Priority: 100,
		})
	}
	return candidates
}

func (s *Service) canAdvertiseDirectWorkerWSS() bool {
	trustMode := normalizeDirectWorkerTLSTrustMode(s.cfg.DataPlane.DirectWorkerTLSTrustMode)
	return !secrets.IsProductionEnv(s.cfg.App.Env) || directWorkerTLSTrustModeIsProductionTrusted(trustMode)
}

func (s *Service) addDirectWorkerTLSTrustMetadata(metadata map[string]any) {
	trustMode := normalizeDirectWorkerTLSTrustMode(s.cfg.DataPlane.DirectWorkerTLSTrustMode)
	metadata["tls_trust_mode"] = trustMode
	metadata["production_trusted"] = directWorkerTLSTrustModeIsProductionTrusted(trustMode)
	metadata["smoke_only"] = trustMode == directWorkerTLSTrustModeSmokeInsecure
	if s.cfg.DataPlane.DirectWorkerTLSCARef != "" {
		metadata["tls_ca_ref"] = s.cfg.DataPlane.DirectWorkerTLSCARef
	}
}

func normalizeDirectWorkerTLSTrustMode(mode string) string {
	switch strings.ToLower(strings.TrimSpace(mode)) {
	case directWorkerTLSTrustModePublicCA:
		return directWorkerTLSTrustModePublicCA
	case directWorkerTLSTrustModePlatformCA:
		return directWorkerTLSTrustModePlatformCA
	default:
		return directWorkerTLSTrustModeSmokeInsecure
	}
}

func directWorkerTLSTrustModeIsProductionTrusted(mode string) bool {
	return mode == directWorkerTLSTrustModePublicCA || mode == directWorkerTLSTrustModePlatformCA
}

func (s *Service) directWorkerWSSURL(workerID string) string {
	template := strings.TrimSpace(s.cfg.DataPlane.DirectWorkerWSSURLTemplate)
	if template == "" || workerID == "" {
		return ""
	}
	return strings.ReplaceAll(template, "{worker_id}", workerID)
}

func signDataPlaneToken(claims DataPlaneTokenClaims, privateKeyPEM string) (string, error) {
	privateKey, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(privateKeyPEM))
	if err != nil {
		return "", fmt.Errorf("parse data-plane private key: %w", err)
	}
	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	signed, err := token.SignedString(privateKey)
	if err != nil {
		return "", fmt.Errorf("sign data-plane token: %w", err)
	}
	return signed, nil
}

func parseDataPlaneToken(tokenValue string, publicKey *rsa.PublicKey) (*DataPlaneTokenClaims, error) {
	claims := &DataPlaneTokenClaims{}
	token, err := jwt.ParseWithClaims(tokenValue, claims, func(token *jwt.Token) (any, error) {
		if token.Method != jwt.SigningMethodRS256 {
			return nil, fmt.Errorf("unexpected data-plane signing method: %s", token.Header["alg"])
		}
		return publicKey, nil
	})
	if err != nil {
		return nil, err
	}
	if !token.Valid {
		return nil, fmt.Errorf("data-plane token invalid")
	}
	return claims, nil
}

func dataPlaneAllowedChannelsFromSession(session RemoteSession) []string {
	channels := []string{
		sessioncontracts.DataPlaneChannelControl,
		sessioncontracts.DataPlaneChannelInput,
		sessioncontracts.DataPlaneChannelRender,
		sessioncontracts.DataPlaneChannelTelemetry,
	}
	metadata := decodeJSONMap(session.Metadata)
	policy, _ := metadata["policy"].(map[string]any)
	if policy != nil {
		if mode, _ := policy["clipboard_mode"].(string); mode != "" && mode != string(ResourceClipboardModeDisabled) {
			channels = append(channels, sessioncontracts.DataPlaneChannelClipboard)
		}
		if mode, _ := policy["file_transfer_mode"].(string); fileTransferAllowsClientToServer(ResourceFileTransferMode(mode)) {
			channels = append(channels, sessioncontracts.DataPlaneChannelFileUpload)
		}
		if mode, _ := policy["file_transfer_mode"].(string); fileTransferAllowsServerToClient(ResourceFileTransferMode(mode)) {
			channels = append(channels, sessioncontracts.DataPlaneChannelFileDownload)
		}
	}
	return channels
}

func (s *Service) attachDataPlaneOffer(result *SessionControlResult) error {
	if result == nil || result.Attachment == nil {
		return nil
	}
	result.GatewayURL = s.cfg.DataPlane.BackendGatewayURL
	offer, err := s.buildDataPlaneOffer(result.Session, *result.Attachment)
	if err != nil {
		return err
	}
	result.DataPlane = offer
	return nil
}