From 8947b711aa31452a97b22a93f28169fe3f990c2b Mon Sep 17 00:00:00 2001
From: Willem van Dreumel <willemvd@users.noreply.github.com>
Date: Mon, 27 Feb 2017 11:10:26 +0100
Subject: [PATCH] Link OAuth2 account to 2FA enabled account (fix #1050)
 (#1052)

* fixes #1050 where linking an account to a 2fa enabled account failed because we forgot to really link the account when 2fa is completed

* handle errors
---
 routers/user/auth.go | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/routers/user/auth.go b/routers/user/auth.go
index fa9f50139..ed8aac0de 100644
--- a/routers/user/auth.go
+++ b/routers/user/auth.go
@@ -221,6 +221,20 @@ func TwoFactorPost(ctx *context.Context, form auth.TwoFactorAuthForm) {
 			return
 		}
 
+		if ctx.Session.Get("linkAccount") != nil {
+			gothUser := ctx.Session.Get("linkAccountGothUser")
+			if gothUser == nil {
+				ctx.Handle(500, "UserSignIn", errors.New("not in LinkAccount session"))
+				return
+			}
+
+			err = models.LinkAccountToUser(u, gothUser.(goth.User))
+			if err != nil {
+				ctx.Handle(500, "UserSignIn", err)
+				return
+			}
+		}
+
 		handleSignIn(ctx, u, remember)
 		return
 	}
@@ -532,8 +546,12 @@ func LinkAccountPostSignIn(ctx *context.Context, signInForm auth.SignInForm) {
 	_, err = models.GetTwoFactorByUID(u.ID)
 	if err != nil {
 		if models.IsErrTwoFactorNotEnrolled(err) {
-			models.LinkAccountToUser(u, gothUser.(goth.User))
-			handleSignIn(ctx, u, signInForm.Remember)
+			err = models.LinkAccountToUser(u, gothUser.(goth.User))
+			if err != nil {
+				ctx.Handle(500, "UserLinkAccount", err)
+			} else {
+				handleSignIn(ctx, u, signInForm.Remember)
+			}
 		} else {
 			ctx.Handle(500, "UserLinkAccount", err)
 		}