db: add tests for two factors (#6099)

* Rename to TwoFactors.Create

* Use GORM to execute queries

* TwoFactor.GetByUserID

* Add tests

* Fix failing tests

* Add MD5 tests

* Add tests for RandomChars

8 files changed, 234 insertions, 107 deletions

diff --git a/internal/db/errors/two_factor.go b/internal/db/errors/two_factor.go
deleted file mode 100644
index c474152d..00000000
--- a/internal/db/errors/two_factor.go
+++ /dev/null
@@ -1,20 +0,0 @@
-// Copyright 2017 The Gogs Authors. All rights reserved.
-// Use of this source code is governed by a MIT-style
-// license that can be found in the LICENSE file.
-
-package errors
-
-import "fmt"
-
-type TwoFactorNotFound struct {
-	UserID int64
-}
-
-func IsTwoFactorNotFound(err error) bool {
-	_, ok := err.(TwoFactorNotFound)
-	return ok
-}
-
-func (err TwoFactorNotFound) Error() string {
-	return fmt.Sprintf("two-factor authentication does not found [user_id: %d]", err.UserID)
-}
diff --git a/internal/db/migrations/migrations.go b/internal/db/migrations/migrations.go
index dce1476d..20ab0a8b 100644
--- a/internal/db/migrations/migrations.go
+++ b/internal/db/migrations/migrations.go
@@ -13,7 +13,7 @@ import (
 	log "unknwon.dev/clog/v2"
 	"xorm.io/xorm"
 
-	"gogs.io/gogs/internal/tool"
+	"gogs.io/gogs/internal/strutil"
 )
 
 const _MIN_DB_VER = 10
@@ -156,10 +156,10 @@ func generateOrgRandsAndSalt(x *xorm.Engine) (err error) {
 	}
 
 	for _, org := range orgs {
-		if org.Rands, err = tool.RandomString(10); err != nil {
+		if org.Rands, err = strutil.RandomChars(10); err != nil {
 			return err
 		}
-		if org.Salt, err = tool.RandomString(10); err != nil {
+		if org.Salt, err = strutil.RandomChars(10); err != nil {
 			return err
 		}
 		if _, err = sess.ID(org.ID).Update(org); err != nil {
diff --git a/internal/db/mocks.go b/internal/db/mocks.go
index 9b1e9028..06bafc75 100644
--- a/internal/db/mocks.go
+++ b/internal/db/mocks.go
@@ -180,9 +180,19 @@ func SetMockReposStore(t *testing.T, mock ReposStore) {
 var _ TwoFactorsStore = (*MockTwoFactorsStore)(nil)
 
 type MockTwoFactorsStore struct {
+	MockCreate        func(userID int64, key, secret string) error
+	MockGetByUserID   func(userID int64) (*TwoFactor, error)
 	MockIsUserEnabled func(userID int64) bool
 }
 
+func (m *MockTwoFactorsStore) Create(userID int64, key, secret string) error {
+	return m.MockCreate(userID, key, secret)
+}
+
+func (m *MockTwoFactorsStore) GetByUserID(userID int64) (*TwoFactor, error) {
+	return m.MockGetByUserID(userID)
+}
+
 func (m *MockTwoFactorsStore) IsUserEnabled(userID int64) bool {
 	return m.MockIsUserEnabled(userID)
 }
diff --git a/internal/db/repo_editor.go b/internal/db/repo_editor.go
index c525bcd6..652d6944 100644
--- a/internal/db/repo_editor.go
+++ b/internal/db/repo_editor.go
@@ -22,6 +22,7 @@ import (
 	"github.com/gogs/git-module"
 
 	"gogs.io/gogs/internal/conf"
+	"gogs.io/gogs/internal/cryptoutil"
 	"gogs.io/gogs/internal/db/errors"
 	"gogs.io/gogs/internal/gitutil"
 	"gogs.io/gogs/internal/osutil"
@@ -56,7 +57,7 @@ func ComposeHookEnvs(opts ComposeHookEnvsOptions) []string {
 		ENV_AUTH_USER_NAME + "=" + opts.AuthUser.Name,
 		ENV_AUTH_USER_EMAIL + "=" + opts.AuthUser.Email,
 		ENV_REPO_OWNER_NAME + "=" + opts.OwnerName,
-		ENV_REPO_OWNER_SALT_MD5 + "=" + tool.MD5(opts.OwnerSalt),
+		ENV_REPO_OWNER_SALT_MD5 + "=" + cryptoutil.MD5(opts.OwnerSalt),
 		ENV_REPO_ID + "=" + com.ToStr(opts.RepoID),
 		ENV_REPO_NAME + "=" + opts.RepoName,
 		ENV_REPO_CUSTOM_HOOKS_PATH + "=" + filepath.Join(opts.RepoPath, "custom_hooks"),
diff --git a/internal/db/two_factor.go b/internal/db/two_factor.go
index e827d59b..ca990454 100644
--- a/internal/db/two_factor.go
+++ b/internal/db/two_factor.go
@@ -7,38 +7,24 @@ package db
 import (
 	"encoding/base64"
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/pquerna/otp/totp"
 	"github.com/unknwon/com"
-	"xorm.io/xorm"
 
 	"gogs.io/gogs/internal/conf"
-	"gogs.io/gogs/internal/db/errors"
-	"gogs.io/gogs/internal/tool"
+	"gogs.io/gogs/internal/cryptoutil"
 )
 
-// TwoFactor represents a two-factor authentication token.
+// TwoFactor is a 2FA token of a user.
 type TwoFactor struct {
 	ID          int64
-	UserID      int64 `xorm:"UNIQUE"`
+	UserID      int64 `xorm:"UNIQUE" gorm:"UNIQUE"`
 	Secret      string
-	Created     time.Time `xorm:"-" json:"-"`
+	Created     time.Time `xorm:"-" gorm:"-" json:"-"`
 	CreatedUnix int64
 }
 
-func (t *TwoFactor) BeforeInsert() {
-	t.CreatedUnix = time.Now().Unix()
-}
-
-func (t *TwoFactor) AfterSet(colName string, _ xorm.Cell) {
-	switch colName {
-	case "created_unix":
-		t.Created = time.Unix(t.CreatedUnix, 0).Local()
-	}
-}
-
 // ValidateTOTP returns true if given passcode is valid for two-factor authentication token.
 // It also returns possible validation error.
 func (t *TwoFactor) ValidateTOTP(passcode string) (bool, error) {
@@ -46,74 +32,13 @@ func (t *TwoFactor) ValidateTOTP(passcode string) (bool, error) {
 	if err != nil {
 		return false, fmt.Errorf("DecodeString: %v", err)
 	}
-	decryptSecret, err := com.AESGCMDecrypt(tool.MD5Bytes(conf.Security.SecretKey), secret)
+	decryptSecret, err := com.AESGCMDecrypt(cryptoutil.MD5Bytes(conf.Security.SecretKey), secret)
 	if err != nil {
 		return false, fmt.Errorf("AESGCMDecrypt: %v", err)
 	}
 	return totp.Validate(passcode, string(decryptSecret)), nil
 }
 
-func generateRecoveryCodes(userID int64) ([]*TwoFactorRecoveryCode, error) {
-	recoveryCodes := make([]*TwoFactorRecoveryCode, 10)
-	for i := 0; i < 10; i++ {
-		code, err := tool.RandomString(10)
-		if err != nil {
-			return nil, fmt.Errorf("RandomString: %v", err)
-		}
-		recoveryCodes[i] = &TwoFactorRecoveryCode{
-			UserID: userID,
-			Code:   strings.ToLower(code[:5] + "-" + code[5:]),
-		}
-	}
-	return recoveryCodes, nil
-}
-
-// NewTwoFactor creates a new two-factor authentication token and recovery codes for given user.
-func NewTwoFactor(userID int64, secret string) error {
-	t := &TwoFactor{
-		UserID: userID,
-	}
-
-	// Encrypt secret
-	encryptSecret, err := com.AESGCMEncrypt(tool.MD5Bytes(conf.Security.SecretKey), []byte(secret))
-	if err != nil {
-		return fmt.Errorf("AESGCMEncrypt: %v", err)
-	}
-	t.Secret = base64.StdEncoding.EncodeToString(encryptSecret)
-
-	recoveryCodes, err := generateRecoveryCodes(userID)
-	if err != nil {
-		return fmt.Errorf("generateRecoveryCodes: %v", err)
-	}
-
-	sess := x.NewSession()
-	defer sess.Close()
-	if err = sess.Begin(); err != nil {
-		return err
-	}
-
-	if _, err = sess.Insert(t); err != nil {
-		return fmt.Errorf("insert two-factor: %v", err)
-	} else if _, err = sess.Insert(recoveryCodes); err != nil {
-		return fmt.Errorf("insert recovery codes: %v", err)
-	}
-
-	return sess.Commit()
-}
-
-// GetTwoFactorByUserID returns two-factor authentication token of given user.
-func GetTwoFactorByUserID(userID int64) (*TwoFactor, error) {
-	t := new(TwoFactor)
-	has, err := x.Where("user_id = ?", userID).Get(t)
-	if err != nil {
-		return nil, err
-	} else if !has {
-		return nil, errors.TwoFactorNotFound{UserID: userID}
-	}
-
-	return t, nil
-}
-
 // DeleteTwoFactor removes two-factor authentication token and recovery codes of given user.
 func DeleteTwoFactor(userID int64) (err error) {
 	sess := x.NewSession()
@@ -152,7 +77,7 @@ func deleteRecoveryCodesByUserID(e Engine, userID int64) error {
 
 // RegenerateRecoveryCodes regenerates new set of recovery codes for given user.
 func RegenerateRecoveryCodes(userID int64) error {
-	recoveryCodes, err := generateRecoveryCodes(userID)
+	recoveryCodes, err := generateRecoveryCodes(userID, 10)
 	if err != nil {
 		return fmt.Errorf("generateRecoveryCodes: %v", err)
 	}
diff --git a/internal/db/two_factors.go b/internal/db/two_factors.go
index 376cc40c..671beb22 100644
--- a/internal/db/two_factors.go
+++ b/internal/db/two_factors.go
@@ -5,24 +5,118 @@
 package db
 
 import (
+	"encoding/base64"
+	"fmt"
+	"strings"
+	"time"
+
 	"github.com/jinzhu/gorm"
+	"github.com/pkg/errors"
 	log "unknwon.dev/clog/v2"
+
+	"gogs.io/gogs/internal/cryptoutil"
+	"gogs.io/gogs/internal/errutil"
+	"gogs.io/gogs/internal/strutil"
 )
 
 // TwoFactorsStore is the persistent interface for 2FA.
 //
 // 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)
 	// IsUserEnabled returns true if the user has enabled 2FA.
 	IsUserEnabled(userID int64) bool
 }
 
 var TwoFactors TwoFactorsStore
 
+// NOTE: This is a GORM create hook.
+func (t *TwoFactor) BeforeCreate() {
+	t.CreatedUnix = gorm.NowFunc().Unix()
+}
+
+// NOTE: This is a GORM query hook.
+func (t *TwoFactor) AfterFind() {
+	t.Created = time.Unix(t.CreatedUnix, 0).Local()
+}
+
+var _ TwoFactorsStore = (*twoFactors)(nil)
+
 type twoFactors struct {
 	*gorm.DB
 }
 
+func (db *twoFactors) Create(userID int64, key, secret string) error {
+	encrypted, err := cryptoutil.AESGCMEncrypt(cryptoutil.MD5Bytes(key), []byte(secret))
+	if err != nil {
+		return errors.Wrap(err, "encrypt secret")
+	}
+	tf := &TwoFactor{
+		UserID: userID,
+		Secret: base64.StdEncoding.EncodeToString(encrypted),
+	}
+
+	recoveryCodes, err := generateRecoveryCodes(userID, 10)
+	if err != nil {
+		return errors.Wrap(err, "generate recovery codes")
+	}
+
+	vals := make([]string, 0, len(recoveryCodes))
+	items := make([]interface{}, 0, len(recoveryCodes)*2)
+	for _, code := range recoveryCodes {
+		vals = append(vals, "(?, ?)")
+		items = append(items, code.UserID, code.Code)
+	}
+
+	return db.Transaction(func(tx *gorm.DB) error {
+		err := tx.Create(tf).Error
+		if err != nil {
+			return err
+		}
+
+		sql := "INSERT INTO two_factor_recovery_code (user_id, code) VALUES " + strings.Join(vals, ", ")
+		return tx.Exec(sql, items...).Error
+	})
+}
+
+var _ errutil.NotFound = (*ErrTwoFactorNotFound)(nil)
+
+type ErrTwoFactorNotFound struct {
+	args errutil.Args
+}
+
+func IsErrTwoFactorNotFound(err error) bool {
+	_, ok := err.(ErrTwoFactorNotFound)
+	return ok
+}
+
+func (err ErrTwoFactorNotFound) Error() string {
+	return fmt.Sprintf("2FA does not found: %v", err.args)
+}
+
+func (ErrTwoFactorNotFound) NotFound() bool {
+	return true
+}
+
+func (db *twoFactors) GetByUserID(userID int64) (*TwoFactor, error) {
+	tf := new(TwoFactor)
+	err := db.Where("user_id = ?", userID).First(tf).Error
+	if err != nil {
+		if gorm.IsRecordNotFoundError(err) {
+			return nil, ErrTwoFactorNotFound{args: errutil.Args{"userID": userID}}
+		}
+		return nil, err
+	}
+	return tf, nil
+}
+
 func (db *twoFactors) IsUserEnabled(userID int64) bool {
 	var count int64
 	err := db.Model(new(TwoFactor)).Where("user_id = ?", userID).Count(&count).Error
@@ -31,3 +125,19 @@ func (db *twoFactors) IsUserEnabled(userID int64) bool {
 	}
 	return count > 0
 }
+
+// generateRecoveryCodes generates N number of recovery codes for 2FA.
+func generateRecoveryCodes(userID int64, n int) ([]*TwoFactorRecoveryCode, error) {
+	recoveryCodes := make([]*TwoFactorRecoveryCode, n)
+	for i := 0; i < n; i++ {
+		code, err := strutil.RandomChars(10)
+		if err != nil {
+			return nil, errors.Wrap(err, "generate random characters")
+		}
+		recoveryCodes[i] = &TwoFactorRecoveryCode{
+			UserID: userID,
+			Code:   strings.ToLower(code[:5] + "-" + code[5:]),
+		}
+	}
+	return recoveryCodes, nil
+}
diff --git a/internal/db/two_factors_test.go b/internal/db/two_factors_test.go
new file mode 100644
index 00000000..e4dea6ca
--- /dev/null
+++ b/internal/db/two_factors_test.go
@@ -0,0 +1,100 @@
+// Copyright 2020 The Gogs Authors. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package db
+
+import (
+	"testing"
+	"time"
+
+	"github.com/jinzhu/gorm"
+	"github.com/stretchr/testify/assert"
+
+	"gogs.io/gogs/internal/errutil"
+)
+
+func Test_twoFactors(t *testing.T) {
+	if testing.Short() {
+		t.Skip()
+	}
+
+	t.Parallel()
+
+	tables := []interface{}{new(TwoFactor), new(TwoFactorRecoveryCode)}
+	db := &twoFactors{
+		DB: initTestDB(t, "twoFactors", tables...),
+	}
+
+	for _, tc := range []struct {
+		name string
+		test func(*testing.T, *twoFactors)
+	}{
+		{"Create", test_twoFactors_Create},
+		{"GetByUserID", test_twoFactors_GetByUserID},
+		{"IsUserEnabled", test_twoFactors_IsUserEnabled},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Cleanup(func() {
+				err := clearTables(db.DB, tables...)
+				if err != nil {
+					t.Fatal(err)
+				}
+			})
+			tc.test(t, db)
+		})
+	}
+}
+
+func test_twoFactors_Create(t *testing.T, db *twoFactors) {
+	// Create a 2FA token
+	err := db.Create(1, "secure-key", "secure-secret")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Get it back and check the Created field
+	tf, err := db.GetByUserID(1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, gorm.NowFunc().Format(time.RFC3339), tf.Created.Format(time.RFC3339))
+
+	// Verify there are 10 recover codes generated
+	var count int64
+	err = db.Model(new(TwoFactorRecoveryCode)).Count(&count).Error
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, int64(10), count)
+}
+
+func test_twoFactors_GetByUserID(t *testing.T, db *twoFactors) {
+	// Create a 2FA token for user 1
+	err := db.Create(1, "secure-key", "secure-secret")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// We should be able to get it back
+	_, err = db.GetByUserID(1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Try to get a non-existent 2FA token
+	_, err = db.GetByUserID(2)
+	expErr := ErrTwoFactorNotFound{args: errutil.Args{"userID": int64(2)}}
+	assert.Equal(t, expErr, err)
+}
+
+func test_twoFactors_IsUserEnabled(t *testing.T, db *twoFactors) {
+	// Create a 2FA token for user 1
+	err := db.Create(1, "secure-key", "secure-secret")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.True(t, db.IsUserEnabled(1))
+	assert.False(t, db.IsUserEnabled(2))
+}
diff --git a/internal/db/user.go b/internal/db/user.go
index 11697bc7..b650fa14 100644
--- a/internal/db/user.go
+++ b/internal/db/user.go
@@ -32,6 +32,7 @@ import (
 	"gogs.io/gogs/internal/conf"
 	"gogs.io/gogs/internal/db/errors"
 	"gogs.io/gogs/internal/errutil"
+	"gogs.io/gogs/internal/strutil"
 	"gogs.io/gogs/internal/tool"
 )
 
@@ -483,9 +484,9 @@ func IsUserExist(uid int64, name string) (bool, error) {
 	return x.Where("id != ?", uid).Get(&User{LowerName: strings.ToLower(name)})
 }
 
-// GetUserSalt returns a ramdom user salt token.
+// GetUserSalt returns a random user salt token.
 func GetUserSalt() (string, error) {
-	return tool.RandomString(10)
+	return strutil.RandomChars(10)
 }
 
 // NewGhostUser creates and returns a fake user for someone who has deleted his/her account.