add rate limiting and 2fa hardening

server/ratelimit.go

@@ -0,0 +1,101 @@
+package server
+
+import (
+	"log"
+	"net"
+	"net/http"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+type rateLimitState struct {
+	count        int
+	resetAt      time.Time
+	blockedUntil time.Time
+}
+
+type rateLimiter struct {
+	mu          sync.Mutex
+	states      map[string]*rateLimitState
+	maxRequests int
+	window      time.Duration
+	blockFor    time.Duration
+}
+
+func newRateLimiter(maxRequests int, window, blockFor time.Duration) *rateLimiter {
+	return &rateLimiter{
+		states:      make(map[string]*rateLimitState),
+		maxRequests: maxRequests,
+		window:      window,
+		blockFor:    blockFor,
+	}
+}
+
+func (limiter *rateLimiter) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !limiter.allow(r) {
+			w.Header().Set("Retry-After", strconvSeconds(limiter.blockFor))
+			http.Error(w, "Too many requests", http.StatusTooManyRequests)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
+func (limiter *rateLimiter) allow(r *http.Request) bool {
+	key := clientIP(r) + ":" + r.URL.Path
+	now := time.Now()
+
+	limiter.mu.Lock()
+	defer limiter.mu.Unlock()
+
+	state, ok := limiter.states[key]
+	if ok && now.Before(state.blockedUntil) {
+		return false
+	}
+
+	if !ok || now.After(state.resetAt) {
+		limiter.states[key] = &rateLimitState{
+			count:   1,
+			resetAt: now.Add(limiter.window),
+		}
+		return true
+	}
+
+	state.count++
+	if state.count > limiter.maxRequests {
+		state.blockedUntil = now.Add(limiter.blockFor)
+		state.resetAt = now.Add(limiter.window)
+		state.count = 0
+		log.Printf("Rate limit triggered for %s on %s", clientIP(r), r.URL.Path)
+		return false
+	}
+
+	return true
+}
+
+func clientIP(r *http.Request) string {
+	if forwardedFor := r.Header.Get("X-Forwarded-For"); forwardedFor != "" {
+		parts := strings.Split(forwardedFor, ",")
+		if ip := strings.TrimSpace(parts[0]); ip != "" {
+			return ip
+		}
+	}
+	if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); realIP != "" {
+		return realIP
+	}
+	if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
+		return host
+	}
+	return r.RemoteAddr
+}
+
+func strconvSeconds(duration time.Duration) string {
+	seconds := int(duration.Seconds())
+	if seconds < 1 {
+		seconds = 1
+	}
+	return strconv.Itoa(seconds)
+}

server/server.go

@@ -10,6 +10,7 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"time"
 )
 
 type Server struct {
@@ -82,20 +83,23 @@ func (this *Server) Run() {
 	//
 	// API
 	//
-	mux.HandleFunc("/api/login", handlers.APILogin)
-	mux.HandleFunc("/api/login/2fa", handlers.APILoginTwoFactor)
-	mux.HandleFunc("/api/refresh", handlers.RefreshToken)
+	loginLimiter := newRateLimiter(10, time.Minute, 5*time.Minute)
+	accountLimiter := newRateLimiter(20, time.Minute, 2*time.Minute)
+
+	mux.Handle("/api/login", loginLimiter.Middleware(http.HandlerFunc(handlers.APILogin)))
+	mux.Handle("/api/login/2fa", loginLimiter.Middleware(http.HandlerFunc(handlers.APILoginTwoFactor)))
+	mux.Handle("/api/refresh", loginLimiter.Middleware(http.HandlerFunc(handlers.RefreshToken)))
 	mux.Handle("/api/logout", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.Logout)))
 	mux.Handle("/api/profile", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.UserInfo)))
-	mux.Handle("/api/2fa/setup", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorSetup)))
-	mux.Handle("/api/2fa/enable", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorEnable)))
-	mux.Handle("/api/2fa/disable", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorDisable)))
-	mux.Handle("/api/2fa/recovery-codes/regenerate", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorRegenerateRecoveryCodes)))
+	mux.Handle("/api/2fa/setup", accountLimiter.Middleware(auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorSetup))))
+	mux.Handle("/api/2fa/enable", loginLimiter.Middleware(auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorEnable))))
+	mux.Handle("/api/2fa/disable", loginLimiter.Middleware(auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorDisable))))
+	mux.Handle("/api/2fa/recovery-codes/regenerate", loginLimiter.Middleware(auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.TwoFactorRegenerateRecoveryCodes))))
 	mux.Handle("/api/userinfo", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.UserInfo)))
-	mux.Handle("/api/account/username", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.AccountUpdateUsername)))
-	mux.Handle("/api/account/password", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.AccountUpdatePassword)))
+	mux.Handle("/api/account/username", accountLimiter.Middleware(auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.AccountUpdateUsername))))
+	mux.Handle("/api/account/password", loginLimiter.Middleware(auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.AccountUpdatePassword))))
 	if this.AllowRegistration {
-		mux.HandleFunc("/api/register", handlers.APIRegister)
+		mux.Handle("/api/register", loginLimiter.Middleware(http.HandlerFunc(handlers.APIRegister)))
 	}
 
 	mux.Handle("/api/item", auth.AuthMiddleware(this.JWTSecret)(http.HandlerFunc(handlers.Item)))