Merge pull request #226 from rancher-sandbox/seamless-upgrade

feat: make agent protection persistent during rolling update

Author: Andreagit97

charts/runtime-enforcer/templates/agent/daemonset.yaml

@@ -9,6 +9,11 @@ spec:
     matchLabels:
       app.kubernetes.io/component: agent
     {{- include "runtime-enforcer.selectorLabels" . | nindent 6 }}
+  updateStrategy: 
+    type: RollingUpdate
+    rollingUpdate:
+      maxUnavailable: 0
+      maxSurge: 1
   template:
     metadata:
       labels:
@@ -59,6 +64,12 @@ spec:
           | default .Chart.AppVersion }}
         imagePullPolicy: {{ .Values.agent.agent.image.pullPolicy }}
         name: agent
+        startupProbe:
+          httpGet:
+            path: /readyz
+            port: 8081
+          initialDelaySeconds: 5
+          periodSeconds: 3
         resources: {{- toYaml .Values.agent.agent.resources | nindent 10 }}
         securityContext: {{- toYaml .Values.agent.agent.containerSecurityContext | nindent 10 }}
         volumeMounts:

cmd/agent/main.go

@@ -2,8 +2,10 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
+	"net/http"
 	"os"
 
 	"github.com/cilium/ebpf"
@@ -32,16 +34,20 @@ type Config struct {
 	enableLearning    bool
 	nriSocketPath     string
 	nriPluginIdx      string
+	probeAddr         string
 }
 
 // +kubebuilder:rbac:groups=security.rancher.io,resources=workloadpolicies,verbs=get;list;watch
 
-func newControllerManager() (manager.Manager, error) {
+func newControllerManager(config Config) (manager.Manager, error) {
 	scheme := runtime.NewScheme()
 	utilruntime.Must(v1alpha1.AddToScheme(scheme))
 	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
 	utilruntime.Must(securityv1alpha1.AddToScheme(scheme))
-	controllerOptions := ctrl.Options{Scheme: scheme}
+	controllerOptions := ctrl.Options{
+		Scheme:                 scheme,
+		HealthProbeBindAddress: config.probeAddr,
+	}
 	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), controllerOptions)
 	if err != nil {
 		return nil, fmt.Errorf("unable to start manager: %w", err)
@@ -55,7 +61,7 @@ func startAgent(ctx context.Context, logger *slog.Logger, config Config) error {
 	//////////////////////
 	// Create controller manager
 	//////////////////////
-	ctrlMgr, err := newControllerManager()
+	ctrlMgr, err := newControllerManager(config)
 	if err != nil {
 		return fmt.Errorf("cannot create manager: %w", err)
 	}
@@ -118,6 +124,12 @@ func startAgent(ctx context.Context, logger *slog.Logger, config Config) error {
 		return fmt.Errorf("failed to add NRI handler to controller manager: %w", err)
 	}
 
+	// controller-runtime doesn't support a separate startup probe, so we use the readiness probe instead.
+	// See https://github.com/kubernetes-sigs/controller-runtime/issues/2644 for more details.
+	if err = ctrlMgr.AddReadyzCheck("resolver readyz", resolver.Ping); err != nil {
+		return fmt.Errorf("failed to add resolver's readiness probe: %w", err)
+	}
+
 	//////////////////////
 	// Create the scraper
 	//////////////////////
@@ -139,7 +151,23 @@ func startAgent(ctx context.Context, logger *slog.Logger, config Config) error {
 	if err != nil {
 		return fmt.Errorf("cannot get workload policy informer: %w", err)
 	}
-	_, _ = workloadPolicyInformer.AddEventHandler(resolver.PolicyEventHandlers())
+	handlerRegistration, err := workloadPolicyInformer.AddEventHandler(resolver.PolicyEventHandlers())
+	if err != nil {
+		return fmt.Errorf("failed to add event handler for workload policy: %w", err)
+	}
+
+	if err = ctrlMgr.AddReadyzCheck("policy readyz", func(_ *http.Request) error {
+		// Instead of informer.HasSynced(), which checks if the internal storage is synced,
+		// we use ResourceEventHandlerRegistration.HasSynced() to ensure that
+		// the event handlers have been synced.
+		if !handlerRegistration.HasSynced() {
+			logger.Warn("workload policy informer has not yet synced")
+			return errors.New("workload policy informer has not yet synced")
+		}
+		return nil
+	}); err != nil {
+		return fmt.Errorf("failed to add NRI handler's readiness probe: %w", err)
+	}
 
 	logger.InfoContext(ctx, "starting manager")
 	if err = ctrlMgr.Start(ctx); err != nil {
@@ -167,6 +195,7 @@ func main() {
 	flag.BoolVar(&config.enableLearning, "enable-learning", false, "Enable learning mode")
 	flag.StringVar(&config.nriSocketPath, "nri-socket-path", "/var/run/nri/nri.sock", "NRI socket path")
 	flag.StringVar(&config.nriPluginIdx, "nri-plugin-index", "00", "NRI plugin index")
+	flag.StringVar(&config.probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
 
 	flag.Parse()
 

internal/nri/handler.go

@@ -24,6 +24,24 @@ type Handler struct {
 	resolver    *resolver.Resolver
 }
 
+func newNRIPlugin(logger *slog.Logger, resolver *resolver.Resolver, opts ...stub.Option) (*plugin, error) {
+	var err error
+	p := &plugin{
+		logger:   logger.With("component", "nri-plugin"),
+		resolver: resolver,
+	}
+
+	p.stub, err = stub.New(p, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create NRI plugin stub: %w", err)
+	}
+	return p, nil
+}
+
+func (p *plugin) Run(ctx context.Context) error {
+	return p.stub.Run(ctx)
+}
+
 func NewNRIHandler(socketPath, pluginIndex string, logger *slog.Logger, r *resolver.Resolver) (*Handler, error) {
 	h := &Handler{
 		socketPath:  socketPath,
@@ -71,24 +89,17 @@ func (h *Handler) checkNRISupport() error {
 }
 
 func (h *Handler) startNRIPlugin(ctx context.Context) error {
-	var err error
-
-	p := &plugin{
-		logger:   h.logger.With("component", "nri-plugin"),
-		resolver: h.resolver,
-	}
-
-	opts := []stub.Option{
+	p, err := newNRIPlugin(
+		h.logger,
+		h.resolver,
 		stub.WithPluginIdx(h.pluginIndex),
 		stub.WithSocketPath(h.socketPath),
-	}
-
-	p.stub, err = stub.New(p, opts...)
+	)
 	if err != nil {
-		return fmt.Errorf("failed to create NRI plugin stub: %w", err)
+		return fmt.Errorf("failed to create NRI plugin: %w", err)
 	}
 
-	err = p.stub.Run(ctx)
+	err = p.Run(ctx)
 	if err != nil {
 		return fmt.Errorf("NRI plugin exited with error: %w", err)
 	}

internal/nri/plugin.go

@@ -93,6 +93,8 @@ func (p *plugin) Synchronize(
 				"error", err)
 		}
 	}
+	// Mark resolver as synchronized, so old agent can be safely removed.
+	p.resolver.NRISynchronized()
 	return nil, nil
 }
 

internal/resolver/nri_api.go

@@ -1,7 +1,9 @@
 package resolver
 
 import (
+	"errors"
 	"fmt"
+	"net/http"
 
 	"github.com/rancher-sandbox/runtime-enforcer/internal/bpf"
 )
@@ -123,3 +125,15 @@ func (r *Resolver) RemovePodContainerFromNri(podID PodID, containerID ContainerI
 
 	return r.cgroupToPolicyMapUpdateFunc(PolicyIDNone, []CgroupID{container.cgID}, bpf.RemoveCgroups)
 }
+
+func (r *Resolver) NRISynchronized() {
+	r.nriSynchronized.Store(true)
+}
+
+func (r *Resolver) Ping(_ *http.Request) error {
+	if !r.nriSynchronized.Load() {
+		r.logger.Warn("NRI handler has not yet synchronized")
+		return errors.New("NRI handler has not yet synchronized")
+	}
+	return nil
+}

internal/resolver/resolver.go

@@ -3,6 +3,7 @@ package resolver
 import (
 	"log/slog"
 	"sync"
+	"sync/atomic"
 
 	"github.com/rancher-sandbox/runtime-enforcer/internal/bpf"
 	"github.com/rancher-sandbox/runtime-enforcer/internal/types/policymode"
@@ -15,8 +16,9 @@ type ContainerName = string
 
 type Resolver struct {
 	// let's see if we can split this unique lock in multiple locks later
-	mu     sync.Mutex
-	logger *slog.Logger
+	mu              sync.Mutex
+	logger          *slog.Logger
+	nriSynchronized atomic.Bool
 	// todo!: we should add a cache with deleted pods/containers so that we can resolve also recently deleted ones
 	podCache        map[PodID]*podState
 	cgroupIDToPodID map[CgroupID]PodID

test/e2e/e2e_test.go

@@ -107,3 +107,9 @@ func TestValidatingAdmissionPolicyPodPolicyLabel(t *testing.T) {
 
 	testEnv.Test(t, getValidatingAdmissionPolicyPodPolicyLabelTest())
 }
+
+func TestRollingUpdate(t *testing.T) {
+	t.Log("test rolling update")
+
+	testEnv.Test(t, getRollingUpdateTest())
+}

test/e2e/enforcement_test.go

@@ -105,7 +105,7 @@ func getEnforcementOnExistingPodsTest() types.Feature {
 						Spec: v1alpha1.WorkloadPolicySpec{
 							Mode: "protect",
 							RulesByContainer: map[string]*v1alpha1.WorkloadPolicyRules{
-								"ubuntu": &v1alpha1.WorkloadPolicyRules{
+								"ubuntu": {
 									Executables: tc.AllowedExecutables,
 								},
 							},
@@ -210,7 +210,7 @@ func getEnforcementOnNewPodsTest() types.Feature {
 						Spec: v1alpha1.WorkloadPolicySpec{
 							Mode: "protect",
 							RulesByContainer: map[string]*v1alpha1.WorkloadPolicyRules{
-								"ubuntu": &v1alpha1.WorkloadPolicyRules{
+								"ubuntu": {
 									Executables: tc.AllowedExecutables,
 								},
 							},