polished up TOTP enrolment

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)
+            fmt.Fprintf(os.Stderr, "WARN: Failed to generate TOTP QR code: %v\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
         }
 
-        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)

admin/views/totp-confirm.html

@@ -19,6 +19,7 @@ code {
     {{end}}
 
     <form action="/admin/account/totp-confirm?totp-name={{.NameEscaped}}" method="POST" id="totp-setup">
+        {{if .QRBase64Image}}
         <img src="data:image/png;base64,{{.QRBase64Image}}" alt="" class="qr-code">
 
         <p>
@@ -29,6 +30,12 @@ code {
         <p>
         If the QR code does not work, you may also enter this secret code:
         </p>
+        {{else}}
+        <p>
+        Paste the below secret code into your authentication app or password manager,
+        then enter your 2FA code below:
+        </p>
+        {{end}}
 
         <p><code>{{.TOTP.Secret}}</code></p>
 

controller/totp.go

@@ -78,7 +78,7 @@ func GetTOTPsForAccount(db *sqlx.DB, accountID string) ([]model.TOTP, error) {
     err := db.Select(
         &totps,
         "SELECT * FROM totp " +
-        "WHERE account=$1 " +
+        "WHERE account=$1 AND confirmed=true " +
         "ORDER BY created_at ASC",
         accountID,
     )
@@ -130,6 +130,15 @@ func GetTOTP(db *sqlx.DB, accountID string, name string) (*model.TOTP, error) {
     return &totp, nil
 }
 
+func ConfirmTOTP(db *sqlx.DB, accountID string, name string) error {
+    _, err := db.Exec(
+        "UPDATE totp SET confirmed=true WHERE account=$1 AND name=$2",
+        accountID,
+        name,
+    )
+    return err
+}
+
 func CreateTOTP(db *sqlx.DB, totp *model.TOTP) error {
     _, err := db.Exec(
         "INSERT INTO totp (account, name, secret) " +
@@ -149,3 +158,8 @@ func DeleteTOTP(db *sqlx.DB, accountID string, name string) error {
     )
     return err
 }
+
+func DeleteUnconfirmedTOTPs(db *sqlx.DB) error {
+    _, err := db.Exec("DELETE FROM totp WHERE confirmed=false")
+    return err
+}

main.go

@@ -215,6 +215,15 @@ func main() {
             code := controller.GenerateTOTP(totp.Secret, 0)
             fmt.Printf("%s\n", code)
             return
+        
+        case "cleanTOTP":
+            err := controller.DeleteUnconfirmedTOTPs(app.DB)
+            if err != nil {
+                fmt.Fprintf(os.Stderr, "FATAL: Failed to clean up TOTP methods: %v\n", err)
+                os.Exit(1)
+            }
+            fmt.Printf("Cleaned up dangling TOTP methods successfully.\n")
+            return
 
         case "createInvite":
             fmt.Printf("Creating invite...\n")
@@ -342,6 +351,7 @@ func main() {
             "listTOTP <username>:\n\tLists an account's TOTP methods.\n" +
             "deleteTOTP <username> <name>:\n\tDeletes an account's TOTP method.\n" +
             "testTOTP <username> <name>:\n\tGenerates the code for an account's TOTP method.\n" +
+            "cleanTOTP:\n\tCleans up unconfirmed (dangling) TOTP methods.\n" +
             "\n" +
             "createInvite:\n\tCreates an invite code to register new accounts.\n" +
             "purgeInvites:\n\tDeletes all available invite codes.\n" +
@@ -381,6 +391,13 @@ func main() {
         os.Exit(1)
     }
 
+    // clean up unconfirmed TOTP methods
+    err = controller.DeleteUnconfirmedTOTPs(app.DB)
+    if err != nil {
+        fmt.Fprintf(os.Stderr, "FATAL: Failed to clean up unconfirmed TOTP methods: %v\n", err)
+        os.Exit(1)
+    }
+
     // start the web server!
     mux := createServeMux(&app)
     fmt.Printf("Now serving at http://%s:%d\n", app.Config.Host, app.Config.Port)

model/totp.go

@@ -9,4 +9,5 @@ type TOTP struct {
     AccountID   string      `json:"accountID" db:"account"`
     Secret      string      `json:"-" db:"secret"`
     CreatedAt   time.Time   `json:"created_at" db:"created_at"`
+    Confirmed   bool        `json:"-" db:"confirmed"`
 }

schema_migration/000-init.sql

@@ -23,12 +23,12 @@ ALTER TABLE arimelody.privilege ADD CONSTRAINT privilege_pk PRIMARY KEY (account
 -- Invites
 CREATE TABLE arimelody.invite (
     code text NOT NULL, 
-    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, 
+    created_at TIMESTAMP NOT NULL DEFAULT current_timestamp, 
     expires_at TIMESTAMP NOT NULL
 );
 ALTER TABLE arimelody.invite ADD CONSTRAINT invite_pk PRIMARY KEY (code);
 
--- Session
+-- Sessions
 CREATE TABLE arimelody.session (
     token TEXT,
     user_agent TEXT NOT NULL,
@@ -40,12 +40,13 @@ CREATE TABLE arimelody.session (
 );
 ALTER TABLE arimelody.session ADD CONSTRAINT session_pk PRIMARY KEY (token);
 
--- TOTPs
+-- TOTP methods
 CREATE TABLE arimelody.totp (
     name TEXT NOT NULL,
     account UUID NOT NULL,
     secret TEXT,
     created_at TIMESTAMP NOT NULL DEFAULT current_timestamp
+    confirmed BOOLEAN DEFAULT false,
 );
 ALTER TABLE arimelody.totp ADD CONSTRAINT totp_pk PRIMARY KEY (account, name);
 

schema_migration/001-pre-versioning.sql

@@ -23,12 +23,12 @@ ALTER TABLE arimelody.privilege ADD CONSTRAINT privilege_pk PRIMARY KEY (account
 -- Invites
 CREATE TABLE arimelody.invite (
     code text NOT NULL, 
-    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, 
+    created_at TIMESTAMP NOT NULL DEFAULT current_timestamp, 
     expires_at TIMESTAMP NOT NULL
 );
 ALTER TABLE arimelody.invite ADD CONSTRAINT invite_pk PRIMARY KEY (code);
 
--- Session
+-- Sessions
 CREATE TABLE arimelody.session (
     token TEXT,
     user_agent TEXT NOT NULL,
@@ -40,12 +40,13 @@ CREATE TABLE arimelody.session (
 );
 ALTER TABLE arimelody.session ADD CONSTRAINT session_pk PRIMARY KEY (token);
 
--- TOTPs
+-- TOTP methods
 CREATE TABLE arimelody.totp (
     name TEXT NOT NULL,
     account UUID NOT NULL,
     secret TEXT,
     created_at TIMESTAMP NOT NULL DEFAULT current_timestamp
+    confirmed BOOLEAN DEFAULT false,
 );
 ALTER TABLE arimelody.totp ADD CONSTRAINT totp_pk PRIMARY KEY (account, name);