diff --git a/cmd/ateapi/internal/controlapi/dialer.go b/cmd/ateapi/internal/controlapi/dialer.go index 30976a41d..4944a2364 100644 --- a/cmd/ateapi/internal/controlapi/dialer.go +++ b/cmd/ateapi/internal/controlapi/dialer.go @@ -19,8 +19,11 @@ import ( "crypto/x509" "errors" "fmt" + "net" "slices" + "strconv" + "github.com/agent-substrate/substrate/internal/atelet" "github.com/agent-substrate/substrate/internal/credbundle" "github.com/agent-substrate/substrate/internal/substratex509" "github.com/spiffe/go-spiffe/v2/bundle/x509bundle" @@ -141,7 +144,7 @@ func (d *AteletDialer) DialForAteletOnNode(nodeName string) (*grpc.ClientConn, e } ateletConn, err := grpc.NewClient( - selectedAtelet.Status.PodIPs[0].IP+":8085", + net.JoinHostPort(selectedAtelet.Status.PodIPs[0].IP, strconv.Itoa(atelet.DefaultPort)), grpc.WithTransportCredentials(creds), grpc.WithStatsHandler(otelgrpc.NewClientHandler()), ) diff --git a/cmd/ateapi/internal/controlapi/dialer_test.go b/cmd/ateapi/internal/controlapi/dialer_test.go index 4c995ccf5..73b9879f9 100644 --- a/cmd/ateapi/internal/controlapi/dialer_test.go +++ b/cmd/ateapi/internal/controlapi/dialer_test.go @@ -36,6 +36,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/tools/cache" + "k8s.io/utils/lru" ) const testAteletSPIFFEID = "spiffe://cluster.local/ns/ate-system/sa/atelet" @@ -133,6 +134,119 @@ func makeLeafCert(t *testing.T, ca *x509.Certificate, caKey *ecdsa.PrivateKey, o return cert } +// newDialerForPods builds an AteletDialer for testing. +func newDialerForPods(t *testing.T, workerPod, ateletPod *corev1.Pod) *AteletDialer { + t.Helper() + + workerIndexer := cache.NewIndexer(cache.MetaNamespaceKeyFunc, cache.Indexers{ + byNamespaceAndName: func(obj any) ([]string, error) { + pod := obj.(*corev1.Pod) + return []string{pod.ObjectMeta.Namespace + "/" + pod.ObjectMeta.Name}, nil + }, + }) + if err := workerIndexer.Add(workerPod); err != nil { + t.Fatalf("adding worker pod to indexer: %v", err) + } + + ateletIndexer := cache.NewIndexer(cache.MetaNamespaceKeyFunc, cache.Indexers{ + byNode: func(obj any) ([]string, error) { + pod := obj.(*corev1.Pod) + return []string{pod.Spec.NodeName}, nil + }, + }) + if err := ateletIndexer.Add(ateletPod); err != nil { + t.Fatalf("adding atelet pod to indexer: %v", err) + } + + return &AteletDialer{ + workerIndexer: workerIndexer, + ateletIndexer: ateletIndexer, + ateletConns: lru.New(16), + dialCredentials: func(string) (credentials.TransportCredentials, error) { + return insecure.NewCredentials(), nil + }, + } +} + +func TestDialForWorkerTarget(t *testing.T) { + tests := []struct { + name string + ateletIP string + wantTarget string + }{ + { + name: "IPv4 atelet", + ateletIP: "10.244.1.7", + wantTarget: "10.244.1.7:8085", + }, + { + name: "IPv6 atelet is bracketed", + ateletIP: "fd00:10:244::7", + wantTarget: "[fd00:10:244::7]:8085", + }, + { + name: "IPv6 loopback is bracketed", + ateletIP: "::1", + wantTarget: "[::1]:8085", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + workerPod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Namespace: "team-a", Name: "worker-1", UID: "worker-uid"}, + Spec: corev1.PodSpec{NodeName: "node-1"}, + } + ateletPod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Namespace: ateletNamespace, Name: "atelet-abc", UID: "atelet-uid"}, + Spec: corev1.PodSpec{NodeName: "node-1"}, + Status: corev1.PodStatus{PodIPs: []corev1.PodIP{{IP: tc.ateletIP}}}, + } + + d := newDialerForPods(t, workerPod, ateletPod) + conn, err := d.DialForWorker("team-a", "worker-1") + if err != nil { + t.Fatalf("DialForWorker returned error: %v", err) + } + t.Cleanup(func() { conn.Close() }) + + if got := conn.Target(); got != tc.wantTarget { + t.Errorf("dial target = %q, want %q", got, tc.wantTarget) + } + }) + } +} + +func TestDialForWorkerErrors(t *testing.T) { + workerPod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Namespace: "team-a", Name: "worker-1", UID: "worker-uid"}, + Spec: corev1.PodSpec{NodeName: "node-1"}, + } + + t.Run("unknown worker pod", func(t *testing.T) { + ateletPod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Namespace: ateletNamespace, Name: "atelet-abc", UID: "atelet-uid"}, + Spec: corev1.PodSpec{NodeName: "node-1"}, + Status: corev1.PodStatus{PodIPs: []corev1.PodIP{{IP: "10.244.1.7"}}}, + } + d := newDialerForPods(t, workerPod, ateletPod) + if _, err := d.DialForWorker("team-a", "no-such-worker"); !errors.Is(err, ErrWorkerPodNotFound) { + t.Fatalf("DialForWorker error = %v, want ErrWorkerPodNotFound", err) + } + }) + + t.Run("atelet without assigned IPs", func(t *testing.T) { + ateletPod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Namespace: ateletNamespace, Name: "atelet-abc", UID: "atelet-uid"}, + Spec: corev1.PodSpec{NodeName: "node-1"}, + } + d := newDialerForPods(t, workerPod, ateletPod) + if _, err := d.DialForWorker("team-a", "worker-1"); err == nil { + t.Fatal("DialForWorker succeeded, want error for atelet with no IPs") + } + }) +} + func TestVerifyAteletServerCert(t *testing.T) { ca, caKey, bundle := makeTestCA(t) otherCA, otherCAKey, _ := makeTestCA(t) diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index ea3f33214..0fcd67c02 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -40,6 +40,7 @@ import ( "github.com/agent-substrate/substrate/internal/ateattr" "github.com/agent-substrate/substrate/internal/ateerrors" "github.com/agent-substrate/substrate/internal/ateinterceptors" + "github.com/agent-substrate/substrate/internal/atelet" "github.com/agent-substrate/substrate/internal/ateompath" "github.com/agent-substrate/substrate/internal/credbundle" "github.com/agent-substrate/substrate/internal/imagecache" @@ -81,7 +82,7 @@ import ( ) var ( - port = pflag.Int("port", 8085, "The port to listen on") + port = pflag.Int("port", atelet.DefaultPort, "The port to listen on") metricsListenAddr = pflag.String("metrics-listen-addr", ":9090", "Address and port the prometheus metrics server should listen on.") grpcServerCredBundle = pflag.String("grpc-server-cred-bundle", "/run/podidentity.podcert.ate.dev/credential-bundle.pem", "Credential bundle atelet presents as its gRPC serving certificate.") diff --git a/cmd/atelet/main_test.go b/cmd/atelet/main_test.go index b08e203d6..1bb492db4 100644 --- a/cmd/atelet/main_test.go +++ b/cmd/atelet/main_test.go @@ -26,6 +26,7 @@ import ( "os" "path/filepath" "slices" + "strconv" "strings" "sync" "syscall" @@ -34,6 +35,7 @@ import ( "github.com/agent-substrate/substrate/internal/ateattr" "github.com/agent-substrate/substrate/internal/ateerrors" + "github.com/agent-substrate/substrate/internal/atelet" "github.com/agent-substrate/substrate/internal/ateompath" "github.com/agent-substrate/substrate/internal/proto/ateletpb" "github.com/agent-substrate/substrate/internal/proto/ateompb" @@ -41,6 +43,7 @@ import ( "github.com/agent-substrate/substrate/internal/serverboot" "github.com/google/go-cmp/cmp" "github.com/klauspost/compress/zstd" + "github.com/spf13/pflag" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" @@ -51,6 +54,17 @@ import ( const testPauseImage = "registry.k8s.io/pause:3.10.2@sha256:f548e0e8e3dc1896ca956272154dde3314e8cc4fde0a57577ee9fa1c63f5baf4" +// TestPortFlagDefault verifies the default value of the --port flag. +func TestPortFlagDefault(t *testing.T) { + f := pflag.Lookup("port") + if f == nil { + t.Fatal("no --port flag registered") + } + if want := strconv.Itoa(atelet.DefaultPort); f.DefValue != want { + t.Errorf("--port default = %q, want %q", f.DefValue, want) + } +} + func TestSnapshotManifestActorMetadata(t *testing.T) { rec := sandboxAssetsRecord{ Atespace: "team-a", diff --git a/internal/atelet/doc.go b/internal/atelet/doc.go new file mode 100644 index 000000000..bde92841e --- /dev/null +++ b/internal/atelet/doc.go @@ -0,0 +1,17 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package atelet contains shared constants and definitions that are used across +// Substrate components related to atelet. +package atelet diff --git a/internal/atelet/port.go b/internal/atelet/port.go new file mode 100644 index 000000000..d89ad3adf --- /dev/null +++ b/internal/atelet/port.go @@ -0,0 +1,20 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package atelet + +const ( + // DefaultPort is the port atelet's gRPC server listens on. + DefaultPort = 8085 +)