Adding Live 2fa capacity

* Some requests may be sensitive enough to require a secondary
two-factor authorization on the spot
* Examples: changing password, changing email address, viewing API
tokens etc
* This creates a core handler that can attach to any Auth-able method
which will require a "twofactorcode" query param before processing

controllers/core/core.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
@@ -93,11 +94,36 @@ func UserLogin() gin.HandlerFunc {
 			return
 		}
 
+		if loginVals.TwoFactor != "" && !checkTwoFactorNotReused(&u.Auth, loginVals.TwoFactor) {
+			fmt.Printf("WARNING: two factor code %v reused for %v\n", loginVals.TwoFactor, u.Uid)
+			c.AbortWithStatusJSON(http.StatusUnauthorized, util.FailMsg{Reason: "2fa reused"})
+			return
+		}
+
 		jwt, maxAge := u.GetJwt()
 		c.SetCookie(JwtHeader, jwt, maxAge, "/v1/sec/", "", true, true)
 	}
 }
 
+func checkTwoFactorNotReused(a *models.Auth, tfCode string) bool {
+	var count int64
+	database.Db.Model(&models.TotpUsage{}).Where("login_uid = ? AND code = ?", a.Uid, tfCode).Count(&count)
+
+	if count > 0 {
+		// We found a token, should reject
+		return false
+	}
+
+	used := models.TotpUsage{
+		LoginUid: a.Uid,
+		Code:     tfCode,
+		Used:     time.Now(),
+	}
+	go database.Db.Create(&used)
+
+	return true
+}
+
 func UserVerify() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		verifyJwt, _ := c.GetQuery("verify")
@@ -239,6 +265,12 @@ func AdminLogin() gin.HandlerFunc {
 			return
 		}
 
+		if loginVals.TwoFactor != "" && !checkTwoFactorNotReused(&a.Auth, loginVals.TwoFactor) {
+			fmt.Printf("WARNING: two factor code %v reused for admin %v\n", loginVals.TwoFactor, a.Uid)
+			c.AbortWithStatusJSON(http.StatusUnauthorized, util.FailMsg{Reason: "2fa reused"})
+			return
+		}
+
 		jwt, maxAge := a.GetJwt()
 		c.SetCookie(JwtHeader, jwt, maxAge, "/v1/adm", "", true, true)
 	}
@@ -285,6 +317,59 @@ func AdminAuth() gin.HandlerFunc {
 	return genericAuth("admin", models.AdminHmac)
 }
 
+// A handler to attach to any method which requires a two-factor check
+// at the time of calling. An example of this might be: changing email,
+// changing password, or other high-sensitivity actions that warrant
+// an extra 2FA check.
+func LiveTwoFactor() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		fmt.Println("Required live 2fa")
+		pif, exists := c.Get("principal")
+		p := pif.(util.PrincipalInfo)
+		if !exists {
+			c.AbortWithStatus(http.StatusUnauthorized)
+			return
+		}
+
+		var a models.Auth
+		fmt.Println(p)
+		if p.Role == "user" {
+			u := models.User{}
+			if err := database.Db.Find(&u, "uid = ?", p.Uid).Error; err != nil {
+				c.AbortWithStatus(http.StatusUnauthorized)
+				return
+			}
+			a = u.Auth
+		} else if p.Role == "admin" {
+			adm := models.Admin{}
+			if err := database.Db.Find(&adm, "uid = ?", p.Uid).Error; err != nil {
+				c.AbortWithStatus(http.StatusUnauthorized)
+				return
+			}
+			a = adm.Auth
+		}
+
+		if a.TwoFactorSecret != "" {
+			tfCode, exists := c.GetQuery("twofactorcode")
+			if !exists || len(tfCode) != 6 {
+				c.AbortWithStatusJSON(http.StatusUnauthorized, util.FailMsg{Reason: "2fa required"})
+				return
+			}
+
+			if err := a.ValidateTwoFactor(tfCode, time.Now()); err != nil {
+				c.AbortWithStatusJSON(http.StatusUnauthorized, util.FailMsg{Reason: err.Error()})
+				return
+			}
+
+			if !checkTwoFactorNotReused(&a, tfCode) {
+				c.AbortWithStatusJSON(http.StatusUnauthorized, util.FailMsg{Reason: "2fa reused"})
+				return
+			}
+		}
+
+	}
+}
+
 func Doot() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		piCtx, exists := c.Get("principal")

main.go

@@ -15,6 +15,7 @@ import (
 func Migrate(g *gorm.DB) {
 	g.AutoMigrate(&models.User{})
 	g.AutoMigrate(&models.Admin{})
+	g.AutoMigrate(&models.TotpUsage{})
 }
 
 func main() {
@@ -36,6 +37,7 @@ func main() {
 	v1Sec := v1.Group("/sec", core.UserAuth())
 
 	v1Sec.GET("/doot", core.Doot())
+	v1Sec.GET("/2fa-doot", core.LiveTwoFactor(), core.Doot())
 
 	// Administrative login
 	v1.POST("/admin", core.AdminLogin())

models/auth.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/pquerna/otp/totp"
 	"golang.org/x/crypto/bcrypt"
 )
@@ -55,6 +56,12 @@ func (a *Auth) CheckPassword(pass string) error {
 	return bcrypt.CompareHashAndPassword([]byte(a.PasswordHash), []byte(pass))
 }
 
+type TotpUsage struct {
+	LoginUid uuid.UUID `gorm:"index"`
+	Used     time.Time
+	Code     string `gorm:"index"`
+}
+
 func (a *Auth) ValidateTwoFactor(tfCode string, stamp time.Time) error {
 	if tfCode == "" && a.TwoFactorSecret != "" {
 		return errors.New("requires 2FA")

models/auth_test.go

@@ -3,6 +3,8 @@ package models
 import (
 	"testing"
 	"time"
+
+	"github.com/yxzzy-wtf/gin-gonic-prepack/database"
 )
 
 func TestBadPasswords(t *testing.T) {
@@ -75,6 +77,8 @@ func TestTwoFactorWhenNotSet(t *testing.T) {
 }
 
 func TestTwoFactor(t *testing.T) {
+	database.InitTestDb()
+
 	a := Auth{}
 	a.TwoFactorSecret = "AAAAAAAAAAAAAAAA"