From 43397b85b98422b47bb499ac2d8d1b93f5b9a910 Mon Sep 17 00:00:00 2001
From: Przemko Robakowski <przemko.robakowski@goteleport.com>
Date: Wed, 13 Nov 2024 23:21:27 +0100
Subject: [PATCH] Respect windows_desktop_labels for dynamic desktops

---
 .../dynamicwindowsv1/service.go               | 47 ++++++++++++++++++-
 lib/services/role.go                          |  1 +
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/lib/auth/dynamicwindows/dynamicwindowsv1/service.go b/lib/auth/dynamicwindows/dynamicwindowsv1/service.go
index 5a42eefe8edca..c2899af0712ff 100644
--- a/lib/auth/dynamicwindows/dynamicwindowsv1/service.go
+++ b/lib/auth/dynamicwindows/dynamicwindowsv1/service.go
@@ -20,6 +20,7 @@ package dynamicwindowsv1
 
 import (
 	"context"
+	"github.com/gravitational/teleport/lib/services"
 	"log/slog"
 
 	"github.com/gravitational/trace"
@@ -108,6 +109,9 @@ func (s *Service) GetDynamicWindowsDesktop(ctx context.Context, request *dynamic
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
+	if err := checkAccess(auth, d); err != nil {
+		return nil, trace.Wrap(err)
+	}
 
 	desktop, ok := d.(*types.DynamicWindowsDesktopV1)
 	if !ok {
@@ -136,6 +140,9 @@ func (s *Service) ListDynamicWindowsDesktops(ctx context.Context, request *dynam
 		NextPageToken: next,
 	}
 	for _, d := range desktops {
+		if err := checkAccess(auth, d); err != nil {
+			continue
+		}
 		desktop, ok := d.(*types.DynamicWindowsDesktopV1)
 		if !ok {
 			return nil, trace.BadParameter("unexpected type %T", d)
@@ -158,6 +165,9 @@ func (s *Service) CreateDynamicWindowsDesktop(ctx context.Context, req *dynamicw
 	if err := auth.CheckAccessToKind(types.KindDynamicWindowsDesktop, types.VerbCreate); err != nil {
 		return nil, trace.Wrap(err)
 	}
+	if err := checkAccess(auth, req.GetDesktop()); err != nil {
+		return nil, trace.Wrap(err)
+	}
 	d, err := s.backend.CreateDynamicWindowsDesktop(ctx, types.DynamicWindowsDesktop(req.Desktop))
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -171,6 +181,10 @@ func (s *Service) CreateDynamicWindowsDesktop(ctx context.Context, req *dynamicw
 	return createdDesktop, nil
 }
 
+func checkAccess(auth *authz.Context, desktop types.DynamicWindowsDesktop) error {
+	return auth.Checker.CheckAccess(desktop, services.AccessState{MFAVerified: true})
+}
+
 // UpdateDynamicWindowsDesktop updates an existing dynamic Windows desktop.
 func (s *Service) UpdateDynamicWindowsDesktop(ctx context.Context, req *dynamicwindowspb.UpdateDynamicWindowsDesktopRequest) (*types.DynamicWindowsDesktopV1, error) {
 	auth, err := s.authorizer.Authorize(ctx)
@@ -183,7 +197,17 @@ func (s *Service) UpdateDynamicWindowsDesktop(ctx context.Context, req *dynamicw
 	if err := auth.CheckAccessToKind(types.KindDynamicWindowsDesktop, types.VerbUpdate); err != nil {
 		return nil, trace.Wrap(err)
 	}
-	d, err := s.backend.UpdateDynamicWindowsDesktop(ctx, req.Desktop)
+	d, err := s.cache.GetDynamicWindowsDesktop(ctx, req.GetDesktop().GetName())
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	if err := checkAccess(auth, d); err != nil {
+		return nil, trace.Wrap(err)
+	}
+	if err := checkAccess(auth, req.GetDesktop()); err != nil {
+		return nil, trace.Wrap(err)
+	}
+	d, err = s.backend.UpdateDynamicWindowsDesktop(ctx, req.Desktop)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -208,7 +232,19 @@ func (s *Service) UpsertDynamicWindowsDesktop(ctx context.Context, req *dynamicw
 	if err := auth.CheckAccessToKind(types.KindDynamicWindowsDesktop, types.VerbCreate, types.VerbUpdate); err != nil {
 		return nil, trace.Wrap(err)
 	}
-	d, err := s.backend.UpsertDynamicWindowsDesktop(ctx, req.Desktop)
+	d, err := s.cache.GetDynamicWindowsDesktop(ctx, req.GetDesktop().GetName())
+	if !trace.IsNotFound(err) {
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		if err := checkAccess(auth, d); err != nil {
+			return nil, trace.Wrap(err)
+		}
+	}
+	if err := checkAccess(auth, req.GetDesktop()); err != nil {
+		return nil, trace.Wrap(err)
+	}
+	d, err = s.backend.UpsertDynamicWindowsDesktop(ctx, req.Desktop)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -233,6 +269,13 @@ func (s *Service) DeleteDynamicWindowsDesktop(ctx context.Context, req *dynamicw
 	if err := auth.CheckAccessToKind(types.KindDynamicWindowsDesktop, types.VerbDelete); err != nil {
 		return nil, trace.Wrap(err)
 	}
+	d, err := s.cache.GetDynamicWindowsDesktop(ctx, req.GetName())
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	if err := checkAccess(auth, d); err != nil {
+		return nil, trace.Wrap(err)
+	}
 	if err := s.backend.DeleteDynamicWindowsDesktop(ctx, req.GetName()); err != nil {
 		return nil, trace.Wrap(err)
 	}
diff --git a/lib/services/role.go b/lib/services/role.go
index 16b1c79287e87..60891a725e3bd 100644
--- a/lib/services/role.go
+++ b/lib/services/role.go
@@ -349,6 +349,7 @@ func validateRoleExpressions(r types.Role) error {
 			{"db_labels", types.KindDatabase},
 			{"db_service_labels", types.KindDatabaseService},
 			{"windows_desktop_labels", types.KindWindowsDesktop},
+			{"windows_desktop_labels", types.KindDynamicWindowsDesktop},
 			{"group_labels", types.KindUserGroup},
 		} {
 			labelMatchers, err := r.GetLabelMatchers(condition.condition, labels.kind)