diff --git a/admin/accounthttp.go b/admin/accounthttp.go index 408b4c5..5e3f4b6 100644 --- a/admin/accounthttp.go +++ b/admin/accounthttp.go @@ -1,6 +1,7 @@ package admin import ( + "database/sql" "fmt" "net/http" "net/url" @@ -190,6 +191,13 @@ func deleteAccountHandler(app *model.AppState) http.Handler { }) } +type totpConfirmData struct { + Session *model.Session + TOTP *model.TOTP + NameEscaped string + QRBase64Image string +} + func totpSetupHandler(app *model.AppState) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodGet { @@ -212,13 +220,6 @@ func totpSetupHandler(app *model.AppState) http.Handler { return } - type totpSetupData struct { - Session *model.Session - TOTP *model.TOTP - NameEscaped string - QRBase64Image string - } - err := r.ParseForm() if err != nil { http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest) @@ -243,7 +244,7 @@ func totpSetupHandler(app *model.AppState) http.Handler { if err != nil { fmt.Printf("WARN: Failed to create TOTP method: %s\n", err) controller.SetSessionError(app.DB, session, "Something went wrong. Please try again.") - err := totpSetupTemplate.Execute(w, totpSetupData{ Session: session }) + err := totpSetupTemplate.Execute(w, totpConfirmData{ Session: session }) if err != nil { fmt.Printf("WARN: Failed to render TOTP setup page: %s\n", err) http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) @@ -254,17 +255,10 @@ func totpSetupHandler(app *model.AppState) http.Handler { qrBase64Image, err := controller.GenerateQRCode( controller.GenerateTOTPURI(session.Account.Username, totp.Secret)) if err != nil { - fmt.Printf("WARN: Failed to generate TOTP setup QR code: %s\n", err) - controller.SetSessionError(app.DB, session, "Something went wrong. Please try again.") - err := totpSetupTemplate.Execute(w, totpSetupData{ Session: session }) - if err != nil { - fmt.Printf("WARN: Failed to render TOTP setup page: %s\n", err) - http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) - } - return + fmt.Fprintf(os.Stderr, "WARN: Failed to generate TOTP QR code: %v\n", err) } - err = totpConfirmTemplate.Execute(w, totpSetupData{ + err = totpConfirmTemplate.Execute(w, totpConfirmData{ Session: session, TOTP: &totp, NameEscaped: url.PathEscape(totp.Name), @@ -284,11 +278,6 @@ func totpConfirmHandler(app *model.AppState) http.Handler { return } - type totpConfirmData struct { - Session *model.Session - TOTP *model.TOTP - } - session := r.Context().Value("session").(*model.Session) err := r.ParseForm() @@ -309,7 +298,7 @@ func totpConfirmHandler(app *model.AppState) http.Handler { totp, err := controller.GetTOTP(app.DB, session.Account.ID, name) if err != nil { - fmt.Printf("WARN: Failed to fetch TOTP method: %s\n", err) + fmt.Printf("WARN: Failed to fetch TOTP method: %v\n", err) controller.SetSessionError(app.DB, session, "Something went wrong. Please try again.") http.Redirect(w, r, "/admin/account", http.StatusFound) return @@ -319,19 +308,39 @@ func totpConfirmHandler(app *model.AppState) http.Handler { return } + qrBase64Image, err := controller.GenerateQRCode( + controller.GenerateTOTPURI(session.Account.Username, totp.Secret)) + if err != nil { + fmt.Fprintf(os.Stderr, "WARN: Failed to generate TOTP QR code: %v\n", err) + } + confirmCode := controller.GenerateTOTP(totp.Secret, 0) if code != confirmCode { confirmCodeOffset := controller.GenerateTOTP(totp.Secret, 1) if code != confirmCodeOffset { - controller.SetSessionError(app.DB, session, "Incorrect TOTP code. Please try again.") + session.Error = sql.NullString{ Valid: true, String: "Incorrect TOTP code. Please try again." } err = totpConfirmTemplate.Execute(w, totpConfirmData{ Session: session, TOTP: totp, + NameEscaped: url.PathEscape(totp.Name), + QRBase64Image: qrBase64Image, }) + if err != nil { + fmt.Fprintf(os.Stderr, "WARN: Failed to render TOTP setup page: %v\n", err) + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + } return } } + err = controller.ConfirmTOTP(app.DB, session.Account.ID, name) + if err != nil { + fmt.Printf("WARN: Failed to confirm TOTP method: %s\n", err) + controller.SetSessionError(app.DB, session, "Something went wrong. Please try again.") + http.Redirect(w, r, "/admin/account", http.StatusFound) + return + } + controller.SetSessionError(app.DB, session, "") controller.SetSessionMessage(app.DB, session, fmt.Sprintf("TOTP method \"%s\" created successfully.", totp.Name)) http.Redirect(w, r, "/admin/account", http.StatusFound) diff --git a/admin/views/totp-confirm.html b/admin/views/totp-confirm.html index b0810e2..7d305ec 100644 --- a/admin/views/totp-confirm.html +++ b/admin/views/totp-confirm.html @@ -19,6 +19,7 @@ code { {{end}}