switching watcher to a sync run function rather than spawning a new goroutine so we can execute as a service within the auth server

Author: eriktate

go.mod

@@ -15,6 +15,7 @@ require (
 	cloud.google.com/go/storage v1.52.0
 	code.dny.dev/ssrf v0.2.0
 	connectrpc.com/connect v1.18.1
+	filippo.io/age v1.2.1
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.9.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/authorization/armauthorization/v2 v2.2.0
@@ -272,7 +273,6 @@ require (
 	cloud.google.com/go/monitoring v1.24.1 // indirect
 	cloud.google.com/go/pubsub v1.47.0 // indirect
 	dario.cat/mergo v1.0.1 // indirect
-	filippo.io/age v1.2.1
 	github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 // indirect
 	github.com/99designs/keyring v1.2.2 // indirect
 	github.com/AdaLogics/go-fuzz-headers v0.0.0-20230811130428-ced1acdcaa24 // indirect

lib/auth/recordingencryption/manager.go

@@ -24,15 +24,17 @@ import (
 	"iter"
 	"log/slog"
 	"slices"
+	"time"
 
 	"filippo.io/age"
 	"github.com/gravitational/trace"
 
+	"github.com/gravitational/teleport"
 	recordingencryptionv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/recordingencryption/v1"
 	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/api/utils/retryutils"
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/cryptosuites"
-	"github.com/gravitational/teleport/lib/events"
 	"github.com/gravitational/teleport/lib/services"
 )
 
@@ -51,10 +53,6 @@ type ManagerConfig struct {
 
 // NewManager returns a new Manager using the given ManagerConfig.
 func NewManager(cfg ManagerConfig) (*Manager, error) {
-	if cfg.Logger == nil {
-		cfg.Logger = slog.Default()
-	}
-
 	if cfg.Backend == nil {
 		return nil, trace.BadParameter("backend is required")
 	}
@@ -63,6 +61,10 @@ func NewManager(cfg ManagerConfig) (*Manager, error) {
 		return nil, trace.BadParameter("key store is required")
 	}
 
+	if cfg.Logger == nil {
+		cfg.Logger = slog.With(teleport.ComponentKey, "encryption-manager")
+	}
+
 	return &Manager{
 		RecordingEncryption: cfg.Backend,
 		keyStore:            cfg.KeyStore,
@@ -78,7 +80,6 @@ type Manager struct {
 
 	logger   *slog.Logger
 	keyStore EncryptionKeyStore
-	uploader events.MultipartUploader
 }
 
 // ensureActiveRecordingEncryption returns the configured RecordingEncryption resource if it exists with active keys. If it does not,
@@ -301,94 +302,130 @@ type RecordingEncryptionResolver interface {
 	ResolveRecordingEncryption(ctx context.Context) (*recordingencryptionv1.RecordingEncryption, error)
 }
 
-// WatchConfig captures required dependencies for building a RecordingEncyprtion watcher that
+// WatchConfig captures required dependencies for building a RecordingEncryption watcher that
 // automatically resolves state.
 type WatchConfig struct {
 	Events        types.Events
 	Resolver      RecordingEncryptionResolver
 	ClusterConfig services.ClusterConfiguration
 	Logger        *slog.Logger
-	LockConfig    backend.RunWhileLockedConfig
+	LockConfig    *backend.RunWhileLockedConfig
+}
+
+// A Watcher watches for changes to the RecordingEncryption resource and resolves the state for the calling
+// auth server.
+type Watcher struct {
+	events        types.Events
+	resolver      RecordingEncryptionResolver
+	clusterConfig services.ClusterConfiguration
+	logger        *slog.Logger
+	lockConfig    *backend.RunWhileLockedConfig
 }
 
-// Watch creates a watcher responsible for responding to changes in the RecordingEncryption
-// resource. This is how auth servers cooperate and ensure there are accessible wrapped keys for each unique
-// keystore configuration in a cluster.
-func Watch(ctx context.Context, cfg WatchConfig) error {
+// NewWatcher returns a new Watcher.
+func NewWatcher(cfg WatchConfig) (*Watcher, error) {
 	switch {
 	case cfg.Events == nil:
-		return trace.BadParameter("events is required")
+		return nil, trace.BadParameter("events is required")
 	case cfg.Resolver == nil:
-		return trace.BadParameter("recording encryption resolver is required")
+		return nil, trace.BadParameter("recording encryption resolver is required")
 	case cfg.ClusterConfig == nil:
-		return trace.BadParameter("cluster config backend is required")
+		return nil, trace.BadParameter("cluster config backend is required")
+	case cfg.LockConfig == nil:
+		return nil, trace.BadParameter("lock config is required")
 	}
 	if cfg.Logger == nil {
-		cfg.Logger = slog.Default()
+		cfg.Logger = slog.With(teleport.ComponentKey, "encryption-watcher")
 	}
 
-	cfg.Logger.DebugContext(ctx, "creating recording_encryption watcher")
-	w, err := cfg.Events.NewWatcher(ctx, types.Watch{
-		Name: "recording_encryption_watcher",
-		Kinds: []types.WatchKind{
-			{
-				Kind: types.KindRecordingEncryption,
-			},
-		},
-	})
-	if err != nil {
-		return trace.Wrap(err)
+	return &Watcher{
+		events:        cfg.Events,
+		resolver:      cfg.Resolver,
+		clusterConfig: cfg.ClusterConfig,
+		logger:        cfg.Logger,
+		lockConfig:    cfg.LockConfig,
+	}, nil
+}
+
+// Watch creates a watcher responsible for responding to changes in the RecordingEncryption resource.
+// This is how auth servers cooperate and ensure there are accessible wrapped keys for each unique keystore
+// configuration in a cluster.
+func (w *Watcher) Run(ctx context.Context) (err error) {
+	jitter := func() {
+		<-time.After(retryutils.SeventhJitter(time.Second * 5))
 	}
 
-	go func() {
+	defer func() {
+		w.logger.InfoContext(ctx, "stopping encryption watcher", "error", err)
+	}()
+
+	for {
+		watch, err := w.events.NewWatcher(ctx, types.Watch{
+			Name: "recording_encryption_watcher",
+			Kinds: []types.WatchKind{
+				{
+					Kind: types.KindRecordingEncryption,
+				},
+			},
+		})
+		if err != nil {
+			w.logger.ErrorContext(ctx, "failed to create watcher, retrying", "error", err)
+			jitter()
+		}
+
+	HandleEvents:
 		for {
+			err := w.handleRecordingEncryptionChange(ctx)
+			if err != nil {
+				w.logger.ErrorContext(ctx, "failed to handle session recording config change", "error", err)
+				jitter()
+				continue
+
+			}
+
 			select {
-			case ev := <-w.Events():
+			case ev := <-watch.Events():
 				if ev.Type != types.OpPut {
 					continue
 				}
-				const retries = 3
-				for tries := range retries {
-					err := handleRecordingEncryptionChange(ctx, cfg)
-					if err == nil {
-						break
-					}
-
-					cfg.Logger.ErrorContext(ctx, "failed to handle session recording config change", "error", err, "remaining_tries", retries-tries-1)
+			case <-watch.Done():
+				if err := watch.Error(); err == nil {
+					return nil
 				}
 
-			case <-w.Done():
-				cfg.Logger.DebugContext(ctx, "no longer watching recording_encryption")
-				return
+				w.logger.ErrorContext(ctx, "watcher failed, retrying", "error", err)
+				jitter()
+				break HandleEvents
+			case <-ctx.Done():
+				watch.Close()
+				return ctx.Err()
 			}
 		}
-	}()
-
-	return nil
+	}
 }
 
 // this helper handles reacting to individual Put events on the RecordingEncryption resource and updates the
 // SessionRecordingConfig with the results, if necessary
-func handleRecordingEncryptionChange(ctx context.Context, cfg WatchConfig) error {
-	return trace.Wrap(backend.RunWhileLocked(ctx, cfg.LockConfig, func(ctx context.Context) error {
-		recConfig, err := cfg.ClusterConfig.GetSessionRecordingConfig(ctx)
+func (w *Watcher) handleRecordingEncryptionChange(ctx context.Context) error {
+	return trace.Wrap(backend.RunWhileLocked(ctx, *w.lockConfig, func(ctx context.Context) error {
+		recConfig, err := w.clusterConfig.GetSessionRecordingConfig(ctx)
 		if err != nil {
 			return trace.Wrap(err, "fetching recording config")
 		}
 
 		if !recConfig.GetEncrypted() {
-			cfg.Logger.DebugContext(ctx, "session recording encryption disabled, skip resolving keys")
+			w.logger.DebugContext(ctx, "session recording encryption disabled, skip resolving keys")
 			return nil
 		}
 
-		encryption, err := cfg.Resolver.ResolveRecordingEncryption(ctx)
+		encryption, err := w.resolver.ResolveRecordingEncryption(ctx)
 		if err != nil {
-			cfg.Logger.ErrorContext(ctx, "failed to resolve recording encryption state", "error", err)
+			w.logger.ErrorContext(ctx, "failed to resolve recording encryption state", "error", err)
 			return trace.Wrap(err, "resolving recording encryption")
 		}
 
 		if recConfig.SetEncryptionKeys(GetAgeEncryptionKeys(encryption.GetSpec().ActiveKeys)) {
-			_, err = cfg.ClusterConfig.UpdateSessionRecordingConfig(ctx, recConfig)
+			_, err = w.clusterConfig.UpdateSessionRecordingConfig(ctx, recConfig)
 			return trace.Wrap(err, "updating encryption keys")
 		}
 

lib/auth/recordingencryption/manager_test.go

@@ -19,7 +19,6 @@ package recordingencryption_test
 import (
 	"context"
 	"crypto"
-	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
@@ -58,11 +57,16 @@ type fakeEncryptionKeyStore struct {
 }
 
 func (f *fakeEncryptionKeyStore) NewEncryptionKeyPair(ctx context.Context, purpose cryptosuites.KeyPurpose) (*types.EncryptionKeyPair, error) {
-	private, err := rsa.GenerateKey(rand.Reader, 2048)
+	decrypter, err := cryptosuites.GenerateDecrypterWithAlgorithm(cryptosuites.RSA2048)
 	if err != nil {
 		return nil, err
 	}
 
+	private, ok := decrypter.(*rsa.PrivateKey)
+	if !ok {
+		return nil, errors.New("expected RSA private key")
+	}
+
 	privatePEM := pem.EncodeToMemory(&pem.Block{
 		Type:  keys.PKCS1PrivateKeyType,
 		Bytes: x509.MarshalPKCS1PrivateKey(private),
@@ -139,7 +143,7 @@ func TestResolveRecordingEncryption(t *testing.T) {
 	require.NoError(t, err)
 	activeKeys := encryption.GetSpec().GetActiveKeys()
 
-	require.Equal(t, 1, len(activeKeys))
+	require.Len(t, activeKeys, 1)
 	firstKey := activeKeys[0]
 
 	// should generate a wrapped key with the initial recording encryption pair
@@ -151,7 +155,7 @@ func TestResolveRecordingEncryption(t *testing.T) {
 	require.NoError(t, err)
 
 	activeKeys = encryption.GetSpec().ActiveKeys
-	require.Equal(t, 2, len(activeKeys))
+	require.Len(t, activeKeys, 2)
 	for _, key := range activeKeys {
 		require.NotNil(t, key.KeyEncryptionPair)
 		if key.KeyEncryptionPair.PrivateKeyType == serviceAType {
@@ -165,7 +169,7 @@ func TestResolveRecordingEncryption(t *testing.T) {
 	encryption, err = serviceB.ResolveRecordingEncryption(ctx)
 	require.NoError(t, err)
 	activeKeys = encryption.GetSpec().ActiveKeys
-	require.Equal(t, 2, len(activeKeys))
+	require.Len(t, activeKeys, 2)
 	for _, key := range activeKeys {
 		require.NotNil(t, key.KeyEncryptionPair)
 		if key.KeyEncryptionPair.PrivateKeyType == serviceAType {
@@ -179,7 +183,7 @@ func TestResolveRecordingEncryption(t *testing.T) {
 	encryption, err = serviceA.ResolveRecordingEncryption(ctx)
 	require.NoError(t, err)
 	activeKeys = encryption.GetSpec().ActiveKeys
-	require.Equal(t, 2, len(activeKeys))
+	require.Len(t, activeKeys, 2)
 	for _, key := range activeKeys {
 		require.NotNil(t, key.KeyEncryptionPair)
 		require.NotNil(t, key.RecordingEncryptionPair)