Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 13 additions & 13 deletions components/ambient-control-plane/cmd/ambient-control-plane/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/ambient-code/platform/components/ambient-control-plane/internal/auth"
"github.com/ambient-code/platform/components/ambient-control-plane/internal/config"
"github.com/ambient-code/platform/components/ambient-control-plane/internal/informer"
"github.com/ambient-code/platform/components/ambient-control-plane/internal/keypair"
"github.com/ambient-code/platform/components/ambient-control-plane/internal/kubeclient"
"github.com/ambient-code/platform/components/ambient-control-plane/internal/reconciler"
"github.com/ambient-code/platform/components/ambient-control-plane/internal/tokenserver"
Expand All @@ -25,8 +26,6 @@ import (
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"
"k8s.io/client-go/rest"
"k8s.io/client-go/tools/clientcmd"
)

var (
Expand Down Expand Up @@ -123,6 +122,12 @@ func runKubeMode(ctx context.Context, cfg *config.ControlPlaneConfig) error {
provisioner := buildNamespaceProvisioner(cfg, provisionerKube)
tokenProvider := buildTokenProvider(cfg, log.Logger)

kp, err := keypair.EnsureKeypairSecret(ctx, provisionerKube, cfg.CPRuntimeNamespace, log.Logger)
if err != nil {
return fmt.Errorf("bootstrapping CP token keypair: %w", err)
}
log.Info().Str("namespace", cfg.CPRuntimeNamespace).Msg("CP token keypair ready")

factory := reconciler.NewSDKClientFactory(cfg.APIServerURL, tokenProvider, log.Logger)
kubeReconcilerCfg := reconciler.KubeReconcilerConfig{
RunnerImage: cfg.RunnerImage,
Expand All @@ -142,6 +147,7 @@ func runKubeMode(ctx context.Context, cfg *config.ControlPlaneConfig) error {
RunnerLogLevel: cfg.RunnerLogLevel,
CPRuntimeNamespace: cfg.CPRuntimeNamespace,
CPTokenURL: cfg.CPTokenURL,
CPTokenPublicKey: string(kp.PublicKeyPEM),
}

conn, err := grpc.NewClient(cfg.GRPCServerAddr, grpc.WithTransportCredentials(grpcCredentials(cfg.GRPCUseTLS)))
Expand Down Expand Up @@ -184,7 +190,7 @@ func runKubeMode(ctx context.Context, cfg *config.ControlPlaneConfig) error {

tsErrCh := make(chan error, 1)
go func() {
tsErrCh <- startTokenServer(ctx, cfg, tokenProvider)
tsErrCh <- startTokenServer(ctx, cfg, tokenProvider, kp)
}()

infErrCh := make(chan error, 1)
Expand All @@ -203,24 +209,18 @@ func runKubeMode(ctx context.Context, cfg *config.ControlPlaneConfig) error {
}
}

func startTokenServer(ctx context.Context, cfg *config.ControlPlaneConfig, tokenProvider auth.TokenProvider) error {
k8sConfig, err := buildK8sRestConfig(cfg.Kubeconfig)
func startTokenServer(ctx context.Context, cfg *config.ControlPlaneConfig, tokenProvider auth.TokenProvider, kp *keypair.KeyPair) error {
privKey, err := keypair.ParsePrivateKey(kp.PrivateKeyPEM)
if err != nil {
return fmt.Errorf("building k8s rest config for token server: %w", err)
return fmt.Errorf("parsing CP token private key: %w", err)
}
ts, err := tokenserver.New(cfg.CPTokenListenAddr, tokenProvider, k8sConfig, log.Logger)
ts, err := tokenserver.New(cfg.CPTokenListenAddr, tokenProvider, privKey, log.Logger)
if err != nil {
return fmt.Errorf("creating token server: %w", err)
}
return ts.Start(ctx)
}

func buildK8sRestConfig(kubeconfig string) (*rest.Config, error) {
if kubeconfig != "" {
return clientcmd.BuildConfigFromFlags("", kubeconfig)
}
return rest.InClusterConfig()
}

func createSessionReconcilers(reconcilerTypes []string, factory *reconciler.SDKClientFactory, kube *kubeclient.KubeClient, projectKube *kubeclient.KubeClient, provisioner kubeclient.NamespaceProvisioner, cfg reconciler.KubeReconcilerConfig, logger zerolog.Logger) []reconciler.Reconciler {
var reconcilers []reconciler.Reconciler
Expand Down
134 changes: 134 additions & 0 deletions components/ambient-control-plane/internal/keypair/bootstrap.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
package keypair

import (
"context"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/base64"
"encoding/pem"
"fmt"

"github.com/ambient-code/platform/components/ambient-control-plane/internal/kubeclient"
"github.com/rs/zerolog"
k8serrors "k8s.io/apimachinery/pkg/api/errors"
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
)

const (
SecretName = "ambient-cp-token-keypair"
privateKeyKey = "private.pem"
publicKeyKey = "public.pem"
rsaKeyBits = 4096
)

type KeyPair struct {
PrivateKeyPEM []byte
PublicKeyPEM []byte
}

func EnsureKeypairSecret(ctx context.Context, kube *kubeclient.KubeClient, namespace string, logger zerolog.Logger) (*KeyPair, error) {
existing, err := kube.GetSecret(ctx, namespace, SecretName)
if err == nil {
return keypairFromSecret(existing)
}
if !k8serrors.IsNotFound(err) {
return nil, fmt.Errorf("checking for keypair secret: %w", err)
}

logger.Info().Str("namespace", namespace).Str("secret", SecretName).Msg("keypair secret not found, generating new RSA keypair")

kp, err := generateKeypair()
if err != nil {
return nil, fmt.Errorf("generating RSA keypair: %w", err)
}

secret := &unstructured.Unstructured{
Object: map[string]interface{}{
"apiVersion": "v1",
"kind": "Secret",
"metadata": map[string]interface{}{
"name": SecretName,
"namespace": namespace,
"labels": map[string]interface{}{
"app": "ambient-control-plane",
"ambient-code.io/managed-by": "ambient-control-plane",
},
},
"type": "Opaque",
"data": map[string]interface{}{
privateKeyKey: base64.StdEncoding.EncodeToString(kp.PrivateKeyPEM),
publicKeyKey: base64.StdEncoding.EncodeToString(kp.PublicKeyPEM),
},
},
}

if _, createErr := kube.CreateSecret(ctx, secret); createErr != nil {
if !k8serrors.IsAlreadyExists(createErr) {
return nil, fmt.Errorf("creating keypair secret: %w", createErr)
}
existing, err = kube.GetSecret(ctx, namespace, SecretName)
if err != nil {
return nil, fmt.Errorf("re-reading keypair secret after race: %w", err)
}
return keypairFromSecret(existing)
}

logger.Info().Str("namespace", namespace).Str("secret", SecretName).Msg("RSA keypair secret created")
return kp, nil
}

func keypairFromSecret(secret *unstructured.Unstructured) (*KeyPair, error) {
data, _, _ := unstructured.NestedMap(secret.Object, "data")

Comment on lines +81 to +83
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Silently ignoring NestedMap error may mask corrupted Secret data.

If the Secret exists but has a malformed structure, the error is discarded and the code proceeds to fail confusingly on subsequent key lookups.

Proposed fix
 func keypairFromSecret(secret *unstructured.Unstructured) (*KeyPair, error) {
-	data, _, _ := unstructured.NestedMap(secret.Object, "data")
+	data, found, err := unstructured.NestedMap(secret.Object, "data")
+	if err != nil || !found {
+		return nil, fmt.Errorf("keypair secret has invalid or missing data field")
+	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
func keypairFromSecret(secret *unstructured.Unstructured) (*KeyPair, error) {
data, _, _ := unstructured.NestedMap(secret.Object, "data")
func keypairFromSecret(secret *unstructured.Unstructured) (*KeyPair, error) {
data, found, err := unstructured.NestedMap(secret.Object, "data")
if err != nil || !found {
return nil, fmt.Errorf("keypair secret has invalid or missing data field")
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@components/ambient-control-plane/internal/keypair/bootstrap.go` around lines
81 - 83, In keypairFromSecret, don't ignore the error returned by
unstructured.NestedMap(secret.Object, "data"); instead capture both the map and
the error (and the existence bool), validate that the call succeeded and return
a descriptive error if it failed or if the "data" field is missing/invalid so
downstream lookups on `data` don't panic—update keypairFromSecret to check the
returned error and existence flag and propagate a clear error when malformed
Secret data is detected.

privB64, ok := data[privateKeyKey].(string)
if !ok || privB64 == "" {
return nil, fmt.Errorf("keypair secret missing %q key", privateKeyKey)
}
pubB64, ok := data[publicKeyKey].(string)
if !ok || pubB64 == "" {
return nil, fmt.Errorf("keypair secret missing %q key", publicKeyKey)
}

privPEM, err := base64.StdEncoding.DecodeString(privB64)
if err != nil {
return nil, fmt.Errorf("decoding private key from secret: %w", err)
}
pubPEM, err := base64.StdEncoding.DecodeString(pubB64)
if err != nil {
return nil, fmt.Errorf("decoding public key from secret: %w", err)
}

return &KeyPair{PrivateKeyPEM: privPEM, PublicKeyPEM: pubPEM}, nil
}

func generateKeypair() (*KeyPair, error) {
privKey, err := rsa.GenerateKey(rand.Reader, rsaKeyBits)
if err != nil {
return nil, fmt.Errorf("generating RSA key: %w", err)
}

privPEM := pem.EncodeToMemory(&pem.Block{
Type: "RSA PRIVATE KEY",
Bytes: x509.MarshalPKCS1PrivateKey(privKey),
})

pubDER, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
if err != nil {
return nil, fmt.Errorf("marshaling public key: %w", err)
}
pubPEM := pem.EncodeToMemory(&pem.Block{
Type: "PUBLIC KEY",
Bytes: pubDER,
})

return &KeyPair{PrivateKeyPEM: privPEM, PublicKeyPEM: pubPEM}, nil
}

func ParsePrivateKey(pemBytes []byte) (*rsa.PrivateKey, error) {
block, _ := pem.Decode(pemBytes)
if block == nil {
return nil, fmt.Errorf("failed to decode PEM block for private key")
}
return x509.ParsePKCS1PrivateKey(block.Bytes)
}
160 changes: 160 additions & 0 deletions components/ambient-control-plane/internal/keypair/bootstrap_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
package keypair

import (
"context"
"crypto/rsa"
"encoding/base64"
"testing"

"github.com/rs/zerolog"
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/client-go/dynamic/fake"

"github.com/ambient-code/platform/components/ambient-control-plane/internal/kubeclient"
)

func newFakeKubeClient(objects ...runtime.Object) *kubeclient.KubeClient {
scheme := runtime.NewScheme()
dynClient := fake.NewSimpleDynamicClient(scheme, objects...)
return kubeclient.NewFromDynamic(dynClient, zerolog.Nop())
}

func TestGenerateKeypair(t *testing.T) {
kp, err := generateKeypair()
if err != nil {
t.Fatalf("generateKeypair() error: %v", err)
}
if len(kp.PrivateKeyPEM) == 0 {
t.Error("PrivateKeyPEM is empty")
}
if len(kp.PublicKeyPEM) == 0 {
t.Error("PublicKeyPEM is empty")
}
}

func TestParsePrivateKey(t *testing.T) {
kp, err := generateKeypair()
if err != nil {
t.Fatalf("generateKeypair() error: %v", err)
}
privKey, err := ParsePrivateKey(kp.PrivateKeyPEM)
if err != nil {
t.Fatalf("ParsePrivateKey() error: %v", err)
}
if privKey == nil {
t.Fatal("ParsePrivateKey() returned nil")
}
if _, ok := interface{}(privKey).(*rsa.PrivateKey); !ok {
t.Error("parsed key is not *rsa.PrivateKey")
}
}

func TestParsePrivateKey_InvalidPEM(t *testing.T) {
_, err := ParsePrivateKey([]byte("not a pem block"))
if err == nil {
t.Error("expected error for invalid PEM, got nil")
}
}

func TestKeypairFromSecret_MissingPrivateKey(t *testing.T) {
secret := &unstructured.Unstructured{
Object: map[string]interface{}{
"apiVersion": "v1",
"kind": "Secret",
"metadata": map[string]interface{}{"name": SecretName, "namespace": "test"},
"data": map[string]interface{}{
publicKeyKey: base64.StdEncoding.EncodeToString([]byte("pub")),
},
},
}
_, err := keypairFromSecret(secret)
if err == nil {
t.Error("expected error for missing private key, got nil")
}
}

func TestKeypairFromSecret_MissingPublicKey(t *testing.T) {
secret := &unstructured.Unstructured{
Object: map[string]interface{}{
"apiVersion": "v1",
"kind": "Secret",
"metadata": map[string]interface{}{"name": SecretName, "namespace": "test"},
"data": map[string]interface{}{
privateKeyKey: base64.StdEncoding.EncodeToString([]byte("priv")),
},
},
}
_, err := keypairFromSecret(secret)
if err == nil {
t.Error("expected error for missing public key, got nil")
}
}

func TestKeypairFromSecret_ValidSecret(t *testing.T) {
kp, err := generateKeypair()
if err != nil {
t.Fatalf("generateKeypair() error: %v", err)
}
secret := &unstructured.Unstructured{
Object: map[string]interface{}{
"apiVersion": "v1",
"kind": "Secret",
"metadata": map[string]interface{}{"name": SecretName, "namespace": "test"},
"data": map[string]interface{}{
privateKeyKey: base64.StdEncoding.EncodeToString(kp.PrivateKeyPEM),
publicKeyKey: base64.StdEncoding.EncodeToString(kp.PublicKeyPEM),
},
},
}
got, err := keypairFromSecret(secret)
if err != nil {
t.Fatalf("keypairFromSecret() error: %v", err)
}
if string(got.PrivateKeyPEM) != string(kp.PrivateKeyPEM) {
t.Error("PrivateKeyPEM mismatch")
}
if string(got.PublicKeyPEM) != string(kp.PublicKeyPEM) {
t.Error("PublicKeyPEM mismatch")
}
}

func TestEnsureKeypairSecret_CreatesWhenMissing(t *testing.T) {
kube := newFakeKubeClient()
ctx := context.Background()

kp, err := EnsureKeypairSecret(ctx, kube, "test-ns", zerolog.Nop())
if err != nil {
t.Fatalf("EnsureKeypairSecret() error: %v", err)
}
if len(kp.PrivateKeyPEM) == 0 || len(kp.PublicKeyPEM) == 0 {
t.Error("returned keypair has empty PEM fields")
}

privKey, err := ParsePrivateKey(kp.PrivateKeyPEM)
if err != nil {
t.Fatalf("generated private key is not parseable: %v", err)
}
if privKey.N.BitLen() != rsaKeyBits {
t.Errorf("key size: got %d, want %d", privKey.N.BitLen(), rsaKeyBits)
}
}

func TestEnsureKeypairSecret_ReturnsExistingWhenPresent(t *testing.T) {
ctx := context.Background()
kube := newFakeKubeClient()

first, err := EnsureKeypairSecret(ctx, kube, "test-ns", zerolog.Nop())
if err != nil {
t.Fatalf("first call error: %v", err)
}

second, err := EnsureKeypairSecret(ctx, kube, "test-ns", zerolog.Nop())
if err != nil {
t.Fatalf("second call error: %v", err)
}

if string(first.PrivateKeyPEM) != string(second.PrivateKeyPEM) {
t.Error("second call returned different private key — should reuse existing Secret")
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ type KubeReconcilerConfig struct {
RunnerLogLevel string
CPRuntimeNamespace string
CPTokenURL string
CPTokenPublicKey string
}

type SimpleKubeReconciler struct {
Expand Down Expand Up @@ -582,6 +583,7 @@ func (r *SimpleKubeReconciler) buildEnv(ctx context.Context, session types.Sessi
envVar("USE_VERTEX", useVertex),
envVar("CLAUDE_CODE_USE_VERTEX", useVertex),
envVar("AMBIENT_CP_TOKEN_URL", r.cfg.CPTokenURL),
envVar("AMBIENT_CP_TOKEN_PUBLIC_KEY", r.cfg.CPTokenPublicKey),
envVar("AMBIENT_GRPC_URL", r.cfg.RunnerGRPCURL),
envVar("AMBIENT_GRPC_ENABLED", boolToStr(r.cfg.RunnerGRPCURL != "")),
envVar("AMBIENT_GRPC_USE_TLS", boolToStr(r.cfg.RunnerGRPCUseTLS)),
Expand Down Expand Up @@ -831,6 +833,7 @@ func (r *SimpleKubeReconciler) buildMCPSidecar() interface{} {
envVar("MCP_BIND_ADDR", fmt.Sprintf(":%d", mcpSidecarPort)),
envVar("AMBIENT_API_URL", r.cfg.MCPAPIServerURL),
envVar("AMBIENT_CP_TOKEN_URL", r.cfg.CPTokenURL),
envVar("AMBIENT_CP_TOKEN_PUBLIC_KEY", r.cfg.CPTokenPublicKey),
},
"resources": map[string]interface{}{
"requests": map[string]interface{}{
Expand Down
Loading
Loading