Add MFA functionality

3 files changed, 96 insertions, 0 deletions

diff --git a/store/sql_user_store.go b/store/sql_user_store.go
index 6062b8a6a..33d1887ad 100644
--- a/store/sql_user_store.go
+++ b/store/sql_user_store.go
@@ -40,6 +40,7 @@ func NewSqlUserStore(sqlStore *SqlStore) UserStore {
 		table.ColMap("NotifyProps").SetMaxSize(2000)
 		table.ColMap("ThemeProps").SetMaxSize(2000)
 		table.ColMap("Locale").SetMaxSize(5)
+		table.ColMap("MfaSecret").SetMaxSize(128)
 		table.SetUniqueTogether("Email", "TeamId")
 		table.SetUniqueTogether("Username", "TeamId")
 	}
@@ -50,6 +51,9 @@ func NewSqlUserStore(sqlStore *SqlStore) UserStore {
 func (us SqlUserStore) UpgradeSchemaIfNeeded() {
 	// ADDED for 2.0 REMOVE for 2.4
 	us.CreateColumnIfNotExists("Users", "Locale", "varchar(5)", "character varying(5)", model.DEFAULT_LOCALE)
+	// ADDED for 2.2 REMOVE for 2.6
+	us.CreateColumnIfNotExists("Users", "MfaActive", "tinyint(1)", "boolean", "0")
+	us.CreateColumnIfNotExists("Users", "MfaSecret", "varchar(128)", "character varying(128)", "")
 }
 
 func (us SqlUserStore) CreateIndexesIfNotExists() {
@@ -141,6 +145,8 @@ func (us SqlUserStore) Update(user *model.User, allowActiveUpdate bool) StoreCha
 			user.LastPingAt = oldUser.LastPingAt
 			user.EmailVerified = oldUser.EmailVerified
 			user.FailedAttempts = oldUser.FailedAttempts
+			user.MfaSecret = oldUser.MfaSecret
+			user.MfaActive = oldUser.MfaActive
 
 			if !allowActiveUpdate {
 				user.Roles = oldUser.Roles
@@ -346,6 +352,50 @@ func (us SqlUserStore) UpdateAuthData(userId, service, authData, email string) S
 	return storeChannel
 }
 
+func (us SqlUserStore) UpdateMfaSecret(userId, secret string) StoreChannel {
+
+	storeChannel := make(StoreChannel)
+
+	go func() {
+		result := StoreResult{}
+
+		updateAt := model.GetMillis()
+
+		if _, err := us.GetMaster().Exec("UPDATE Users SET MfaSecret = :Secret, UpdateAt = :UpdateAt WHERE Id = :UserId", map[string]interface{}{"Secret": secret, "UpdateAt": updateAt, "UserId": userId}); err != nil {
+			result.Err = model.NewLocAppError("SqlUserStore.UpdateMfaSecret", "store.sql_user.update_mfa_secret.app_error", nil, "id="+userId+", "+err.Error())
+		} else {
+			result.Data = userId
+		}
+
+		storeChannel <- result
+		close(storeChannel)
+	}()
+
+	return storeChannel
+}
+
+func (us SqlUserStore) UpdateMfaActive(userId string, active bool) StoreChannel {
+
+	storeChannel := make(StoreChannel)
+
+	go func() {
+		result := StoreResult{}
+
+		updateAt := model.GetMillis()
+
+		if _, err := us.GetMaster().Exec("UPDATE Users SET MfaActive = :Active, UpdateAt = :UpdateAt WHERE Id = :UserId", map[string]interface{}{"Active": active, "UpdateAt": updateAt, "UserId": userId}); err != nil {
+			result.Err = model.NewLocAppError("SqlUserStore.UpdateMfaActive", "store.sql_user.update_mfa_active.app_error", nil, "id="+userId+", "+err.Error())
+		} else {
+			result.Data = userId
+		}
+
+		storeChannel <- result
+		close(storeChannel)
+	}()
+
+	return storeChannel
+}
+
 func (us SqlUserStore) Get(id string) StoreChannel {
 
 	storeChannel := make(StoreChannel)
diff --git a/store/sql_user_store_test.go b/store/sql_user_store_test.go
index 8f2366136..dcd2440ac 100644
--- a/store/sql_user_store_test.go
+++ b/store/sql_user_store_test.go
@@ -502,3 +502,47 @@ func TestUserUnreadCount(t *testing.T) {
 		t.Fatal("should have 3 unread messages")
 	}
 }
+
+func TestUserStoreUpdateMfaSecret(t *testing.T) {
+	Setup()
+
+	u1 := model.User{}
+	u1.TeamId = model.NewId()
+	u1.Email = model.NewId()
+	Must(store.User().Save(&u1))
+
+	time.Sleep(100 * time.Millisecond)
+
+	if err := (<-store.User().UpdateMfaSecret(u1.Id, "12345")).Err; err != nil {
+		t.Fatal(err)
+	}
+
+	// should pass, no update will occur though
+	if err := (<-store.User().UpdateMfaSecret("junk", "12345")).Err; err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestUserStoreUpdateMfaActive(t *testing.T) {
+	Setup()
+
+	u1 := model.User{}
+	u1.TeamId = model.NewId()
+	u1.Email = model.NewId()
+	Must(store.User().Save(&u1))
+
+	time.Sleep(100 * time.Millisecond)
+
+	if err := (<-store.User().UpdateMfaActive(u1.Id, true)).Err; err != nil {
+		t.Fatal(err)
+	}
+
+	if err := (<-store.User().UpdateMfaActive(u1.Id, false)).Err; err != nil {
+		t.Fatal(err)
+	}
+
+	// should pass, no update will occur though
+	if err := (<-store.User().UpdateMfaActive("junk", true)).Err; err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/store/store.go b/store/store.go
index 7ec5ac3a5..323595ffb 100644
--- a/store/store.go
+++ b/store/store.go
@@ -117,6 +117,8 @@ type UserStore interface {
 	UpdateUserAndSessionActivity(userId string, sessionId string, time int64) StoreChannel
 	UpdatePassword(userId, newPassword string) StoreChannel
 	UpdateAuthData(userId, service, authData, email string) StoreChannel
+	UpdateMfaSecret(userId, secret string) StoreChannel
+	UpdateMfaActive(userId string, active bool) StoreChannel
 	Get(id string) StoreChannel
 	GetProfiles(teamId string) StoreChannel
 	GetByEmail(teamId string, email string) StoreChannel