From 262b4a9b127b3a64e9860f47cfba3735769193cf Mon Sep 17 00:00:00 2001
From: Kristine Kunapuli <49738148+kkunapuli@users.noreply.github.com>
Date: Thu, 3 Oct 2024 11:11:46 -0400
Subject: [PATCH] chore: check task config policies against slots and max_slots
 (#10015)

---
 master/internal/api_command.go                 |  2 +-
 .../configpolicy/task_config_policy.go         | 17 +++++++++++------
 .../task_config_policy_intg_test.go            | 18 ++++++++++++++++++
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/master/internal/api_command.go b/master/internal/api_command.go
index 30e40d178bc..3af94d5227d 100644
--- a/master/internal/api_command.go
+++ b/master/internal/api_command.go
@@ -153,7 +153,7 @@ func (a *apiServer) getCommandLaunchParams(ctx context.Context, req *protoComman
 	// Check submitted config against task config policies.
 	valid, err := configpolicy.CheckNTSCConstraints(ctx, int(cmdSpec.Metadata.WorkspaceID), config, a.m.rm)
 	if !valid {
-		return nil, nil, err
+		return nil, nil, status.Errorf(codes.InvalidArgument, "failed constraint check: %v", err)
 	}
 
 	token, err := getTaskSessionToken(ctx, userModel)
diff --git a/master/internal/configpolicy/task_config_policy.go b/master/internal/configpolicy/task_config_policy.go
index 667b0a7c770..4eb62aed894 100644
--- a/master/internal/configpolicy/task_config_policy.go
+++ b/master/internal/configpolicy/task_config_policy.go
@@ -55,15 +55,20 @@ func CheckNTSCConstraints(
 	if err == nil && constraints.PriorityLimit != nil && workloadConfig.Resources.Priority != nil {
 		if !priorityWithinLimit(*workloadConfig.Resources.Priority, *constraints.PriorityLimit, smallerHigher) {
 			return false, fmt.Errorf("requested priority [%d] exceeds limit set by admin [%d]: %w",
-				*constraints.PriorityLimit, *workloadConfig.Resources.Priority, errPriorityConstraintFailure)
+				*workloadConfig.Resources.Priority, *constraints.PriorityLimit, errPriorityConstraintFailure)
 		}
 	}
 
-	if constraints.ResourceConstraints != nil && constraints.ResourceConstraints.MaxSlots != nil &&
-		workloadConfig.Resources.MaxSlots != nil {
-		if *constraints.ResourceConstraints.MaxSlots < *workloadConfig.Resources.MaxSlots {
-			return false, fmt.Errorf("requested resources.max_slots [%d] exceeds limit set by admin [%d]: %w",
-				*constraints.ResourceConstraints.MaxSlots, *workloadConfig.Resources.MaxSlots, errResourceConstraintFailure)
+	if constraints.ResourceConstraints != nil && constraints.ResourceConstraints.MaxSlots != nil {
+		if workloadConfig.Resources.MaxSlots != nil {
+			if *constraints.ResourceConstraints.MaxSlots < *workloadConfig.Resources.MaxSlots {
+				return false, fmt.Errorf("requested resources.max_slots [%d] exceeds limit set by admin [%d]: %w",
+					*workloadConfig.Resources.MaxSlots, *constraints.ResourceConstraints.MaxSlots, errResourceConstraintFailure)
+			}
+		}
+		if *constraints.ResourceConstraints.MaxSlots < workloadConfig.Resources.Slots {
+			return false, fmt.Errorf("requested resources.slots [%d] exceeds limit set by admin [%d]: %w",
+				workloadConfig.Resources.Slots, *constraints.ResourceConstraints.MaxSlots, errResourceConstraintFailure)
 		}
 	}
 
diff --git a/master/internal/configpolicy/task_config_policy_intg_test.go b/master/internal/configpolicy/task_config_policy_intg_test.go
index b6e520b797f..cfe8bb8db4f 100644
--- a/master/internal/configpolicy/task_config_policy_intg_test.go
+++ b/master/internal/configpolicy/task_config_policy_intg_test.go
@@ -111,6 +111,24 @@ func TestValidateNTSCConstraints(t *testing.T) {
 		require.ErrorIs(t, err, errResourceConstraintFailure)
 	})
 
+	t.Run("exceeds slots - not ok", func(t *testing.T) {
+		constraints := DefaultConstraints()
+		w := model.Workspace{Name: uuid.NewString(), UserID: user.ID}
+		_, err := db.Bun().NewInsert().Model(&w).Exec(context.Background())
+		require.NoError(t, err)
+		addConstraints(t, user, &w.ID, *constraints)
+
+		resourceManager := mocks.ResourceManager{}
+		resourceManager.On("SmallerValueIsHigherPriority", mock.Anything).Return(true, nil)
+
+		config := defaultConfig()
+		config.Resources.Slots = *config.Resources.MaxSlots
+		config.Resources.MaxSlots = nil // ensure only slots is set
+		_, err = CheckNTSCConstraints(context.Background(), w.ID, config, &resourceManager)
+		require.Error(t, err)
+		require.ErrorIs(t, err, errResourceConstraintFailure)
+	})
+
 	t.Run("rm priority not supported - ok", func(t *testing.T) {
 		w := addWorkspacePriorityLimit(t, user, wkspPriorityLimit)
 		rm1 := mocks.ResourceManager{}