feat: add support for TOTP multi-factor authentication

app/app.go

@@ -24,6 +24,7 @@ type App struct {
 	AccountStore      data.AccountStore
 	RefreshTokenStore data.RefreshTokenStore
 	KeyStore          data.KeyStore
+	TOTPCache         data.TOTPCache
 	Actives           data.Actives
 	Reporter          ops.ErrorReporter
 	OauthProviders    map[string]oauth.Provider
@@ -69,10 +70,12 @@ func NewApp(cfg *Config, logger logrus.FieldLogger) (*App, error) {
 		return nil, errors.Wrap(err, "NewBlobStore")
 	}
 
+	encryptedBlobStore := data.NewEncryptedBlobStore(blobStore, cfg.DBEncryptionKey)
+
 	keyStore := data.NewRotatingKeyStore()
 	if cfg.IdentitySigningKey == nil {
 		m := data.NewKeyStoreRotater(
-			data.NewEncryptedBlobStore(blobStore, cfg.DBEncryptionKey),
+			encryptedBlobStore,
 			cfg.AccessTokenTTL,
 			logger,
 		)
@@ -84,6 +87,8 @@ func NewApp(cfg *Config, logger logrus.FieldLogger) (*App, error) {
 		keyStore.Rotate(cfg.IdentitySigningKey)
 	}
 
+	totpCache := data.NewTOTPCache(encryptedBlobStore)
+
 	var actives data.Actives
 	if redis != nil {
 		actives = dataRedis.NewActives(
@@ -121,6 +126,7 @@ func NewApp(cfg *Config, logger logrus.FieldLogger) (*App, error) {
 		AccountStore:      accountStore,
 		RefreshTokenStore: tokenStore,
 		KeyStore:          keyStore,
+		TOTPCache:         totpCache,
 		Actives:           actives,
 		Reporter:          errorReporter,
 		OauthProviders:    oauthProviders,

app/data/account_store.go

@@ -25,6 +25,8 @@ type AccountStore interface {
 	SetPassword(id int, p []byte) (bool, error)
 	UpdateUsername(id int, u string) (bool, error)
 	SetLastLogin(id int) (bool, error)
+	SetTOTPSecret(id int, secret []byte) (bool, error)
+	DeleteTOTPSecret(id int) (bool, error)
 }
 
 func NewAccountStore(db sqlx.Ext) (AccountStore, error) {

app/data/blob_store.go

@@ -17,6 +17,9 @@ type BlobStore interface {
 
 	// WriteNX will write the blob into the store only if the name does not exist.
 	WriteNX(name string, blob []byte) (bool, error)
+
+	// Write will write the blob into the store
+	Write(name string, blob []byte) (bool, error)
 }
 
 func NewBlobStore(interval time.Duration, redis *redis.Client, db *sqlx.DB, reporter ops.ErrorReporter) (BlobStore, error) {

app/data/encrypted_blob_store.go

@@ -30,3 +30,11 @@ func (bs *EncryptedBlobStore) WriteNX(name string, blob []byte) (bool, error) {
 	}
 	return bs.store.WriteNX(name, encryptedBlob)
 }
+
+func (bs *EncryptedBlobStore) Write(name string, blob []byte) (bool, error) {
+	encryptedBlob, err := compat.Encrypt(blob, bs.encryptionKey)
+	if err != nil {
+		return false, err
+	}
+	return bs.store.Write(name, encryptedBlob)
+}

app/data/mock/account_store.go

@@ -1,6 +1,7 @@
 package mock
 
 import (
+	"database/sql"
 	"fmt"
 	"strings"
 	"time"
@@ -204,6 +205,28 @@ func (s *accountStore) SetLastLogin(id int) (bool, error) {
 	return true, nil
 }
 
+func (s *accountStore) SetTOTPSecret(id int, secret []byte) (bool, error) {
+	account := s.accountsByID[id]
+	if account == nil {
+		return false, nil
+	}
+
+	account.TOTPSecret = sql.NullString{String: string(secret), Valid: true}
+
+	return true, nil
+}
+
+func (s *accountStore) DeleteTOTPSecret(id int) (bool, error) {
+	account := s.accountsByID[id]
+	if account == nil {
+		return false, nil
+	}
+
+	account.TOTPSecret = sql.NullString{}
+
+	return true, nil
+}
+
 // i think this works? i want to avoid accidentally giving callers the ability
 // to reach into the memory map and modify things or see changes without relying
 // on the store api.

app/data/mock/blob_store.go

@@ -1,7 +1,9 @@
 package mock
 
-import "time"
-import "sync"
+import (
+	"sync"
+	"time"
+)
 
 type BlobStore struct {
 	blobs    map[string][]byte
@@ -39,3 +41,11 @@ func (bs *BlobStore) WriteNX(name string, blob []byte) (bool, error) {
 	bs.blobs[name] = blob
 	return true, nil
 }
+
+func (bs *BlobStore) Write(name string, blob []byte) (bool, error) {
+	bs.mutex.Lock()
+	defer bs.mutex.Unlock()
+
+	bs.blobs[name] = blob
+	return true, nil
+}

app/data/mysql/account_store.go

@@ -138,6 +138,16 @@ func (db *AccountStore) SetLastLogin(id int) (bool, error) {
 	return ok(result, err)
 }
 
+func (db *AccountStore) SetTOTPSecret(id int, secret []byte) (bool, error) {
+	result, err := db.Exec("UPDATE accounts SET totp_secret = ? WHERE id = ?", secret, id)
+	return ok(result, err)
+}
+
+func (db *AccountStore) DeleteTOTPSecret(id int) (bool, error) {
+	result, err := db.Exec("UPDATE accounts SET totp_secret = NULL WHERE id = ?", id)
+	return ok(result, err)
+}
+
 func ok(result sql.Result, err error) (bool, error) {
 	if err != nil {
 		return false, err

app/data/mysql/migrations.go

@@ -1,7 +1,9 @@
 package mysql
 
-import "github.com/jmoiron/sqlx"
-import "github.com/go-sql-driver/mysql"
+import (
+	"github.com/go-sql-driver/mysql"
+	"github.com/jmoiron/sqlx"
+)
 
 // MigrateDB is committed to doing the work necessary to converge the database
 // in a safe, production-grade fashion. This will mean conditional logic as it
@@ -12,6 +14,7 @@ func MigrateDB(db *sqlx.DB) error {
 		createAccounts,
 		createOauthAccounts,
 		createAccountLastLoginAtField,
+		createAccountTOTPFields,
 	}
 	for _, m := range migrations {
 		if err := m(db); err != nil {
@@ -69,3 +72,15 @@ func createAccountLastLoginAtField(db *sqlx.DB) error {
 	}
 	return err
 }
+
+func createAccountTOTPFields(db *sqlx.DB) error {
+	_, err := db.Exec(`
+        ALTER TABLE accounts ADD totp_secret VARCHAR(255) DEFAULT NULL
+    `)
+	if mysqlError, ok := err.(*mysql.MySQLError); ok {
+		if mysqlError.Number == 1060 { // 1060 = Duplicate column name
+			err = nil
+		}
+	}
+	return err
+}

app/data/postgres/account_store.go

@@ -164,6 +164,16 @@ func (db *AccountStore) SetLastLogin(id int) (bool, error) {
 	return ok(result, err)
 }
 
+func (db *AccountStore) SetTOTPSecret(id int, secret []byte) (bool, error) {
+	result, err := db.Exec("UPDATE accounts SET totp_secret = $1 WHERE id = $2", secret, id)
+	return ok(result, err)
+}
+
+func (db *AccountStore) DeleteTOTPSecret(id int) (bool, error) {
+	result, err := db.Exec("UPDATE accounts SET totp_secret = NULL WHERE id = $1", id)
+	return ok(result, err)
+}
+
 func ok(result sql.Result, err error) (bool, error) {
 	if err != nil {
 		return false, err

app/data/postgres/migrations.go

@@ -12,6 +12,7 @@ func MigrateDB(db *sqlx.DB) error {
 		createOauthAccounts,
 		createAccountLastLoginAtField,
 		caseInsensitiveUsername,
+		createAccountTOTPFields,
 	}
 	for _, m := range migrations {
 		if err := m(db); err != nil {
@@ -20,6 +21,7 @@ func MigrateDB(db *sqlx.DB) error {
 	}
 	return nil
 }
+
 func migrateAccounts(db *sqlx.DB) error {
 	_, err := db.Exec(`
         CREATE TABLE IF NOT EXISTS accounts (
@@ -68,3 +70,10 @@ func caseInsensitiveUsername(db *sqlx.DB) error {
     `)
 	return err
 }
+
+func createAccountTOTPFields(db *sqlx.DB) error {
+	_, err := db.Exec(`
+        ALTER TABLE accounts ADD COLUMN IF NOT EXISTS totp_secret TEXT DEFAULT NULL
+    `)
+	return err
+}

app/data/redis/blob_store.go

@@ -31,3 +31,11 @@ func (s *BlobStore) Read(name string) ([]byte, error) {
 func (s *BlobStore) WriteNX(name string, blob []byte) (bool, error) {
 	return s.Client.SetNX(context.TODO(), name, blob, s.TTL).Result()
 }
+
+func (s *BlobStore) Write(name string, blob []byte) (bool, error) {
+	res, err := s.Client.Set(context.TODO(), name, blob, s.TTL).Result()
+	if res != "OK" {
+		return false, err
+	}
+	return true, nil
+}

app/data/sqlite3/account_store.go

@@ -138,6 +138,16 @@ func (db *AccountStore) SetLastLogin(id int) (bool, error) {
 	return ok(result, err)
 }
 
+func (db *AccountStore) SetTOTPSecret(id int, secret []byte) (bool, error) {
+	result, err := db.Exec("UPDATE accounts SET totp_secret = ? WHERE id = ?", secret, id)
+	return ok(result, err)
+}
+
+func (db *AccountStore) DeleteTOTPSecret(id int) (bool, error) {
+	result, err := db.Exec("UPDATE accounts SET totp_secret = NULL WHERE id = ?", id)
+	return ok(result, err)
+}
+
 func ok(result sql.Result, err error) (bool, error) {
 	if err != nil {
 		return false, err

app/data/sqlite3/blob_store.go

@@ -51,3 +51,12 @@ func (s *BlobStore) WriteNX(name string, blob []byte) (bool, error) {
 	}
 	return true, nil
 }
+
+func (s *BlobStore) Write(name string, blob []byte) (bool, error) {
+	expiresAt := time.Now().Add(s.TTL)
+	_, err := s.DB.Exec("INSERT or REPLACE INTO blobs (name, blob, expires_at) VALUES (?, ?, ?)", name, blob, expiresAt, blob, expiresAt, name)
+	if err != nil {
+		return false, err
+	}
+	return true, nil
+}

app/data/sqlite3/migrations.go

@@ -19,6 +19,7 @@ func MigrateDB(db *sqlx.DB) error {
 		createOauthAccounts,
 		createAccountLastLoginAtField,
 		caseInsensitiveUsername,
+		createAccountTOTPFields,
 	}
 	for _, m := range migrations {
 		if err := m(db); err != nil {
@@ -137,3 +138,13 @@ func caseInsensitiveUsername(db *sqlx.DB) error {
     `)
 	return err
 }
+
+func createAccountTOTPFields(db *sqlx.DB) error {
+	_, err := db.Exec(`
+        ALTER TABLE accounts ADD totp_secret VARCHAR(255) DEFAULT NULL
+    `)
+	if isDuplicateError(err) {
+		return nil
+	}
+	return err
+}

app/data/testers/account_store_testers.go

@@ -17,6 +17,7 @@ var AccountStoreTesters = []func(*testing.T, data.AccountStore){
 	testArchiveWithOauth,
 	testRequireNewPassword,
 	testSetPassword,
+	testSetAndDeleteTOTP,
 	testUpdateUsername,
 	testAddOauthAccount,
 	testFindByOauthAccount,
@@ -193,6 +194,37 @@ func testSetPassword(t *testing.T, store data.AccountStore) {
 	assert.Equal(t, 1, getOpenConnectionCount(store))
 }
 
+func testSetAndDeleteTOTP(t *testing.T, store data.AccountStore) {
+	account, err := store.Create("[email protected]", []byte("password"))
+	require.NoError(t, err)
+	assert.False(t, account.TOTPEnabled())
+	assert.False(t, account.TOTPSecret.Valid)
+
+	//Check set
+	ok, err := store.SetTOTPSecret(account.ID, []byte("secret"))
+	assert.True(t, ok)
+	require.NoError(t, err)
+
+	after, err := store.Find(account.ID)
+	require.NoError(t, err)
+	assert.Equal(t, "secret", after.TOTPSecret.String)
+	assert.True(t, after.TOTPEnabled())
+	assert.True(t, after.TOTPSecret.Valid)
+
+	//Check delete
+	ok, err = store.DeleteTOTPSecret(account.ID)
+	assert.True(t, ok)
+	require.NoError(t, err)
+
+	after, err = store.Find(account.ID)
+	require.NoError(t, err)
+	assert.False(t, after.TOTPEnabled())
+	assert.False(t, after.TOTPSecret.Valid)
+
+	// Assert that db connections are released to pool
+	assert.Equal(t, 1, getOpenConnectionCount(store))
+}
+
 func testUpdateUsername(t *testing.T, store data.AccountStore) {
 	other, err := store.Create("other", []byte("other"))
 	require.NoError(t, err)

app/data/testers/blob_store_testers.go

@@ -11,6 +11,7 @@ import (
 var BlobStoreTesters = []func(*testing.T, data.BlobStore){
 	testRead,
 	testWriteNX,
+	testWrite,
 }
 
 func testRead(t *testing.T, bs data.BlobStore) {
@@ -36,3 +37,13 @@ func testWriteNX(t *testing.T, bs data.BlobStore) {
 	assert.NoError(t, err)
 	assert.False(t, set)
 }
+
+func testWrite(t *testing.T, bs data.BlobStore) {
+	set, err := bs.Write("key", []byte("first"))
+	assert.NoError(t, err)
+	assert.Equal(t, true, set)
+
+	set, err = bs.Write("key", []byte("second"))
+	assert.NoError(t, err)
+	assert.Equal(t, true, set)
+}

app/data/totp_cache.go

@@ -0,0 +1,35 @@
+package data
+
+import (
+	"fmt"
+
+	"github.com/pkg/errors"
+)
+
+type TOTPCache struct {
+	ebs *EncryptedBlobStore
+}
+
+func NewTOTPCache(ebs *EncryptedBlobStore) TOTPCache {
+	return TOTPCache{
+		ebs: ebs,
+	}
+}
+
+func (t *TOTPCache) CacheTOTPSecret(accountID int, secret []byte) error {
+	keyName := fmt.Sprintf("totp:%d", accountID)
+	_, err := t.ebs.Write(keyName, secret)
+	if err != nil {
+		return errors.Wrap(err, "CacheTOTPSecret")
+	}
+	return nil
+}
+
+func (t *TOTPCache) LoadTOTPSecret(accountID int) ([]byte, error) {
+	keyName := fmt.Sprintf("totp:%d", accountID)
+	val, err := t.ebs.Read(keyName)
+	if err != nil {
+		return nil, errors.Wrap(err, "LoadTOTPSecret")
+	}
+	return val, nil
+}