lavender/server/otp.go

package server

import (
	"bytes"
	"encoding/base64"
	auth2 "github.com/1f349/lavender/auth"
	"github.com/1f349/lavender/database"
	"github.com/1f349/lavender/pages"
	"github.com/julienschmidt/httprouter"
	"github.com/skip2/go-qrcode"
	"github.com/xlzd/gotp"
	"html/template"
	"image/png"
	"net/http"
	"time"
)

func (h *httpServer) editOtpPost(rw http.ResponseWriter, req *http.Request, _ httprouter.Params, auth auth2.UserAuth) {
	if req.Method == http.MethodPost && req.FormValue("remove") == "1" {
		if !req.Form.Has("code") {
			// render page
			pages.RenderPageTemplate(rw, "remove-otp", map[string]any{
				"ServiceName": h.conf.ServiceName,
			})
			return
		}

		otpInput := req.Form.Get("code")
		err := h.authOtp.VerifyOtpCode(req.Context(), auth.Subject, otpInput)
		if err != nil {
			http.Error(rw, "Invalid OTP code", http.StatusBadRequest)
			return
		}

		if h.DbTx(rw, func(tx *database.Queries) error {
			return tx.DeleteOtp(req.Context(), auth.Subject)
		}) {
			return
		}

		http.Redirect(rw, req, "/", http.StatusFound)
		return
	}

	var digits int
	switch req.FormValue("digits") {
	case "6":
		digits = 6
	case "7":
		digits = 7
	case "8":
		digits = 8
	default:
		http.Error(rw, "400 Bad Request: Invalid number of digits for OTP code", http.StatusBadRequest)
		return
	}

	secret := req.FormValue("secret")
	if !gotp.IsSecretValid(secret) {
		http.Error(rw, "400 Bad Request: Invalid secret", http.StatusBadRequest)
		return
	}

	if secret == "" {
		// get user email
		var email string
		if h.DbTx(rw, func(tx *database.Queries) error {
			var err error
			email, err = tx.GetUserEmail(req.Context(), auth.Subject)
			return err
		}) {
			return
		}

		secret = gotp.RandomSecret(64)
		if secret == "" {
			http.Error(rw, "500 Internal Server Error: failed to generate OTP secret", http.StatusInternalServerError)
			return
		}
		totp := gotp.NewTOTP(secret, digits, 30, nil)
		otpUri := totp.ProvisioningUri(email, h.conf.OtpIssuer)
		code, err := qrcode.New(otpUri, qrcode.Medium)
		if err != nil {
			http.Error(rw, "500 Internal Server Error: failed to generate QR code", http.StatusInternalServerError)
			return
		}
		qrImg := code.Image(60 * 4)
		qrBounds := qrImg.Bounds()
		qrWidth := qrBounds.Dx()

		qrBuf := new(bytes.Buffer)
		if png.Encode(qrBuf, qrImg) != nil {
			http.Error(rw, "500 Internal Server Error: failed to generate PNG image of QR code", http.StatusInternalServerError)
			return
		}

		// render page
		pages.RenderPageTemplate(rw, "edit-otp", map[string]any{
			"ServiceName": h.conf.ServiceName,
			"OtpQr":       template.URL("data:qrImg/png;base64," + base64.StdEncoding.EncodeToString(qrBuf.Bytes())),
			"QrWidth":     qrWidth,
			"OtpUrl":      otpUri,
			"OtpSecret":   secret,
			"OtpDigits":   digits,
		})
		return
	}

	totp := gotp.NewTOTP(secret, digits, 30, nil)

	if !verifyTotp(totp, req.FormValue("code")) {
		http.Error(rw, "400 Bad Request: invalid OTP code", http.StatusBadRequest)
		return
	}

	if h.DbTx(rw, func(tx *database.Queries) error {
		return tx.SetOtp(req.Context(), database.SetOtpParams{
			Subject:   auth.Subject,
			OtpSecret: secret,
			OtpDigits: int64(digits),
		})
	}) {
		return
	}

	http.Redirect(rw, req, "/", http.StatusFound)
}

func verifyTotp(totp *gotp.TOTP, code string) bool {
	t := time.Now()
	if totp.VerifyTime(code, t) {
		return true
	}
	if totp.VerifyTime(code, t.Add(-30*time.Second)) {
		return true
	}
	return totp.VerifyTime(code, t.Add(30*time.Second))
}