security: remediate vulnerabilities and issues from codebase audit

backend/handlers/auth.go

@@ -4,6 +4,9 @@ import (
 	"crypto/rand"
 	"encoding/base64"
 	"fmt"
+	"strings"
+	"crypto/subtle"
+	"sync"
 
 	"git.elijahkuntz.com/Elijah/drive/config"
 	"git.elijahkuntz.com/Elijah/drive/database"
@@ -17,6 +20,18 @@ type AuthHandler struct {
 	DB     *database.DB
 	Config *config.Config
 	Guard  *middleware.BruteForceGuard
+	
+	tempTOTP map[string]string
+	totpMu   sync.Mutex
+}
+
+func NewAuthHandler(db *database.DB, cfg *config.Config, guard *middleware.BruteForceGuard) *AuthHandler {
+	return &AuthHandler{
+		DB:       db,
+		Config:   cfg,
+		Guard:    guard,
+		tempTOTP: make(map[string]string),
+	}
 }
 
 type SetupRequest struct {
@@ -173,6 +188,26 @@ func (h *AuthHandler) Refresh(c *fiber.Ctx) error {
 	})
 }
 
+// GetDownloadToken issues a short-lived token for file downloads
+// GET /api/auth/download-token
+func (h *AuthHandler) GetDownloadToken(c *fiber.Ctx) error {
+	userID := c.Locals("userID")
+	username := c.Locals("username")
+
+	token, err := middleware.GenerateDownloadToken(
+		int(userID.(float64)),
+		username.(string),
+		h.Config.JWTSecret,
+	)
+	if err != nil {
+		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate download token"})
+	}
+
+	return c.JSON(fiber.Map{
+		"token": token,
+	})
+}
+
 // Enable2FA generates a TOTP secret and returns the provisioning URI.
 // POST /api/auth/2fa/enable
 func (h *AuthHandler) Enable2FA(c *fiber.Ctx) error {
@@ -186,11 +221,10 @@ func (h *AuthHandler) Enable2FA(c *fiber.Ctx) error {
 		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate TOTP secret"})
 	}
 
-	// Store the secret but don't enable 2FA until confirmed
-	_, err = h.DB.Exec("UPDATE users SET totp_secret = ? WHERE username = ?", key.Secret(), username)
-	if err != nil {
-		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to save TOTP secret"})
-	}
+	// Store the secret temporarily in memory
+	h.totpMu.Lock()
+	h.tempTOTP[username] = key.Secret()
+	h.totpMu.Unlock()
 
 	return c.JSON(fiber.Map{
 		"secret": key.Secret(),
@@ -210,21 +244,27 @@ func (h *AuthHandler) Confirm2FA(c *fiber.Ctx) error {
 		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request"})
 	}
 
-	var secret string
-	err := h.DB.QueryRow("SELECT totp_secret FROM users WHERE username = ?", username).Scan(&secret)
-	if err != nil {
-		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
+	h.totpMu.Lock()
+	secret, exists := h.tempTOTP[username]
+	h.totpMu.Unlock()
+
+	if !exists {
+		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "no pending 2FA setup found"})
 	}
 
 	if !totp.Validate(body.Code, secret) {
 		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid TOTP code"})
 	}
 
-	_, err = h.DB.Exec("UPDATE users SET totp_enabled = 1 WHERE username = ?", username)
+	_, err := h.DB.Exec("UPDATE users SET totp_enabled = 1, totp_secret = ? WHERE username = ?", secret, username)
 	if err != nil {
 		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to enable 2FA"})
 	}
 
+	h.totpMu.Lock()
+	delete(h.tempTOTP, username)
+	h.totpMu.Unlock()
+
 	h.DB.AddAuditLog("2fa_enabled", fmt.Sprintf("User '%s' enabled 2FA", username), c.IP())
 
 	return c.JSON(fiber.Map{"message": "2FA enabled successfully"})
@@ -235,7 +275,25 @@ func (h *AuthHandler) Confirm2FA(c *fiber.Ctx) error {
 func (h *AuthHandler) Disable2FA(c *fiber.Ctx) error {
 	username := c.Locals("username").(string)
 
-	_, err := h.DB.Exec("UPDATE users SET totp_enabled = 0, totp_secret = '' WHERE username = ?", username)
+	var body struct {
+		Code string `json:"code"`
+	}
+	if err := c.BodyParser(&body); err != nil {
+		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request"})
+	}
+
+	var secret string
+	err := h.DB.QueryRow("SELECT totp_secret FROM users WHERE username = ?", username).Scan(&secret)
+	if err != nil {
+		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
+	}
+
+	if !totp.Validate(body.Code, secret) {
+		h.Guard.RecordFailure(c.IP())
+		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid TOTP code"})
+	}
+
+	_, err = h.DB.Exec("UPDATE users SET totp_enabled = 0, totp_secret = '' WHERE username = ?", username)
 	if err != nil {
 		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to disable 2FA"})
 	}
@@ -294,7 +352,8 @@ func hashPassword(password string) (string, error) {
 	}
 
 	hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)
-	encoded := fmt.Sprintf("%s$%s",
+	// Output PHC format
+	encoded := fmt.Sprintf("$argon2id$v=19$m=65536,t=2,p=4$%s$%s",
 		base64.RawStdEncoding.EncodeToString(salt),
 		base64.RawStdEncoding.EncodeToString(hash),
 	)
@@ -302,7 +361,23 @@ func hashPassword(password string) (string, error) {
 }
 
 func verifyPassword(password, encoded string) bool {
-	parts := splitN(encoded, "$", 2)
+	if strings.HasPrefix(encoded, "$argon2id$") {
+		// PHC Format
+		parts := strings.Split(encoded, "$")
+		if len(parts) != 6 {
+			return false
+		}
+		salt, err := base64.RawStdEncoding.DecodeString(parts[4])
+		if err != nil { return false }
+		expectedHash, err := base64.RawStdEncoding.DecodeString(parts[5])
+		if err != nil { return false }
+
+		hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)
+		return subtle.ConstantTimeCompare(hash, expectedHash) == 1
+	}
+
+	// Legacy format: base64(salt)$base64(hash)
+	parts := strings.SplitN(encoded, "$", 2)
 	if len(parts) != 2 {
 		return false
 	}
@@ -318,38 +393,5 @@ func verifyPassword(password, encoded string) bool {
 	}
 
 	hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)
-
-	if len(hash) != len(expectedHash) {
-		return false
-	}
-
-	// Constant-time comparison to prevent timing attacks
-	var diff byte
-	for i := range hash {
-		diff |= hash[i] ^ expectedHash[i]
-	}
-	return diff == 0
-}
-
-func splitN(s, sep string, n int) []string {
-	result := make([]string, 0, n)
-	for i := 0; i < n-1; i++ {
-		idx := indexOf(s, sep)
-		if idx < 0 {
-			break
-		}
-		result = append(result, s[:idx])
-		s = s[idx+len(sep):]
-	}
-	result = append(result, s)
-	return result
-}
-
-func indexOf(s, sub string) int {
-	for i := 0; i <= len(s)-len(sub); i++ {
-		if s[i:i+len(sub)] == sub {
-			return i
-		}
-	}
-	return -1
+	return subtle.ConstantTimeCompare(hash, expectedHash) == 1
 }