feat: add validation for ng/cz input (#206)

Signed-off-by: Amber Xue <ambermingxin@nvidia.com>

Author: ambermingxin

cmd/fleetint/enroll.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"regexp"
 	"strings"
 	"time"
 
@@ -35,6 +36,9 @@ var (
 )
 
 const defaultEnrollTimeout = time.Minute
+const reservedUnassignedName = "Unassigned"
+
+var enrollMetadataNamePattern = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9 ._-]*$`)
 
 // resolveToken returns the SAK token from --token, --token-file, or stdin.
 func resolveToken(cliContext *cli.Context) (string, error) {
@@ -84,9 +88,20 @@ func resolveToken(cliContext *cli.Context) (string, error) {
 func enrollCommand(cliContext *cli.Context) error {
 	baseEndpoint := cliContext.String("endpoint")
 	force := cliContext.Bool("force")
+	nodeGroup, err := validatedOptionalMetadataFlagValue(cliContext, "node-group", "Node group")
+	if err != nil {
+		return err
+	}
+	computeZone, err := validatedOptionalMetadataFlagValue(cliContext, "compute-zone", "Compute zone")
+	if err != nil {
+		return err
+	}
+	if err := validateReservedPairMetadata(nodeGroup, computeZone); err != nil {
+		return err
+	}
 	metadata := &enrollment.EnrollMetadata{
-		NodeGroup:   optionalFlagValue(cliContext, "node-group"),
-		ComputeZone: optionalFlagValue(cliContext, "compute-zone"),
+		NodeGroup:   nodeGroup,
+		ComputeZone: computeZone,
 	}
 
 	sakToken, err := resolveToken(cliContext)
@@ -124,10 +139,40 @@ func enrollCommand(cliContext *cli.Context) error {
 	return performEnrollWorkflow(ctx, baseEndpoint, sakToken, cfg, metadata)
 }
 
-func optionalFlagValue(cliContext *cli.Context, name string) *string {
+func validatedOptionalMetadataFlagValue(cliContext *cli.Context, name, fieldName string) (*string, error) {
 	if !cliContext.IsSet(name) {
-		return nil
+		return nil, nil
 	}
 	value := strings.TrimSpace(cliContext.String(name))
-	return &value
+	// Explicit empty means clear/unassign for reserved assignment fields.
+	if value == "" {
+		return &value, nil
+	}
+	if strings.EqualFold(value, reservedUnassignedName) {
+		return nil, fmt.Errorf("%s name %q is reserved; use empty value to clear assignment", fieldName, reservedUnassignedName)
+	}
+	if len(value) > 255 {
+		return nil, fmt.Errorf("%s name must be 255 characters or fewer", fieldName)
+	}
+	if !enrollMetadataNamePattern.MatchString(value) {
+		return nil, fmt.Errorf("%s name must start with a letter and contain only letters, numbers, spaces, hyphens, underscores, or periods", fieldName)
+	}
+	return &value, nil
+}
+
+func validateReservedPairMetadata(nodeGroup, computeZone *string) error {
+	nodeGroupSet := nodeGroup != nil
+	computeZoneSet := computeZone != nil
+
+	if !nodeGroupSet && !computeZoneSet {
+		return nil
+	}
+	if nodeGroupSet && computeZoneSet {
+		nodeGroupEmpty := *nodeGroup == ""
+		computeZoneEmpty := *computeZone == ""
+		if nodeGroupEmpty == computeZoneEmpty {
+			return nil
+		}
+	}
+	return fmt.Errorf("--node-group and --compute-zone must be both omitted, both empty, or both non-empty")
 }

cmd/fleetint/enroll_test.go

@@ -18,14 +18,17 @@ package main
 import (
 	"bytes"
 	"context"
+	"flag"
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/urfave/cli"
 
 	"github.com/NVIDIA/fleet-intelligence-agent/internal/config"
 	"github.com/NVIDIA/fleet-intelligence-agent/internal/enrollment"
@@ -288,42 +291,108 @@ func TestEnrollCommandPassesOptionalMetadata(t *testing.T) {
 	require.NoError(t, err)
 }
 
-func TestEnrollCommandTreatsExplicitEmptyMetadataAsClear(t *testing.T) {
-	useMissingFleetintEnvFile(t)
+func TestValidatedOptionalMetadataFlagValueAllowsExplicitEmpty(t *testing.T) {
+	flagSet := flag.NewFlagSet("enroll", flag.ContinueOnError)
+	flagSet.String("node-group", "", "")
+	require.NoError(t, flagSet.Set("node-group", ""))
+	cliContext := cli.NewContext(cli.NewApp(), flagSet, nil)
 
-	originalRunPrecheck := runPrecheck
-	originalEnrollWorkflow := performEnrollWorkflow
-	t.Cleanup(func() {
-		runPrecheck = originalRunPrecheck
-		performEnrollWorkflow = originalEnrollWorkflow
+	value, err := validatedOptionalMetadataFlagValue(cliContext, "node-group", "Node group")
+	require.NoError(t, err)
+	require.NotNil(t, value)
+	require.Empty(t, *value)
+}
+
+func TestValidatedOptionalMetadataFlagValueRejectsReservedUnassignedName(t *testing.T) {
+	flagSet := flag.NewFlagSet("enroll", flag.ContinueOnError)
+	flagSet.String("node-group", "", "")
+	require.NoError(t, flagSet.Set("node-group", "unassigned"))
+	cliContext := cli.NewContext(cli.NewApp(), flagSet, nil)
+
+	_, err := validatedOptionalMetadataFlagValue(cliContext, "node-group", "Node group")
+	require.ErrorContains(t, err, `Node group name "Unassigned" is reserved; use empty value to clear assignment`)
+}
+
+func TestEnrollCommandRejectsInvalidMetadataCharacters(t *testing.T) {
+	app := App()
+	app.Writer = &bytes.Buffer{}
+
+	err := app.Run([]string{
+		"fleetint", "enroll",
+		"--endpoint", "https://example.com",
+		"--token", "token",
+		"--compute-zone", "@bad-zone",
 	})
+	require.ErrorContains(t, err, "Compute zone name must start with a letter")
+}
 
-	runPrecheck = func() (precheck.Result, error) {
-		return precheck.Result{
-			Checks: []precheck.Check{
-				{Name: "gpu-present", Message: "ok", Passed: true},
-			},
-		}, nil
+func TestEnrollCommandRejectsOverlongMetadataNames(t *testing.T) {
+	app := App()
+	app.Writer = &bytes.Buffer{}
+	longName := strings.Repeat("a", 256)
+
+	err := app.Run([]string{
+		"fleetint", "enroll",
+		"--endpoint", "https://example.com",
+		"--token", "token",
+		"--node-group", longName,
+	})
+	require.ErrorContains(t, err, "Node group name must be 255 characters or fewer")
+}
+
+func TestValidateReservedPairMetadata(t *testing.T) {
+	strPtr := func(v string) *string { return &v }
+	tests := []struct {
+		name        string
+		nodeGroup   *string
+		computeZone *string
+		wantErr     bool
+	}{
+		{name: "both omitted", nodeGroup: nil, computeZone: nil, wantErr: false},
+		{name: "both empty", nodeGroup: strPtr(""), computeZone: strPtr(""), wantErr: false},
+		{name: "both non-empty", nodeGroup: strPtr("ng-a"), computeZone: strPtr("cz-a"), wantErr: false},
+		{name: "node-group only", nodeGroup: strPtr("ng-a"), computeZone: nil, wantErr: true},
+		{name: "compute-zone only", nodeGroup: nil, computeZone: strPtr("cz-a"), wantErr: true},
+		{name: "node-group empty compute-zone non-empty", nodeGroup: strPtr(""), computeZone: strPtr("cz-a"), wantErr: true},
+		{name: "node-group non-empty compute-zone empty", nodeGroup: strPtr("ng-a"), computeZone: strPtr(""), wantErr: true},
 	}
 
-	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config, metadata *enrollment.EnrollMetadata) error {
-		require.NotNil(t, metadata)
-		require.NotNil(t, metadata.NodeGroup)
-		require.Empty(t, *metadata.NodeGroup)
-		require.NotNil(t, metadata.ComputeZone)
-		require.Empty(t, *metadata.ComputeZone)
-		return nil
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			err := validateReservedPairMetadata(tc.nodeGroup, tc.computeZone)
+			if tc.wantErr {
+				require.ErrorContains(t, err, "--node-group and --compute-zone must be both omitted, both empty, or both non-empty")
+				return
+			}
+			require.NoError(t, err)
+		})
 	}
+}
 
+func TestEnrollCommandRejectsMixedReservedPairValues(t *testing.T) {
 	app := App()
 	app.Writer = &bytes.Buffer{}
 
 	err := app.Run([]string{
 		"fleetint", "enroll",
 		"--endpoint", "https://example.com",
 		"--token", "token",
-		"--node-group=",
-		"--compute-zone=",
+		"--node-group", "",
+		"--compute-zone", "cz-a",
 	})
-	require.NoError(t, err)
+	require.ErrorContains(t, err, "--node-group and --compute-zone must be both omitted, both empty, or both non-empty")
+}
+
+func TestEnrollCommandRejectsReservedUnassignedName(t *testing.T) {
+	app := App()
+	app.Writer = &bytes.Buffer{}
+
+	err := app.Run([]string{
+		"fleetint", "enroll",
+		"--endpoint", "https://example.com",
+		"--token", "token",
+		"--node-group", "Unassigned",
+		"--compute-zone", "cz-a",
+	})
+	require.ErrorContains(t, err, `Node group name "Unassigned" is reserved; use empty value to clear assignment`)
 }

third_party/fleet-intelligence-sdk/pkg/metadata/metadata.go

@@ -21,6 +21,9 @@ const (
 	columnKey             = "key"
 	columnValue           = "value"
 )
+const setMetadataUpsertQuery = `
+INSERT INTO gpud_metadata (key, value) VALUES (?, ?)
+ON CONFLICT(key) DO UPDATE SET value = excluded.value`
 
 // CreateTableMetadata creates the table for the metadata.
 func CreateTableMetadata(ctx context.Context, dbRW *sql.DB) error {
@@ -51,28 +54,10 @@ const (
 
 // SetMetadata sets the value of a metadata entry.
 // If the metadata entry is not found, it is created.
-// If the metadata entry is found and the value is the same, it is not updated.
-// If the metadata entry is found and the value is different, it is updated.
+// If the metadata entry is found, it is updated in-place.
 func SetMetadata(ctx context.Context, dbRW *sql.DB, key string, value string) error {
-	prev, err := ReadMetadata(ctx, dbRW, key)
-	if err != nil {
-		return err
-	}
-
-	if prev == value {
-		return nil
-	}
-
 	start := time.Now()
-	if prev == "" {
-		// the "name" is not in the table, so we need to insert it
-		_, err = dbRW.ExecContext(ctx, fmt.Sprintf(`
-INSERT INTO %s (%s, %s) VALUES (?, ?)`, tableNameGPUdMetadata, columnKey, columnValue), key, value)
-	} else {
-		// the "name" is already in the table, so we need to update it
-		_, err = dbRW.ExecContext(ctx, fmt.Sprintf(`
-UPDATE %s SET %s = ? WHERE %s = ?`, tableNameGPUdMetadata, columnValue, columnKey), value, key)
-	}
+	_, err := dbRW.ExecContext(ctx, setMetadataUpsertQuery, key, value)
 	pkgmetricsrecorder.RecordSQLiteInsertUpdate(time.Since(start).Seconds())
 
 	return err

third_party/fleet-intelligence-sdk/pkg/metadata/metadata_test.go

@@ -126,3 +126,23 @@ func TestSetAndReadMetadata(t *testing.T) {
 	_, err = ReadMetadata(canceledCtx, dbRO, testKey)
 	assert.Error(t, err)
 }
+
+func TestSetMetadata_AllowsTransitionFromEmptyToNonEmpty(t *testing.T) {
+	t.Parallel()
+	dbRW, dbRO, cleanup := sqlite.OpenTestDB(t)
+	defer cleanup()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	require.NoError(t, CreateTableMetadata(ctx, dbRW))
+
+	const key = "nodegroup"
+	require.NoError(t, SetMetadata(ctx, dbRW, key, "group-a"))
+	require.NoError(t, SetMetadata(ctx, dbRW, key, ""))
+	require.NoError(t, SetMetadata(ctx, dbRW, key, "group-b"))
+
+	value, err := ReadMetadata(ctx, dbRO, key)
+	require.NoError(t, err)
+	assert.Equal(t, "group-b", value)
+}