codevault/internal/api/middleware.go

package api

import (
	"context"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/metric"
	"go.uber.org/zap"

	"github.com/kanootoko/codevault/internal/db"
	"github.com/kanootoko/codevault/internal/observability/logging"
	"github.com/kanootoko/codevault/internal/observability/metrics"
	"github.com/kanootoko/codevault/internal/ratelimit"
)

type contextKey string

const loginContextKey contextKey = "login"
const ipContextKey contextKey = "ip"

// LoginFromContext returns the Basic Auth login stashed by BasicAuthMiddleware,
// if any. Strict-server handlers only see context.Context, not *gin.Context, so
// this is how they recover the authenticated login.
func LoginFromContext(ctx context.Context) (string, bool) {
	login, ok := ctx.Value(loginContextKey).(string)
	return login, ok
}

// IPFromContext returns the client IP stashed by InstrumentMiddleware, if any.
// Strict-server handlers only see context.Context, not *gin.Context, so this is
// how they recover it.
func IPFromContext(ctx context.Context) (string, bool) {
	ip, ok := ctx.Value(ipContextKey).(string)
	return ip, ok
}

// BasicAuthMiddleware checks the request's HTTP Basic Auth credentials against
// the DB-backed user store. It shares the same per-IP failure limiter as
// RateLimitMiddleware, so repeated bad credentials get throttled the same way
// repeated failed lookups do, instead of the caller being able to hammer the
// user store indefinitely.
func BasicAuthMiddleware(users *db.UsersRepository, limiter *ratelimit.Limiter, vm metrics.ValueMetrics) gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		if !limiter.Allowed(ip) {
			limiter.RecordFailure(ip)
			vm.SetFailedTotal.Add(c.Request.Context(), 1, metric.WithAttributes(attribute.String("reason", "rate_limited")))
			c.Header("WWW-Authenticate", `Basic realm="codevault"`)
			c.AbortWithStatus(http.StatusUnauthorized)
			return
		}

		login, password, ok := c.Request.BasicAuth()
		if !ok || users.Verify(c.Request.Context(), login, password) != nil {
			limiter.RecordFailure(ip)
			vm.SetFailedTotal.Add(c.Request.Context(), 1, metric.WithAttributes(attribute.String("reason", "unauthorized")))
			c.Header("WWW-Authenticate", `Basic realm="codevault"`)
			c.AbortWithStatus(http.StatusUnauthorized)
			return
		}
		c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), loginContextKey, login))
		c.Next()
	}
}

// RateLimitMiddleware short-circuits to 404 once the client IP has accumulated
// too many failures, without touching the handler (and thus the database). A
// short-circuited request counts as a failure itself, same as a genuine 404 from
// the handler.
func RateLimitMiddleware(limiter *ratelimit.Limiter, vm metrics.ValueMetrics) gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		if !limiter.Allowed(ip) {
			limiter.RecordFailure(ip)
			vm.GetFailedTotal.Add(c.Request.Context(), 1, metric.WithAttributes(attribute.String("reason", "rate_limited")))
			c.AbortWithStatus(http.StatusNotFound)
			return
		}

		c.Next()

		if c.Writer.Status() == http.StatusNotFound {
			limiter.RecordFailure(ip)
		}
	}
}

// InstrumentMiddleware records request count and duration metrics per route. Also injects logger to context.
func InstrumentMiddleware(route string, m metrics.HTTPMetrics, logger *zap.Logger) gin.HandlerFunc {
	return func(c *gin.Context) {
		ctx := context.WithValue(c.Request.Context(), ipContextKey, c.ClientIP())
		ctx = logging.CtxWithLogger(ctx, logger)
		c.Request = c.Request.WithContext(ctx)

		start := time.Now()
		c.Next()
		duration := time.Since(start).Seconds()

		attrs := metric.WithAttributes(
			attribute.String("route", route),
			attribute.Int("status", c.Writer.Status()),
		)
		m.RequestsTotal.Add(c.Request.Context(), 1, attrs)
		m.RequestDuration.Record(c.Request.Context(), duration, metric.WithAttributes(attribute.String("route", route)))
	}
}