db: use `context` and go-mockgen for `TwoFactorsStore` (#7045)

1 files changed, 18 insertions, 17 deletions

diff --git a/internal/db/two_factors.go b/internal/db/two_factors.go
index 935f66db..2dba2a02 100644
--- a/internal/db/two_factors.go
+++ b/internal/db/two_factors.go
@@ -5,6 +5,7 @@
 package db
 
 import (
+	"context"
 	"encoding/base64"
 	"fmt"
 	"strings"
@@ -23,21 +24,21 @@ import (
 //
 // NOTE: All methods are sorted in alphabetical order.
 type TwoFactorsStore interface {
-	// Create creates a new 2FA token and recovery codes for given user.
-	// The "key" is used to encrypt and later decrypt given "secret",
-	// which should be configured in site-level and change of the "key"
-	// will break all existing 2FA tokens.
-	Create(userID int64, key, secret string) error
-	// GetByUserID returns the 2FA token of given user.
-	// It returns ErrTwoFactorNotFound when not found.
-	GetByUserID(userID int64) (*TwoFactor, error)
+	// Create creates a new 2FA token and recovery codes for given user. The "key"
+	// is used to encrypt and later decrypt given "secret", which should be
+	// configured in site-level and change of the "key" will break all existing 2FA
+	// tokens.
+	Create(ctx context.Context, userID int64, key, secret string) error
+	// GetByUserID returns the 2FA token of given user. It returns
+	// ErrTwoFactorNotFound when not found.
+	GetByUserID(ctx context.Context, userID int64) (*TwoFactor, error)
 	// IsUserEnabled returns true if the user has enabled 2FA.
-	IsUserEnabled(userID int64) bool
+	IsUserEnabled(ctx context.Context, userID int64) bool
 }
 
 var TwoFactors TwoFactorsStore
 
-// NOTE: This is a GORM create hook.
+// BeforeCreate implements the GORM create hook.
 func (t *TwoFactor) BeforeCreate(tx *gorm.DB) error {
 	if t.CreatedUnix == 0 {
 		t.CreatedUnix = tx.NowFunc().Unix()
@@ -45,7 +46,7 @@ func (t *TwoFactor) BeforeCreate(tx *gorm.DB) error {
 	return nil
 }
 
-// NOTE: This is a GORM query hook.
+// AfterFind implements the GORM query hook.
 func (t *TwoFactor) AfterFind(_ *gorm.DB) error {
 	t.Created = time.Unix(t.CreatedUnix, 0).Local()
 	return nil
@@ -57,7 +58,7 @@ type twoFactors struct {
 	*gorm.DB
 }
 
-func (db *twoFactors) Create(userID int64, key, secret string) error {
+func (db *twoFactors) Create(ctx context.Context, userID int64, key, secret string) error {
 	encrypted, err := cryptoutil.AESGCMEncrypt(cryptoutil.MD5Bytes(key), []byte(secret))
 	if err != nil {
 		return errors.Wrap(err, "encrypt secret")
@@ -72,7 +73,7 @@ func (db *twoFactors) Create(userID int64, key, secret string) error {
 		return errors.Wrap(err, "generate recovery codes")
 	}
 
-	return db.Transaction(func(tx *gorm.DB) error {
+	return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
 		err := tx.Create(tf).Error
 		if err != nil {
 			return err
@@ -101,9 +102,9 @@ func (ErrTwoFactorNotFound) NotFound() bool {
 	return true
 }
 
-func (db *twoFactors) GetByUserID(userID int64) (*TwoFactor, error) {
+func (db *twoFactors) GetByUserID(ctx context.Context, userID int64) (*TwoFactor, error) {
 	tf := new(TwoFactor)
-	err := db.Where("user_id = ?", userID).First(tf).Error
+	err := db.WithContext(ctx).Where("user_id = ?", userID).First(tf).Error
 	if err != nil {
 		if err == gorm.ErrRecordNotFound {
 			return nil, ErrTwoFactorNotFound{args: errutil.Args{"userID": userID}}
@@ -113,9 +114,9 @@ func (db *twoFactors) GetByUserID(userID int64) (*TwoFactor, error) {
 	return tf, nil
 }
 
-func (db *twoFactors) IsUserEnabled(userID int64) bool {
+func (db *twoFactors) IsUserEnabled(ctx context.Context, userID int64) bool {
 	var count int64
-	err := db.Model(new(TwoFactor)).Where("user_id = ?", userID).Count(&count).Error
+	err := db.WithContext(ctx).Model(new(TwoFactor)).Where("user_id = ?", userID).Count(&count).Error
 	if err != nil {
 		log.Error("Failed to count two factors [user_id: %d]: %v", userID, err)
 	}