Bookra/apps/backend/internal/httpx/middleware.go

package httpx

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

func SecurityHeaders() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Header("X-Content-Type-Options", "nosniff")
		c.Header("X-Frame-Options", "DENY")
		c.Header("Referrer-Policy", "strict-origin-when-cross-origin")
		c.Header("Content-Security-Policy", "default-src 'self'; connect-src 'self' https:; img-src 'self' data: https:; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; font-src 'self' https://fonts.gstatic.com; script-src 'self'; base-uri 'self'; form-action 'self'")
		c.Next()
	}
}

type visitor struct {
	limiter  *rate.Limiter
	lastSeen time.Time
}

type RateLimiter struct {
	mu       sync.Mutex
	visitors map[string]*visitor
	limit    rate.Limit
	burst    int
}

func NewRateLimiter(limit rate.Limit, burst int) *RateLimiter {
	return &RateLimiter{
		visitors: make(map[string]*visitor),
		limit:    limit,
		burst:    burst,
	}
}

func (r *RateLimiter) Middleware() gin.HandlerFunc {
	go r.cleanupLoop()

	return func(c *gin.Context) {
		ip := c.ClientIP()
		if ip == "" {
			ip = "unknown"
		}
		limiter := r.getVisitor(ip)
		if !limiter.Allow() {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "rate_limited"})
			return
		}
		c.Next()
	}
}

func (r *RateLimiter) getVisitor(key string) *rate.Limiter {
	r.mu.Lock()
	defer r.mu.Unlock()

	entry, ok := r.visitors[key]
	if !ok {
		entry = &visitor{
			limiter:  rate.NewLimiter(r.limit, r.burst),
			lastSeen: time.Now(),
		}
		r.visitors[key] = entry
		return entry.limiter
	}

	entry.lastSeen = time.Now()
	return entry.limiter
}

func (r *RateLimiter) cleanupLoop() {
	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		r.mu.Lock()
		cutoff := time.Now().Add(-15 * time.Minute)
		for key, entry := range r.visitors {
			if entry.lastSeen.Before(cutoff) {
				delete(r.visitors, key)
			}
		}
		r.mu.Unlock()
	}
}