From b91b6e7ce01c3c9f23847fbe342dbef89392f7b9 Mon Sep 17 00:00:00 2001 From: ari melody Date: Sun, 26 Jan 2025 20:37:20 +0000 Subject: [PATCH] polished up TOTP enrolment --- admin/accounthttp.go | 57 ++++++++++++++----------- admin/views/totp-confirm.html | 7 +++ controller/totp.go | 16 ++++++- main.go | 17 ++++++++ model/totp.go | 1 + schema_migration/000-init.sql | 7 +-- schema_migration/001-pre-versioning.sql | 7 +-- 7 files changed, 81 insertions(+), 31 deletions(-) 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}}
+ {{if .QRBase64Image}}

@@ -29,6 +30,12 @@ code {

If the QR code does not work, you may also enter this secret code:

+ {{else}} +

+ Paste the below secret code into your authentication app or password manager, + then enter your 2FA code below: +

+ {{end}}

{{.TOTP.Secret}}

diff --git a/controller/totp.go b/controller/totp.go index dbbeec7..88f6bc3 100644 --- a/controller/totp.go +++ b/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 +} diff --git a/main.go b/main.go index 3b25ac4..7e8e06f 100644 --- a/main.go +++ b/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 :\n\tLists an account's TOTP methods.\n" + "deleteTOTP :\n\tDeletes an account's TOTP method.\n" + "testTOTP :\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) diff --git a/model/totp.go b/model/totp.go index 8d8422f..cfad10a 100644 --- a/model/totp.go +++ b/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"` } diff --git a/schema_migration/000-init.sql b/schema_migration/000-init.sql index ff5c1af..8751180 100644 --- a/schema_migration/000-init.sql +++ b/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); diff --git a/schema_migration/001-pre-versioning.sql b/schema_migration/001-pre-versioning.sql index cd0c061..37de432 100644 --- a/schema_migration/001-pre-versioning.sql +++ b/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);