Add EC2 joining for Windows Desktop Service (#16439)

lib/auth/join_ec2.go

@@ -253,6 +253,20 @@ func dbExists(ctx context.Context, presence services.Presence, hostID string) (b
 	return false, nil
 }
 
+func desktopServiceExists(ctx context.Context, presence services.Presence, hostID string) (bool, error) {
+	svcs, err := presence.GetWindowsDesktopServices(ctx)
+	if err != nil {
+		return false, trace.Wrap(err)
+	}
+
+	for _, wds := range svcs {
+		if wds.GetName() == hostID {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
 // checkInstanceUnique makes sure the instance which sent the request has not
 // already joined the cluster with the same role. Tokens should be limited to
 // only allow the roles which will actually be used by all expected instances so
@@ -279,6 +293,8 @@ func (a *Server) checkInstanceUnique(ctx context.Context, req *types.RegisterUsi
 		instanceExists, err = appExists(ctx, a, req.HostID)
 	case types.RoleDatabase:
 		instanceExists, err = dbExists(ctx, a, req.HostID)
+	case types.RoleWindowsDesktop:
+		instanceExists, err = desktopServiceExists(ctx, a, req.HostID)
 	default:
 		return trace.BadParameter("unsupported role: %q", req.Role)
 	}
@@ -296,11 +312,12 @@ func (a *Server) checkInstanceUnique(ctx context.Context, req *types.RegisterUsi
 
 // checkEC2JoinRequest checks register requests which use EC2 Simplified Node
 // Joining. This method checks that:
-// 1. The given Instance Identity Document has a valid signature (signed by AWS).
-// 2. A node has not already joined the cluster from this EC2 instance (to
-//    prevent re-use of a stolen Instance Identity Document).
-// 3. The signed instance attributes match one of the allow rules for the
-//    corresponding token.
+//  1. The given Instance Identity Document has a valid signature (signed by AWS).
+//  2. A node has not already joined the cluster from this EC2 instance (to
+//     prevent re-use of a stolen Instance Identity Document).
+//  3. The signed instance attributes match one of the allow rules for the
+//     corresponding token.
+//
 // If the request does not include an Instance Identity Document, and the
 // token does not include any allow rules, this method returns nil and the
 // normal token checking logic resumes.

lib/auth/join_ec2_test.go

@@ -592,9 +592,10 @@ func TestHostUniqueCheck(t *testing.T) {
 				types.RoleKube,
 				types.RoleDatabase,
 				types.RoleApp,
+				types.RoleWindowsDesktop,
 			},
 			Allow: []*types.TokenRule{
-				&types.TokenRule{
+				{
 					AWSAccount: instance1.account,
 					AWSRegions: []string{instance1.region},
 				},
@@ -703,6 +704,20 @@ func TestHostUniqueCheck(t *testing.T) {
 				require.NoError(t, err)
 			},
 		},
+		{
+			role: types.RoleWindowsDesktop,
+			upserter: func(name string) {
+				wds, err := types.NewWindowsDesktopServiceV3(instance1.account+"-"+instance1.instanceID,
+					types.WindowsDesktopServiceSpecV3{
+						Addr:            "localhost:3028",
+						TeleportVersion: "10.2.2",
+					})
+				require.NoError(t, err)
+
+				_, err = a.UpsertWindowsDesktopService(context.Background(), wds)
+				require.NoError(t, err)
+			},
+		},
 	}
 
 	ctx = context.WithValue(ctx, ec2ClientKey{}, ec2ClientRunning{})