cv-site/internal/middleware/csrf.go

package middleware

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
)

const (
	csrfTokenLength = 32
	csrfCookieName  = "csrf_token"
	csrfFormField   = "csrf_token"
	csrfTokenTTL    = 24 * time.Hour
)

// csrfTokenEntry stores token with expiration
type csrfTokenEntry struct {
	token      string
	expiresAt  time.Time
}

// CSRFProtection provides CSRF token generation and validation
type CSRFProtection struct {
	mu     sync.RWMutex
	tokens map[string]*csrfTokenEntry // map[token]entry
}

// NewCSRFProtection creates a new CSRF protection instance
func NewCSRFProtection() *CSRFProtection {
	csrf := &CSRFProtection{
		tokens: make(map[string]*csrfTokenEntry),
	}

	// Cleanup expired tokens every hour
	go csrf.cleanup()

	return csrf
}

// Middleware provides CSRF protection for state-changing operations
// GET requests: Generate and set CSRF token
// POST/PUT/DELETE: Validate CSRF token
func (c *CSRFProtection) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Only validate on state-changing methods
		if r.Method == http.MethodPost || r.Method == http.MethodPut || r.Method == http.MethodDelete {
			if !c.validateToken(r) {
				log.Printf("SECURITY: CSRF validation failed from IP %s", getClientIP(r))

				// Check if HTMX request
				isHTMX := r.Header.Get("HX-Request") != ""

				if isHTMX {
					w.Header().Set("Content-Type", "text/html; charset=utf-8")
					w.WriteHeader(http.StatusForbidden)
					_, _ = w.Write([]byte(`<div class="alert alert-error">
						<h3>Security Error</h3>
						<p>Invalid security token. Please refresh the page and try again.</p>
					</div>`))
				} else {
					http.Error(w, "CSRF validation failed", http.StatusForbidden)
				}
				return
			}
		}

		next.ServeHTTP(w, r)
	})
}

// generateToken creates a new CSRF token
func (c *CSRFProtection) generateToken() (string, error) {
	bytes := make([]byte, csrfTokenLength)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}

	token := base64.URLEncoding.EncodeToString(bytes)

	// Store token with expiration
	c.mu.Lock()
	c.tokens[token] = &csrfTokenEntry{
		token:     token,
		expiresAt: time.Now().Add(csrfTokenTTL),
	}
	c.mu.Unlock()

	return token, nil
}

// GetToken retrieves or generates a CSRF token for the request
// This should be called when rendering forms
func (c *CSRFProtection) GetToken(w http.ResponseWriter, r *http.Request) (string, error) {
	// Check if token exists in cookie
	cookie, err := r.Cookie(csrfCookieName)
	if err == nil && cookie.Value != "" {
		// Validate existing token
		c.mu.RLock()
		entry, exists := c.tokens[cookie.Value]
		c.mu.RUnlock()

		if exists && time.Now().Before(entry.expiresAt) {
			// Token is valid, return it
			return cookie.Value, nil
		}
	}

	// Generate new token
	token, err := c.generateToken()
	if err != nil {
		return "", fmt.Errorf("failed to generate CSRF token: %w", err)
	}

	// Set cookie
	http.SetCookie(w, &http.Cookie{
		Name:     csrfCookieName,
		Value:    token,
		Path:     "/",
		HttpOnly: true,
		Secure:   r.TLS != nil, // Only set Secure flag if using HTTPS
		SameSite: http.SameSiteStrictMode,
		MaxAge:   int(csrfTokenTTL.Seconds()),
	})

	return token, nil
}

// validateToken validates the CSRF token from the request
func (c *CSRFProtection) validateToken(r *http.Request) bool {
	// Get token from form
	var formToken string

	// Try form value first
	if err := r.ParseForm(); err == nil {
		formToken = r.FormValue(csrfFormField)
	}

	// If not in form, try header (for AJAX requests)
	if formToken == "" {
		formToken = r.Header.Get("X-CSRF-Token")
	}

	if formToken == "" {
		log.Printf("CSRF: No token in request")
		return false
	}

	// Get token from cookie
	cookie, err := r.Cookie(csrfCookieName)
	if err != nil || cookie.Value == "" {
		log.Printf("CSRF: No token in cookie")
		return false
	}

	// Tokens must match
	if formToken != cookie.Value {
		log.Printf("CSRF: Token mismatch")
		return false
	}

	// Validate token exists and is not expired
	c.mu.RLock()
	entry, exists := c.tokens[formToken]
	c.mu.RUnlock()

	if !exists {
		log.Printf("CSRF: Token not found in store")
		return false
	}

	if time.Now().After(entry.expiresAt) {
		log.Printf("CSRF: Token expired")
		return false
	}

	return true
}

// cleanup removes expired tokens periodically
func (c *CSRFProtection) cleanup() {
	ticker := time.NewTicker(1 * time.Hour)
	defer ticker.Stop()

	for range ticker.C {
		c.mu.Lock()
		now := time.Now()
		for token, entry := range c.tokens {
			if now.After(entry.expiresAt) {
				delete(c.tokens, token)
			}
		}
		c.mu.Unlock()
	}
}

// Note: getClientIP is defined in security_logger.go