Add basic network policy for Ingress to Workers

cmd/atecontroller/internal/controllers/networkpolicy_controller.go

@@ -0,0 +1,132 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package controllers
+
+import (
+	"context"
+	"fmt"
+
+	networkingv1 "k8s.io/api/networking/v1"
+	k8errors "k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/runtime"
+	metav1ac "k8s.io/client-go/applyconfigurations/meta/v1"
+	networkingv1ac "k8s.io/client-go/applyconfigurations/networking/v1"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+
+	"github.com/agent-substrate/substrate/internal/resources"
+	atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
+)
+
+const (
+	networkPolicyFieldOwner = "ate-networkpolicy"
+	ateSystemNamespace      = "ate-system"
+	atenetRouterAppName     = "atenet-router"
+)
+
+type NetworkPolicyReconciler struct {
+	client.Client
+	Scheme *runtime.Scheme
+}
+
+//+kubebuilder:rbac:groups=ate.dev,resources=workerpools,verbs=get;list;watch
+//+kubebuilder:rbac:groups=networking.k8s.io,resources=networkpolicies,verbs=get;list;watch;create;update;patch;delete
+
+func (r *NetworkPolicyReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
+	log := log.FromContext(ctx)
+
+	wp := &atev1alpha1.WorkerPool{}
+	if err := r.Get(ctx, req.NamespacedName, wp); err != nil {
+		if k8errors.IsNotFound(err) {
+			return ctrl.Result{}, nil
+		}
+		return ctrl.Result{}, fmt.Errorf("failed to get worker pool %q: %w", req.NamespacedName, err)
+	}
+
+	if !wp.GetDeletionTimestamp().IsZero() {
+		log.Info("WorkerPool is being deleted, NetworkPolicy will be GC'd via OwnerReference",
+			"namespace", wp.Namespace,
+			"name", wp.Name)
+		return ctrl.Result{}, nil
+	}
+
+	if err := r.reconcileImpl(ctx, wp); err != nil {
+		log.Error(err, "Failed to reconcile NetworkPolicy")
+		return ctrl.Result{}, err
+	}
+
+	return ctrl.Result{}, nil
+}
+
+func (r *NetworkPolicyReconciler) reconcileImpl(ctx context.Context, wp *atev1alpha1.WorkerPool) error {
+	log := log.FromContext(ctx)
+
+	npAC := buildNetworkPolicyApplyConfig(wp)
+
+	if err := r.Apply(ctx, npAC, client.FieldOwner(networkPolicyFieldOwner), client.ForceOwnership); err != nil {
+		return fmt.Errorf("failed to apply NetworkPolicy %s:%s: %w", *npAC.Namespace, *npAC.Name, err)
+	}
+	log.Info("reconcileImpl done",
+		"namespace", *npAC.Namespace,
+		"name", *npAC.Name)
+
+	return nil
+}
+
+func buildNetworkPolicyApplyConfig(wp *atev1alpha1.WorkerPool) *networkingv1ac.NetworkPolicyApplyConfiguration {
+	np := networkingv1ac.NetworkPolicy(resources.NetworkPolicyName(wp.Name), wp.Namespace).
+		WithLabels(map[string]string{
+			"ate.dev/worker-pool": wp.Name,
+		}).
+		WithOwnerReferences(metav1ac.OwnerReference().
+			WithAPIVersion(atev1alpha1.GroupVersion.String()).
+			WithKind("WorkerPool").
+			WithName(wp.Name).
+			WithUID(wp.UID).
+			WithController(true).
+			WithBlockOwnerDeletion(true))
+
+	// Ingress policy: only accept connections from the atenet-router, all ports.
+	np.
+		WithSpec(networkingv1ac.NetworkPolicySpec().
+			WithPodSelector(metav1ac.LabelSelector().
+				WithMatchLabels(map[string]string{"ate.dev/worker-pool": wp.Name})).
+			WithPolicyTypes(networkingv1.PolicyTypeIngress).
+			WithIngress(
+				networkingv1ac.NetworkPolicyIngressRule().
+					WithFrom(
+						networkingv1ac.NetworkPolicyPeer().
+							WithNamespaceSelector(metav1ac.LabelSelector().
+								WithMatchLabels(map[string]string{"kubernetes.io/metadata.name": ateSystemNamespace})).
+							WithPodSelector(metav1ac.LabelSelector().
+								WithMatchLabels(map[string]string{"app": atenetRouterAppName})),
+					),
+			))
+
+	// Egress is left unmanaged by Kubernetes NetworkPolicy for now while waiting for
+	// further progress on Egress API designs and because some egress rules may not be
+	// expressible at the Kubernetes layer.
+
+	return np
+}
+
+func (r *NetworkPolicyReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		Named("networkpolicy").
+		For(&atev1alpha1.WorkerPool{}).
+		Owns(&networkingv1.NetworkPolicy{}).
+		Complete(r)
+}

cmd/atecontroller/internal/controllers/networkpolicy_controller_test.go

@@ -0,0 +1,92 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package controllers
+
+import (
+	"context"
+	"testing"
+
+	networkingv1 "k8s.io/api/networking/v1"
+	"k8s.io/apimachinery/pkg/types"
+
+	"github.com/agent-substrate/substrate/internal/resources"
+)
+
+func TestWorkerPoolCreatesNetworkPolicy(t *testing.T) {
+	ctx := t.Context()
+	wp := makeWorkerPool("test-netpolicy-create", "default", 2, "ateom:v1")
+	if err := k8sClient.Create(ctx, wp); err != nil {
+		t.Fatalf("create WorkerPool: %v", err)
+	}
+	deleteOnCleanup(t, wp)
+
+	eventually(t, func(ctx context.Context) (bool, error) {
+		npName := resources.NetworkPolicyName(wp.Name)
+		np := &networkingv1.NetworkPolicy{}
+		err := k8sClient.Get(ctx, types.NamespacedName{Name: npName, Namespace: wp.Namespace}, np)
+		if err != nil {
+			return false, nil
+		}
+
+		// Verify OwnerReference
+		if len(np.OwnerReferences) == 0 || np.OwnerReferences[0].Name != wp.Name {
+			return false, nil
+		}
+
+		// Verify metadata label matches the worker pool
+		if np.Labels == nil || np.Labels["ate.dev/worker-pool"] != wp.Name {
+			return false, nil
+		}
+
+		// Verify PodSelector matches the worker pool
+		if np.Spec.PodSelector.MatchLabels == nil || np.Spec.PodSelector.MatchLabels["ate.dev/worker-pool"] != wp.Name {
+			return false, nil
+		}
+
+		// Verify PolicyTypes contains Ingress
+		hasIngress := false
+		for _, pt := range np.Spec.PolicyTypes {
+			if pt == networkingv1.PolicyTypeIngress {
+				hasIngress = true
+			}
+		}
+		if !hasIngress {
+			return false, nil
+		}
+
+		// Verify Ingress Rules (Allow only ingress from ATE router)
+		if len(np.Spec.Ingress) != 1 {
+			return false, nil
+		}
+		ingressRule := np.Spec.Ingress[0]
+		if len(ingressRule.From) != 1 {
+			return false, nil
+		}
+		fromPeer := ingressRule.From[0]
+		if fromPeer.NamespaceSelector == nil || fromPeer.NamespaceSelector.MatchLabels["kubernetes.io/metadata.name"] != ateSystemNamespace {
+			return false, nil
+		}
+		if fromPeer.PodSelector == nil || fromPeer.PodSelector.MatchLabels["app"] != atenetRouterAppName {
+			return false, nil
+		}
+
+		// Verify Egress Rules are unmanaged (empty)
+		if len(np.Spec.Egress) != 0 {
+			return false, nil
+		}
+
+		return true, nil
+	})
+}

cmd/atecontroller/internal/controllers/workerpool_controller_test.go

@@ -94,6 +94,14 @@ func TestMain(m *testing.M) {
 		os.Exit(1)
 	}
 
+	if err := (&NetworkPolicyReconciler{
+		Client: mgr.GetClient(),
+		Scheme: mgr.GetScheme(),
+	}).SetupWithManager(mgr); err != nil {
+		fmt.Fprintf(os.Stderr, "netpolicy controller setup failed: %v\n", err)
+		os.Exit(1)
+	}
+
 	mgrCtx, mgrCancel := context.WithCancel(context.Background())
 	go func() {
 		_ = mgr.Start(mgrCtx)

cmd/atecontroller/main.go

@@ -91,6 +91,14 @@ func main() {
 		os.Exit(1)
 	}
 
+	if err = (&controllers.NetworkPolicyReconciler{
+		Client: mgr.GetClient(),
+		Scheme: mgr.GetScheme(),
+	}).SetupWithManager(mgr); err != nil {
+		setupLog.Error(err, "unable to create controller", "controller", "NetPolicy")
+		os.Exit(1)
+	}
+
 	if err = (&controllers.ActorTemplateReconciler{
 		Client:    mgr.GetClient(),
 		Scheme:    mgr.GetScheme(),