Encryption: Improve the DX of encryption operations within database transactions (#41654)

* Move user oauth info encryption away from db transaction

* Add encryption methods with support for db session reusability

pkg/services/login/authinfoservice/database.go

@@ -80,72 +80,73 @@ func (s *Implementation) GetAuthInfo(query *models.GetAuthInfoQuery) error {
 }
 
 func (s *Implementation) SetAuthInfo(cmd *models.SetAuthInfoCommand) error {
+	authUser := &models.UserAuth{
+		UserId:     cmd.UserId,
+		AuthModule: cmd.AuthModule,
+		AuthId:     cmd.AuthId,
+		Created:    getTime(),
+	}
+
+	if cmd.OAuthToken != nil {
+		secretAccessToken, err := s.encryptAndEncode(cmd.OAuthToken.AccessToken)
+		if err != nil {
+			return err
+		}
+		secretRefreshToken, err := s.encryptAndEncode(cmd.OAuthToken.RefreshToken)
+		if err != nil {
+			return err
+		}
+		secretTokenType, err := s.encryptAndEncode(cmd.OAuthToken.TokenType)
+		if err != nil {
+			return err
+		}
+
+		authUser.OAuthAccessToken = secretAccessToken
+		authUser.OAuthRefreshToken = secretRefreshToken
+		authUser.OAuthTokenType = secretTokenType
+		authUser.OAuthExpiry = cmd.OAuthToken.Expiry
+	}
+
 	return s.SQLStore.WithTransactionalDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
-		authUser := &models.UserAuth{
-			UserId:     cmd.UserId,
-			AuthModule: cmd.AuthModule,
-			AuthId:     cmd.AuthId,
-			Created:    getTime(),
-		}
-
-		if cmd.OAuthToken != nil {
-			secretAccessToken, err := s.encryptAndEncode(cmd.OAuthToken.AccessToken)
-			if err != nil {
-				return err
-			}
-			secretRefreshToken, err := s.encryptAndEncode(cmd.OAuthToken.RefreshToken)
-			if err != nil {
-				return err
-			}
-			secretTokenType, err := s.encryptAndEncode(cmd.OAuthToken.TokenType)
-			if err != nil {
-				return err
-			}
-
-			authUser.OAuthAccessToken = secretAccessToken
-			authUser.OAuthRefreshToken = secretRefreshToken
-			authUser.OAuthTokenType = secretTokenType
-			authUser.OAuthExpiry = cmd.OAuthToken.Expiry
-		}
-
 		_, err := sess.Insert(authUser)
 		return err
 	})
 }
 
 func (s *Implementation) UpdateAuthInfo(cmd *models.UpdateAuthInfoCommand) error {
+	authUser := &models.UserAuth{
+		UserId:     cmd.UserId,
+		AuthModule: cmd.AuthModule,
+		AuthId:     cmd.AuthId,
+		Created:    getTime(),
+	}
+
+	if cmd.OAuthToken != nil {
+		secretAccessToken, err := s.encryptAndEncode(cmd.OAuthToken.AccessToken)
+		if err != nil {
+			return err
+		}
+		secretRefreshToken, err := s.encryptAndEncode(cmd.OAuthToken.RefreshToken)
+		if err != nil {
+			return err
+		}
+		secretTokenType, err := s.encryptAndEncode(cmd.OAuthToken.TokenType)
+		if err != nil {
+			return err
+		}
+
+		authUser.OAuthAccessToken = secretAccessToken
+		authUser.OAuthRefreshToken = secretRefreshToken
+		authUser.OAuthTokenType = secretTokenType
+		authUser.OAuthExpiry = cmd.OAuthToken.Expiry
+	}
+
+	cond := &models.UserAuth{
+		UserId:     cmd.UserId,
+		AuthModule: cmd.AuthModule,
+	}
+
 	return s.SQLStore.WithTransactionalDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
-		authUser := &models.UserAuth{
-			UserId:     cmd.UserId,
-			AuthModule: cmd.AuthModule,
-			AuthId:     cmd.AuthId,
-			Created:    getTime(),
-		}
-
-		if cmd.OAuthToken != nil {
-			secretAccessToken, err := s.encryptAndEncode(cmd.OAuthToken.AccessToken)
-			if err != nil {
-				return err
-			}
-			secretRefreshToken, err := s.encryptAndEncode(cmd.OAuthToken.RefreshToken)
-			if err != nil {
-				return err
-			}
-			secretTokenType, err := s.encryptAndEncode(cmd.OAuthToken.TokenType)
-			if err != nil {
-				return err
-			}
-
-			authUser.OAuthAccessToken = secretAccessToken
-			authUser.OAuthRefreshToken = secretRefreshToken
-			authUser.OAuthTokenType = secretTokenType
-			authUser.OAuthExpiry = cmd.OAuthToken.Expiry
-		}
-
-		cond := &models.UserAuth{
-			UserId:     cmd.UserId,
-			AuthModule: cmd.AuthModule,
-		}
 		upd, err := sess.Update(authUser, cond)
 		s.logger.Debug("Updated user_auth", "user_id", cmd.UserId, "auth_module", cmd.AuthModule, "rows", upd)
 		return err

pkg/services/secrets/database/database.go

@@ -5,12 +5,10 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/grafana/grafana/pkg/services/secrets"
-
-	"xorm.io/xorm"
-
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/secrets"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"xorm.io/xorm"
 )
 
 const dataKeysTable = "data_keys"

pkg/services/secrets/fakes/fake_store.go

@@ -36,7 +36,7 @@ func (f FakeSecretsStore) CreateDataKey(_ context.Context, dataKey secrets.DataK
 	return nil
 }
 
-func (f FakeSecretsStore) CreateDataKeyWithDBSession(ctx context.Context, dataKey secrets.DataKey, sess *xorm.Session) error {
+func (f FakeSecretsStore) CreateDataKeyWithDBSession(_ context.Context, dataKey secrets.DataKey, _ *xorm.Session) error {
 	f.store[dataKey.Name] = &dataKey
 	return nil
 }

pkg/services/secrets/manager/manager.go

@@ -14,6 +14,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/secrets"
 	grafana "github.com/grafana/grafana/pkg/services/secrets/defaultprovider"
 	"github.com/grafana/grafana/pkg/setting"
+	"xorm.io/xorm"
 )
 
 const (
@@ -59,6 +60,10 @@ type dataKeyCacheItem struct {
 var b64 = base64.RawStdEncoding
 
 func (s *SecretsService) Encrypt(ctx context.Context, payload []byte, opt secrets.EncryptionOptions) ([]byte, error) {
+	return s.EncryptWithDBSession(ctx, payload, opt, nil)
+}
+
+func (s *SecretsService) EncryptWithDBSession(ctx context.Context, payload []byte, opt secrets.EncryptionOptions, sess *xorm.Session) ([]byte, error) {
 	// Use legacy encryption service if envelopeEncryptionFeatureToggle toggle is off
 	if !s.settings.IsFeatureToggleEnabled(envelopeEncryptionFeatureToggle) {
 		return s.enc.Encrypt(ctx, payload, setting.SecretKey)
@@ -71,7 +76,7 @@ func (s *SecretsService) Encrypt(ctx context.Context, payload []byte, opt secret
 	dataKey, err := s.dataKey(ctx, keyName)
 	if err != nil {
 		if errors.Is(err, secrets.ErrDataKeyNotFound) {
-			dataKey, err = s.newDataKey(ctx, keyName, scope)
+			dataKey, err = s.newDataKey(ctx, keyName, scope, sess)
 			if err != nil {
 				return nil, err
 			}
@@ -137,9 +142,13 @@ func (s *SecretsService) Decrypt(ctx context.Context, payload []byte) ([]byte, e
 }
 
 func (s *SecretsService) EncryptJsonData(ctx context.Context, kv map[string]string, opt secrets.EncryptionOptions) (map[string][]byte, error) {
+	return s.EncryptJsonDataWithDBSession(ctx, kv, opt, nil)
+}
+
+func (s *SecretsService) EncryptJsonDataWithDBSession(ctx context.Context, kv map[string]string, opt secrets.EncryptionOptions, sess *xorm.Session) (map[string][]byte, error) {
 	encrypted := make(map[string][]byte)
 	for key, value := range kv {
-		encryptedData, err := s.Encrypt(ctx, []byte(value), opt)
+		encryptedData, err := s.EncryptWithDBSession(ctx, []byte(value), opt, sess)
 		if err != nil {
 			return nil, err
 		}
@@ -185,7 +194,7 @@ func newRandomDataKey() ([]byte, error) {
 }
 
 // newDataKey creates a new random DEK, caches it and returns its value
-func (s *SecretsService) newDataKey(ctx context.Context, name string, scope string) ([]byte, error) {
+func (s *SecretsService) newDataKey(ctx context.Context, name string, scope string, sess *xorm.Session) ([]byte, error) {
 	// 1. Create new DEK
 	dataKey, err := newRandomDataKey()
 	if err != nil {
@@ -203,13 +212,20 @@ func (s *SecretsService) newDataKey(ctx context.Context, name string, scope stri
 	}
 
 	// 3. Store its encrypted value in db
-	err = s.store.CreateDataKey(ctx, secrets.DataKey{
+	dek := secrets.DataKey{
 		Active:        true, // TODO: right now we never mark a key as deactivated
 		Name:          name,
 		Provider:      s.currentProvider,
 		EncryptedData: encrypted,
 		Scope:         scope,
-	})
+	}
+
+	if sess == nil {
+		err = s.store.CreateDataKey(ctx, dek)
+	} else {
+		err = s.store.CreateDataKeyWithDBSession(ctx, dek, sess)
+	}
+
 	if err != nil {
 		return nil, err
 	}

pkg/services/secrets/secrets.go

@@ -9,10 +9,18 @@ import (
 // Service is an envelope encryption service in charge of encrypting/decrypting secrets.
 // It is a replacement for encryption.Service
 type Service interface {
+	// Encrypt MUST NOT be used within database transactions, it may cause database locks.
+	// For those specific use cases where the encryption operation cannot be moved outside
+	// the database transaction, look at database-specific methods present at the specific
+	// implementation present at manager.SecretsService.
 	Encrypt(ctx context.Context, payload []byte, opt EncryptionOptions) ([]byte, error)
 	Decrypt(ctx context.Context, payload []byte) ([]byte, error)
+
+	// EncryptJsonData MUST NOT be used within database transactions.
+	// Look at Encrypt method comment for further details.
 	EncryptJsonData(ctx context.Context, kv map[string]string, opt EncryptionOptions) (map[string][]byte, error)
 	DecryptJsonData(ctx context.Context, sjd map[string][]byte) (map[string]string, error)
+
 	GetDecryptedValue(ctx context.Context, sjd map[string][]byte, key, fallback string) string
 }