ipn/{ipnlocal,store},kube/kubeclient: store TLS cert and key pair to a Secret in a single operation. (#15147)

To avoid duplicate issuances/slowness while the state Secret
contains a mismatched cert and key.

Updates tailscale/tailscale#15134
Updates tailscale/corp#24795

Signed-off-by: Irbe Krumina <irbe@tailscale.com>

ipn/ipnlocal/cert.go

@@ -250,15 +250,13 @@ type certStore interface {
 	// for now. If they're expired, it returns errCertExpired.
 	// If they don't exist, it returns ipn.ErrStateNotExist.
 	Read(domain string, now time.Time) (*TLSCertKeyPair, error)
-	// WriteCert writes the cert for domain.
-	WriteCert(domain string, cert []byte) error
-	// WriteKey writes the key for domain.
-	WriteKey(domain string, key []byte) error
 	// ACMEKey returns the value previously stored via WriteACMEKey.
 	// It is a PEM encoded ECDSA key.
 	ACMEKey() ([]byte, error)
 	// WriteACMEKey stores the provided PEM encoded ECDSA key.
 	WriteACMEKey([]byte) error
+	// WriteTLSCertAndKey writes the cert and key for domain.
+	WriteTLSCertAndKey(domain string, cert, key []byte) error
 }
 
 var errCertExpired = errors.New("cert expired")
@@ -344,6 +342,13 @@ func (f certFileStore) WriteKey(domain string, key []byte) error {
 	return atomicfile.WriteFile(keyFile(f.dir, domain), key, 0600)
 }
 
+func (f certFileStore) WriteTLSCertAndKey(domain string, cert, key []byte) error {
+	if err := f.WriteKey(domain, key); err != nil {
+		return err
+	}
+	return f.WriteCert(domain, cert)
+}
+
 // certStateStore implements certStore by storing the cert & key files in an ipn.StateStore.
 type certStateStore struct {
 	ipn.StateStore
@@ -384,6 +389,27 @@ func (s certStateStore) WriteACMEKey(key []byte) error {
 	return ipn.WriteState(s.StateStore, ipn.StateKey(acmePEMName), key)
 }
 
+// TLSCertKeyWriter is an interface implemented by state stores that can write the TLS
+// cert and key in a single atomic operation. Currently this is only implemented
+// by the kubestore.StoreKube.
+type TLSCertKeyWriter interface {
+	WriteTLSCertAndKey(domain string, cert, key []byte) error
+}
+
+// WriteTLSCertAndKey writes the TLS cert and key for domain to the current
+// LocalBackend's StateStore.
+func (s certStateStore) WriteTLSCertAndKey(domain string, cert, key []byte) error {
+	// If we're using a store that supports atomic writes, use that.
+	if aw, ok := s.StateStore.(TLSCertKeyWriter); ok {
+		return aw.WriteTLSCertAndKey(domain, cert, key)
+	}
+	// Otherwise fall back to separate writes for cert and key.
+	if err := s.WriteKey(domain, key); err != nil {
+		return err
+	}
+	return s.WriteCert(domain, cert)
+}
+
 // TLSCertKeyPair is a TLS public and private key, and whether they were obtained
 // from cache or freshly obtained.
 type TLSCertKeyPair struct {
@@ -546,9 +572,6 @@ func (b *LocalBackend) getCertPEM(ctx context.Context, cs certStore, logf logger
 	if err := encodeECDSAKey(&privPEM, certPrivKey); err != nil {
 		return nil, err
 	}
-	if err := cs.WriteKey(domain, privPEM.Bytes()); err != nil {
-		return nil, err
-	}
 
 	csr, err := certRequest(certPrivKey, domain, nil)
 	if err != nil {
@@ -570,7 +593,7 @@ func (b *LocalBackend) getCertPEM(ctx context.Context, cs certStore, logf logger
 			return nil, err
 		}
 	}
-	if err := cs.WriteCert(domain, certPEM.Bytes()); err != nil {
+	if err := cs.WriteTLSCertAndKey(domain, certPEM.Bytes(), privPEM.Bytes()); err != nil {
 		return nil, err
 	}
 	b.domainRenewed(domain)

ipn/ipnlocal/cert_test.go

@@ -86,13 +86,9 @@ func TestCertStoreRoundTrip(t *testing.T) {
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			if err := test.store.WriteCert(testDomain, testCert); err != nil {
-				t.Fatalf("WriteCert: unexpected error: %v", err)
+			if err := test.store.WriteTLSCertAndKey(testDomain, testCert, testKey); err != nil {
+				t.Fatalf("WriteTLSCertAndKey: unexpected error: %v", err)
 			}
-			if err := test.store.WriteKey(testDomain, testKey); err != nil {
-				t.Fatalf("WriteKey: unexpected error: %v", err)
-			}
-
 			kp, err := test.store.Read(testDomain, testNow)
 			if err != nil {
 				t.Fatalf("Read: unexpected error: %v", err)

ipn/store/kubestore/store_kube.go

@@ -18,6 +18,7 @@ import (
 	"tailscale.com/kube/kubeapi"
 	"tailscale.com/kube/kubeclient"
 	"tailscale.com/types/logger"
+	"tailscale.com/util/mak"
 )
 
 const (
@@ -83,10 +84,22 @@ func (s *Store) ReadState(id ipn.StateKey) ([]byte, error) {
 
 // WriteState implements the StateStore interface.
 func (s *Store) WriteState(id ipn.StateKey, bs []byte) (err error) {
+	return s.updateStateSecret(map[string][]byte{string(id): bs})
+}
+
+// WriteTLSCertAndKey writes a TLS cert and key to domain.crt, domain.key fields of a Tailscale Kubernetes node's state
+// Secret.
+func (s *Store) WriteTLSCertAndKey(domain string, cert, key []byte) error {
+	return s.updateStateSecret(map[string][]byte{domain + ".crt": cert, domain + ".key": key})
+}
+
+func (s *Store) updateStateSecret(data map[string][]byte) (err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer func() {
 		if err == nil {
-			s.memory.WriteState(ipn.StateKey(sanitizeKey(id)), bs)
+			for id, bs := range data {
+				s.memory.WriteState(ipn.StateKey(id), bs)
+			}
 		}
 		if err != nil {
 			if err := s.client.Event(ctx, eventTypeWarning, reasonTailscaleStateUpdateFailed, err.Error()); err != nil {
@@ -99,9 +112,9 @@ func (s *Store) WriteState(id ipn.StateKey, bs []byte) (err error) {
 		}
 		cancel()
 	}()
-
 	secret, err := s.client.GetSecret(ctx, s.secretName)
 	if err != nil {
+		// If the Secret does not exist, create it with the required data.
 		if kubeclient.IsNotFoundErr(err) {
 			return s.client.CreateSecret(ctx, &kubeapi.Secret{
 				TypeMeta: kubeapi.TypeMeta{
@@ -111,40 +124,53 @@ func (s *Store) WriteState(id ipn.StateKey, bs []byte) (err error) {
 				ObjectMeta: kubeapi.ObjectMeta{
 					Name: s.secretName,
 				},
-				Data: map[string][]byte{
-					sanitizeKey(id): bs,
-				},
+				Data: func(m map[string][]byte) map[string][]byte {
+					d := make(map[string][]byte, len(m))
+					for key, val := range m {
+						d[sanitizeKey(key)] = val
+					}
+					return d
+				}(data),
 			})
 		}
 		return err
 	}
 	if s.canPatch {
-		if len(secret.Data) == 0 { // if user has pre-created a blank Secret
-			m := []kubeclient.JSONPatch{
+		var m []kubeclient.JSONPatch
+		// If the user has pre-created a Secret with no data, we need to ensure the top level /data field.
+		if len(secret.Data) == 0 {
+			m = []kubeclient.JSONPatch{
 				{
-					Op:    "add",
-					Path:  "/data",
-					Value: map[string][]byte{sanitizeKey(id): bs},
+					Op:   "add",
+					Path: "/data",
+					Value: func(m map[string][]byte) map[string][]byte {
+						d := make(map[string][]byte, len(m))
+						for key, val := range m {
+							d[sanitizeKey(key)] = val
+						}
+						return d
+					}(data),
 				},
 			}
-			if err := s.client.JSONPatchResource(ctx, s.secretName, kubeclient.TypeSecrets, m); err != nil {
-				return fmt.Errorf("error patching Secret %s with a /data field: %v", s.secretName, err)
+			// If the Secret has data, patch it with the new data.
+		} else {
+			for key, val := range data {
+				m = append(m, kubeclient.JSONPatch{
+					Op:    "add",
+					Path:  "/data/" + sanitizeKey(key),
+					Value: val,
+				})
 			}
-			return nil
-		}
-		m := []kubeclient.JSONPatch{
-			{
-				Op:    "add",
-				Path:  "/data/" + sanitizeKey(id),
-				Value: bs,
-			},
 		}
 		if err := s.client.JSONPatchResource(ctx, s.secretName, kubeclient.TypeSecrets, m); err != nil {
-			return fmt.Errorf("error patching Secret %s with /data/%s field: %v", s.secretName, sanitizeKey(id), err)
+			return fmt.Errorf("error patching Secret %s: %w", s.secretName, err)
 		}
 		return nil
 	}
-	secret.Data[sanitizeKey(id)] = bs
+	// No patch permissions, use UPDATE instead.
+	for key, val := range data {
+		mak.Set(&secret.Data, sanitizeKey(key), val)
+	}
 	if err := s.client.UpdateSecret(ctx, secret); err != nil {
 		return err
 	}
@@ -172,9 +198,9 @@ func (s *Store) loadState() (err error) {
 	return nil
 }
 
-func sanitizeKey(k ipn.StateKey) string {
-	// The only valid characters in a Kubernetes secret key are alphanumeric, -,
-	// _, and .
+// sanitizeKey converts any value that can be converted to a string into a valid Kubernetes secret key.
+// Valid characters are alphanumeric, -, _, and .
+func sanitizeKey[T ~string](k T) string {
 	return strings.Map(func(r rune) rune {
 		if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '-' || r == '_' || r == '.' {
 			return r

ipn/store/kubestore/store_kube_test.go

@@ -0,0 +1,183 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package kubestore
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"tailscale.com/ipn"
+	"tailscale.com/ipn/store/mem"
+	"tailscale.com/kube/kubeapi"
+	"tailscale.com/kube/kubeclient"
+)
+
+func TestUpdateStateSecret(t *testing.T) {
+	tests := []struct {
+		name       string
+		initial    map[string][]byte
+		updates    map[string][]byte
+		wantData   map[string][]byte
+		allowPatch bool
+	}{
+		{
+			name: "basic_update",
+			initial: map[string][]byte{
+				"existing": []byte("old"),
+			},
+			updates: map[string][]byte{
+				"foo": []byte("bar"),
+			},
+			wantData: map[string][]byte{
+				"existing": []byte("old"),
+				"foo":      []byte("bar"),
+			},
+			allowPatch: true,
+		},
+		{
+			name: "update_existing",
+			initial: map[string][]byte{
+				"foo": []byte("old"),
+			},
+			updates: map[string][]byte{
+				"foo": []byte("new"),
+			},
+			wantData: map[string][]byte{
+				"foo": []byte("new"),
+			},
+			allowPatch: true,
+		},
+		{
+			name: "multiple_updates",
+			initial: map[string][]byte{
+				"keep": []byte("keep"),
+			},
+			updates: map[string][]byte{
+				"foo": []byte("bar"),
+				"baz": []byte("qux"),
+			},
+			wantData: map[string][]byte{
+				"keep": []byte("keep"),
+				"foo":  []byte("bar"),
+				"baz":  []byte("qux"),
+			},
+			allowPatch: true,
+		},
+		{
+			name: "create_new_secret",
+			updates: map[string][]byte{
+				"foo": []byte("bar"),
+			},
+			wantData: map[string][]byte{
+				"foo": []byte("bar"),
+			},
+			allowPatch: true,
+		},
+		{
+			name: "patch_denied",
+			initial: map[string][]byte{
+				"foo": []byte("old"),
+			},
+			updates: map[string][]byte{
+				"foo": []byte("new"),
+			},
+			wantData: map[string][]byte{
+				"foo": []byte("new"),
+			},
+			allowPatch: false,
+		},
+		{
+			name: "sanitize_keys",
+			initial: map[string][]byte{
+				"clean-key": []byte("old"),
+			},
+			updates: map[string][]byte{
+				"dirty@key": []byte("new"),
+				"also/bad":  []byte("value"),
+				"good.key":  []byte("keep"),
+			},
+			wantData: map[string][]byte{
+				"clean-key": []byte("old"),
+				"dirty_key": []byte("new"),
+				"also_bad":  []byte("value"),
+				"good.key":  []byte("keep"),
+			},
+			allowPatch: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			secret := tt.initial // track current state
+			client := &kubeclient.FakeClient{
+				GetSecretImpl: func(ctx context.Context, name string) (*kubeapi.Secret, error) {
+					if secret == nil {
+						return nil, &kubeapi.Status{Code: 404}
+					}
+					return &kubeapi.Secret{Data: secret}, nil
+				},
+				CheckSecretPermissionsImpl: func(ctx context.Context, name string) (bool, bool, error) {
+					return tt.allowPatch, true, nil
+				},
+				CreateSecretImpl: func(ctx context.Context, s *kubeapi.Secret) error {
+					secret = s.Data
+					return nil
+				},
+				UpdateSecretImpl: func(ctx context.Context, s *kubeapi.Secret) error {
+					secret = s.Data
+					return nil
+				},
+				JSONPatchResourceImpl: func(ctx context.Context, name, resourceType string, patches []kubeclient.JSONPatch) error {
+					if !tt.allowPatch {
+						return &kubeapi.Status{Reason: "Forbidden"}
+					}
+					if secret == nil {
+						secret = make(map[string][]byte)
+					}
+					for _, p := range patches {
+						if p.Op == "add" && p.Path == "/data" {
+							secret = p.Value.(map[string][]byte)
+						} else if p.Op == "add" && strings.HasPrefix(p.Path, "/data/") {
+							key := strings.TrimPrefix(p.Path, "/data/")
+							secret[key] = p.Value.([]byte)
+						}
+					}
+					return nil
+				},
+			}
+
+			s := &Store{
+				client:     client,
+				canPatch:   tt.allowPatch,
+				secretName: "test-secret",
+				memory:     mem.Store{},
+			}
+
+			err := s.updateStateSecret(tt.updates)
+			if err != nil {
+				t.Errorf("updateStateSecret() error = %v", err)
+				return
+			}
+
+			// Verify secret data
+			if diff := cmp.Diff(secret, tt.wantData); diff != "" {
+				t.Errorf("secret data mismatch (-got +want):\n%s", diff)
+			}
+
+			// Verify memory store was updated
+			for k, v := range tt.updates {
+				got, err := s.memory.ReadState(ipn.StateKey(k))
+				if err != nil {
+					t.Errorf("reading from memory store: %v", err)
+					continue
+				}
+				if !cmp.Equal(got, v) {
+					t.Errorf("memory store key %q = %v, want %v", k, got, v)
+				}
+			}
+		})
+	}
+}

kube/kubeclient/fake_client.go

@@ -15,6 +15,9 @@ var _ Client = &FakeClient{}
 type FakeClient struct {
 	GetSecretImpl              func(context.Context, string) (*kubeapi.Secret, error)
 	CheckSecretPermissionsImpl func(ctx context.Context, name string) (bool, bool, error)
+	CreateSecretImpl           func(context.Context, *kubeapi.Secret) error
+	UpdateSecretImpl           func(context.Context, *kubeapi.Secret) error
+	JSONPatchResourceImpl      func(context.Context, string, string, []JSONPatch) error
 }
 
 func (fc *FakeClient) CheckSecretPermissions(ctx context.Context, name string) (bool, bool, error) {
@@ -33,8 +36,12 @@ func (fc *FakeClient) Event(context.Context, string, string, string) error {
 	return nil
 }
 
-func (fc *FakeClient) JSONPatchResource(context.Context, string, string, []JSONPatch) error {
-	return nil
+func (fc *FakeClient) JSONPatchResource(ctx context.Context, resource, name string, patches []JSONPatch) error {
+	return fc.JSONPatchResourceImpl(ctx, resource, name, patches)
+}
+func (fc *FakeClient) UpdateSecret(ctx context.Context, secret *kubeapi.Secret) error {
+	return fc.UpdateSecretImpl(ctx, secret)
+}
+func (fc *FakeClient) CreateSecret(ctx context.Context, secret *kubeapi.Secret) error {
+	return fc.CreateSecretImpl(ctx, secret)
 }
-func (fc *FakeClient) UpdateSecret(context.Context, *kubeapi.Secret) error { return nil }
-func (fc *FakeClient) CreateSecret(context.Context, *kubeapi.Secret) error { return nil }