diff --git a/README.md b/README.md index a7176ce..ac9702e 100644 --- a/README.md +++ b/README.md @@ -309,6 +309,234 @@ pod_template: ``` +#### Restricted Kubernetes egress + +Restricted assignments require an upgraded Kubernetes worker with explicit local +opt-in. Docker, Direct, Command, and Kubernetes workers without opt-in reject +them before claiming work. An absent/null `network_policy` remains unrestricted; +a present policy with zero rules denies all egress, including DNS. Host rules +allow all ports to their destinations; they are not HTTPS-only rules. + +File configuration: + +```yaml +backend: + kubernetes: + network_policy: + enabled: false + dns: + server_ips: [] + namespace: kube-system + pod_labels: + k8s-app: kube-dns +``` + +The Helm equivalents are `kubernetesBackend.networkPolicy.enabled` and +`kubernetesBackend.networkPolicy.dns.{serverIPs,namespace,podLabels}`. No CLI +switch enables this feature. Leave it disabled until the prerequisites below +are verified. + +**Install and verify the baseline first.** The chart always renders a persistent +native egress-only default deny, even with support disabled or `rbac.create=false`. +Non-Helm operators must install the equivalent in the task namespace: + +```yaml +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: oz-restricted-baseline + namespace: agents +spec: + podSelector: + matchLabels: + warp-network-restricted: "true" + policyTypes: [Egress] + egress: [] +``` + +The baseline does not select unrestricted tasks or change ingress. Successful +Helm installation or policy creation proves API acceptance, not dataplane +enforcement. Verify blocking on your CNI before opting in, including on newly +scheduled nodes. The CNI must continue enforcing a policy while it has a +`deletionTimestamp` and the consumer finalizer remains, including across CNI +agent restarts. Revalidate on CNI upgrades. The worker does not run per-task +connectivity probes or attest CNI health. + +Other native, Cilium, and administrative policies are additive and must not widen +restricted-task egress. Admission must preserve reserved task labels and the +network boundary. All `warp-` task label keys are reserved; `extra_labels` and +Helm `extraLabels` reject them. Worker-owned `warp-` annotations override +conflicting extra annotations. + +**DNS and providers.** A definite absence of the Cilium API selects native +`NetworkPolicy`; a present `cilium.io/v2` `ciliumnetworkpolicies` API selects +Cilium. Discovery errors and Cilium permission/API errors fail closed, without +native fallback. The CRD alone does not prove that Cilium is enforcing. + +- Native rules snapshot every A/AAAA answer once per launch and allow exact + `/32` or `/128` destinations. Resolution errors or empty answers prevent Job + creation. Existing snapshots never refresh. Long-lived/debug-idle executions + can lose reachability or retain stale allowed addresses as DNS changes. +- Cilium rules use exact `toFQDNs.matchName` and literal `toCIDR` grants, with DNS + proxy rules limited to allowed query names. Cilium tracks changing answers and + handles CNAME chains according to its DNS policy semantics; the worker does not + add unrelated hostnames or worker-side snapshots. +- Hostnames sharing an allowed IP share IP-level access. This is not strict + hostname, HTTP authority, or data-loss-prevention isolation. Service DNAT and + Cilium-managed endpoint identity can deny a named IP; no broad cluster/world + exception is added to fix this. + +Explicit `dns.server_ips` take precedence. Otherwise only an actual in-cluster +client reads the worker Pod's `/etc/resolv.conf` nameservers. Out-of-cluster or +explicit-kubeconfig workers must supply the target cluster's trusted resolver +addresses. Loopback, unspecified, malformed, and zone-qualified addresses fail. +There is no public-resolver fallback or worker Service lookup. Empty/IP-only +policies need no DNS grant. Hostname policies require valid resolver addresses, +DNS namespace, and nonempty Pod labels. + +Native policies grant UDP/TCP 53 only to exact resolver IPs and a combined +namespace/Pod selector. Cilium uses that DNS Pod selector and its restricted +query proxy. Ordinary `ClusterFirst` Pod networking is the zero-configuration +case. NodeLocal DNS, custom forwarders, service meshes, secondary interfaces, and +different task DNS views require explicit topology configuration and live CNI +validation. Set the resolver/selector overrides and task `pod_template.dnsPolicy` +/ `dnsConfig` as appropriate; the worker does not overwrite those Pod settings +or synthesize Cilium `host`/broad DNS allowances. Inspecting +`kubectl -n kube-system get service kube-dns` can help diagnose a cluster, but its +ClusterIP is not necessarily the resolver tasks use. + +**Security boundary.** The worker does not generate host networking, privileged +containers, or `NET_ADMIN`/`SYS_ADMIN` grants. Explicit operator PodSpec/container +settings remain supported and can bypass enforcement; `hostNetwork` semantics +are CNI-dependent. Root sidecar-copy init containers remain supported and root +UID alone is not privileged mode. Keep task ServiceAccounts separate from the +worker's policy-management account. Task credentials should not permit creating +Pods, changing labels/policies, reaching privileged daemons, or escaping to host +networking. Node/runtime image pulls are not task-Pod egress. + +**Shared policies and cleanup.** Policies are immutable and named +`warp-np-<52-character SHA-256 base32 digest>` from canonical final enforcement +content, not run IDs, metadata, or source descriptions. Equal content shares one +policy in the selected namespace/resource kind. Native DNS changes create new +snapshots; they never edit an existing policy's rules. + +The policy is created with its consumer finalizer and first durable registration +before the ordinary unsuspended Job is posted. Later consumers register through +UID/resourceVersion CAS. The versioned `warp-network-policy-consumers` annotation +records full run/execution IDs, expected Job name, launcher incarnation, and +optional Job UID. The canonical annotation is diagnostic, not proof of content. +The complete annotations must fit below 256 KiB; exhaustion fails before a Job +is created rather than evicting registrations. + +Startup, 30-second sweeps, and terminal notifications reconcile both provider +kinds. A consumer is released only after its Job cannot restart and all owned +regular/init/native-sidecar/ephemeral containers have stopped. Pending, Unknown, +missing statuses, live sidecars, or unexplained matching Jobs/Pods retain +protection. Job cancellation/expiry is asynchronous; worker shutdown preserves +live Jobs, policies, and registrations. `cleanup: false` retains terminal +inspection Jobs/Pods, but does not retain their policies after execution stops. + +The last stopped consumer triggers UID/resourceVersion-fenced deletion and +removal of only `warp.dev/network-policy-consumers`. New consumers cannot join a +terminating UID; an identical launch waits for actual removal within its +two-minute provisioning budget or fails. It does not bypass a stuck policy with +a new selector. Policy/DNS/lifecycle failures after claim are terminal +`network_policy` task failures, not capacity requeues. + +A launcher submits at most one Job POST per registration. Lost responses are +recovered only through matching API identity. A crash before posting, or a Job +that disappears before UID binding, can leave an irreducibly ambiguous +registration. Age, disconnect, cancellation, and one missing-Job read are never +permission to expire it. + +For a stuck reservation, an operator must: + +1. Fence every original launcher that could still post that execution's Job. + Establish that no create request remains in flight. +2. Inspect the policy UID, complete ledger, expected Job names, full + run/execution/policy-UID annotations, Job UIDs, and owned/matching Pods. + Confirm actual termination or fence affected/partitioned nodes; force deletion + alone does not prove processes stopped. +3. Take a fresh policy read. CAS-remove only entries proven stale, preserving all + other consumers. For example, a JSON Patch must test `/metadata/uid` and + `/metadata/resourceVersion` before replacing the ledger annotation with the + verified reduced version-1 ledger. Retry from a fresh read on conflict. +4. Let an enabled reconciler perform normal zero-consumer finalization. + +Unknown ledger schemas, malformed metadata, permission errors, operator-mutated +rules, or unexplained live objects retain protection and emit policy UID and +consumer diagnostics. Without an enabled reconciler or necessary rights, a +policy can remain terminating indefinitely. Direct finalizer removal is a +break-glass action outside this guarantee. The worker does not repair +force-removed enforcement or resurrect a deleted shared policy. + +**Coordinated rollout and rollback.** There is no capability negotiation: old +workers silently ignore the new JSON field. Publish the upgraded worker/chart +first. Keep policy-bearing dispatch disabled for the target pool, upgrade every +connected replica to an immutable image, verify the baseline and CNI/DNS +prerequisites, then explicitly enable local support. Only then deploy/enable +server assignment plumbing under the existing `WorkerNetworkPoliciesEnabled` +gate. That gate is not assumed to be pool-scoped. Mixed old/new workers with +policy-bearing dispatch are unsupported and insecure. Turning the gate off +intentionally returns new runs to nil-policy dispatch; it is not an isolation +workaround. Before rollback to old binaries, stop policy-bearing assignments. +Drain restricted Jobs, reservations, and policies before removing policy RBAC, +the baseline/chart, or reconciliation. Preserve in-use ledger schemas during +upgrades. + +**Measured validation matrix (not a universal support claim).** + +- Kubernetes **v1.32.2**, Cilium **1.18.3**, Cilium provider, single-node kind: + controlled external IPv4 allowed/forbidden endpoints and controlled CoreDNS + records passed setup/main probes, empty-policy DNS denial, two-consumer + immutable reuse, and final + cleanup with retained terminal Jobs/Pods. Terminating-policy UID/spec and + allow/deny enforcement survived an authorized Job POST delayed until after + DELETE, Cilium agent restart, and fresh reconciliation. Existing controller + expiry/TTL and retained-Job tests also passed. +- Kubernetes **v1.32.2**, Calico **3.29.3**, native provider: **BLOCKED / NOT + VERIFIED**. This sandbox kernel lacks iptables `rpfilter` support; Felix could + not program iptables rules. The documented eBPF/VXLAN alternative also failed: + its connect-time BPF program could not load (`Invalid argument`). With the + optional connect-time load balancer disabled and packet-path NAT enabled, + `calico-node` became Ready, but the unrestricted external control failed and + node-side cluster-DNS UDP writes returned `invalid argument`. Controller + expiry/TTL tests passed; native network enforcement and terminating-policy + retention did not. A compatible isolated native-CNI cluster is required before + native acceptance. Agent readiness/API acceptance do not constitute proof. +- Default kindnet is non-enforcing and is not a supported validation result. + Dual-stack/IPv6 dataplane, NodeLocal/custom DNS, additional Kubernetes/CNI + versions, and multi-node propagation remain **NOT VERIFIED**. + +The opt-in integration tests require an explicitly disposable kubeconfig, +`kubectl`, two externally reachable HTTP endpoints on port 8080 with distinct IPs, +and trusted DNS mapping the supplied names to those endpoints. They delete test +namespaces and restart selected CNI agent Pods; never point them at production. +The recorded Cilium run used Docker-network endpoints `172.18.0.12` and +`172.18.0.13`, with CoreDNS records `allowed.remote3452.test` and +`blocked.remote3452.test`. To reproduce in an equivalent disposable cluster: + +```bash +OZ_K8S_TEST_KUBECONFIG=/path/to/disposable-kubeconfig \ +OZ_K8S_TEST_CNI_VERSION=cilium-1.18.3 \ +OZ_K8S_TEST_PROVIDER=cilium \ +OZ_K8S_TEST_ALLOWED_IP=172.18.0.12 \ +OZ_K8S_TEST_FORBIDDEN_IP=172.18.0.13 \ +OZ_K8S_TEST_ALLOWED_HOST=allowed.remote3452.test \ +OZ_K8S_TEST_FORBIDDEN_HOST=blocked.remote3452.test \ +OZ_K8S_TEST_DNS_IPS=10.96.0.10 \ +OZ_K8S_TEST_CNI_NAMESPACE=kube-system \ +OZ_K8S_TEST_CNI_POD_SELECTOR=k8s-app=cilium \ +go test -tags=integration ./internal/worker \ + -run 'Test(KubernetesNetworkPolicyWithCNI|KubernetesNetworkPolicyTerminatingWithCNI|JobExpirationWithController)' \ + -count=1 -v +``` + +Use `OZ_K8S_TEST_PROVIDER=native` on an enforcing native cluster without the +Cilium API, and supply that CNI's version and agent selector. Out-of-cluster +native hostname compilation additionally needs a route to the configured DNS +resolver. A skipped fixture is not an enforcement result. + ### Helm Chart This repo includes a namespace-scoped Helm chart at `charts/oz-agent-worker`. diff --git a/charts/oz-agent-worker/templates/configmap.yaml b/charts/oz-agent-worker/templates/configmap.yaml index 2d5a3d5..817e6c3 100644 --- a/charts/oz-agent-worker/templates/configmap.yaml +++ b/charts/oz-agent-worker/templates/configmap.yaml @@ -17,6 +17,12 @@ data: backend: kubernetes: namespace: {{ default .Release.Namespace .Values.kubernetesBackend.namespace | quote }} + network_policy: + enabled: {{ .Values.kubernetesBackend.networkPolicy.enabled }} + dns: + server_ips: {{ .Values.kubernetesBackend.networkPolicy.dns.serverIPs | toJson }} + namespace: {{ .Values.kubernetesBackend.networkPolicy.dns.namespace | quote }} + pod_labels: {{ .Values.kubernetesBackend.networkPolicy.dns.podLabels | toJson }} {{- if .Values.kubernetesBackend.defaultImage }} default_image: {{ .Values.kubernetesBackend.defaultImage | quote }} {{- end }} @@ -44,6 +50,11 @@ data: {{ .Values.kubernetesBackend.teardownCommand | indent 10 }} {{- end }} {{- with .Values.kubernetesBackend.extraLabels }} + {{- range $key, $_ := . }} + {{- if hasPrefix "warp-" $key }} + {{- fail (printf "kubernetesBackend.extraLabels key %q is reserved for Warp" $key) }} + {{- end }} + {{- end }} extra_labels: {{ toYaml . | indent 10 }} {{- end }} diff --git a/charts/oz-agent-worker/templates/networkpolicy.yaml b/charts/oz-agent-worker/templates/networkpolicy.yaml new file mode 100644 index 0000000..39846a8 --- /dev/null +++ b/charts/oz-agent-worker/templates/networkpolicy.yaml @@ -0,0 +1,13 @@ +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: {{ include "oz-agent-worker.fullname" . }}-restricted-baseline + namespace: {{ default .Release.Namespace .Values.kubernetesBackend.namespace }} + labels: + {{- include "oz-agent-worker.labels" . | nindent 4 }} +spec: + podSelector: + matchLabels: + warp-network-restricted: "true" + policyTypes: [Egress] + egress: [] diff --git a/charts/oz-agent-worker/templates/role.yaml b/charts/oz-agent-worker/templates/role.yaml index 687aa24..d10780f 100644 --- a/charts/oz-agent-worker/templates/role.yaml +++ b/charts/oz-agent-worker/templates/role.yaml @@ -19,4 +19,12 @@ rules: - apiGroups: [""] resources: ["events"] verbs: ["list"] + {{- if .Values.kubernetesBackend.networkPolicy.enabled }} + - apiGroups: ["networking.k8s.io"] + resources: ["networkpolicies"] + verbs: ["create", "get", "list", "patch", "delete"] + - apiGroups: ["cilium.io"] + resources: ["ciliumnetworkpolicies"] + verbs: ["create", "get", "list", "patch", "delete"] + {{- end }} {{- end }} diff --git a/charts/oz-agent-worker/values.yaml b/charts/oz-agent-worker/values.yaml index ce66e2f..9339876 100644 --- a/charts/oz-agent-worker/values.yaml +++ b/charts/oz-agent-worker/values.yaml @@ -93,6 +93,13 @@ kubernetesBackend: teardownCommand: "" extraLabels: {} extraAnnotations: {} + networkPolicy: + enabled: false + dns: + serverIPs: [] + namespace: kube-system + podLabels: + k8s-app: kube-dns # Hard wall-clock cap on how long a single task Job (and its pod) may run # before Kubernetes terminates it and marks the Job failed with a # DeadlineExceeded condition. This is a safety backstop against tasks whose diff --git a/internal/config/config.go b/internal/config/config.go index e4276ad..a682ccd 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -30,6 +30,17 @@ type FileConfig struct { Backend BackendConfig `yaml:"backend"` } +type NetworkPolicyConfig struct { + Enabled bool `yaml:"enabled"` + DNS NetworkPolicyDNSConfig `yaml:"dns"` +} + +type NetworkPolicyDNSConfig struct { + ServerIPs []string `yaml:"server_ips"` + Namespace string `yaml:"namespace"` + PodLabels map[string]string `yaml:"pod_labels"` +} + // BackendConfig contains the backend selection. // At most one backend field may be non-nil; configuring multiple backends simultaneously is an error. type BackendConfig struct { @@ -75,10 +86,11 @@ type DirectConfig struct { // KubernetesConfig holds Kubernetes-backend-specific configuration. type KubernetesConfig struct { - Namespace string `yaml:"namespace"` - Kubeconfig string `yaml:"kubeconfig"` - DefaultImage string `yaml:"default_image" validate:"omitempty,no_whitespace"` - ImagePullPolicy string `yaml:"image_pull_policy" validate:"omitempty,oneof=Always Never IfNotPresent"` + NetworkPolicy NetworkPolicyConfig `yaml:"network_policy"` + Namespace string `yaml:"namespace"` + Kubeconfig string `yaml:"kubeconfig"` + DefaultImage string `yaml:"default_image" validate:"omitempty,no_whitespace"` + ImagePullPolicy string `yaml:"image_pull_policy" validate:"omitempty,oneof=Always Never IfNotPresent"` // UseImageVolumes is nil when support should be detected by the startup preflight. UseImageVolumes *bool `yaml:"use_image_volumes"` PreflightImage string `yaml:"preflight_image" validate:"omitempty,no_whitespace"` @@ -173,6 +185,13 @@ func Load(path string) (*FileConfig, error) { if err := configValidator.Struct(cfg); err != nil { return nil, fmt.Errorf("invalid config: %w", formatValidationErrors(err)) } + if cfg.Backend.Kubernetes != nil { + for key := range cfg.Backend.Kubernetes.ExtraLabels { + if strings.HasPrefix(key, "warp-") { + return nil, fmt.Errorf("backend.kubernetes.extra_labels key %q is reserved for Warp", key) + } + } + } return &cfg, nil } diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index 61656a6..2e0f567 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -139,6 +139,7 @@ const ( TaskFailureReasonInvalidImage TaskFailureReason = "invalid_image" TaskFailureReasonActiveDeadline TaskFailureReason = "active_deadline" TaskFailureReasonCleanup TaskFailureReason = "cleanup" + TaskFailureReasonNetworkPolicy TaskFailureReason = "network_policy" ) var ( @@ -183,6 +184,7 @@ var ( TaskFailureReasonInvalidImage, TaskFailureReasonActiveDeadline, TaskFailureReasonCleanup, + TaskFailureReasonNetworkPolicy, } ) diff --git a/internal/types/messages.go b/internal/types/messages.go index 6aad91b..dfede37 100644 --- a/internal/types/messages.go +++ b/internal/types/messages.go @@ -24,6 +24,16 @@ type WebSocketMessage struct { Data json.RawMessage `json:"data,omitempty"` } +type NetworkPolicy struct { + Rules []NetworkPolicyRule `json:"rules"` +} + +type NetworkPolicyRule struct { + Hosts []string `json:"hosts"` + Action string `json:"action"` + Description string `json:"description"` +} + // FailureDetails contains backend-specific structured diagnostics. type FailureDetails struct { Kubernetes *KubernetesFailureDetails `json:"kubernetes,omitempty"` @@ -139,6 +149,7 @@ type TaskAssignmentMessage struct { // backends size the task container/pod from it; omitted when the run has no explicit // runner instance shape, in which case the worker keeps its default sizing. InstanceShape *InstanceShape `json:"instance_shape,omitempty"` + NetworkPolicy *NetworkPolicy `json:"network_policy,omitempty"` } // TaskClaimedMessage is sent from worker to server after successfully claiming a task @@ -170,11 +181,14 @@ type TaskFailedMessage struct { FailureDetails *FailureDetails `json:"failure_details,omitempty"` } -// TaskRejectedMessage is sent from worker to server when the worker cannot accept the task -// (e.g. at maximum concurrency). The server should keep the task queued rather than marking it failed. +// TaskRejectedMessage is sent before claim when the worker cannot accept a task. +// Uncoded rejections, such as capacity, keep the task queued. Codes network_policy_unsupported +// and network_policy_invalid are terminal failures for the matching execution. type TaskRejectedMessage struct { - TaskID string `json:"task_id"` - Reason string `json:"reason"` + TaskID string `json:"task_id"` + Reason string `json:"reason"` + ExecutionID string `json:"execution_id,omitempty"` + Code string `json:"code,omitempty"` } // TaskCancellationMessage is sent from server to worker to cancel an active task. diff --git a/internal/worker/backend.go b/internal/worker/backend.go index 9b4057c..fd766d2 100644 --- a/internal/worker/backend.go +++ b/internal/worker/backend.go @@ -165,6 +165,7 @@ type TaskParams struct { // Containerized backends apply it as CPU/memory limits (Docker) or resource // requests/limits (Kubernetes). Backends that cannot enforce a shape (direct) ignore it. InstanceShape *types.InstanceShape + NetworkPolicy *types.NetworkPolicy // SetupEvents reports worker-observed setup phase durations to warp-server. // It may be nil, and all of its methods are safe to call on a nil reporter. diff --git a/internal/worker/kubernetes.go b/internal/worker/kubernetes.go index c9a3685..c380153 100644 --- a/internal/worker/kubernetes.go +++ b/internal/worker/kubernetes.go @@ -7,9 +7,11 @@ import ( "io" "sort" "strings" + "sync" "time" "github.com/rs/zerolog" + "github.com/warpdotdev/oz-agent-worker/internal/config" "github.com/warpdotdev/oz-agent-worker/internal/log" "github.com/warpdotdev/oz-agent-worker/internal/metrics" "github.com/warpdotdev/oz-agent-worker/internal/tasktelemetry/tracing" @@ -22,6 +24,8 @@ import ( "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/watch" + "k8s.io/client-go/discovery" + "k8s.io/client-go/dynamic" "k8s.io/client-go/kubernetes" "k8s.io/client-go/rest" "k8s.io/client-go/tools/clientcmd" @@ -90,6 +94,7 @@ func kubernetesTaskOwnedEnvVars(taskID string) []string { // KubernetesBackendConfig holds configuration specific to the Kubernetes backend. type KubernetesBackendConfig struct { + NetworkPolicy config.NetworkPolicyConfig WorkerID string Namespace string Kubeconfig string @@ -131,9 +136,16 @@ func terminatedExitCode(terminated *corev1.ContainerStateTerminated) int { // KubernetesBackend executes tasks in Kubernetes Jobs. type KubernetesBackend struct { - config KubernetesBackendConfig - clientset kubernetes.Interface - useImageVolumes bool + config KubernetesBackendConfig + clientset kubernetes.Interface + useImageVolumes bool + networkDynamic dynamic.Interface + networkDiscovery discovery.DiscoveryInterface + networkInCluster bool + networkLaunchMu sync.Mutex + networkLaunches map[[2]string]chan struct{} + networkReconcile chan struct{} + networkStop context.CancelFunc } func (b *KubernetesBackend) Capabilities() backendCapabilities { @@ -146,6 +158,11 @@ func (b *KubernetesBackend) PreservesTasksOnShutdown() bool { // NewKubernetesBackend creates a new Kubernetes backend and validates startup requirements. func NewKubernetesBackend(ctx context.Context, config KubernetesBackendConfig) (*KubernetesBackend, error) { + for key := range config.ExtraLabels { + if strings.HasPrefix(key, "warp-") { + return nil, fmt.Errorf("task label %q is reserved for Warp", key) + } + } if config.Namespace == "" { config.Namespace = defaultKubernetesNamespace } @@ -157,7 +174,7 @@ func NewKubernetesBackend(ctx context.Context, config KubernetesBackendConfig) ( config.PreflightImage = kubernetesStartupPreflightImage } - clientConfig, err := loadKubernetesClientConfig(config.Kubeconfig) + clientConfig, inCluster, err := loadKubernetesClientConfig(config.Kubeconfig) if err != nil { return nil, fmt.Errorf("failed to build Kubernetes client config: %w", err) } @@ -174,12 +191,29 @@ func NewKubernetesBackend(ctx context.Context, config KubernetesBackendConfig) ( log.Debugf(ctx, "Connected to Kubernetes API server version %s", serverVersion.String()) backend := &KubernetesBackend{ - config: config, - clientset: clientset, + config: config, + clientset: clientset, + networkInCluster: inCluster, } if err := backend.configureImageVolumes(ctx); err != nil { return nil, err } + if config.NetworkPolicy.Enabled { + policyClientConfig := rest.CopyConfig(clientConfig) + policyClientConfig.Timeout = kubernetesAPIRequestTimeout + backend.networkDynamic, err = dynamic.NewForConfig(policyClientConfig) + if err != nil { + return nil, fmt.Errorf("create network policy client: %w", err) + } + backend.networkDiscovery, err = discovery.NewDiscoveryClientForConfig(policyClientConfig) + if err != nil { + return nil, fmt.Errorf("create network policy discovery client: %w", err) + } + backend.networkReconcile = make(chan struct{}, 1) + reconcileCtx, stop := context.WithCancel(ctx) + backend.networkStop = stop + go backend.runNetworkPolicyReconciler(reconcileCtx) + } return backend, nil } @@ -371,9 +405,19 @@ func (b *KubernetesBackend) ExecuteTask(ctx context.Context, params *TaskParams) log.Infof(ctx, "Creating Kubernetes Job %s in namespace %s", jobName, b.config.Namespace) doneJobCreate := params.SetupEvents.startPhase(ctx, SetupEventJobCreate) - createdJob, err := b.clientset.BatchV1().Jobs(b.config.Namespace).Create(ctx, job, metav1.CreateOptions{}) + var createdJob *batchv1.Job + var err error + if params.NetworkPolicy != nil { + createdJob, err = b.createRestrictedJob(ctx, params, job) + defer b.notifyNetworkPolicyReconciler() + } else { + createdJob, err = b.clientset.BatchV1().Jobs(b.config.Namespace).Create(ctx, job, metav1.CreateOptions{}) + } if err != nil { doneJobCreate(true) + if params.NetworkPolicy != nil { + return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonNetworkPolicy, err)) + } failure := newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonJobCreate, fmt.Errorf("failed to create Kubernetes Job: %w", err)) return executeError(withFailureDetails(failure, buildKubernetesFailureDetails("", job, nil, nil, nil))) } @@ -583,6 +627,7 @@ func (b *KubernetesBackend) ExecuteTask(ctx context.Context, params *TaskParams) // Job whose pod is still pending outlives the cancellation and boots an agent against a run // the server has already finished. A Job that no longer exists is not an error. func (b *KubernetesBackend) CancelTask(ctx context.Context, params *CancelParams) error { + defer b.notifyNetworkPolicyReconciler() if params == nil { return fmt.Errorf("cancel params are required") } @@ -612,6 +657,9 @@ func (b *KubernetesBackend) CancelTask(ctx context.Context, params *CancelParams // worker pod disruption (for example Karpenter node rotation), deleting Jobs // here would kill otherwise healthy running sessions. func (b *KubernetesBackend) Shutdown(ctx context.Context) { + if b.networkStop != nil { + b.networkStop() + } log.Infof(ctx, "Preserving Kubernetes task Jobs during worker shutdown") } @@ -702,18 +750,20 @@ func (b *KubernetesBackend) watchTaskPods(ctx context.Context, executionID strin }) } -func loadKubernetesClientConfig(kubeconfigPath string) (*rest.Config, error) { +func loadKubernetesClientConfig(kubeconfigPath string) (*rest.Config, bool, error) { if kubeconfigPath != "" { - return clientcmd.BuildConfigFromFlags("", kubeconfigPath) + config, err := clientcmd.BuildConfigFromFlags("", kubeconfigPath) + return config, false, err } if inCluster, err := rest.InClusterConfig(); err == nil { - return inCluster, nil + return inCluster, true, nil } loader := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( clientcmd.NewDefaultClientConfigLoadingRules(), &clientcmd.ConfigOverrides{}, ) - return loader.ClientConfig() + config, err := loader.ClientConfig() + return config, false, err } func (b *KubernetesBackend) basePodSpec() corev1.PodSpec { diff --git a/internal/worker/kubernetes_network_policy.go b/internal/worker/kubernetes_network_policy.go new file mode 100644 index 0000000..f796c2b --- /dev/null +++ b/internal/worker/kubernetes_network_policy.go @@ -0,0 +1,486 @@ +package worker + +import ( + "context" + "crypto/sha256" + "encoding/base32" + "encoding/json" + "fmt" + "net" + "net/netip" + "os" + "slices" + "strings" + "sync/atomic" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/config" + "github.com/warpdotdev/oz-agent-worker/internal/types" + networkingv1 "k8s.io/api/networking/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + k8stypes "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/validation" + "k8s.io/client-go/dynamic" + "k8s.io/client-go/rest" +) + +const ( + networkRestrictedLabel = "warp-network-restricted" + networkPolicyIDLabel = "warp-network-policy-id" + networkManagedLabel = "warp-network-managed" + networkCanonicalAnnotation = "warp-network-policy-canonical" + networkConsumersAnnotation = "warp-network-policy-consumers" + networkPolicyUIDAnnotation = "warp-network-policy-uid" + networkRunAnnotation = "warp-run-id" + networkExecutionAnnotation = "warp-execution-id" + networkPolicyFinalizer = "warp.dev/network-policy-consumers" + networkPolicyPlaceholder = "POLICY_ID" +) + +var ciliumNetworkPolicyResource = schema.GroupVersionResource{ + Group: "cilium.io", Version: "v2", Resource: "ciliumnetworkpolicies", +} + +func (w *Worker) admitNetworkPolicy(policy *types.NetworkPolicy) (string, error) { + if err := validateNetworkPolicy(policy); err != nil { + return "network_policy_invalid", err + } + backend, ok := w.backend.(*KubernetesBackend) + if !ok || !backend.config.NetworkPolicy.Enabled { + return "network_policy_unsupported", fmt.Errorf("restricted executions require the Kubernetes backend with backend.kubernetes.network_policy.enabled: true and an operator-verified enforcing baseline") + } + return "", nil +} + +func validateNetworkPolicy(policy *types.NetworkPolicy) error { + if policy == nil { + return nil + } + for i, rule := range policy.Rules { + if rule.Action != "ALLOW" || len(rule.Hosts) == 0 { + return fmt.Errorf("network policy rule %d requires action ALLOW and nonempty hosts", i) + } + for _, host := range rule.Hosts { + if ip, err := netip.ParseAddr(host); err == nil { + if ip.Zone() == "" && ip.String() == host { + continue + } + return fmt.Errorf("network policy IP %q must be canonical and have no zone", host) + } + if host == "" || len(host) > 253 || host != strings.ToLower(host) { + return fmt.Errorf("network policy host %q must be a canonical literal hostname or IP", host) + } + for _, label := range strings.Split(host, ".") { + if len(label) == 0 || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' { + return fmt.Errorf("invalid network policy hostname %q", host) + } + for _, char := range label { + if (char < 'a' || char > 'z') && (char < '0' || char > '9') && char != '-' { + return fmt.Errorf("invalid network policy hostname %q", host) + } + } + } + } + } + return nil +} + +func effectiveNetworkPolicyDNS(cfg config.NetworkPolicyDNSConfig, inCluster bool, resolvConf []byte) (config.NetworkPolicyDNSConfig, error) { + if cfg.Namespace == "" { + cfg.Namespace = "kube-system" + } + if cfg.PodLabels == nil { + cfg.PodLabels = map[string]string{"k8s-app": "kube-dns"} + } + if len(validation.IsDNS1123Label(cfg.Namespace)) != 0 || len(cfg.PodLabels) == 0 { + return cfg, fmt.Errorf("network_policy.dns requires a valid namespace and nonempty pod_labels") + } + for key, value := range cfg.PodLabels { + if len(validation.IsQualifiedName(key)) != 0 || len(validation.IsValidLabelValue(value)) != 0 { + return cfg, fmt.Errorf("invalid network_policy.dns.pod_labels entry %q", key) + } + } + ips := cfg.ServerIPs + if len(ips) == 0 { + if !inCluster { + return cfg, fmt.Errorf("out-of-cluster workers must set backend.kubernetes.network_policy.dns.server_ips") + } + for line := range strings.SplitSeq(string(resolvConf), "\n") { + fields := strings.Fields(line) + if len(fields) >= 2 && fields[0] == "nameserver" { + ips = append(ips, fields[1]) + } + } + } + var usable []string + for _, value := range ips { + ip, err := netip.ParseAddr(value) + if err != nil || ip.Zone() != "" || ip.IsLoopback() || ip.IsUnspecified() { + return cfg, fmt.Errorf("unusable DNS resolver %q; set backend.kubernetes.network_policy.dns.server_ips to trusted task resolvers", value) + } + usable = append(usable, ip.Unmap().String()) + } + slices.Sort(usable) + cfg.ServerIPs = slices.Compact(usable) + if len(cfg.ServerIPs) == 0 { + return cfg, fmt.Errorf("no usable in-cluster DNS nameservers; set backend.kubernetes.network_policy.dns.server_ips") + } + return cfg, nil +} + +func (b *KubernetesBackend) networkPolicyDNS() (config.NetworkPolicyDNSConfig, error) { + cfg := b.config.NetworkPolicy.DNS + var data []byte + if len(cfg.ServerIPs) == 0 && b.networkInCluster { + var err error + data, err = os.ReadFile("/etc/resolv.conf") + if err != nil { + return cfg, fmt.Errorf("read cluster DNS: %w; set backend.kubernetes.network_policy.dns.server_ips", err) + } + } + return effectiveNetworkPolicyDNS(cfg, b.networkInCluster, data) +} + +type networkHostResolver interface { + LookupNetIP(context.Context, string, string) ([]netip.Addr, error) +} + +func networkDNSResolver(ips []string) *net.Resolver { + var next atomic.Uint64 + return &net.Resolver{PreferGo: true, StrictErrors: true, Dial: func(ctx context.Context, network, _ string) (net.Conn, error) { + ip := ips[(next.Add(1)-1)%uint64(len(ips))] + dialer := net.Dialer{Timeout: 5 * time.Second} + return dialer.DialContext(ctx, network, net.JoinHostPort(ip, "53")) + }} +} + +type networkPolicyProvider struct { + backend *KubernetesBackend + cilium bool +} + +func (b *KubernetesBackend) discoverNetworkPolicyProvider() (networkPolicyProvider, error) { + provider := networkPolicyProvider{backend: b} + if b.networkDiscovery == nil { + return provider, fmt.Errorf("network policy discovery client is unavailable") + } + resources, err := b.networkDiscovery.ServerResourcesForGroupVersion("cilium.io/v2") + if apierrors.IsNotFound(err) { + return provider, nil + } + if err != nil { + return provider, fmt.Errorf("discover Cilium policies (native fallback requires definite API absence): %w", err) + } + if resources == nil { + return provider, fmt.Errorf("cilium discovery returned no resource list") + } + for _, resource := range resources.APIResources { + if resource.Name == ciliumNetworkPolicyResource.Resource { + provider.cilium = true + if b.networkDynamic == nil { + return provider, fmt.Errorf("cilium dynamic client is unavailable") + } + break + } + } + return provider, nil +} + +func compileNetworkPolicy(ctx context.Context, policy *types.NetworkPolicy, cilium bool, dns config.NetworkPolicyDNSConfig, resolver networkHostResolver) (*unstructured.Unstructured, error) { + if policy == nil { + return nil, fmt.Errorf("restricted launch requires a present policy") + } + if err := validateNetworkPolicy(policy); err != nil { + return nil, err + } + var names, cidrs []string + for _, rule := range policy.Rules { + for _, host := range rule.Hosts { + if ip, err := netip.ParseAddr(host); err == nil { + cidrs = append(cidrs, exactNetworkCIDR(ip)) + } else { + names = append(names, host) + } + } + } + slices.Sort(names) + names = slices.Compact(names) + if len(names) > 0 { + var err error + dns, err = effectiveNetworkPolicyDNS(dns, false, nil) + if err != nil { + return nil, err + } + } + if !cilium { + for _, name := range names { + if resolver == nil { + return nil, fmt.Errorf("no resolver for network policy host %s", name) + } + lookupCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + ips, err := resolver.LookupNetIP(lookupCtx, "ip", name) + cancel() + if err != nil || len(ips) == 0 { + return nil, fmt.Errorf("resolve network policy host %s: no complete DNS answer (%v)", name, err) + } + for _, ip := range ips { + if !ip.IsValid() || ip.Zone() != "" { + return nil, fmt.Errorf("invalid DNS answer for %s", name) + } + cidrs = append(cidrs, exactNetworkCIDR(ip)) + } + } + } + slices.Sort(cidrs) + cidrs = slices.Compact(cidrs) + selector := map[string]any{"matchLabels": map[string]any{ + networkRestrictedLabel: "true", networkPolicyIDLabel: networkPolicyPlaceholder, + }} + spec := map[string]any{} + var egress []any + kind, apiVersion := "NetworkPolicy", "networking.k8s.io/v1" + if cilium { + kind, apiVersion = "CiliumNetworkPolicy", "cilium.io/v2" + spec["endpointSelector"] = selector + spec["enableDefaultDeny"] = map[string]any{"egress": true} + if len(cidrs) > 0 { + egress = append(egress, map[string]any{"toCIDR": stringAnySlice(cidrs)}) + } + if len(names) > 0 { + var fqdns, queries []any + for _, name := range names { + fqdns = append(fqdns, map[string]any{"matchName": name}) + queries = append(queries, map[string]any{"matchName": name}) + } + egress = append(egress, map[string]any{"toFQDNs": fqdns}) + dnsLabels := map[string]any{"k8s:io.kubernetes.pod.namespace": dns.Namespace} + for key, value := range dns.PodLabels { + dnsLabels["k8s:"+key] = value + } + egress = append(egress, map[string]any{ + "toEndpoints": []any{map[string]any{"matchLabels": dnsLabels}}, + "toPorts": []any{map[string]any{ + "ports": []any{map[string]any{"port": "53", "protocol": "UDP"}, map[string]any{"port": "53", "protocol": "TCP"}}, + "rules": map[string]any{"dns": queries}, + }}, + }) + } + } else { + spec["podSelector"] = selector + spec["policyTypes"] = []any{"Egress"} + if len(cidrs) > 0 { + var peers []any + for _, cidr := range cidrs { + peers = append(peers, map[string]any{"ipBlock": map[string]any{"cidr": cidr}}) + } + egress = append(egress, map[string]any{"to": peers}) + } + if len(names) > 0 { + var peers []any + for _, ip := range dns.ServerIPs { + addr, _ := netip.ParseAddr(ip) + peers = append(peers, map[string]any{"ipBlock": map[string]any{"cidr": exactNetworkCIDR(addr)}}) + } + podLabels := make(map[string]any, len(dns.PodLabels)) + for key, value := range dns.PodLabels { + podLabels[key] = value + } + peers = append(peers, map[string]any{ + "namespaceSelector": map[string]any{"matchLabels": map[string]any{"kubernetes.io/metadata.name": dns.Namespace}}, + "podSelector": map[string]any{"matchLabels": podLabels}, + }) + egress = append(egress, map[string]any{"to": peers, "ports": []any{ + map[string]any{"port": int64(53), "protocol": "UDP"}, + map[string]any{"port": int64(53), "protocol": "TCP"}, + }}) + } + } + if cilium && egress == nil { + egress = []any{} + } + if len(egress) > 0 || cilium { + spec["egress"] = egress + } + object := &unstructured.Unstructured{Object: map[string]any{"apiVersion": apiVersion, "kind": kind, "spec": spec}} + canonical, err := canonicalNetworkSpec(object, networkPolicyPlaceholder) + if err != nil { + return nil, err + } + sum := sha256.Sum256(canonical) + id := strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(sum[:])) + selector["matchLabels"].(map[string]any)[networkPolicyIDLabel] = id + object.SetName("warp-np-" + id) + object.SetLabels(map[string]string{networkManagedLabel: "true", networkPolicyIDLabel: id}) + object.SetAnnotations(map[string]string{networkCanonicalAnnotation: string(canonical)}) + object.SetFinalizers([]string{networkPolicyFinalizer}) + return object, nil +} + +func exactNetworkCIDR(ip netip.Addr) string { + ip = ip.Unmap() + return netip.PrefixFrom(ip, ip.BitLen()).String() +} + +func stringAnySlice(values []string) []any { + result := make([]any, len(values)) + for i, value := range values { + result[i] = value + } + return result +} + +func canonicalNetworkSpec(object *unstructured.Unstructured, expectedID string) ([]byte, error) { + spec, found, err := unstructured.NestedMap(object.Object, "spec") + if err != nil || !found { + return nil, fmt.Errorf("policy spec is absent or malformed") + } + selectorKey := "podSelector" + if object.GetKind() == "CiliumNetworkPolicy" { + selectorKey = "endpointSelector" + } + id, _, err := unstructured.NestedString(spec, selectorKey, "matchLabels", networkPolicyIDLabel) + if err != nil || id != expectedID { + return nil, fmt.Errorf("policy generated selector does not match %q", expectedID) + } + if err := unstructured.SetNestedField(spec, networkPolicyPlaceholder, selectorKey, "matchLabels", networkPolicyIDLabel); err != nil { + return nil, err + } + if object.GetKind() == "NetworkPolicy" { + var typed networkingv1.NetworkPolicySpec + if err := runtime.DefaultUnstructuredConverter.FromUnstructured(spec, &typed); err != nil { + return nil, err + } + encoded, err := json.Marshal(typed) + if err != nil { + return nil, err + } + if err := json.Unmarshal(encoded, &spec); err != nil { + return nil, err + } + } + return json.Marshal(canonicalNetworkSets(spec)) +} + +func canonicalNetworkSets(value any) any { + switch value := value.(type) { + case map[string]any: + for key, item := range value { + value[key] = canonicalNetworkSets(item) + } + return value + case []any: + items := make(map[string]any, len(value)) + var keys []string + for _, item := range value { + item = canonicalNetworkSets(item) + key, _ := json.Marshal(item) + if _, exists := items[string(key)]; !exists { + keys = append(keys, string(key)) + items[string(key)] = item + } + } + slices.Sort(keys) + result := make([]any, 0, len(keys)) + for _, key := range keys { + result = append(result, items[key]) + } + return result + default: + return value + } +} + +func (p networkPolicyProvider) ciliumClient() dynamic.ResourceInterface { + return p.backend.networkDynamic.Resource(ciliumNetworkPolicyResource).Namespace(p.backend.config.Namespace) +} + +func nativePolicyObject(policy *networkingv1.NetworkPolicy) (*unstructured.Unstructured, error) { + object, err := runtime.DefaultUnstructuredConverter.ToUnstructured(policy) + if err != nil { + return nil, err + } + result := &unstructured.Unstructured{Object: object} + result.SetKind("NetworkPolicy") + return result, nil +} + +func (p networkPolicyProvider) get(ctx context.Context, name string) (*unstructured.Unstructured, error) { + if p.cilium { + return p.ciliumClient().Get(ctx, name, metav1.GetOptions{}) + } + object, err := p.backend.clientset.NetworkingV1().NetworkPolicies(p.backend.config.Namespace).Get(ctx, name, metav1.GetOptions{}) + if err != nil { + return nil, err + } + return nativePolicyObject(object) +} + +func (p networkPolicyProvider) create(ctx context.Context, object *unstructured.Unstructured) error { + body, err := json.Marshal(object) + if err != nil { + return err + } + return p.writeRequest("POST").SetHeader("Content-Type", "application/json"). + Body(body).Do(ctx).Error() +} + +func (p networkPolicyProvider) writeRequest(method string) *rest.Request { + // Replayed writes can hide a committed reservation behind AlreadyExists or a + // resourceVersion conflict. + request := p.backend.clientset.NetworkingV1().RESTClient().Verb(method). + SetHeader("Accept", "application/json").MaxRetries(0) + if p.cilium { + return request.AbsPath("apis", ciliumNetworkPolicyResource.Group, ciliumNetworkPolicyResource.Version). + Namespace(p.backend.config.Namespace).Resource(ciliumNetworkPolicyResource.Resource) + } + return request.Namespace(p.backend.config.Namespace).Resource("networkpolicies") +} + +func (p networkPolicyProvider) patchMetadata(ctx context.Context, object *unstructured.Unstructured) error { + patch, err := json.Marshal(map[string]any{"metadata": map[string]any{ + "uid": object.GetUID(), "resourceVersion": object.GetResourceVersion(), + "annotations": object.GetAnnotations(), "finalizers": object.GetFinalizers(), + }}) + if err != nil { + return err + } + return p.writeRequest("PATCH").Name(object.GetName()). + SetHeader("Content-Type", string(k8stypes.MergePatchType)).Body(patch).Do(ctx).Error() +} + +func (p networkPolicyProvider) delete(ctx context.Context, object *unstructured.Unstructured) error { + uid, rv := object.GetUID(), object.GetResourceVersion() + options := metav1.DeleteOptions{Preconditions: &metav1.Preconditions{UID: &uid, ResourceVersion: &rv}} + if p.cilium { + return p.ciliumClient().Delete(ctx, object.GetName(), options) + } + return p.backend.clientset.NetworkingV1().NetworkPolicies(p.backend.config.Namespace).Delete(ctx, object.GetName(), options) +} + +func (p networkPolicyProvider) list(ctx context.Context) ([]unstructured.Unstructured, error) { + options := metav1.ListOptions{LabelSelector: networkManagedLabel + "=true"} + if p.cilium { + list, err := p.ciliumClient().List(ctx, options) + if err != nil { + return nil, err + } + return list.Items, nil + } + list, err := p.backend.clientset.NetworkingV1().NetworkPolicies(p.backend.config.Namespace).List(ctx, options) + if err != nil { + return nil, err + } + var result []unstructured.Unstructured + for i := range list.Items { + object, err := nativePolicyObject(&list.Items[i]) + if err != nil { + return nil, err + } + result = append(result, *object) + } + return result, nil +} diff --git a/internal/worker/kubernetes_network_policy_integration_test.go b/internal/worker/kubernetes_network_policy_integration_test.go new file mode 100644 index 0000000..3ff51b0 --- /dev/null +++ b/internal/worker/kubernetes_network_policy_integration_test.go @@ -0,0 +1,388 @@ +//go:build integration + +package worker + +import ( + "context" + "fmt" + "net/netip" + "os" + "os/exec" + "reflect" + "strings" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/config" + "github.com/warpdotdev/oz-agent-worker/internal/types" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + networkingv1 "k8s.io/api/networking/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/util/wait" + "k8s.io/client-go/dynamic" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/tools/clientcmd" +) + +type liveNetworkFixture struct { + backend *KubernetesBackend + kubeconfig string + provider networkPolicyProvider + ctx context.Context + allowedIP, forbiddenIP, allowedHost, forbiddenHost string +} + +func newLiveNetworkFixture(t *testing.T) *liveNetworkFixture { + t.Helper() + kubeconfig := os.Getenv("OZ_K8S_TEST_KUBECONFIG") + if kubeconfig == "" { + t.Skip("set OZ_K8S_TEST_KUBECONFIG to an explicit disposable-cluster kubeconfig") + } + if _, err := exec.LookPath("kubectl"); err != nil { + t.Fatal("kubectl is required for live container probes") + } + required := []string{"OZ_K8S_TEST_CNI_VERSION", "OZ_K8S_TEST_ALLOWED_IP", "OZ_K8S_TEST_FORBIDDEN_IP", "OZ_K8S_TEST_ALLOWED_HOST", "OZ_K8S_TEST_FORBIDDEN_HOST", "OZ_K8S_TEST_DNS_IPS"} + for _, name := range required { + if os.Getenv(name) == "" { + t.Skipf("set %s for the controlled external endpoint/DNS fixture; API acceptance is not CNI proof", name) + } + } + cfg, err := clientcmd.BuildConfigFromFlags("", kubeconfig) + if err != nil { + t.Fatal(err) + } + cfg.Timeout = 30 * time.Second + client, err := kubernetes.NewForConfig(cfg) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + t.Cleanup(cancel) + namespace, err := client.CoreV1().Namespaces().Create(ctx, &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{GenerateName: "oz-network-test-"}}, metav1.CreateOptions{}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := client.CoreV1().Namespaces().Delete(cleanupCtx, namespace.Name, metav1.DeleteOptions{}); err != nil { + t.Errorf("delete disposable namespace: %v", err) + } + }) + dynamicClient, err := dynamic.NewForConfig(cfg) + if err != nil { + t.Fatal(err) + } + backend := &KubernetesBackend{clientset: client, networkDiscovery: client.Discovery(), networkDynamic: dynamicClient, + config: KubernetesBackendConfig{Namespace: namespace.Name, NoCleanup: true, NetworkPolicy: config.NetworkPolicyConfig{Enabled: true, + DNS: config.NetworkPolicyDNSConfig{ServerIPs: strings.Split(os.Getenv("OZ_K8S_TEST_DNS_IPS"), ",")}}}} + provider, err := backend.discoverNetworkPolicyProvider() + if err != nil { + t.Fatal(err) + } + wantProvider := os.Getenv("OZ_K8S_TEST_PROVIDER") + if wantProvider != "native" && wantProvider != "cilium" { + t.Fatal("set OZ_K8S_TEST_PROVIDER to native or cilium") + } + if provider.cilium != (wantProvider == "cilium") { + t.Fatalf("selected provider=%v, expected %s", provider.cilium, wantProvider) + } + f := &liveNetworkFixture{backend: backend, kubeconfig: kubeconfig, provider: provider, ctx: ctx, + allowedIP: os.Getenv("OZ_K8S_TEST_ALLOWED_IP"), forbiddenIP: os.Getenv("OZ_K8S_TEST_FORBIDDEN_IP"), + allowedHost: os.Getenv("OZ_K8S_TEST_ALLOWED_HOST"), forbiddenHost: os.Getenv("OZ_K8S_TEST_FORBIDDEN_HOST")} + for _, ip := range []string{f.allowedIP, f.forbiddenIP} { + if _, err := netip.ParseAddr(ip); err != nil { + t.Fatal(err) + } + } + if f.allowedIP == f.forbiddenIP { + t.Fatal("controlled endpoints must have distinct addresses") + } + if err := validateNetworkPolicy(allowNetworkPolicy(f.allowedHost, f.forbiddenHost)); err != nil { + t.Fatal(err) + } + version, err := client.Discovery().ServerVersion() + if err != nil { + t.Fatal(err) + } + t.Logf("Kubernetes=%s CNI=%s provider=%s allowed=%s forbidden=%s", version, os.Getenv("OZ_K8S_TEST_CNI_VERSION"), wantProvider, f.allowedIP, f.forbiddenIP) + _, err = client.NetworkingV1().NetworkPolicies(namespace.Name).Create(ctx, &networkingv1.NetworkPolicy{ + ObjectMeta: metav1.ObjectMeta{Name: "restricted-baseline"}, + Spec: networkingv1.NetworkPolicySpec{PodSelector: metav1.LabelSelector{MatchLabels: map[string]string{networkRestrictedLabel: "true"}}, PolicyTypes: []networkingv1.PolicyType{networkingv1.PolicyTypeEgress}, Egress: []networkingv1.NetworkPolicyEgressRule{}}, + }, metav1.CreateOptions{}) + if err != nil { + t.Fatal(err) + } + return f +} + +func (f *liveNetworkFixture) probeScript(allow bool) string { + fetch := func(host string) string { return fmt.Sprintf("wget -q -T 3 -O /dev/null 'http://%s:8080/'", host) } + var script []string + if allow { + script = append(script, fetch(f.allowedHost)+" || exit 12", fetch(f.allowedIP)+" || exit 11") + } else { + script = append(script, fetch(f.allowedIP)+" && exit 13") + } + script = append(script, fetch(f.forbiddenIP)+" && exit 14", fetch(f.forbiddenHost)+" && exit 15", "exit 0") + return strings.Join(script, "\n") +} + +func (f *liveNetworkFixture) job(run string, probe string, hold bool) *batchv1.Job { + job := testRestrictedJob(run, "execution") + job.Namespace = f.backend.config.Namespace + command := probe + if hold { + command = strings.TrimSuffix(probe, "exit 0") + "while [ ! -f /tmp/finish ]; do sleep 1; done" + } + job.Spec.Template.Spec.InitContainers = []corev1.Container{{Name: "setup", Image: "busybox:1.36", Command: []string{"sh", "-c", probe}}} + job.Spec.Template.Spec.Containers[0].Command = []string{"sh", "-c", command} + return job +} + +func (f *liveNetworkFixture) awaitPod(t *testing.T, job *batchv1.Job) *corev1.Pod { + t.Helper() + var result *corev1.Pod + err := wait.PollUntilContextTimeout(f.ctx, 250*time.Millisecond, 2*time.Minute, true, func(ctx context.Context) (bool, error) { + pods, err := f.backend.clientset.CoreV1().Pods(f.backend.config.Namespace).List(ctx, metav1.ListOptions{LabelSelector: batchv1.ControllerUidLabel + "=" + string(job.UID)}) + if err != nil { + return false, err + } + for i := range pods.Items { + pod := &pods.Items[i] + if pod.Status.Phase == corev1.PodFailed { + return false, fmt.Errorf("probe failed: Pod %s status %+v", pod.Name, pod.Status) + } + if podReady(pod) || pod.Status.Phase == corev1.PodSucceeded { + result = pod + return true, nil + } + } + return false, nil + }) + if err != nil { + t.Fatal(err) + } + return result +} + +func (f *liveNetworkFixture) exec(t *testing.T, pod *corev1.Pod, script string) { + t.Helper() + ctx, cancel := context.WithTimeout(f.ctx, 30*time.Second) + defer cancel() + output, err := exec.CommandContext(ctx, "kubectl", "--kubeconfig", f.kubeconfig, "--namespace", f.backend.config.Namespace, "exec", pod.Name, "-c", "task", "--", "sh", "-c", script).CombinedOutput() + if err != nil { + t.Fatalf("probe exec: %v %s", err, output) + } +} + +func (f *liveNetworkFixture) awaitComplete(t *testing.T, job *batchv1.Job) { + t.Helper() + if err := wait.PollUntilContextTimeout(f.ctx, 250*time.Millisecond, 2*time.Minute, true, func(ctx context.Context) (bool, error) { + current, err := f.backend.getTaskJob(ctx, job.Name) + if err != nil { + return false, err + } + if jobFailed(current) { + return false, fmt.Errorf("probe Job failed: %+v", current.Status) + } + return jobComplete(current), nil + }); err != nil { + t.Fatal(err) + } +} + +func TestKubernetesNetworkPolicyWithCNI(t *testing.T) { + f := newLiveNetworkFixture(t) + // Establish both external endpoints are reachable without restrictions. + control := f.job("control", fmt.Sprintf("wget -q -T 3 -O /dev/null http://%s:8080/ && wget -q -T 3 -O /dev/null http://%s:8080/", f.allowedIP, f.forbiddenIP), false) + control, err := f.backend.clientset.BatchV1().Jobs(f.backend.config.Namespace).Create(f.ctx, control, metav1.CreateOptions{}) + if err != nil { + t.Fatal(err) + } + f.awaitComplete(t, control) + if control.Labels[networkRestrictedLabel] != "" { + t.Fatal("unrestricted control acquired reserved selector") + } + baseline := f.job("baseline", f.probeScript(false), false) + baseline.Spec.Template.Labels = map[string]string{networkRestrictedLabel: "true"} + baseline, err = f.backend.clientset.BatchV1().Jobs(f.backend.config.Namespace).Create(f.ctx, baseline, metav1.CreateOptions{}) + if err != nil { + t.Fatal(err) + } + f.awaitComplete(t, baseline) + emptyProbe := strings.TrimSuffix(f.probeScript(false), "exit 0") + fmt.Sprintf("timeout 4 nslookup '%s' '%s' && exit 16\nexit 0", f.allowedHost, strings.Split(os.Getenv("OZ_K8S_TEST_DNS_IPS"), ",")[0]) + empty, err := f.backend.createRestrictedJob(f.ctx, &TaskParams{TaskID: "empty", ExecutionID: "execution", NetworkPolicy: &types.NetworkPolicy{}}, f.job("empty", emptyProbe, false)) + if err != nil { + t.Fatal(err) + } + f.awaitComplete(t, empty) + emptyPolicies, err := f.provider.list(f.ctx) + if err != nil || len(emptyPolicies) != 1 { + t.Fatalf("empty policy: %v %v", emptyPolicies, err) + } + if err := f.provider.reconcile(f.ctx, &emptyPolicies[0]); err != nil { + t.Fatal(err) + } + policy := allowNetworkPolicy(f.allowedHost) + var jobs []*batchv1.Job + var pods []*corev1.Pod + for _, run := range []string{"one", "two"} { + job, err := f.backend.createRestrictedJob(f.ctx, &TaskParams{TaskID: run, ExecutionID: "execution", NetworkPolicy: policy}, f.job(run, f.probeScript(true), true)) + if err != nil { + t.Fatal(err) + } + jobs = append(jobs, job) + pods = append(pods, f.awaitPod(t, job)) + } + objects, err := f.provider.list(f.ctx) + if err != nil || len(objects) != 1 { + t.Fatalf("expected shared policy: %v %v", objects, err) + } + original := objects[0].DeepCopy() + _, ledger := requireNetworkLedger(t, f.provider, original.GetName()) + if len(ledger.Consumers) != 2 { + t.Fatalf("consumers=%+v", ledger) + } + for i, job := range jobs { + f.exec(t, pods[i], "touch /tmp/finish") + f.awaitComplete(t, job) + if err := f.provider.reconcile(f.ctx, original); err != nil { + t.Fatal(err) + } + if i == 0 { + f.exec(t, pods[1], f.probeScript(true)) + _, remaining := requireNetworkLedger(t, f.provider, original.GetName()) + if len(remaining.Consumers) != 1 { + t.Fatal("released active consumer") + } + } + } + if _, err := f.provider.get(f.ctx, original.GetName()); !apierrors.IsNotFound(err) { + t.Fatalf("policy retained after stopped executions: %v", err) + } + for _, job := range jobs { + if _, err := f.backend.getTaskJob(f.ctx, job.Name); err != nil { + t.Fatalf("NoCleanup lost terminal Job: %v", err) + } + } + t.Log("shared enforcement, setup/main probes, retained Jobs and zero-consumer deletion passed") +} + +func TestKubernetesNetworkPolicyTerminatingWithCNI(t *testing.T) { + f := newLiveNetworkFixture(t) + job, err := f.backend.createRestrictedJob(f.ctx, &TaskParams{TaskID: "terminating", ExecutionID: "execution", NetworkPolicy: allowNetworkPolicy(f.allowedHost)}, f.job("terminating", f.probeScript(true), true)) + if err != nil { + t.Fatal(err) + } + pod := f.awaitPod(t, job) + objects, err := f.provider.list(f.ctx) + if err != nil || len(objects) != 1 { + t.Fatalf("policy: %v %v", objects, err) + } + policy := objects[0].DeepCopy() + pending := testNetworkConsumer("delayed", "execution") + intended, err := compileNetworkPolicy(f.ctx, allowNetworkPolicy(f.allowedHost), f.provider.cilium, f.backend.config.NetworkPolicy.DNS, networkDNSResolver(f.backend.config.NetworkPolicy.DNS.ServerIPs)) + if err != nil { + t.Fatal(err) + } + if _, err := f.provider.register(f.ctx, intended, pending); err != nil { + t.Fatal(err) + } + policy, err = f.provider.get(f.ctx, policy.GetName()) + if err != nil { + t.Fatal(err) + } + if err := f.provider.delete(f.ctx, policy); err != nil { + t.Fatal(err) + } + f.exec(t, pod, f.probeScript(true)) + if err := f.provider.reconcile(f.ctx, policy); err != nil { + t.Fatal(err) + } + delayed := f.job("delayed", f.probeScript(true), true) + applyNetworkExecutionMetadata(delayed, pending, policy) + applyNetworkExecutionMetadata(&delayed.Spec.Template, pending, policy) + delayed, err = f.backend.clientset.BatchV1().Jobs(f.backend.config.Namespace).Create(f.ctx, delayed, metav1.CreateOptions{}) + if err != nil { + t.Fatal(err) + } + delayedPod := f.awaitPod(t, delayed) + if err := f.provider.bind(f.ctx, policy, pending, delayed.UID); err != nil { + t.Fatal(err) + } + current, err := f.provider.get(f.ctx, policy.GetName()) + if err != nil || current.GetUID() != policy.GetUID() || current.GetDeletionTimestamp() == nil { + t.Fatalf("terminating UID lost: %v %v", current, err) + } + want, _, _ := unstructured.NestedMap(policy.Object, "spec") + got, _, _ := unstructured.NestedMap(current.Object, "spec") + if !reflect.DeepEqual(want, got) { + t.Fatal("terminating policy spec changed") + } + cniNS, selector := os.Getenv("OZ_K8S_TEST_CNI_NAMESPACE"), os.Getenv("OZ_K8S_TEST_CNI_POD_SELECTOR") + if cniNS == "" || selector == "" { + t.Fatal("CNI restart proof requires OZ_K8S_TEST_CNI_NAMESPACE and OZ_K8S_TEST_CNI_POD_SELECTOR") + } + agents, err := f.backend.clientset.CoreV1().Pods(cniNS).List(f.ctx, metav1.ListOptions{LabelSelector: selector}) + if err != nil || len(agents.Items) == 0 { + t.Fatalf("CNI agents: %v %v", agents, err) + } + for _, agent := range agents.Items { + uid := agent.UID + if err := f.backend.clientset.CoreV1().Pods(cniNS).Delete(f.ctx, agent.Name, metav1.DeleteOptions{Preconditions: &metav1.Preconditions{UID: &uid}}); err != nil { + t.Fatal(err) + } + } + f.exec(t, pod, f.probeScript(true)) + if err := wait.PollUntilContextTimeout(f.ctx, time.Second, 3*time.Minute, true, func(ctx context.Context) (bool, error) { + current, err := f.backend.clientset.CoreV1().Pods(cniNS).List(ctx, metav1.ListOptions{LabelSelector: selector}) + if err != nil { + return false, err + } + if len(current.Items) != len(agents.Items) { + return false, nil + } + for _, agent := range current.Items { + for _, old := range agents.Items { + if old.UID == agent.UID { + return false, nil + } + } + if !podReady(&agent) { + return false, nil + } + } + return true, nil + }); err != nil { + t.Fatal(err) + } + f.exec(t, pod, f.probeScript(true)) + f.exec(t, pod, "touch /tmp/finish") + f.awaitComplete(t, job) + if err := f.provider.reconcile(f.ctx, policy); err != nil { + t.Fatal(err) + } + f.exec(t, delayedPod, f.probeScript(true)) + f.exec(t, delayedPod, "touch /tmp/finish") + f.awaitComplete(t, delayed) + if err := f.provider.reconcile(f.ctx, policy); err != nil { + t.Fatal(err) + } + if _, err := f.provider.get(f.ctx, policy.GetName()); !apierrors.IsNotFound(err) { + t.Fatalf("terminating policy not finalized: %v", err) + } + t.Log("terminating UID retained enforcement before/during/after CNI restart and fresh API reconciliation") +} + +func podReady(pod *corev1.Pod) bool { + for _, condition := range pod.Status.Conditions { + if condition.Type == corev1.PodReady && condition.Status == corev1.ConditionTrue { + return true + } + } + return false +} diff --git a/internal/worker/kubernetes_network_policy_lifecycle.go b/internal/worker/kubernetes_network_policy_lifecycle.go new file mode 100644 index 0000000..bb27c67 --- /dev/null +++ b/internal/worker/kubernetes_network_policy_lifecycle.go @@ -0,0 +1,639 @@ +package worker + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base32" + "encoding/json" + "fmt" + "net/netip" + "slices" + "strings" + "time" + + "github.com/google/uuid" + "github.com/warpdotdev/oz-agent-worker/internal/log" + batchv1 "k8s.io/api/batch/v1" + 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/apis/meta/v1/unstructured" + k8stypes "k8s.io/apimachinery/pkg/types" +) + +var networkLauncherIncarnation = uuid.NewString() + +type networkPolicyConsumer struct { + RunID string `json:"run_id"` + ExecutionID string `json:"execution_id"` + JobName string `json:"job_name"` + Launcher string `json:"launcher"` + JobUID k8stypes.UID `json:"job_uid,omitempty"` +} + +type networkPolicyLedger struct { + Version int `json:"version"` + Consumers []networkPolicyConsumer `json:"consumers"` +} + +func readNetworkLedger(object *unstructured.Unstructured) (networkPolicyLedger, error) { + var ledger networkPolicyLedger + raw := object.GetAnnotations()[networkConsumersAnnotation] + if err := json.Unmarshal([]byte(raw), &ledger); err != nil { + return ledger, fmt.Errorf("policy %s UID %s has a malformed consumer ledger: %w", object.GetName(), object.GetUID(), err) + } + if ledger.Version != 1 { + return ledger, fmt.Errorf("policy %s UID %s has unsupported consumer ledger version %d", object.GetName(), object.GetUID(), ledger.Version) + } + seen := make(map[[2]string]bool) + for _, entry := range ledger.Consumers { + key := [2]string{entry.RunID, entry.ExecutionID} + if entry.RunID == "" || entry.ExecutionID == "" || entry.Launcher == "" || + entry.JobName != kubernetesTaskJobName(entry.RunID, entry.ExecutionID) || seen[key] { + return ledger, fmt.Errorf("policy %s UID %s has an invalid/duplicate consumer %q/%q", object.GetName(), object.GetUID(), entry.RunID, entry.ExecutionID) + } + seen[key] = true + } + return ledger, nil +} + +func writeNetworkLedger(object *unstructured.Unstructured, ledger networkPolicyLedger) error { + data, err := json.Marshal(ledger) + if err != nil { + return err + } + annotations := object.GetAnnotations() + if annotations == nil { + annotations = make(map[string]string) + } + annotations[networkConsumersAnnotation] = string(data) + size := 0 + for key, value := range annotations { + size += len(key) + len(value) + } + if size >= 256*1024 { + return fmt.Errorf("network policy annotations exceed the 256 KiB limit; wait for consumers to stop") + } + object.SetAnnotations(annotations) + return nil +} + +func validateManagedNetworkPolicy(object *unstructured.Unstructured) (networkPolicyLedger, error) { + id := object.GetLabels()[networkPolicyIDLabel] + if object.GetUID() == "" || object.GetResourceVersion() == "" || + object.GetLabels()[networkManagedLabel] != "true" || len(object.GetOwnerReferences()) != 0 || + !slices.Contains(object.GetFinalizers(), networkPolicyFinalizer) { + return networkPolicyLedger{}, fmt.Errorf("policy %s UID %s lacks managed identity, immutable ownership, or consumer finalizer", object.GetName(), object.GetUID()) + } + canonical, err := canonicalNetworkSpec(object, id) + if err != nil { + return networkPolicyLedger{}, err + } + sum := sha256.Sum256(canonical) + digest := strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(sum[:])) + if id != digest || object.GetName() != "warp-np-"+digest { + return networkPolicyLedger{}, fmt.Errorf("policy %s UID %s live enforcement content does not match its fingerprint", object.GetName(), object.GetUID()) + } + return readNetworkLedger(object) +} + +func validateMatchingNetworkPolicy(object, intended *unstructured.Unstructured) (networkPolicyLedger, error) { + ledger, err := validateManagedNetworkPolicy(object) + if err != nil { + return ledger, err + } + id := intended.GetLabels()[networkPolicyIDLabel] + actual, err := canonicalNetworkSpec(object, id) + if err != nil { + return ledger, err + } + expected, err := canonicalNetworkSpec(intended, id) + if err != nil { + return ledger, err + } + if !bytes.Equal(actual, expected) { + return ledger, fmt.Errorf("network policy %s conflicts with intended enforcement content", object.GetName()) + } + return ledger, nil +} + +func pauseNetworkRetry(ctx context.Context) error { + timer := time.NewTimer(time.Second) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func (p networkPolicyProvider) register(ctx context.Context, intended *unstructured.Unstructured, entry networkPolicyConsumer) (*unstructured.Unstructured, error) { + intended.SetNamespace(p.backend.config.Namespace) + if err := writeNetworkLedger(intended, networkPolicyLedger{Version: 1, Consumers: []networkPolicyConsumer{entry}}); err != nil { + return nil, err + } + for ctx.Err() == nil { + object, err := p.get(ctx, intended.GetName()) + insertedByCreate := false + if apierrors.IsNotFound(err) { + createErr := p.create(ctx, intended) + // Recover the original create, including an accepted request whose response was lost. + object, err = p.get(ctx, intended.GetName()) + if err != nil { + return nil, fmt.Errorf("cannot establish policy creation result (%v): %w", createErr, err) + } + insertedByCreate = createErr == nil || (!apierrors.IsAlreadyExists(createErr) && isTransientKubernetesAPIError(createErr)) + } + if err != nil { + if isTransientKubernetesAPIError(err) && pauseNetworkRetry(ctx) == nil { + continue + } + return nil, err + } + ledger, err := validateMatchingNetworkPolicy(object, intended) + if err != nil { + return nil, err + } + if insertedByCreate { + for _, existing := range ledger.Consumers { + if existing == entry { + return object, nil + } + } + return nil, fmt.Errorf("created policy lacks this launcher's exact reservation") + } + if object.GetDeletionTimestamp() != nil { + if err := pauseNetworkRetry(ctx); err != nil { + return nil, fmt.Errorf("policy %s UID %s is terminating; existing consumers must stop before reuse: %w", object.GetName(), object.GetUID(), err) + } + continue + } + for _, existing := range ledger.Consumers { + if existing.RunID == entry.RunID && existing.ExecutionID == entry.ExecutionID { + return nil, fmt.Errorf("execution %s/%s already has a durable launch registration; it cannot authorize another Job POST", entry.RunID, entry.ExecutionID) + } + } + ledger.Consumers = append(ledger.Consumers, entry) + if err := writeNetworkLedger(object, ledger); err != nil { + return nil, err + } + err = p.patchMetadata(ctx, object) + if err == nil { + return object, nil + } + if !apierrors.IsConflict(err) && isTransientKubernetesAPIError(err) { + recovered, readErr := p.get(ctx, intended.GetName()) + if readErr != nil { + return nil, fmt.Errorf("cannot recover registration response: %w", readErr) + } + recoveredLedger, readErr := validateMatchingNetworkPolicy(recovered, intended) + if readErr != nil || recovered.GetUID() != object.GetUID() { + return nil, fmt.Errorf("policy changed during registration recovery: %v", readErr) + } + for _, existing := range recoveredLedger.Consumers { + if existing == entry { + return recovered, nil + } + } + // An unresolved patch may still succeed; do not issue another launch. + return nil, fmt.Errorf("registration outcome is ambiguous for %s/%s: %w", entry.RunID, entry.ExecutionID, err) + } + if !apierrors.IsConflict(err) { + return nil, err + } + if err := pauseNetworkRetry(ctx); err != nil { + return nil, err + } + } + return nil, ctx.Err() +} + +func networkExecutionMatches(meta metav1.Object, entry networkPolicyConsumer, policy *unstructured.Unstructured) bool { + annotations, labels := meta.GetAnnotations(), meta.GetLabels() + return annotations[networkRunAnnotation] == entry.RunID && + annotations[networkExecutionAnnotation] == entry.ExecutionID && + annotations[networkPolicyUIDAnnotation] == string(policy.GetUID()) && + labels[networkRestrictedLabel] == "true" && + labels[networkPolicyIDLabel] == policy.GetLabels()[networkPolicyIDLabel] +} + +func (p networkPolicyProvider) bind(ctx context.Context, policy *unstructured.Unstructured, entry networkPolicyConsumer, jobUID k8stypes.UID) error { + for ctx.Err() == nil { + object, err := p.get(ctx, policy.GetName()) + if err != nil { + return err + } + if object.GetUID() != policy.GetUID() { + return fmt.Errorf("policy UID changed before Job binding") + } + ledger, err := validateMatchingNetworkPolicy(object, policy) + if err != nil { + return err + } + found := false + for i, existing := range ledger.Consumers { + if existing.RunID == entry.RunID && existing.ExecutionID == entry.ExecutionID { + if existing.Launcher != entry.Launcher || (existing.JobUID != "" && existing.JobUID != jobUID) { + return fmt.Errorf("consumer identity changed before Job binding") + } + found = true + ledger.Consumers[i].JobUID = jobUID + } + } + if !found || jobUID == "" { + return fmt.Errorf("consumer registration or Job UID missing at binding") + } + if err := writeNetworkLedger(object, ledger); err != nil { + return err + } + err = p.patchMetadata(ctx, object) + if err == nil { + return nil + } + if !apierrors.IsConflict(err) && !isTransientKubernetesAPIError(err) { + return err + } + if err := pauseNetworkRetry(ctx); err != nil { + return err + } + } + return ctx.Err() +} + +func (b *KubernetesBackend) createRestrictedJob(ctx context.Context, params *TaskParams, job *batchv1.Job) (*batchv1.Job, error) { + ctx, cancel := context.WithTimeout(ctx, 2*time.Minute) + defer cancel() + key := [2]string{params.TaskID, taskExecutionID(params)} + for { + if err := ctx.Err(); err != nil { + return nil, err + } + b.networkLaunchMu.Lock() + pending, exists := b.networkLaunches[key] + if !exists { + if b.networkLaunches == nil { + b.networkLaunches = make(map[[2]string]chan struct{}) + } + pending = make(chan struct{}) + b.networkLaunches[key] = pending + b.networkLaunchMu.Unlock() + defer func() { + b.networkLaunchMu.Lock() + delete(b.networkLaunches, key) + close(pending) + b.networkLaunchMu.Unlock() + }() + break + } + b.networkLaunchMu.Unlock() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-pending: + } + } + if !b.config.NetworkPolicy.Enabled { + return nil, fmt.Errorf("backend.kubernetes.network_policy.enabled is false") + } + if err := validateNetworkPolicy(params.NetworkPolicy); err != nil { + return nil, err + } + provider, err := b.discoverNetworkPolicyProvider() + if err != nil { + return nil, err + } + dns := b.config.NetworkPolicy.DNS + needsDNS := false + for _, rule := range params.NetworkPolicy.Rules { + for _, host := range rule.Hosts { + if _, err := netip.ParseAddr(host); err != nil { + needsDNS = true + } + } + } + var resolver networkHostResolver + if needsDNS { + dns, err = b.networkPolicyDNS() + if err != nil { + return nil, err + } + resolver = networkDNSResolver(dns.ServerIPs) + } + intended, err := compileNetworkPolicy(ctx, params.NetworkPolicy, provider.cilium, dns, resolver) + if err != nil { + return nil, err + } + entry := networkPolicyConsumer{RunID: params.TaskID, ExecutionID: taskExecutionID(params), JobName: job.Name, Launcher: networkLauncherIncarnation} + policy, err := provider.register(ctx, intended, entry) + if err != nil { + return nil, fmt.Errorf("register restricted execution %s/%s: %w", entry.RunID, entry.ExecutionID, err) + } + applyNetworkExecutionMetadata(job, entry, policy) + applyNetworkExecutionMetadata(&job.Spec.Template, entry, policy) + if err := ctx.Err(); err != nil { + b.releaseRejectedNetworkLaunch(ctx, provider, policy, entry) + return nil, err + } + // A timeout is not permission to repeat this POST or discard its reservation. + created := &batchv1.Job{} + createErr := b.clientset.BatchV1().RESTClient().Post().Namespace(b.config.Namespace). + Resource("jobs").Body(job).MaxRetries(0).Do(ctx).Into(created) + if createErr != nil { + if apierrors.IsInvalid(createErr) || apierrors.IsBadRequest(createErr) || apierrors.IsForbidden(createErr) || apierrors.IsUnauthorized(createErr) { + b.releaseRejectedNetworkLaunch(ctx, provider, policy, entry) + return nil, fmt.Errorf("job launch rejected before admission: %w", createErr) + } + created, err = b.getTaskJob(ctx, job.Name) + if err != nil { + return nil, fmt.Errorf("job launch outcome unresolved; reservation retained (POST: %v; GET: %w)", createErr, err) + } + } + if created == nil || created.Name != entry.JobName || !networkExecutionMatches(created, entry, policy) { + return nil, fmt.Errorf("job identity conflicts with restricted execution; reservation retained") + } + if !networkExecutionMatches(&created.Spec.Template, entry, policy) { + stopCtx, stopCancel := context.WithTimeout(context.WithoutCancel(ctx), kubernetesCleanupTimeout) + defer stopCancel() + err := fmt.Errorf("job Pod template conflicts with restricted execution") + stopErr := b.requestJobExecutionStop(stopCtx, created, err) + return nil, fmt.Errorf("%w (stop request: %v); reservation retained", err, stopErr) + } + if err := provider.bind(ctx, policy, entry, created.UID); err != nil { + stopCtx, stopCancel := context.WithTimeout(context.WithoutCancel(ctx), kubernetesCleanupTimeout) + defer stopCancel() + stopErr := b.requestJobExecutionStop(stopCtx, created, err) + return nil, fmt.Errorf("bind restricted Job UID: %w (stop request: %v); reservation retained", err, stopErr) + } + return created, nil +} + +func applyNetworkExecutionMetadata(meta metav1.Object, entry networkPolicyConsumer, policy *unstructured.Unstructured) { + labels := meta.GetLabels() + if labels == nil { + labels = make(map[string]string) + } + labels[networkRestrictedLabel] = "true" + labels[networkPolicyIDLabel] = policy.GetLabels()[networkPolicyIDLabel] + meta.SetLabels(labels) + annotations := meta.GetAnnotations() + if annotations == nil { + annotations = make(map[string]string) + } + annotations[networkRunAnnotation] = entry.RunID + annotations[networkExecutionAnnotation] = entry.ExecutionID + annotations[networkPolicyUIDAnnotation] = string(policy.GetUID()) + meta.SetAnnotations(annotations) +} + +func (b *KubernetesBackend) releaseRejectedNetworkLaunch(ctx context.Context, provider networkPolicyProvider, policy *unstructured.Unstructured, entry networkPolicyConsumer) { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), kubernetesCleanupTimeout) + defer cancel() + _, err := retryKubernetesAPI(cleanupCtx, kubernetesAPIBackoff(), func() (struct{}, error) { + object, err := provider.get(cleanupCtx, policy.GetName()) + if err != nil { + return struct{}{}, err + } + if object.GetUID() != policy.GetUID() { + return struct{}{}, fmt.Errorf("policy UID changed while releasing a rejected launch") + } + ledger, err := validateMatchingNetworkPolicy(object, policy) + if err != nil { + return struct{}{}, err + } + ledger.Consumers = slices.DeleteFunc(ledger.Consumers, func(existing networkPolicyConsumer) bool { return existing == entry }) + if err := writeNetworkLedger(object, ledger); err != nil { + return struct{}{}, err + } + return struct{}{}, provider.patchMetadata(cleanupCtx, object) + }) + if err != nil && !apierrors.IsNotFound(err) { + log.Errorf(ctx, "Retaining rejected launch %s/%s on policy UID %s after cleanup error: %v", entry.RunID, entry.ExecutionID, policy.GetUID(), err) + } + b.notifyNetworkPolicyReconciler() +} + +func networkPodStopped(pod *corev1.Pod) bool { + if pod.Status.Phase != corev1.PodSucceeded && pod.Status.Phase != corev1.PodFailed { + return false + } + statuses := make(map[string]corev1.ContainerStatus) + for _, status := range append(append(slices.Clone(pod.Status.ContainerStatuses), pod.Status.InitContainerStatuses...), pod.Status.EphemeralContainerStatuses...) { + if status.State.Terminated == nil || status.State.Running != nil || status.State.Waiting != nil { + return false + } + statuses[status.Name] = status + } + for _, container := range append(slices.Clone(pod.Spec.Containers), pod.Spec.InitContainers...) { + if _, ok := statuses[container.Name]; !ok { + return false + } + } + for _, container := range pod.Spec.EphemeralContainers { + if _, ok := statuses[container.Name]; !ok { + return false + } + } + return len(pod.Spec.Containers) > 0 +} + +func (b *KubernetesBackend) networkConsumerStopped(ctx context.Context, policy *unstructured.Unstructured, entry networkPolicyConsumer) (bool, error) { + job, err := b.getTaskJob(ctx, entry.JobName) + if err != nil && !apierrors.IsNotFound(err) { + return false, err + } + if err == nil { + if job.UID != entry.JobUID || !networkExecutionMatches(job, entry, policy) || !networkExecutionMatches(&job.Spec.Template, entry, policy) { + return false, fmt.Errorf("job %s conflicts with consumer %s/%s UID %s", entry.JobName, entry.RunID, entry.ExecutionID, entry.JobUID) + } + if !jobComplete(job) && !jobFailed(job) { + return false, nil + } + } + pods, err := b.clientset.CoreV1().Pods(b.config.Namespace).List(ctx, metav1.ListOptions{ + LabelSelector: batchv1.ControllerUidLabel + "=" + string(entry.JobUID), + }) + if err != nil { + return false, err + } + for i := range pods.Items { + pod := &pods.Items[i] + owner := metav1.GetControllerOf(pod) + if owner == nil || owner.UID != entry.JobUID || owner.Kind != "Job" || + !networkExecutionMatches(pod, entry, policy) || !networkPodStopped(pod) { + return false, nil + } + } + return true, nil +} + +func (b *KubernetesBackend) networkMatchingObjectsStopped(ctx context.Context, policy *unstructured.Unstructured) (bool, error) { + selector := networkRestrictedLabel + "=true," + networkPolicyIDLabel + "=" + policy.GetLabels()[networkPolicyIDLabel] + jobs, err := b.clientset.BatchV1().Jobs(b.config.Namespace).List(ctx, metav1.ListOptions{LabelSelector: selector}) + if err != nil { + return false, err + } + for i := range jobs.Items { + job := &jobs.Items[i] + entry := networkPolicyConsumer{RunID: job.Annotations[networkRunAnnotation], ExecutionID: job.Annotations[networkExecutionAnnotation]} + priorPolicy := policy.DeepCopy() + priorPolicy.SetUID(k8stypes.UID(job.Annotations[networkPolicyUIDAnnotation])) + if entry.RunID == "" || entry.ExecutionID == "" || job.Name != kubernetesTaskJobName(entry.RunID, entry.ExecutionID) || + job.UID == "" || priorPolicy.GetUID() == "" || + !networkExecutionMatches(job, entry, priorPolicy) || !networkExecutionMatches(&job.Spec.Template, entry, priorPolicy) || + (!jobComplete(job) && !jobFailed(job)) { + return false, fmt.Errorf("unexplained or active matching Job %s blocks policy %s UID %s cleanup", job.Name, policy.GetName(), policy.GetUID()) + } + } + pods, err := b.clientset.CoreV1().Pods(b.config.Namespace).List(ctx, metav1.ListOptions{LabelSelector: selector}) + if err != nil { + return false, err + } + for i := range pods.Items { + pod := &pods.Items[i] + entry := networkPolicyConsumer{RunID: pod.Annotations[networkRunAnnotation], ExecutionID: pod.Annotations[networkExecutionAnnotation]} + owner := metav1.GetControllerOf(pod) + if entry.RunID == "" || entry.ExecutionID == "" || pod.Annotations[networkPolicyUIDAnnotation] == "" || + owner == nil || owner.Kind != "Job" || owner.UID == "" || owner.Name != kubernetesTaskJobName(entry.RunID, entry.ExecutionID) || + pod.Labels[batchv1.ControllerUidLabel] != string(owner.UID) || + !networkPodStopped(pod) { + return false, fmt.Errorf("unexplained or active matching Pod %s blocks policy %s UID %s cleanup", pod.Name, policy.GetName(), policy.GetUID()) + } + } + return true, nil +} + +func (p networkPolicyProvider) reconcile(ctx context.Context, observed *unstructured.Unstructured) error { + for ctx.Err() == nil { + object, err := p.get(ctx, observed.GetName()) + if apierrors.IsNotFound(err) { + return nil + } + if err != nil { + return err + } + if object.GetUID() != observed.GetUID() { + return fmt.Errorf("policy %s UID was replaced during reconciliation", observed.GetName()) + } + ledger, err := validateManagedNetworkPolicy(object) + if err != nil { + return err + } + changed := false + for i, entry := range ledger.Consumers { + if entry.JobUID == "" { + job, getErr := p.backend.getTaskJob(ctx, entry.JobName) + if apierrors.IsNotFound(getErr) { + continue + } + if getErr != nil { + return getErr + } + if !networkExecutionMatches(job, entry, object) || !networkExecutionMatches(&job.Spec.Template, entry, object) || job.UID == "" { + return fmt.Errorf("unbound consumer %s/%s conflicts with Job %s; reservation retained", entry.RunID, entry.ExecutionID, entry.JobName) + } + ledger.Consumers[i].JobUID = job.UID + changed = true + break + } + stopped, err := p.backend.networkConsumerStopped(ctx, object, entry) + if err != nil { + return err + } + if stopped { + ledger.Consumers = slices.Delete(ledger.Consumers, i, i+1) + changed = true + break + } + } + if changed { + if err := writeNetworkLedger(object, ledger); err != nil { + return err + } + err = p.patchMetadata(ctx, object) + } else if len(ledger.Consumers) > 0 { + return nil + } else { + stopped, scanErr := p.backend.networkMatchingObjectsStopped(ctx, object) + if scanErr != nil { + return scanErr + } + if !stopped { + return nil + } + if object.GetDeletionTimestamp() == nil { + err = p.delete(ctx, object) + } else { + // Terminating policies cannot accept registrations; recheck the ledger + // and objects on this fresh read before removing only our finalizer. + object.SetFinalizers(slices.DeleteFunc(object.GetFinalizers(), func(value string) bool { return value == networkPolicyFinalizer })) + err = p.patchMetadata(ctx, object) + if err == nil { + return nil + } + } + } + if apierrors.IsNotFound(err) { + return nil + } + if err != nil && !apierrors.IsConflict(err) && !isTransientKubernetesAPIError(err) { + return err + } + if err != nil { + if err := pauseNetworkRetry(ctx); err != nil { + return err + } + } + } + return ctx.Err() +} + +func (b *KubernetesBackend) reconcileNetworkPolicies(ctx context.Context) { + for _, cilium := range []bool{false, true} { + if cilium && b.networkDynamic == nil { + continue + } + provider := networkPolicyProvider{backend: b, cilium: cilium} + objects, err := provider.list(ctx) + if cilium && apierrors.IsNotFound(err) { + continue + } + if err != nil { + log.Errorf(ctx, "Cannot list managed network policies (cilium=%t); consumers retained: %v", cilium, err) + continue + } + for i := range objects { + object := &objects[i] + _, err := retryKubernetesAPI(ctx, kubernetesAPIBackoff(), func() (struct{}, error) { + return struct{}{}, provider.reconcile(ctx, object) + }) + if err != nil { + log.Errorf(ctx, "Network policy %s UID %s cleanup retained consumers %s: %v", object.GetName(), object.GetUID(), object.GetAnnotations()[networkConsumersAnnotation], err) + } + } + } +} + +func (b *KubernetesBackend) runNetworkPolicyReconciler(ctx context.Context) { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), kubernetesCleanupTimeout) + b.reconcileNetworkPolicies(cleanupCtx) + cancel() + select { + case <-ctx.Done(): + return + case <-ticker.C: + case <-b.networkReconcile: + } + } +} + +func (b *KubernetesBackend) notifyNetworkPolicyReconciler() { + select { + case b.networkReconcile <- struct{}{}: + default: + } +} diff --git a/internal/worker/kubernetes_network_policy_lifecycle_test.go b/internal/worker/kubernetes_network_policy_lifecycle_test.go new file mode 100644 index 0000000..f7c1b67 --- /dev/null +++ b/internal/worker/kubernetes_network_policy_lifecycle_test.go @@ -0,0 +1,670 @@ +package worker + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "reflect" + "slices" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/warpdotdev/oz-agent-worker/internal/config" + "github.com/warpdotdev/oz-agent-worker/internal/types" + jsonpatch "gopkg.in/evanphx/json-patch.v4" + batchv1 "k8s.io/api/batch/v1" + 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/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/runtime" + k8stypes "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/dynamic" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/rest" +) + +type networkPolicyAPIFixture struct { + mu sync.Mutex + objects map[string]*unstructured.Unstructured + version int + jobPosts int + policyCreates int + policyPatches int + beforeJobPost func() + loseJobResponse bool + losePolicyResponse bool + losePatchResponse bool + denyPolicy bool + cilium bool + server *httptest.Server +} + +func TestKubernetesNetworkPolicyRegistrationRecovery(t *testing.T) { + for _, cilium := range []bool{false, true} { + for _, write := range []string{"create", "registration"} { + t.Run(fmt.Sprintf("cilium=%v/%s", cilium, write), func(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + fixture.cilium = cilium + provider := networkPolicyProvider{backend: backend, cilium: cilium} + intended := compileTestNetworkPolicy(t, &types.NetworkPolicy{}, cilium, nil) + fixture.losePolicyResponse = write == "create" + first, err := provider.register(context.Background(), intended.DeepCopy(), testNetworkConsumer("one", "execution")) + if err != nil { + t.Fatal(err) + } + fixture.losePatchResponse = write == "registration" + second, err := provider.register(context.Background(), intended.DeepCopy(), testNetworkConsumer("two", "execution")) + if err != nil || second.GetUID() != first.GetUID() { + t.Fatalf("recover registration: %v %v", second, err) + } + _, ledger := requireNetworkLedger(t, provider, first.GetName()) + if len(ledger.Consumers) != 2 || fixture.jobPosts != 0 || fixture.policyCreates != 1 || fixture.policyPatches != 1 { + t.Fatalf("registrations=%+v creates=%d patches=%d jobs=%d", ledger, fixture.policyCreates, fixture.policyPatches, fixture.jobPosts) + } + other := testNetworkConsumer("one", "execution") + other.Launcher = "replacement-process" + if _, err := provider.register(context.Background(), intended.DeepCopy(), other); err == nil { + t.Fatal("replacement launcher reused an ambiguous reservation") + } + resource := "networkpolicies" + if cilium { + resource = "ciliumnetworkpolicies" + } + fixture.mu.Lock() + fixture.objects[resource+"/"+first.GetName()].SetUID("replacement-policy") + fixture.mu.Unlock() + if err := provider.bind(context.Background(), first, testNetworkConsumer("one", "execution"), "job-uid"); err == nil { + t.Fatal("bound against replacement policy UID") + } + }) + } + } +} + +func TestKubernetesNetworkPolicyUnsafeLedger(t *testing.T) { + for _, kind := range []string{"unknown schema", "annotation exhaustion"} { + t.Run(kind, func(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + provider := networkPolicyProvider{backend: backend} + intended := compileTestNetworkPolicy(t, &types.NetworkPolicy{}, false, nil) + policy, err := provider.register(context.Background(), intended.DeepCopy(), testNetworkConsumer("one", "execution")) + if err != nil { + t.Fatal(err) + } + fixture.mu.Lock() + stored := fixture.objects["networkpolicies/"+policy.GetName()] + annotations := stored.GetAnnotations() + if kind == "unknown schema" { + annotations[networkConsumersAnnotation] = `{"version":99,"consumers":[]}` + } else { + size := 0 + for key, value := range annotations { + size += len(key) + len(value) + } + annotations["filler"] = strings.Repeat("x", 256*1024-size-len("filler")-100) + } + stored.SetAnnotations(annotations) + fixture.mu.Unlock() + _, err = backend.createRestrictedJob(context.Background(), &TaskParams{TaskID: "two", ExecutionID: "execution", NetworkPolicy: &types.NetworkPolicy{}}, testRestrictedJob("two", "execution")) + if err == nil || fixture.jobPosts != 0 { + t.Fatalf("unsafe ledger launched: %v POSTs=%d", err, fixture.jobPosts) + } + if err := provider.reconcile(context.Background(), policy); kind == "unknown schema" && err == nil { + t.Fatal("collector accepted unknown schema") + } + current, err := provider.get(context.Background(), policy.GetName()) + if err != nil || !slices.Contains(current.GetFinalizers(), networkPolicyFinalizer) { + t.Fatalf("unsafe ledger lost protection: %v", err) + } + }) + } +} + +func TestKubernetesNetworkPolicyUnexplainedObjects(t *testing.T) { + for _, conflict := range []string{"live object without registration", "reused Job name"} { + t.Run(conflict, func(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + params := &TaskParams{TaskID: "run", ExecutionID: "execution", NetworkPolicy: &types.NetworkPolicy{}} + job, err := backend.createRestrictedJob(context.Background(), params, testRestrictedJob("run", "execution")) + if err != nil { + t.Fatal(err) + } + provider := networkPolicyProvider{backend: backend} + intended := compileTestNetworkPolicy(t, params.NetworkPolicy, false, nil) + policy, ledger := requireNetworkLedger(t, provider, intended.GetName()) + if conflict == "live object without registration" { + ledger.Consumers = nil + if err := writeNetworkLedger(policy, ledger); err != nil { + t.Fatal(err) + } + if err := provider.patchMetadata(context.Background(), policy); err != nil { + t.Fatal(err) + } + } else { + fixture.mu.Lock() + fixture.objects["jobs/"+job.Name].SetUID("other-job") + fixture.mu.Unlock() + } + if err := provider.reconcile(context.Background(), policy); err == nil { + t.Fatal("unexplained consumer did not block cleanup") + } + requireNetworkLedger(t, provider, policy.GetName()) + }) + } +} + +func newNetworkPolicyAPIFixture(t *testing.T) (*networkPolicyAPIFixture, *KubernetesBackend) { + t.Helper() + fixture := &networkPolicyAPIFixture{objects: make(map[string]*unstructured.Unstructured)} + fixture.server = httptest.NewServer(http.HandlerFunc(fixture.serve)) + t.Cleanup(fixture.server.Close) + cfg := &rest.Config{Host: fixture.server.URL, Timeout: time.Second * 5, ContentConfig: rest.ContentConfig{ContentType: "application/json", AcceptContentTypes: "application/json"}} + client, err := kubernetes.NewForConfig(cfg) + if err != nil { + t.Fatal(err) + } + backend := &KubernetesBackend{config: KubernetesBackendConfig{Namespace: "agents", NoCleanup: true, NetworkPolicy: config.NetworkPolicyConfig{Enabled: true}}, + clientset: client, networkDiscovery: client.Discovery()} + backend.networkDynamic, err = dynamic.NewForConfig(cfg) + if err != nil { + t.Fatal(err) + } + return fixture, backend +} + +func fixtureStatus(w http.ResponseWriter, status int, reason metav1.StatusReason) { + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(metav1.Status{TypeMeta: metav1.TypeMeta{APIVersion: "v1", Kind: "Status"}, Status: "Failure", Code: int32(status), Reason: reason, Message: string(reason)}) +} + +func (f *networkPolicyAPIFixture) serve(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/apis/cilium.io/v2" { + if f.cilium { + _ = json.NewEncoder(w).Encode(metav1.APIResourceList{GroupVersion: "cilium.io/v2", APIResources: []metav1.APIResource{{Name: "ciliumnetworkpolicies", Namespaced: true}}}) + return + } + fixtureStatus(w, 404, metav1.StatusReasonNotFound) + return + } + parts := strings.Split(strings.Trim(r.URL.Path, "/"), "/") + resourceIndex := slices.Index(parts, "namespaces") + 2 + if resourceIndex < 2 || resourceIndex >= len(parts) { + fixtureStatus(w, 404, metav1.StatusReasonNotFound) + return + } + resource, name := parts[resourceIndex], "" + if len(parts) > resourceIndex+1 { + name = parts[resourceIndex+1] + } + body, _ := io.ReadAll(r.Body) + if r.Method == http.MethodPost && resource == "jobs" && f.beforeJobPost != nil { + f.beforeJobPost() + } + f.mu.Lock() + defer f.mu.Unlock() + key := resource + "/" + name + current := f.objects[key] + switch r.Method { + case http.MethodGet: + if name == "" { + selector, err := labels.Parse(r.URL.Query().Get("labelSelector")) + if err != nil { + fixtureStatus(w, 400, metav1.StatusReasonBadRequest) + return + } + items := []any{} + for objectKey, object := range f.objects { + if strings.HasPrefix(objectKey, resource+"/") && selector.Matches(labels.Set(object.GetLabels())) { + items = append(items, object.Object) + } + } + kind := map[string]string{"pods": "PodList", "jobs": "JobList", "networkpolicies": "NetworkPolicyList", "ciliumnetworkpolicies": "CiliumNetworkPolicyList"}[resource] + apiVersion := "v1" + if resource == "jobs" { + apiVersion = "batch/v1" + } + if resource == "networkpolicies" { + apiVersion = "networking.k8s.io/v1" + } + if resource == "ciliumnetworkpolicies" { + apiVersion = "cilium.io/v2" + } + _ = json.NewEncoder(w).Encode(map[string]any{"apiVersion": apiVersion, "kind": kind, "items": items}) + return + } + if current == nil { + fixtureStatus(w, 404, metav1.StatusReasonNotFound) + return + } + _ = json.NewEncoder(w).Encode(current) + case http.MethodPost: + if resource == "networkpolicies" || resource == "ciliumnetworkpolicies" { + f.policyCreates++ + if f.denyPolicy { + fixtureStatus(w, 403, metav1.StatusReasonForbidden) + return + } + } + if resource == "jobs" { + f.jobPosts++ + } + object := &unstructured.Unstructured{} + if err := json.Unmarshal(body, &object.Object); err != nil { + fixtureStatus(w, 400, metav1.StatusReasonBadRequest) + return + } + key = resource + "/" + object.GetName() + if f.objects[key] != nil { + fixtureStatus(w, 409, metav1.StatusReasonAlreadyExists) + return + } + f.version++ + object.SetUID(k8stypes.UID(uuid.NewString())) + object.SetResourceVersion(strconv.Itoa(f.version)) + f.objects[key] = object + if resource != "jobs" && f.losePolicyResponse { + f.losePolicyResponse = false + w.Header().Set("Retry-After", "0") + fixtureStatus(w, 504, metav1.StatusReasonTimeout) + return + } + if resource == "jobs" && f.loseJobResponse { + w.Header().Set("Retry-After", "0") + fixtureStatus(w, 504, metav1.StatusReasonTimeout) + return + } + w.WriteHeader(201) + _ = json.NewEncoder(w).Encode(object) + case http.MethodPatch: + f.policyPatches++ + var patch unstructured.Unstructured + _ = json.Unmarshal(body, &patch.Object) + if current == nil { + fixtureStatus(w, 404, metav1.StatusReasonNotFound) + return + } + if patch.GetUID() != current.GetUID() || patch.GetResourceVersion() != current.GetResourceVersion() { + fixtureStatus(w, 409, metav1.StatusReasonConflict) + return + } + raw, _ := json.Marshal(current) + merged, err := jsonpatch.MergePatch(raw, body) + if err != nil { + fixtureStatus(w, 400, metav1.StatusReasonBadRequest) + return + } + object := &unstructured.Unstructured{} + _ = json.Unmarshal(merged, &object.Object) + f.version++ + object.SetResourceVersion(strconv.Itoa(f.version)) + f.objects[key] = object + if object.GetDeletionTimestamp() != nil && len(object.GetFinalizers()) == 0 { + delete(f.objects, key) + } + if f.losePatchResponse { + f.losePatchResponse = false + w.Header().Set("Retry-After", "0") + fixtureStatus(w, 504, metav1.StatusReasonTimeout) + return + } + _ = json.NewEncoder(w).Encode(object) + case http.MethodDelete: + if current == nil { + fixtureStatus(w, 404, metav1.StatusReasonNotFound) + return + } + var options metav1.DeleteOptions + _ = json.Unmarshal(body, &options) + if options.Preconditions == nil || options.Preconditions.UID == nil || *options.Preconditions.UID != current.GetUID() || + (resource != "jobs" && (options.Preconditions.ResourceVersion == nil || *options.Preconditions.ResourceVersion != current.GetResourceVersion())) { + fixtureStatus(w, 409, metav1.StatusReasonConflict) + return + } + f.version++ + if len(current.GetFinalizers()) == 0 { + delete(f.objects, key) + } else { + now := metav1.Now() + current.SetDeletionTimestamp(&now) + current.SetResourceVersion(strconv.Itoa(f.version)) + } + _ = json.NewEncoder(w).Encode(metav1.Status{Status: "Success"}) + } +} + +func testRestrictedJob(run, execution string) *batchv1.Job { + zero := int32(0) + return &batchv1.Job{ObjectMeta: metav1.ObjectMeta{Name: kubernetesTaskJobName(run, execution), Namespace: "agents"}, + Spec: batchv1.JobSpec{BackoffLimit: &zero, Template: corev1.PodTemplateSpec{Spec: corev1.PodSpec{ + RestartPolicy: corev1.RestartPolicyNever, Containers: []corev1.Container{{Name: "task", Image: "busybox:1.36"}}, + }}}} +} + +func testNetworkConsumer(run, execution string) networkPolicyConsumer { + return networkPolicyConsumer{RunID: run, ExecutionID: execution, JobName: kubernetesTaskJobName(run, execution), Launcher: networkLauncherIncarnation} +} + +func requireNetworkLedger(t *testing.T, provider networkPolicyProvider, name string) (*unstructured.Unstructured, networkPolicyLedger) { + t.Helper() + object, err := provider.get(context.Background(), name) + if err != nil { + t.Fatal(err) + } + ledger, err := validateManagedNetworkPolicy(object) + if err != nil { + t.Fatal(err) + } + return object, ledger +} + +func TestKubernetesRestrictedLaunchOrdering(t *testing.T) { + for _, ambiguous := range []bool{false, true} { + t.Run(fmt.Sprint("ambiguous=", ambiguous), func(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + fixture.loseJobResponse = ambiguous + provider := networkPolicyProvider{backend: backend} + intended := compileTestNetworkPolicy(t, &types.NetworkPolicy{}, false, nil) + fixture.beforeJobPost = func() { + object, ledger := requireNetworkLedger(t, provider, intended.GetName()) + if len(ledger.Consumers) != 1 || ledger.Consumers[0].JobUID != "" || object.GetDeletionTimestamp() != nil { + t.Error("Job POST preceded durable protected registration") + } + } + params := &TaskParams{TaskID: "run", ExecutionID: "execution", NetworkPolicy: &types.NetworkPolicy{}} + job, err := backend.createRestrictedJob(context.Background(), params, testRestrictedJob("run", "execution")) + if err != nil || job.UID == "" || (job.Spec.Suspend != nil && *job.Spec.Suspend) { + t.Fatalf("launch: %v %v", job, err) + } + _, ledger := requireNetworkLedger(t, provider, intended.GetName()) + if len(ledger.Consumers) != 1 || ledger.Consumers[0].JobUID != job.UID || fixture.jobPosts != 1 { + t.Fatalf("binding=%+v POSTs=%d", ledger, fixture.jobPosts) + } + if _, err := backend.createRestrictedJob(context.Background(), params, testRestrictedJob("run", "execution")); err == nil || fixture.jobPosts != 1 { + t.Fatal("duplicate registration authorized another POST") + } + }) + } + for _, cilium := range []bool{false, true} { + t.Run(fmt.Sprintf("policy denied cilium=%v", cilium), func(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + fixture.denyPolicy = true + fixture.cilium = cilium + _, err := backend.createRestrictedJob(context.Background(), &TaskParams{TaskID: "run", ExecutionID: "execution", NetworkPolicy: &types.NetworkPolicy{}}, testRestrictedJob("run", "execution")) + if err == nil || fixture.jobPosts != 0 || fixture.policyCreates != 1 { + t.Fatal("policy failure launched a Job") + } + }) + } +} + +func TestKubernetesNetworkPolicySharedReuse(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + provider := networkPolicyProvider{backend: backend} + intended := compileTestNetworkPolicy(t, allowNetworkPolicy("192.0.2.1"), false, nil) + var wg sync.WaitGroup + for _, run := range []string{"one", "two"} { + wg.Go(func() { + if _, err := provider.register(context.Background(), intended.DeepCopy(), testNetworkConsumer(run, "execution")); err != nil { + t.Error(err) + } + }) + } + wg.Wait() + object, ledger := requireNetworkLedger(t, provider, intended.GetName()) + if len(ledger.Consumers) != 2 || len(object.GetOwnerReferences()) != 0 { + t.Fatalf("shared ledger=%+v", ledger) + } + fixture.mu.Lock() + count := len(fixture.objects) + fixture.mu.Unlock() + if count != 1 { + t.Fatalf("created %d policies", count) + } + for _, mutation := range []string{"spec", "selector", "finalizer", "ledger"} { + t.Run(mutation, func(t *testing.T) { + modified := object.DeepCopy() + switch mutation { + case "spec": + _ = unstructured.SetNestedSlice(modified.Object, []any{map[string]any{}}, "spec", "egress") + case "selector": + _ = unstructured.SetNestedField(modified.Object, "wrong", "spec", "podSelector", "matchLabels", networkPolicyIDLabel) + case "finalizer": + modified.SetFinalizers(nil) + case "ledger": + annotations := modified.GetAnnotations() + annotations[networkConsumersAnnotation] = `{"version":99}` + modified.SetAnnotations(annotations) + } + if _, err := validateMatchingNetworkPolicy(modified, intended); err == nil { + t.Fatal("unsafe existing policy accepted") + } + }) + } +} + +func TestKubernetesNetworkPolicyRegistrationFinalizerRace(t *testing.T) { + _, backend := newNetworkPolicyAPIFixture(t) + provider := networkPolicyProvider{backend: backend} + intended := compileTestNetworkPolicy(t, &types.NetworkPolicy{}, false, nil) + first := testNetworkConsumer("one", "execution") + object, err := provider.register(context.Background(), intended.DeepCopy(), first) + if err != nil { + t.Fatal(err) + } + ledger, _ := readNetworkLedger(object) + ledger.Consumers = nil + if err := writeNetworkLedger(object, ledger); err != nil { + t.Fatal(err) + } + if err := provider.patchMetadata(context.Background(), object); err != nil { + t.Fatal(err) + } + empty, err := provider.get(context.Background(), object.GetName()) + if err != nil { + t.Fatal(err) + } + pending, err := provider.register(context.Background(), intended.DeepCopy(), testNetworkConsumer("two", "execution")) + if err != nil { + t.Fatal(err) + } + if err := provider.delete(context.Background(), empty); !apierrors.IsConflict(err) { + t.Fatalf("stale collector DELETE=%v", err) + } + latest, err := provider.get(context.Background(), pending.GetName()) + if err != nil { + t.Fatal(err) + } + if err := provider.delete(context.Background(), latest); err != nil { + t.Fatal(err) + } + if err := provider.reconcile(context.Background(), latest); err != nil { + t.Fatal(err) + } + terminating, ledger := requireNetworkLedger(t, provider, latest.GetName()) + if terminating.GetDeletionTimestamp() == nil || len(ledger.Consumers) != 1 { + t.Fatal("pending launch did not retain terminating policy") + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if _, err := provider.register(ctx, intended.DeepCopy(), testNetworkConsumer("three", "execution")); err == nil { + t.Fatal("joined terminating UID") + } + // The earlier registration still authorizes its one delayed Job POST. + job := testRestrictedJob("two", "execution") + applyNetworkExecutionMetadata(job, ledger.Consumers[0], terminating) + applyNetworkExecutionMetadata(&job.Spec.Template, ledger.Consumers[0], terminating) + created, err := backend.clientset.BatchV1().Jobs("agents").Create(context.Background(), job, metav1.CreateOptions{}) + if err != nil { + t.Fatal(err) + } + if err := provider.bind(context.Background(), terminating, ledger.Consumers[0], created.UID); err != nil { + t.Fatal(err) + } + if err := provider.reconcile(context.Background(), terminating); err != nil { + t.Fatal(err) + } + requireNetworkLedger(t, provider, terminating.GetName()) +} + +func TestKubernetesNetworkPolicyCleanupAndRecovery(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + provider := networkPolicyProvider{backend: backend} + intended := compileTestNetworkPolicy(t, &types.NetworkPolicy{}, false, nil) + entry := testNetworkConsumer("run", "execution") + policy, err := provider.register(context.Background(), intended.DeepCopy(), entry) + if err != nil { + t.Fatal(err) + } + if err := provider.reconcile(context.Background(), policy); err != nil { + t.Fatal(err) + } + _, ledger := requireNetworkLedger(t, provider, policy.GetName()) + if len(ledger.Consumers) != 1 || ledger.Consumers[0].JobUID != "" { + t.Fatal("missing Job expired reservation") + } + job := testRestrictedJob("run", "execution") + applyNetworkExecutionMetadata(job, entry, policy) + applyNetworkExecutionMetadata(&job.Spec.Template, entry, policy) + created, err := backend.clientset.BatchV1().Jobs("agents").Create(context.Background(), job, metav1.CreateOptions{}) + if err != nil { + t.Fatal(err) + } + if err := provider.reconcile(context.Background(), policy); err != nil { + t.Fatal(err) + } + _, ledger = requireNetworkLedger(t, provider, policy.GetName()) + if ledger.Consumers[0].JobUID != created.UID { + t.Fatal("replacement did not recover matching Job UID") + } + fixture.mu.Lock() + stored := fixture.objects["jobs/"+job.Name] + _ = unstructured.SetNestedSlice(stored.Object, []any{map[string]any{"type": "Complete", "status": "True"}}, "status", "conditions") + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "pod", Namespace: "agents", + Labels: map[string]string{batchv1.ControllerUidLabel: string(created.UID)}, + OwnerReferences: []metav1.OwnerReference{{Kind: "Job", Name: job.Name, UID: created.UID, Controller: ptrTrue()}}}, + Spec: job.Spec.Template.Spec, Status: corev1.PodStatus{Phase: corev1.PodRunning}} + applyNetworkExecutionMetadata(pod, entry, policy) + raw, _ := runtime.DefaultUnstructuredConverter.ToUnstructured(pod) + fixture.objects["pods/pod"] = &unstructured.Unstructured{Object: raw} + fixture.mu.Unlock() + if err := provider.reconcile(context.Background(), policy); err != nil { + t.Fatal(err) + } + requireNetworkLedger(t, provider, policy.GetName()) + fixture.mu.Lock() + pod.Status.Phase = corev1.PodSucceeded + pod.Status.ContainerStatuses = []corev1.ContainerStatus{{Name: "task", State: corev1.ContainerState{Terminated: &corev1.ContainerStateTerminated{}}}} + raw, _ = runtime.DefaultUnstructuredConverter.ToUnstructured(pod) + fixture.objects["pods/pod"] = &unstructured.Unstructured{Object: raw} + fixture.mu.Unlock() + if err := provider.reconcile(context.Background(), policy); err != nil { + t.Fatal(err) + } + if _, err := provider.get(context.Background(), policy.GetName()); !apierrors.IsNotFound(err) { + t.Fatalf("terminal retained Job/Pod kept policy: %v", err) + } + fixture.mu.Lock() + defer fixture.mu.Unlock() + if fixture.objects["pods/pod"] == nil || fixture.objects["jobs/"+job.Name] == nil { + t.Fatal("NoCleanup removed inspection objects") + } +} + +func ptrTrue() *bool { value := true; return &value } + +func TestKubernetesNetworkPolicyPodTermination(t *testing.T) { + pod := &corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "task"}}, InitContainers: []corev1.Container{{Name: "sidecar"}}, + EphemeralContainers: []corev1.EphemeralContainer{{EphemeralContainerCommon: corev1.EphemeralContainerCommon{Name: "debug"}}}}, + Status: corev1.PodStatus{Phase: corev1.PodSucceeded}} + terminated := func(name string) corev1.ContainerStatus { + return corev1.ContainerStatus{Name: name, State: corev1.ContainerState{Terminated: &corev1.ContainerStateTerminated{}}} + } + pod.Status.ContainerStatuses = []corev1.ContainerStatus{terminated("task")} + pod.Status.InitContainerStatuses = []corev1.ContainerStatus{terminated("sidecar")} + if networkPodStopped(pod) { + t.Fatal("missing ephemeral status counted stopped") + } + pod.Status.EphemeralContainerStatuses = []corev1.ContainerStatus{terminated("debug")} + if !networkPodStopped(pod) { + t.Fatal("all containers stopped") + } + pod.Status.InitContainerStatuses[0].State = corev1.ContainerState{Running: &corev1.ContainerStateRunning{}} + if networkPodStopped(pod) { + t.Fatal("running native sidecar counted stopped") + } + pod.Status.InitContainerStatuses[0] = terminated("sidecar") + pod.Status.Phase = corev1.PodUnknown + if networkPodStopped(pod) { + t.Fatal("Unknown phase counted stopped") + } +} + +func TestNetworkPolicyTemplateAndMetadata(t *testing.T) { + _, backend := newNetworkPolicyAPIFixture(t) + privileged := true + security := &corev1.SecurityContext{Privileged: &privileged, Capabilities: &corev1.Capabilities{Add: []corev1.Capability{"NET_ADMIN", "SYS_ADMIN"}}} + template := &corev1.PodSpec{HostNetwork: true, DNSPolicy: corev1.DNSNone, DNSConfig: &corev1.PodDNSConfig{Nameservers: []string{"192.0.2.53"}}, + Containers: []corev1.Container{{Name: "task", SecurityContext: security}}, + InitContainers: []corev1.Container{{Name: "operator-init", Image: "operator-image", SecurityContext: security}}} + backend.config.PodTemplate = template.DeepCopy() + spec := backend.buildTaskPodSpec(nil, nil, corev1.Container{Name: "task", Image: "task-image"}) + if !spec.HostNetwork || spec.DNSPolicy != template.DNSPolicy || !reflect.DeepEqual(spec.DNSConfig, template.DNSConfig) || + !reflect.DeepEqual(spec.Containers[0].SecurityContext, security) || !reflect.DeepEqual(spec.InitContainers, template.InitContainers) { + t.Fatalf("operator Pod settings lost: %+v", spec) + } + job := testRestrictedJob("run", "execution") + job.Spec.Template.Spec = spec + job.Annotations = map[string]string{networkRunAnnotation: "operator", networkPolicyUIDAnnotation: "operator", "other": "preserved", "warp-custom": "preserved"} + job.Spec.Template.Annotations = copyStringMap(job.Annotations) + created, err := backend.createRestrictedJob(context.Background(), &TaskParams{TaskID: "run", ExecutionID: "execution", NetworkPolicy: &types.NetworkPolicy{}}, job) + if err != nil { + t.Fatal(err) + } + for _, meta := range []metav1.Object{created, &created.Spec.Template} { + if meta.GetAnnotations()[networkRunAnnotation] != "run" || meta.GetAnnotations()[networkExecutionAnnotation] != "execution" || + meta.GetAnnotations()[networkPolicyUIDAnnotation] == "" || meta.GetAnnotations()[networkPolicyUIDAnnotation] == "operator" || + meta.GetAnnotations()["other"] != "preserved" || meta.GetAnnotations()["warp-custom"] != "preserved" || meta.GetLabels()[networkRestrictedLabel] != "true" { + t.Fatalf("metadata=%+v %+v", meta.GetLabels(), meta.GetAnnotations()) + } + } + if _, err := NewKubernetesBackend(context.Background(), KubernetesBackendConfig{ExtraLabels: map[string]string{"warp-network-restricted": "true"}}); err == nil { + t.Fatal("reserved operator labels accepted") + } +} + +func TestKubernetesRestrictedLaunchDuplicateCancellation(t *testing.T) { + fixture, backend := newNetworkPolicyAPIFixture(t) + entered, release := make(chan struct{}), make(chan struct{}) + fixture.beforeJobPost = func() { close(entered); <-release } + params := &TaskParams{TaskID: "run", ExecutionID: "execution", NetworkPolicy: &types.NetworkPolicy{}} + first := make(chan error, 1) + go func() { + _, err := backend.createRestrictedJob(context.Background(), params, testRestrictedJob("run", "execution")) + first <- err + }() + <-entered + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + _, err := backend.createRestrictedJob(ctx, params, testRestrictedJob("run", "execution")) + close(release) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("duplicate wait ignored cancellation: %v", err) + } + if err := <-first; err != nil { + t.Fatal(err) + } + if fixture.jobPosts != 1 { + t.Fatalf("duplicate launched %d requests", fixture.jobPosts) + } +} diff --git a/internal/worker/kubernetes_network_policy_test.go b/internal/worker/kubernetes_network_policy_test.go new file mode 100644 index 0000000..1e2cfdf --- /dev/null +++ b/internal/worker/kubernetes_network_policy_test.go @@ -0,0 +1,266 @@ +package worker + +import ( + "context" + "encoding/json" + "errors" + "net/netip" + "reflect" + "strings" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/config" + "github.com/warpdotdev/oz-agent-worker/internal/types" + "golang.org/x/sync/semaphore" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + discoveryfake "k8s.io/client-go/discovery/fake" + dynamicfake "k8s.io/client-go/dynamic/fake" + "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" +) + +func TestNetworkPolicyAssignmentJSON(t *testing.T) { + for _, field := range []string{"", `,"network_policy":null`, `,"network_policy":{}`, `,"network_policy":{"rules":[]}`, `,"network_policy":{"rules":null}`} { + var assignment types.TaskAssignmentMessage + if err := json.Unmarshal([]byte(`{"task_id":"run"`+field+`}`), &assignment); err != nil { + t.Fatal(err) + } + wantPresent := field != "" && field != `,"network_policy":null` + if (assignment.NetworkPolicy != nil) != wantPresent { + t.Fatalf("policy presence for %s = %v", field, assignment.NetworkPolicy) + } + encoded, err := json.Marshal(assignment) + if err != nil || strings.Contains(string(encoded), `"network_policy"`) != wantPresent { + t.Fatalf("wire presence changed: %s %v", encoded, err) + } + } + policy := &types.NetworkPolicy{Rules: []types.NetworkPolicyRule{{Action: "ALLOW", Hosts: []string{"example.com", "192.0.2.1", "2001:db8::1"}, Description: "preserved"}}} + raw, err := json.Marshal(types.TaskAssignmentMessage{NetworkPolicy: policy}) + if err != nil { + t.Fatal(err) + } + var assignment types.TaskAssignmentMessage + if err := json.Unmarshal(raw, &assignment); err != nil || !reflect.DeepEqual(policy, assignment.NetworkPolicy) || !strings.Contains(string(raw), `"action":"ALLOW"`) { + t.Fatalf("wire contract: %s %v", raw, err) + } +} + +func TestNetworkPolicyPreclaimRejection(t *testing.T) { + for _, tc := range []struct { + name string + backend Backend + policy *types.NetworkPolicy + code string + }{ + {"unsupported", &recordingBackend{}, &types.NetworkPolicy{}, "network_policy_unsupported"}, + {"disabled", &KubernetesBackend{}, &types.NetworkPolicy{}, "network_policy_unsupported"}, + {"unknown action", &recordingBackend{}, allowNetworkPolicy("example.com"), "network_policy_invalid"}, + } { + t.Run(tc.name, func(t *testing.T) { + if tc.name == "unknown action" { + tc.policy.Rules[0].Action = "DENY" + } + w := &Worker{ctx: context.Background(), config: Config{OneShot: true}, backend: tc.backend, + outbound: newOutboundQueue(4), oneShot: newOneShotState(), taskSemaphore: semaphore.NewWeighted(1)} + w.handleTaskAssignment(&types.TaskAssignmentMessage{TaskID: "run", ExecutionID: "execution", NetworkPolicy: tc.policy}, time.Now()) + messages := drainMessages(t, w.outbound.messages) + if len(messages) != 1 || messages[0].Type != types.MessageTypeTaskRejected { + t.Fatalf("messages=%v", messages) + } + var rejected types.TaskRejectedMessage + if err := json.Unmarshal(messages[0].Data, &rejected); err != nil || rejected.Code != tc.code || rejected.ExecutionID != "execution" { + t.Fatalf("rejection=%+v err=%v", rejected, err) + } + if !w.oneShot.tryAccept() || !w.taskSemaphore.TryAcquire(1) { + t.Fatal("rejected policy consumed one-shot acceptance or a concurrency slot") + } + w.taskSemaphore.Release(1) + }) + } +} + +func TestNetworkPolicyHostValidation(t *testing.T) { + for _, host := range []string{"*", "*.example.com", "https://example.com", "192.0.2.0/24", "Example.com", "example.com.", "fe80::1%eth0", " bad", "-bad.example", "example..com"} { + if err := validateNetworkPolicy(allowNetworkPolicy(host)); err == nil { + t.Errorf("accepted invalid host %q", host) + } + } + if err := validateNetworkPolicy(&types.NetworkPolicy{Rules: []types.NetworkPolicyRule{{Action: "ALLOW"}}}); err == nil { + t.Fatal("accepted empty hosts") + } +} + +func TestNetworkPolicyDNSConfiguration(t *testing.T) { + cfg, err := effectiveNetworkPolicyDNS(config.NetworkPolicyDNSConfig{}, true, []byte("search cluster.local\nnameserver 10.96.0.10\nnameserver 2001:db8::53\nnameserver 10.96.0.10\n")) + if err != nil || !reflect.DeepEqual(cfg.ServerIPs, []string{"10.96.0.10", "2001:db8::53"}) || + cfg.Namespace != "kube-system" || cfg.PodLabels["k8s-app"] != "kube-dns" { + t.Fatalf("discovery: %+v %v", cfg, err) + } + override := config.NetworkPolicyDNSConfig{ServerIPs: []string{"192.0.2.53"}, Namespace: "dns", PodLabels: map[string]string{"app": "resolver"}} + got, err := effectiveNetworkPolicyDNS(override, false, []byte("nameserver 127.0.0.1")) + if err != nil || !reflect.DeepEqual(got, override) { + t.Fatalf("explicit override: %+v %v", got, err) + } + for _, tc := range []struct { + inCluster bool + resolv string + ips []string + }{ + {false, "nameserver 10.96.0.10", nil}, + {true, "", nil}, + {true, "nameserver invalid", nil}, + {true, "nameserver 127.0.0.1", nil}, + {true, "nameserver ::", nil}, + {false, "", []string{"fe80::1%eth0"}}, + } { + if _, err := effectiveNetworkPolicyDNS(config.NetworkPolicyDNSConfig{ServerIPs: tc.ips}, tc.inCluster, []byte(tc.resolv)); err == nil { + t.Errorf("accepted unusable DNS configuration: %+v", tc) + } + } +} + +func allowNetworkPolicy(hosts ...string) *types.NetworkPolicy { + return &types.NetworkPolicy{Rules: []types.NetworkPolicyRule{{Action: "ALLOW", Hosts: hosts}}} +} + +type fixedNetworkResolver struct { + ips []netip.Addr + err error + calls int +} + +func (r *fixedNetworkResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + r.calls++ + return r.ips, r.err +} + +func compileTestNetworkPolicy(t *testing.T, policy *types.NetworkPolicy, cilium bool, resolver networkHostResolver) *unstructured.Unstructured { + t.Helper() + object, err := compileNetworkPolicy(context.Background(), policy, cilium, config.NetworkPolicyDNSConfig{ServerIPs: []string{"10.96.0.10"}}, resolver) + if err != nil { + t.Fatal(err) + } + return object +} + +func TestNativeNetworkPolicyTranslation(t *testing.T) { + resolver := &fixedNetworkResolver{ips: []netip.Addr{netip.MustParseAddr("192.0.2.1"), netip.MustParseAddr("2001:db8::1"), netip.MustParseAddr("192.0.2.1")}} + object := compileTestNetworkPolicy(t, allowNetworkPolicy("example.com", "example.com", "198.51.100.1"), false, resolver) + spec, _, _ := unstructured.NestedMap(object.Object, "spec") + raw, _ := json.Marshal(spec) + for _, required := range []string{`192.0.2.1/32`, `2001:db8::1/128`, `198.51.100.1/32`, `10.96.0.10/32`, `"protocol":"TCP"`, `"protocol":"UDP"`, `"namespaceSelector"`, `"podSelector"`, `"policyTypes":["Egress"]`} { + if !strings.Contains(string(raw), required) { + t.Errorf("missing %s from %s", required, raw) + } + } + if resolver.calls != 1 || strings.Contains(string(raw), `"ingress"`) { + t.Fatalf("unexpected refresh/ingress: calls=%d spec=%s", resolver.calls, raw) + } + for _, r := range []*fixedNetworkResolver{{err: errors.New("DNS failed")}, {}} { + if _, err := compileNetworkPolicy(context.Background(), allowNetworkPolicy("example.com"), false, config.NetworkPolicyDNSConfig{ServerIPs: []string{"10.96.0.10"}}, r); err == nil { + t.Fatal("accepted incomplete DNS snapshot") + } + } +} + +func TestCiliumNetworkPolicyTranslation(t *testing.T) { + resolver := &fixedNetworkResolver{err: errors.New("must not resolve")} + object := compileTestNetworkPolicy(t, allowNetworkPolicy("example.com", "192.0.2.1", "2001:db8::1"), true, resolver) + spec, _, _ := unstructured.NestedMap(object.Object, "spec") + raw, _ := json.Marshal(spec) + for _, required := range []string{`"toFQDNs":[{"matchName":"example.com"}]`, `"toCIDR":["192.0.2.1/32","2001:db8::1/128"]`, + `"rules":{"dns":[{"matchName":"example.com"}]}`, `"enableDefaultDeny":{"egress":true}`} { + if !strings.Contains(string(raw), required) { + t.Errorf("missing %s from %s", required, raw) + } + } + if resolver.calls != 0 || strings.Contains(string(raw), "toEntities") || strings.Contains(string(raw), "matchPattern") || strings.Contains(string(raw), "ingress") { + t.Fatalf("Cilium widened or resolved rules: %s", raw) + } + for _, cilium := range []bool{false, true} { + empty, err := compileNetworkPolicy(context.Background(), &types.NetworkPolicy{}, cilium, config.NetworkPolicyDNSConfig{}, nil) + if err != nil { + t.Fatal(err) + } + spec, _, _ := unstructured.NestedMap(empty.Object, "spec") + if grants, present := spec["egress"]; present && len(grants.([]any)) > 0 { + t.Fatalf("empty policy emitted egress grants: %v", spec) + } + if cilium && spec["egress"] == nil { + t.Fatal("Cilium schema requires an explicit empty egress list") + } + } +} + +func TestKubernetesNetworkPolicyFingerprint(t *testing.T) { + a := compileTestNetworkPolicy(t, allowNetworkPolicy("198.51.100.1", "192.0.2.1", "192.0.2.1"), false, nil) + b := compileTestNetworkPolicy(t, allowNetworkPolicy("192.0.2.1", "198.51.100.1"), false, nil) + if a.GetName() != b.GetName() || len(a.GetName()) != 60 || len(a.GetLabels()[networkPolicyIDLabel]) != 52 { + t.Fatalf("unstable fingerprints: %s / %s", a.GetName(), b.GetName()) + } + for _, cilium := range []bool{false, true} { + policy := allowNetworkPolicy("example.com") + resolver := &fixedNetworkResolver{ips: []netip.Addr{netip.MustParseAddr("192.0.2.1")}} + first := compileTestNetworkPolicy(t, policy, cilium, resolver) + policy.Rules[0].Description = "non-enforcing" + second := compileTestNetworkPolicy(t, policy, cilium, resolver) + if first.GetName() != second.GetName() { + t.Fatal("description changed digest") + } + second.SetNamespace("another") + second.SetFinalizers([]string{"other/finalizer"}) + second.SetAnnotations(map[string]string{"other": "metadata"}) + expected, _ := canonicalNetworkSpec(first, first.GetLabels()[networkPolicyIDLabel]) + actual, _ := canonicalNetworkSpec(second, second.GetLabels()[networkPolicyIDLabel]) + if string(expected) != string(actual) { + t.Fatal("metadata changed content") + } + resolver.ips = []netip.Addr{netip.MustParseAddr("198.51.100.1")} + changed := compileTestNetworkPolicy(t, policy, cilium, resolver) + if (first.GetName() == changed.GetName()) != cilium { + t.Fatal("native snapshot/Cilium identity mismatch") + } + if err := unstructured.SetNestedField(second.Object, "altered", "spec", map[bool]string{false: "podSelector", true: "endpointSelector"}[cilium], "matchLabels", networkPolicyIDLabel); err != nil { + t.Fatal(err) + } + if _, err := canonicalNetworkSpec(second, first.GetLabels()[networkPolicyIDLabel]); err == nil { + t.Fatal("altered actual selector accepted") + } + } +} + +func TestKubernetesNetworkPolicyProviderSelection(t *testing.T) { + for _, tc := range []struct { + name string + resources []metav1.APIResource + err error + cilium bool + fail bool + }{ + {name: "present", resources: []metav1.APIResource{{Name: "ciliumnetworkpolicies", Namespaced: true}}, cilium: true}, + {name: "absent", err: apierrors.NewNotFound(schema.GroupResource{Group: "cilium.io", Resource: "v2"}, "")}, + {name: "resource absent"}, + {name: "denied", err: apierrors.NewForbidden(schema.GroupResource{Resource: "discovery"}, "", errors.New("denied")), fail: true}, + {name: "timeout", err: context.DeadlineExceeded, fail: true}, + {name: "partial", err: errors.New("partial group discovery"), fail: true}, + } { + t.Run(tc.name, func(t *testing.T) { + client := fake.NewClientset() + discovery := client.Discovery().(*discoveryfake.FakeDiscovery) + discovery.Resources = []*metav1.APIResourceList{{GroupVersion: "cilium.io/v2", APIResources: tc.resources}} + if tc.err != nil { + discovery.PrependReactor("get", "resource", func(k8stesting.Action) (bool, runtime.Object, error) { return true, nil, tc.err }) + } + backend := &KubernetesBackend{clientset: client, networkDiscovery: discovery, networkDynamic: dynamicfake.NewSimpleDynamicClient(runtime.NewScheme())} + provider, err := backend.discoverNetworkPolicyProvider() + if (err != nil) != tc.fail || provider.cilium != tc.cilium { + t.Fatalf("provider=%+v err=%v", provider, err) + } + }) + } +} diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 3012519..6ffb61d 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -625,6 +625,18 @@ func (w *Worker) rememberCompletedLocked(key taskExecution, cancellationAccepted } func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage, receivedAt time.Time) { + if assignment.NetworkPolicy != nil { + code, err := w.admitNetworkPolicy(assignment.NetworkPolicy) + if err != nil { + metrics.RecordTaskRejected(code) + if sendErr := w.sendRejection(types.TaskRejectedMessage{ + TaskID: assignment.TaskID, ExecutionID: assignment.ExecutionID, Code: code, Reason: err.Error(), + }); sendErr != nil { + log.Errorf(w.ctx, "Failed to send network policy rejection: %v", sendErr) + } + return + } + } key := taskExecution{assignment.TaskID, assignment.ExecutionID} log.Infof(w.ctx, "Received task assignment: taskID=%s, title=%s", assignment.TaskID, assignment.Task.Title) taskCtx, session, err := w.config.Tracing.Start(w.ctx, w.config.ServerRootURL, assignment, receivedAt) @@ -843,6 +855,7 @@ func (w *Worker) prepareTaskParams(assignment *types.TaskAssignmentMessage) *Tas DockerImage: dockerImage, Sidecars: sidecars, InstanceShape: assignment.InstanceShape, + NetworkPolicy: assignment.NetworkPolicy, SetupEvents: setupEvents, } } @@ -1107,11 +1120,13 @@ func (w *Worker) sendTaskCancelled(taskID, executionID, message string) error { } func (w *Worker) sendTaskRejected(taskID, reason string) error { - rejectedMsg := types.TaskRejectedMessage{ + return w.sendRejection(types.TaskRejectedMessage{ TaskID: taskID, Reason: reason, - } + }) +} +func (w *Worker) sendRejection(rejectedMsg types.TaskRejectedMessage) error { data, err := json.Marshal(rejectedMsg) if err != nil { return fmt.Errorf("failed to marshal task rejected message: %w", err) diff --git a/main.go b/main.go index 0ca4660..082cc20 100644 --- a/main.go +++ b/main.go @@ -220,10 +220,12 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { unschedulableTimeout *time.Duration podTemplate *corev1.PodSpec preflightResources *corev1.ResourceRequirements + networkPolicy config.NetworkPolicyConfig ) if fileConfig != nil && fileConfig.Backend.Kubernetes != nil { kc := fileConfig.Backend.Kubernetes + networkPolicy = kc.NetworkPolicy namespace = kc.Namespace kubeconfig = kc.Kubeconfig defaultImage = kc.DefaultImage @@ -278,6 +280,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { } wc.Kubernetes = &worker.KubernetesBackendConfig{ + NetworkPolicy: networkPolicy, WorkerID: workerID, Namespace: namespace, Kubeconfig: kubeconfig, diff --git a/network_policy_config_test.go b/network_policy_config_test.go new file mode 100644 index 0000000..6263ba4 --- /dev/null +++ b/network_policy_config_test.go @@ -0,0 +1,145 @@ +package main + +import ( + "bytes" + "io" + "os" + "os/exec" + "path/filepath" + "reflect" + "strings" + "testing" + + "github.com/warpdotdev/oz-agent-worker/internal/config" + "gopkg.in/yaml.v3" +) + +func TestNetworkPolicyConfigAndBackendAdmission(t *testing.T) { + for _, block := range []string{"", " network_policy:\n dns:\n namespace: kube-system\n", " network_policy:\n enabled: true\n dns:\n server_ips: [192.0.2.53, '2001:db8::53']\n namespace: trusted-dns\n pod_labels: {app: resolver}\n"} { + t.Run(strings.ReplaceAll(block, "\n", " "), func(t *testing.T) { + resetCLIForTest() + t.Cleanup(resetCLIForTest) + path := filepath.Join(t.TempDir(), "worker.yaml") + raw := "worker_id: worker\nbackend:\n kubernetes:\n namespace: agents\n" + block + if err := os.WriteFile(path, []byte(raw), 0600); err != nil { + t.Fatal(err) + } + loaded, err := config.Load(path) + if err != nil { + t.Fatal(err) + } + merged, err := mergeConfig(loaded) + if err != nil { + t.Fatal(err) + } + policy := merged.Kubernetes.NetworkPolicy + if policy.Enabled != strings.Contains(block, "enabled: true") || !reflect.DeepEqual(policy, loaded.Backend.Kubernetes.NetworkPolicy) { + t.Fatalf("network policy configuration lost: %+v", policy) + } + if policy.Enabled && (!reflect.DeepEqual(policy.DNS.ServerIPs, []string{"192.0.2.53", "2001:db8::53"}) || policy.DNS.Namespace != "trusted-dns" || policy.DNS.PodLabels["app"] != "resolver") { + t.Fatalf("DNS configuration lost: %+v", policy) + } + }) + } + path := filepath.Join(t.TempDir(), "reserved.yaml") + if err := os.WriteFile(path, []byte("backend:\n kubernetes:\n extra_labels: {warp-custom: forbidden}\n"), 0600); err != nil { + t.Fatal(err) + } + if _, err := config.Load(path); err == nil || !strings.Contains(err.Error(), "reserved") { + t.Fatalf("reserved label accepted: %v", err) + } +} + +func TestNetworkPolicyHelm(t *testing.T) { + helm, err := exec.LookPath("helm") + if err != nil { + t.Skip("helm is required for chart acceptance tests") + } + for _, tc := range []struct { + name string + args []string + enabled bool + noRole bool + }{ + {name: "default"}, + {name: "disabled", args: []string{"--set", "kubernetesBackend.networkPolicy.enabled=false"}}, + {name: "enabled", args: []string{"--set", "kubernetesBackend.networkPolicy.enabled=true"}, enabled: true}, + {name: "external RBAC", args: []string{"--set", "kubernetesBackend.networkPolicy.enabled=true", "--set", "rbac.create=false"}, enabled: true, noRole: true}, + {name: "DNS override", args: []string{"--set", "kubernetesBackend.networkPolicy.enabled=true", "--set", "kubernetesBackend.networkPolicy.dns.serverIPs[0]=192.0.2.53", "--set", "kubernetesBackend.networkPolicy.dns.namespace=trusted-dns", "--set", "kubernetesBackend.networkPolicy.dns.podLabels.app=resolver"}, enabled: true}, + } { + t.Run(tc.name, func(t *testing.T) { + args := append([]string{"template", "oz-agent-worker", "./charts/oz-agent-worker", "--namespace", "agents", "--set", "worker.workerId=ci-worker", "--set", "image.tag=ci"}, tc.args...) + raw, err := exec.Command(helm, args...).CombinedOutput() + if err != nil { + t.Fatalf("helm: %v %s", err, raw) + } + decoder := yaml.NewDecoder(bytes.NewReader(raw)) + baselines, roles, policyRules := 0, 0, 0 + for { + var object map[string]any + err := decoder.Decode(&object) + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + switch object["kind"] { + case "NetworkPolicy": + baselines++ + spec := object["spec"].(map[string]any) + want := map[string]any{"podSelector": map[string]any{"matchLabels": map[string]any{"warp-network-restricted": "true"}}, "policyTypes": []any{"Egress"}, "egress": []any{}} + if !reflect.DeepEqual(spec, want) { + t.Fatalf("baseline=%v", spec) + } + case "Role": + roles++ + for _, rawRule := range object["rules"].([]any) { + rule := rawRule.(map[string]any) + resource := rule["resources"].([]any)[0].(string) + if resource == "networkpolicies" || resource == "ciliumnetworkpolicies" { + policyRules++ + if !reflect.DeepEqual(rule["verbs"], []any{"create", "get", "list", "patch", "delete"}) { + t.Fatalf("policy RBAC=%v", rule) + } + } + if resource == "services" { + t.Fatal("DNS service permissions added") + } + } + case "ClusterRole": + t.Fatal("cluster-scoped policy rights added") + case "ConfigMap": + data := object["data"].(map[string]any)["config.yaml"].(string) + path := filepath.Join(t.TempDir(), "rendered.yaml") + if err := os.WriteFile(path, []byte(data), 0600); err != nil { + t.Fatal(err) + } + cfg, err := config.Load(path) + if err != nil { + t.Fatal(err) + } + policy := cfg.Backend.Kubernetes.NetworkPolicy + if policy.Enabled != tc.enabled { + t.Fatalf("enabled=%v", policy.Enabled) + } + if tc.name == "DNS override" && (policy.DNS.Namespace != "trusted-dns" || policy.DNS.PodLabels["app"] != "resolver" || !reflect.DeepEqual(policy.DNS.ServerIPs, []string{"192.0.2.53"})) { + t.Fatalf("DNS=%+v", policy.DNS) + } + } + } + wantRules := 0 + if tc.enabled && !tc.noRole { + wantRules = 2 + } + if baselines != 1 || policyRules != wantRules || (tc.noRole && roles != 0) { + t.Fatalf("baseline=%d role=%d policy rules=%d", baselines, roles, policyRules) + } + }) + } + cmd := exec.Command(helm, "template", "worker", "./charts/oz-agent-worker", "--set", "worker.workerId=ci-worker", "--set", "image.tag=ci", "--set", "kubernetesBackend.extraLabels.warp-custom=forbidden") + raw, err := cmd.CombinedOutput() + if err == nil || !strings.Contains(string(raw), "reserved for Warp") { + t.Fatalf("Helm accepted reserved label: %s %v", raw, err) + } +}