implement providing 2FA token on PDS account login screen

Signed-off-by: Will Andrews <did:plc:dadhhalkfcq3gucaq25hjqon>

server/handle_account_signin.go

···
 
 import (
 	"errors"
+	"fmt"
 	"strings"
+	"time"
 
 	"github.com/bluesky-social/indigo/atproto/syntax"
 	"github.com/gorilla/sessions"
···
 )
 
 type OauthSigninInput struct {
-	Username    string `form:"username"`
-	Password    string `form:"password"`
-	QueryParams string `form:"query_params"`
+	Username        string `form:"username"`
+	Password        string `form:"password"`
+	AuthFactorToken string `form:"token"`
+	QueryParams     string `form:"query_params"`
 }
 
 func (s *Server) getSessionRepoOrErr(e echo.Context) (*models.RepoActor, *sessions.Session, error) {
···
 func getFlashesFromSession(e echo.Context, sess *sessions.Session) map[string]any {
 	defer sess.Save(e.Request(), e.Response())
 	return map[string]any{
-		"errors":    sess.Flashes("error"),
-		"successes": sess.Flashes("success"),
+		"errors":        sess.Flashes("error"),
+		"successes":     sess.Flashes("success"),
+		"tokenrequired": sess.Flashes("tokenrequired"),
 	}
 }
 
···
 		idtype = "email"
 	}
 
+	queryParams := ""
+	if req.QueryParams != "" {
+		queryParams = fmt.Sprintf("?%s", req.QueryParams)
+	}
+
 	// TODO: we should make this a helper since we do it for the base create_session as well
 	var repo models.RepoActor
 	var err error
···
 			sess.AddFlash("Something went wrong!", "error")
 		}
 		sess.Save(e.Request(), e.Response())
-		return e.Redirect(303, "/account/signin")
+		return e.Redirect(303, "/account/signin"+queryParams)
 	}
 
 	if err := bcrypt.CompareHashAndPassword([]byte(repo.Password), []byte(req.Password)); err != nil {
···
 			sess.AddFlash("Something went wrong!", "error")
 		}
 		sess.Save(e.Request(), e.Response())
-		return e.Redirect(303, "/account/signin")
+		return e.Redirect(303, "/account/signin"+queryParams)
+	}
+
+	// if repo requires auth factor token and one hasn't been provided, return error prompting for one
+	if repo.EmailAuthFactor && req.AuthFactorToken == "" {
+		err = s.createAndSendAuthCode(ctx, repo)
+		if err != nil {
+			sess.AddFlash("Something went wrong!", "error")
+			sess.Save(e.Request(), e.Response())
+			return e.Redirect(303, "/account/signin"+queryParams)
+		}
+
+		sess.AddFlash("requires 2FA token", "tokenrequired")
+		sess.Save(e.Request(), e.Response())
+		return e.Redirect(303, "/account/signin"+queryParams)
+	}
+
+	// if auth factor is required, now check that the one provided is valid
+	if repo.EmailAuthFactor {
+		if repo.AuthCode == nil || repo.AuthCodeExpiresAt == nil {
+			err = s.createAndSendAuthCode(ctx, repo)
+			if err != nil {
+				sess.AddFlash("Something went wrong!", "error")
+				sess.Save(e.Request(), e.Response())
+				return e.Redirect(303, "/account/signin"+queryParams)
+			}
+
+			sess.AddFlash("requires 2FA token", "tokenrequired")
+			sess.Save(e.Request(), e.Response())
+			return e.Redirect(303, "/account/signin"+queryParams)
+		}
+
+		if *repo.AuthCode != req.AuthFactorToken {
+			return helpers.InvalidTokenError(e)
+		}
+
+		if time.Now().UTC().After(*repo.AuthCodeExpiresAt) {
+			return helpers.ExpiredTokenError(e)
+		}
 	}
 
 	sess.Options = &sessions.Options{
···
 		return err
 	}
 
-	if req.QueryParams != "" {
-		return e.Redirect(303, "/oauth/authorize?"+req.QueryParams)
+	if queryParams != "" {
+		return e.Redirect(303, "/oauth/authorize"+queryParams)
 	} else {
 		return e.Redirect(303, "/account")
 	}

server/handle_server_create_session.go

···
 		return helpers.ServerError(e, nil)
 	}
 
+	if err := bcrypt.CompareHashAndPassword([]byte(repo.Password), []byte(req.Password)); err != nil {
+		if err != bcrypt.ErrMismatchedHashAndPassword {
+			logger.Error("erorr comparing hash and password", "error", err)
+		}
+		return helpers.InputError(e, to.StringPtr("InvalidRequest"))
+	}
+
 	// if repo requires auth factor token and one hasn't been provided, return error prompting for one
 	if repo.EmailAuthFactor && (req.AuthFactorToken == nil || *req.AuthFactorToken == "") {
 		err = s.createAndSendAuthCode(ctx, repo)
···
 		if repo.AuthCode == nil || repo.AuthCodeExpiresAt == nil {
 			err = s.createAndSendAuthCode(ctx, repo)
 			if err != nil {
-				s.logger.Error("sending auth code", "error", err)
+				logger.Error("sending auth code", "error", err)
 				return helpers.ServerError(e, nil)
 			}
 
···
 		if time.Now().UTC().After(*repo.AuthCodeExpiresAt) {
 			return helpers.ExpiredTokenError(e)
 		}
-	}
-
-	if err := bcrypt.CompareHashAndPassword([]byte(repo.Password), []byte(req.Password)); err != nil {
-		if err != bcrypt.ErrMismatchedHashAndPassword {
-			logger.Error("erorr comparing hash and password", "error", err)
-		}
-		return helpers.InputError(e, to.StringPtr("InvalidRequest"))
 	}
 
 	sess, err := s.createSession(ctx, &repo.Repo)

server/templates/signin.html

···
           type="password"
           placeholder="Password"
         />
+        {{ if .flashes.tokenrequired  }}
+        <br />
+        <input name="token" id="token" placeholder="Enter your 2FA token" />
+        {{ end }}
         <input name="query_params" type="hidden" value="{{ .QueryParams }}" />
         <button class="primary" type="submit" value="Login">Login</button>
       </form>