fix(validator): roll back cluster-admin RBAC on prep failure (#2128)

Signed-off-by: Nathan Hensley <nhensley@nvidia.com>
Co-authored-by: Mark Chmarny <mchmarny@users.noreply.github.com>

Author: njhensley

pkg/validator/job/rbac.go

@@ -110,9 +110,15 @@ func EnsureRBAC(ctx context.Context, clientset kubernetes.Interface, namespace,
 	return nil
 }
 
-// CleanupRBAC removes the per-run ServiceAccount and ClusterRoleBinding.
+// CleanupRBAC removes the per-run ClusterRoleBinding and ServiceAccount.
 // Ignores NotFound errors (idempotent). Call once at end of validation run.
 //
+// The ClusterRoleBinding is deleted BEFORE the ServiceAccount: revoking the
+// binding immediately closes the cluster-admin escalation window, whereas an
+// orphaned ServiceAccount with no binding carries no privileges. If the binding
+// delete fails, the ServiceAccount delete still runs so we do not leave both
+// behind.
+//
 // When both deletes fail, the returned StructuredError wraps the joined
 // underlying errors via stderrors.Join so callers can inspect individual
 // failures with errors.Is / errors.As.
@@ -122,15 +128,15 @@ func CleanupRBAC(ctx context.Context, clientset kubernetes.Interface, namespace,
 
 	var errs []error
 
-	if err := clientset.CoreV1().ServiceAccounts(namespace).Delete(ctx, saName, metav1.DeleteOptions{}); err != nil {
+	if err := clientset.RbacV1().ClusterRoleBindings().Delete(ctx, crbName, metav1.DeleteOptions{}); err != nil {
 		if !apierrors.IsNotFound(err) {
-			errs = append(errs, errors.Wrap(errors.ErrCodeInternal, "failed to delete ServiceAccount", err))
+			errs = append(errs, errors.Wrap(errors.ErrCodeInternal, "failed to delete ClusterRoleBinding", err))
 		}
 	}
 
-	if err := clientset.RbacV1().ClusterRoleBindings().Delete(ctx, crbName, metav1.DeleteOptions{}); err != nil {
+	if err := clientset.CoreV1().ServiceAccounts(namespace).Delete(ctx, saName, metav1.DeleteOptions{}); err != nil {
 		if !apierrors.IsNotFound(err) {
-			errs = append(errs, errors.Wrap(errors.ErrCodeInternal, "failed to delete ClusterRoleBinding", err))
+			errs = append(errs, errors.Wrap(errors.ErrCodeInternal, "failed to delete ServiceAccount", err))
 		}
 	}
 

pkg/validator/job/rbac_test.go

@@ -16,9 +16,16 @@ package job
 
 import (
 	"context"
+	stderrors "errors"
 	"testing"
 
+	"github.com/NVIDIA/aicr/pkg/errors"
+	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	k8sfake "k8s.io/client-go/kubernetes/fake"
+	clienttesting "k8s.io/client-go/testing"
 )
 
 func TestEnsureRBAC(t *testing.T) {
@@ -168,3 +175,71 @@ func TestCleanupRBACNotFound(t *testing.T) {
 		t.Fatalf("CleanupRBAC() on nonexistent resources should not error, got: %v", err)
 	}
 }
+
+// TestCleanupRBACDeletesClusterRoleBindingBeforeServiceAccount guards the
+// ordering fix: the cluster-admin ClusterRoleBinding must be revoked BEFORE the
+// ServiceAccount is deleted, so the privilege-escalation window closes first. A
+// reactor records delete order without handling the delete (returns handled =
+// false to fall through to the tracker). Reverting the reorder in CleanupRBAC
+// flips the recorded order and fails this test.
+func TestCleanupRBACDeletesClusterRoleBindingBeforeServiceAccount(t *testing.T) {
+	cs := k8sfake.NewSimpleClientset()
+
+	var order []string
+	cs.PrependReactor("delete", "*", func(action clienttesting.Action) (bool, runtime.Object, error) {
+		order = append(order, action.GetResource().Resource)
+		return false, nil, nil // fall through to the default tracker
+	})
+
+	if err := CleanupRBAC(context.Background(), cs, "ns", "run"); err != nil {
+		t.Fatalf("CleanupRBAC() error = %v, want nil", err)
+	}
+
+	want := []string{"clusterrolebindings", "serviceaccounts"}
+	if len(order) != len(want) {
+		t.Fatalf("delete order = %v, want %v", order, want)
+	}
+	for i := range want {
+		if order[i] != want[i] {
+			t.Fatalf("delete order = %v, want %v", order, want)
+		}
+	}
+}
+
+// TestCleanupRBACClusterRoleBindingDeleteFailureSurfaced proves CleanupRBAC
+// fails closed AND does not short-circuit: a ClusterRoleBinding delete error
+// (anything other than NotFound) is surfaced as an ErrCodeInternal error rather
+// than swallowed, so a leaked cluster-admin binding cannot masquerade as a
+// clean teardown — and cleanup still proceeds to delete the ServiceAccount so a
+// binding failure does not strand the SA behind an early return.
+func TestCleanupRBACClusterRoleBindingDeleteFailureSurfaced(t *testing.T) {
+	const ns, runID = "ns", "run"
+	saName := ServiceAccountName(runID)
+
+	// Seed the ServiceAccount so its post-cleanup absence proves the delete ran.
+	cs := k8sfake.NewSimpleClientset(&corev1.ServiceAccount{
+		ObjectMeta: metav1.ObjectMeta{Name: saName, Namespace: ns},
+	})
+
+	wantCause := stderrors.New("apiserver unavailable")
+	cs.PrependReactor("delete", "clusterrolebindings", func(clienttesting.Action) (bool, runtime.Object, error) {
+		return true, nil, wantCause
+	})
+
+	err := CleanupRBAC(context.Background(), cs, ns, runID)
+	if err == nil {
+		t.Fatal("CleanupRBAC() error = nil, want error when ClusterRoleBinding delete fails")
+	}
+	if !stderrors.Is(err, errors.New(errors.ErrCodeInternal, "")) {
+		t.Errorf("CleanupRBAC() error = %v, want ErrCodeInternal", err)
+	}
+	if !stderrors.Is(err, wantCause) {
+		t.Errorf("CleanupRBAC() error = %v, want wrapped underlying cause", err)
+	}
+
+	// Despite the binding-delete failure, the ServiceAccount delete must still
+	// have run — proving CleanupRBAC does not return early on the first error.
+	if _, getErr := cs.CoreV1().ServiceAccounts(ns).Get(context.Background(), saName, metav1.GetOptions{}); !apierrors.IsNotFound(getErr) {
+		t.Errorf("ServiceAccount %q still present after cleanup (Get err = %v); cleanup returned early on binding failure", saName, getErr)
+	}
+}

pkg/validator/validator.go

@@ -133,13 +133,12 @@ func (v *Validator) prepareCluster(
 	ctx context.Context,
 	validationInput *v1.ValidationInput,
 	snap *snapshotter.Snapshot,
-) (*clusterState, error) {
+) (cs *clusterState, err error) {
 
 	// Use PropagateOrWrap so a coded inner error (e.g. an invalid kubeconfig
 	// classified as a deterministic config error) survives instead of being
 	// blanket-relabeled ErrCodeInternal, which would mask it as retryable.
 	var clientset kubernetes.Interface
-	var err error
 	kubeconfig := strings.TrimSpace(v.Kubeconfig)
 	switch {
 	case kubeconfig == "":
@@ -166,6 +165,32 @@ func (v *Validator) prepareCluster(
 		return nil, errors.PropagateOrWrap(rbacErr, errors.ErrCodeInternal, "failed to ensure RBAC")
 	}
 
+	// Privileged RBAC (the per-run cluster-admin ClusterRoleBinding) now exists.
+	// Register an immediate rollback so any later failure in prepareCluster
+	// revokes it before returning, instead of leaking a privileged identity
+	// until manual cleanup. On the success path err is nil and the binding is
+	// retained — the caller's deferClusterCleanup owns success-path teardown, so
+	// this defer must not double-clean.
+	//
+	//nolint:contextcheck // rollbackRBAC uses a fresh context: parent may be canceled
+	defer func() {
+		if err != nil {
+			if rollbackErr := v.rollbackRBAC(clientset); rollbackErr != nil {
+				// The privileged binding could not be revoked after a
+				// preparation failure. Fold the rollback failure into the
+				// returned error so the operator sees BOTH the original cause
+				// and the leaked cluster-admin binding — the prep error alone
+				// would hide that manual cleanup is now required. Keep it a
+				// coded StructuredError wrapping the joined causes so callers
+				// can still match ErrCodeInternal and inspect either error.
+				err = errors.WrapWithContext(errors.ErrCodeInternal,
+					"preparation failed and RBAC rollback failed; cluster-admin binding may be orphaned",
+					stderrors.Join(err, rollbackErr),
+					map[string]any{"runID": v.RunID, "namespace": v.Namespace})
+			}
+		}
+	}()
+
 	if cmErr := v.ensureDataConfigMaps(ctx, clientset, snap, validationInput); cmErr != nil {
 		return nil, errors.PropagateOrWrap(cmErr, errors.ErrCodeInternal, "failed to create data ConfigMaps")
 	}
@@ -183,27 +208,62 @@ func (v *Validator) prepareCluster(
 	}, nil
 }
 
-// deferClusterCleanup registers deferred cleanup for RBAC and data ConfigMaps.
-// Both cleanup steps share a single deadline so a stalled apiserver cannot
-// extend total post-run blocking time to 2 * K8sCleanupTimeout. Cleanup
-// failures are surfaced at structured-log level so operators see when
-// resources may have been orphaned in the validator namespace.
-func (v *Validator) deferClusterCleanup(clientset kubernetes.Interface) {
+// deferClusterCleanup performs success-path teardown of RBAC and data
+// ConfigMaps. Both cleanup steps share a single deadline so a stalled apiserver
+// cannot extend total post-run blocking time to 2 * K8sCleanupTimeout.
+//
+// RBAC is privileged (a per-run cluster-admin ClusterRoleBinding), so a failure
+// to revoke it is returned to the caller and promoted into the run's error —
+// fail closed, never leak cluster-admin silently. ConfigMap cleanup is not
+// privileged, so its failure stays warning-only and does not fail the run.
+func (v *Validator) deferClusterCleanup(clientset kubernetes.Interface) error {
 	if !v.Cleanup {
-		return
+		return nil
 	}
 	//nolint:contextcheck // Fresh context: parent may be canceled during cleanup
 	cleanupCtx, cancel := context.WithTimeout(context.Background(), defaults.K8sCleanupTimeout)
 	defer cancel()
 
+	var rbacErr error
 	if cleanupErr := job.CleanupRBAC(cleanupCtx, clientset, v.Namespace, v.RunID); cleanupErr != nil {
-		slog.Warn("failed to cleanup RBAC; resources may be orphaned",
+		slog.Error("failed to cleanup RBAC; cluster-admin binding may be orphaned",
 			"runID", v.RunID, "namespace", v.Namespace, "error", cleanupErr)
+		rbacErr = errors.PropagateOrWrap(cleanupErr, errors.ErrCodeInternal, "failed to revoke privileged RBAC")
 	}
 	if cmErr := v.cleanupDataConfigMaps(cleanupCtx, clientset); cmErr != nil {
 		slog.Warn("failed to cleanup ConfigMaps; resources may be orphaned",
 			"runID", v.RunID, "namespace", v.Namespace, "error", cmErr)
 	}
+	return rbacErr
+}
+
+// rollbackRBAC revokes the per-run RBAC created earlier in prepareCluster when a
+// later preparation step fails. It uses a fresh bounded context because the
+// caller's ctx may already be canceled — the very condition that can trigger the
+// failure. Revoking the cluster-admin ClusterRoleBinding closes the privilege
+// escalation window immediately; the surrounding prepareCluster call still
+// returns its error, so the run fails closed regardless of this rollback's
+// outcome. Respects v.Cleanup for parity with the success-path teardown: a
+// caller that disabled cleanup has opted into managing teardown manually, so a
+// disabled-cleanup run performs no rollback and returns nil.
+//
+// Returns the CleanupRBAC error (still logged) so the caller can fold a failed
+// revocation into the run's error and surface that cluster-admin may be
+// orphaned; returns nil when cleanup is disabled or the revocation succeeds.
+func (v *Validator) rollbackRBAC(clientset kubernetes.Interface) error {
+	if !v.Cleanup {
+		return nil
+	}
+	//nolint:contextcheck // Fresh context: parent may be canceled during rollback
+	cleanupCtx, cancel := context.WithTimeout(context.Background(), defaults.K8sCleanupTimeout)
+	defer cancel()
+
+	if cleanupErr := job.CleanupRBAC(cleanupCtx, clientset, v.Namespace, v.RunID); cleanupErr != nil {
+		slog.Error("failed to roll back RBAC after preparation failure; cluster-admin binding may be orphaned",
+			"runID", v.RunID, "namespace", v.Namespace, "error", cleanupErr)
+		return cleanupErr
+	}
+	return nil
 }
 
 // ValidatePhases runs the specified phases sequentially and returns one
@@ -215,7 +275,7 @@ func (v *Validator) ValidatePhases(
 	phases []Phase,
 	validationInput *v1.ValidationInput,
 	snap *snapshotter.Snapshot,
-) ([]*PhaseResult, error) {
+) (results []*PhaseResult, err error) {
 
 	if len(phases) == 0 {
 		phases = PhaseOrder
@@ -226,15 +286,15 @@ func (v *Validator) ValidatePhases(
 	// Lower any nccl-benchmark-runtime-ref into its inline carrier by reading the
 	// referenced template from the --data tree. Fails fast on a bad ref before
 	// deploying any Jobs.
-	if err := v.resolveBenchmarkRuntimeRef(ctx, validationInput); err != nil {
-		return nil, err
+	if refErr := v.resolveBenchmarkRuntimeRef(ctx, validationInput); refErr != nil {
+		return nil, refErr
 	}
 
 	// Pre-flight: evaluate the top-level and readiness-phase constraints
 	// against the snapshot. Fails fast before deploying any Jobs if
 	// prerequisites aren't met.
-	if err := checkReadiness(validationInput, snap); err != nil {
-		return nil, err
+	if readyErr := checkReadiness(validationInput, snap); readyErr != nil {
+		return nil, readyErr
 	}
 
 	cat, err := catalog.LoadWithDataProvider(ctx, v.dataProvider, v.Version, v.Commit)
@@ -260,9 +320,18 @@ func (v *Validator) ValidatePhases(
 		return nil, err
 	}
 	defer close(cs.stopCh)
-	defer v.deferClusterCleanup(cs.clientset) //nolint:contextcheck // cleanup uses fresh context
+	// Promote a privileged (RBAC) cleanup failure into the run's error, but only
+	// when there is no prior real error — a genuine phase failure takes
+	// precedence over a cleanup problem.
+	//
+	//nolint:contextcheck // deferClusterCleanup uses a fresh context: parent may be canceled
+	defer func() {
+		if cleanupErr := v.deferClusterCleanup(cs.clientset); cleanupErr != nil && err == nil {
+			err = cleanupErr
+		}
+	}()
 
-	results, err := v.runPhases(ctx, func(phase Phase) (*PhaseResult, error) {
+	results, err = v.runPhases(ctx, func(phase Phase) (*PhaseResult, error) {
 		return v.runPhase(ctx, cs.clientset, cs.factory, cat, phase, validationInput)
 	}, cat, phases)
 	if err != nil {
@@ -324,19 +393,19 @@ func (v *Validator) ValidatePhase(
 	phase Phase,
 	validationInput *v1.ValidationInput,
 	snap *snapshotter.Snapshot,
-) (*PhaseResult, error) {
+) (result *PhaseResult, err error) {
 
 	// Lower any nccl-benchmark-runtime-ref into its inline carrier before the
 	// phase runs (or is skipped), so a bad ref fails fast even offline.
-	if err := v.resolveBenchmarkRuntimeRef(ctx, validationInput); err != nil {
-		return nil, err
+	if refErr := v.resolveBenchmarkRuntimeRef(ctx, validationInput); refErr != nil {
+		return nil, refErr
 	}
 
 	// Readiness pre-flight — before the no-cluster short-circuit, matching
 	// ValidatePhases: constraints are evaluated inline against the snapshot
 	// even in test mode.
-	if err := checkReadiness(validationInput, snap); err != nil {
-		return nil, err
+	if readyErr := checkReadiness(validationInput, snap); readyErr != nil {
+		return nil, readyErr
 	}
 
 	cat, err := catalog.LoadWithDataProvider(ctx, v.dataProvider, v.Version, v.Commit)
@@ -361,7 +430,16 @@ func (v *Validator) ValidatePhase(
 		return nil, err
 	}
 	defer close(cs.stopCh)
-	defer v.deferClusterCleanup(cs.clientset) //nolint:contextcheck // cleanup uses fresh context
+	// Promote a privileged (RBAC) cleanup failure into the run's error, but only
+	// when there is no prior real error — a genuine phase failure takes
+	// precedence over a cleanup problem.
+	//
+	//nolint:contextcheck // deferClusterCleanup uses a fresh context: parent may be canceled
+	defer func() {
+		if cleanupErr := v.deferClusterCleanup(cs.clientset); cleanupErr != nil && err == nil {
+			err = cleanupErr
+		}
+	}()
 
 	return v.runPhase(ctx, cs.clientset, cs.factory, cat, phase, validationInput)
 }