addressing some more PR feedback

Author: eriktate

lib/auth/recordingencryption/manager.go

@@ -40,16 +40,16 @@ import (
 	"github.com/gravitational/teleport/lib/services"
 )
 
-// EncryptionKeyStore provides methods for interacting with encryption keys.
-type EncryptionKeyStore interface {
+// KeyStore provides methods for interacting with encryption keys.
+type KeyStore interface {
 	NewEncryptionKeyPair(ctx context.Context, purpose cryptosuites.KeyPurpose) (*types.EncryptionKeyPair, error)
 	GetDecrypter(ctx context.Context, keyPair *types.EncryptionKeyPair) (crypto.Decrypter, error)
 }
 
 // ManagerConfig captures all of the dependencies required to instantiate a Manager.
 type ManagerConfig struct {
 	Backend  services.RecordingEncryption
-	KeyStore EncryptionKeyStore
+	KeyStore KeyStore
 	Logger   *slog.Logger
 }
 
@@ -74,32 +74,30 @@ func NewManager(cfg ManagerConfig) (*Manager, error) {
 	}, nil
 }
 
-// A Manager wraps a services.RecordingEncryption and EncryptionKeyStore in order to provide more complex operations
+// A Manager wraps a services.RecordingEncryption and KeyStore in order to provide more complex operations
 // than the CRUD methods exposed by services.RecordingEncryption. It primarily handles resolving RecordingEncryption
 // state and searching for accessible decryption keys.
 type Manager struct {
 	services.RecordingEncryption
 
 	logger   *slog.Logger
-	keyStore EncryptionKeyStore
+	keyStore KeyStore
 }
 
 // ensureActiveRecordingEncryption returns the configured RecordingEncryption resource if it exists with active keys. If it does not,
 // then the resource will be created or updated with a new active keypair. The bool return value indicates whether or not
 // a new pair was provisioned.
 func (m *Manager) ensureActiveRecordingEncryption(ctx context.Context) (*recordingencryptionv1.RecordingEncryption, bool, error) {
-	upsert := m.UpdateRecordingEncryption
+	persistFn := m.UpdateRecordingEncryption
 	encryption, err := m.GetRecordingEncryption(ctx)
 	if err != nil {
 		if !trace.IsNotFound(err) {
 			return encryption, false, trace.Wrap(err)
 		}
 		encryption = &recordingencryptionv1.RecordingEncryption{
-			Spec: &recordingencryptionv1.RecordingEncryptionSpec{
-				ActiveKeys: nil,
-			},
+			Spec: &recordingencryptionv1.RecordingEncryptionSpec{},
 		}
-		upsert = m.CreateRecordingEncryption
+		persistFn = m.CreateRecordingEncryption
 	}
 
 	activeKeys := encryption.GetSpec().ActiveKeys
@@ -133,7 +131,7 @@ func (m *Manager) ensureActiveRecordingEncryption(ctx context.Context) (*recordi
 		},
 	}
 	encryption.Spec.ActiveKeys = []*recordingencryptionv1.WrappedKey{&wrappedKey}
-	encryption, err = upsert(ctx, encryption)
+	encryption, err = persistFn(ctx, encryption)
 	if err != nil {
 		return encryption, false, trace.Wrap(err)
 	}
@@ -178,13 +176,19 @@ func (m *Manager) getRecordingEncryptionKeyPair(ctx context.Context, keys []*rec
 }
 
 // ResolveRecordingEncryption examines the current state of the RescordingEncryption resource and advances it to the
-// next state on behalf of the current auth server. When no active recording encryption key pairs exist, the first
-// pair will be generated and wrapped using a new key encryption pair generated by the Manager's keystore. When at
-// least one active keypair exists but none are accessible to the Manager's keystore, a new key encryption pair will
-// be generated and saved without a key encryption key pair. This is an unfulfilled key that some other instance of
-// Manager on another auth server will need to fulfill asynchrnously. If at least one active key is accessible ot the
-// Manager's keystore, then unfulfilled keys (identified by missing recording encryption key pairs) will be fulfilled
-// using their public keys. If there are no unfulfilled keys present, then nothing should be done.
+// next state on behalf of the current auth server.
+//
+// When no active recording encryption key pairs exist, the first pair will be generated and wrapped using a new key
+// encryption pair generated by the Manager's keystore.
+//
+// When at least one active keypair exists but none are accessible to the Manager's keystore, a new key encryption pair
+// will be generated and saved without a key encryption pair. This is an unfulfilled key that some other instance of
+// Manager on another auth server will need to fulfill asynchronously.
+//
+// If at least one active key is accessible to the Manager's keystore, then unfulfilled keys (identified by missing
+// recording encryption key pairs) will be fulfilled using their public key encryption keys.
+//
+// If there are no unfulfilled keys present, then nothing should be done.
 func (m *Manager) ResolveRecordingEncryption(ctx context.Context) (*recordingencryptionv1.RecordingEncryption, error) {
 	encryption, generatedKey, err := m.ensureActiveRecordingEncryption(ctx)
 	if err != nil {
@@ -308,16 +312,16 @@ func GetAgeEncryptionKeys(keys []*recordingencryptionv1.WrappedKey) iter.Seq[*ty
 	}
 }
 
-// RecordingEncryptionResolver resolves RecordingEncryption state
-type RecordingEncryptionResolver interface {
+// Resolver resolves RecordingEncryption state
+type Resolver interface {
 	ResolveRecordingEncryption(ctx context.Context) (*recordingencryptionv1.RecordingEncryption, error)
 }
 
 // WatchConfig captures required dependencies for building a RecordingEncryption watcher that
 // automatically resolves state.
 type WatchConfig struct {
 	Events        types.Events
-	Resolver      RecordingEncryptionResolver
+	Resolver      Resolver
 	ClusterConfig services.ClusterConfiguration
 	Logger        *slog.Logger
 	LockConfig    *backend.RunWhileLockedConfig
@@ -327,7 +331,7 @@ type WatchConfig struct {
 // auth server.
 type Watcher struct {
 	events        types.Events
-	resolver      RecordingEncryptionResolver
+	resolver      Resolver
 	clusterConfig services.ClusterConfiguration
 	logger        *slog.Logger
 	lockConfig    *backend.RunWhileLockedConfig
@@ -362,8 +366,14 @@ func NewWatcher(cfg WatchConfig) (*Watcher, error) {
 // 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))
+	// jitter returns a bool specifiying whether or not execution should continue
+	jitter := func() bool {
+		select {
+		case <-time.After(retryutils.SeventhJitter(time.Second * 5)):
+			return true
+		case <-ctx.Done():
+			return false
+		}
 	}
 
 	defer func() {
@@ -381,16 +391,20 @@ func (w *Watcher) Run(ctx context.Context) (err error) {
 		})
 		if err != nil {
 			w.logger.ErrorContext(ctx, "failed to create watcher, retrying", "error", err)
-			jitter()
+			if !jitter() {
+				return nil
+			}
 			continue
 		}
+		defer watch.Close()
 
 	HandleEvents:
 		for {
 			err := w.handleRecordingEncryptionChange(ctx)
 			if err != nil {
-				w.logger.ErrorContext(ctx, "failed to handle session recording config change", "error", err)
-				jitter()
+				if !jitter() {
+					return nil
+				}
 				continue
 
 			}
@@ -406,11 +420,12 @@ func (w *Watcher) Run(ctx context.Context) (err error) {
 				}
 
 				w.logger.ErrorContext(ctx, "watcher failed, retrying", "error", err)
-				jitter()
+				if !jitter() {
+					return nil
+				}
 				break HandleEvents
 			case <-ctx.Done():
-				watch.Close()
-				return ctx.Err()
+				return nil
 			}
 		}
 	}